"""
Audio Augmentation Module for Music Classification

Project: Music Classification Dataset
Author: Molla Samser
Company: RSK World
Designer & Tester: Rima Khatun
Website: https://rskworld.in
Email: help@rskworld.in, support@rskworld.in
Phone: +91 93305 39277
Description: Advanced audio augmentation techniques for data augmentation
License: Educational Purpose Only
"""

import numpy as np
import librosa
import warnings
warnings.filterwarnings('ignore')


class AudioAugmenter:
    """
    Audio augmentation class for data augmentation
    
    Author: Molla Samser (RSK World)
    Website: https://rskworld.in
    Email: help@rskworld.in
    """
    
    def __init__(self, sample_rate=22050):
        """
        Initialize AudioAugmenter
        
        Args:
            sample_rate: Sample rate of audio files
        """
        self.sample_rate = sample_rate
    
    def add_noise(self, audio: np.ndarray, noise_factor: float = 0.005) -> np.ndarray:
        """
        Add random noise to audio
        
        Args:
            audio: Input audio array
            noise_factor: Amount of noise to add
            
        Returns:
            Augmented audio with noise
        """
        noise = np.random.randn(len(audio))
        augmented_audio = audio + noise_factor * noise
        return augmented_audio.astype(audio.dtype)
    
    def time_stretch(self, audio: np.ndarray, rate: float = 1.0) -> np.ndarray:
        """
        Time stretch audio without changing pitch
        
        Args:
            audio: Input audio array
            rate: Stretch factor (>1.0 speeds up, <1.0 slows down)
            
        Returns:
            Time-stretched audio
        """
        return librosa.effects.time_stretch(audio, rate=rate)
    
    def pitch_shift(self, audio: np.ndarray, n_steps: int = 0) -> np.ndarray:
        """
        Shift pitch of audio
        
        Args:
            audio: Input audio array
            n_steps: Number of semitones to shift
            
        Returns:
            Pitch-shifted audio
        """
        return librosa.effects.pitch_shift(
            audio, 
            sr=self.sample_rate, 
            n_steps=n_steps
        )
    
    def time_shift(self, audio: np.ndarray, shift_max: float = 0.2) -> np.ndarray:
        """
        Shift audio in time
        
        Args:
            audio: Input audio array
            shift_max: Maximum shift as fraction of audio length
            
        Returns:
            Time-shifted audio
        """
        shift = int(np.random.uniform(-shift_max, shift_max) * len(audio))
        return np.roll(audio, shift)
    
    def change_volume(self, audio: np.ndarray, factor: float = 1.0) -> np.ndarray:
        """
        Change volume of audio
        
        Args:
            audio: Input audio array
            factor: Volume factor (>1.0 louder, <1.0 quieter)
            
        Returns:
            Volume-adjusted audio
        """
        return audio * factor
    
    def add_reverb(self, audio: np.ndarray, reverb_factor: float = 0.1) -> np.ndarray:
        """
        Add simple reverb effect
        
        Args:
            audio: Input audio array
            reverb_factor: Amount of reverb
            
        Returns:
            Audio with reverb
        """
        delay = int(0.1 * self.sample_rate)
        reverb = np.zeros(len(audio) + delay)
        reverb[:len(audio)] = audio
        reverb[delay:] += audio * reverb_factor
        return reverb[:len(audio)]
    
    def random_augmentation(self, audio: np.ndarray) -> np.ndarray:
        """
        Apply random augmentation
        
        Args:
            audio: Input audio array
            
        Returns:
            Randomly augmented audio
        """
        augmentation = np.random.choice([
            'noise', 'time_stretch', 'pitch_shift', 
            'time_shift', 'volume', 'reverb'
        ])
        
        if augmentation == 'noise':
            return self.add_noise(audio, np.random.uniform(0.001, 0.01))
        elif augmentation == 'time_stretch':
            return self.time_stretch(audio, np.random.uniform(0.8, 1.2))
        elif augmentation == 'pitch_shift':
            return self.pitch_shift(audio, np.random.randint(-3, 4))
        elif augmentation == 'time_shift':
            return self.time_shift(audio, 0.2)
        elif augmentation == 'volume':
            return self.change_volume(audio, np.random.uniform(0.7, 1.3))
        elif augmentation == 'reverb':
            return self.add_reverb(audio, np.random.uniform(0.05, 0.15))
        
        return audio
    
    def augment_batch(self, audio_list: list, augmentations_per_sample: int = 2) -> list:
        """
        Augment a batch of audio samples
        
        Args:
            audio_list: List of audio arrays
            augmentations_per_sample: Number of augmented versions per sample
            
        Returns:
            List of original + augmented audio samples
        """
        augmented_list = []
        
        for audio in audio_list:
            # Add original
            augmented_list.append(audio)
            
            # Add augmented versions
            for _ in range(augmentations_per_sample):
                aug_audio = self.random_augmentation(audio)
                augmented_list.append(aug_audio)
        
        return augmented_list


# Example usage
if __name__ == "__main__":
    """
    Demo script for AudioAugmenter
    
    Author: Molla Samser
    Company: RSK World
    Website: https://rskworld.in
    """
    augmenter = AudioAugmenter()
    
    # Create dummy audio
    dummy_audio = np.random.randn(44100)
    
    print("Audio Augmentation Demo")
    print("Author: Molla Samser | RSK World")
    print("Website: https://rskworld.in")
    print("\nAugmentation techniques available:")
    print("- Add noise")
    print("- Time stretch")
    print("- Pitch shift")
    print("- Time shift")
    print("- Volume change")
    print("- Reverb")
    
    # Test augmentation
    noisy = augmenter.add_noise(dummy_audio)
    print(f"\nOriginal shape: {dummy_audio.shape}")
    print(f"Augmented shape: {noisy.shape}")
    
    print("\n© 2026 RSK World - Molla Samser")