"""
Baby Cry AI - Advanced Data Augmentation Pipeline
Advanced augmentation with time stretching, pitch shifting, noise addition, etc.
"""

import os
import sys
import numpy as np
import librosa
import soundfile as sf
from pathlib import Path
from collections import Counter
import random
import warnings
warnings.filterwarnings('ignore')

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from audio_processor import AudioProcessor


class AdvancedDataAugmenter:
    """Advanced data augmentation for audio files"""
    
    def __init__(self, data_dir="../data", output_dir=None):
        """
        Initialize augmenter
        
        Args:
            data_dir: Path to data directory
            output_dir: Path to output directory (default: data_dir/../data_augmented)
        """
        self.data_dir = Path(data_dir)
        if output_dir:
            self.output_dir = Path(output_dir)
        else:
            self.output_dir = self.data_dir.parent / "data_augmented"
        
        self.processor = AudioProcessor()
        self.generated_count = 0
    
    def pitch_shift(self, y, sr, semitones):
        """Shift pitch by semitones"""
        try:
            return librosa.effects.pitch_shift(y, sr=sr, n_steps=semitones)
        except Exception:
            return y
    
    def time_stretch(self, y, rate):
        """Time stretch audio (rate > 1 speeds up, < 1 slows down)"""
        try:
            return librosa.effects.time_stretch(y, rate=rate)
        except Exception:
            return y
    
    def add_noise(self, y, snr_db=20):
        """Add white noise at specified SNR (dB)"""
        try:
            # Calculate signal power
            signal_power = np.mean(y ** 2)
            
            # Avoid division by zero
            if signal_power < 1e-10:
                return y
            
            # Calculate noise power for desired SNR
            snr_linear = 10 ** (snr_db / 10)
            noise_power = signal_power / snr_linear
            
            # Generate noise
            noise = np.random.normal(0, np.sqrt(noise_power), len(y))
            
            return y + noise
        except Exception:
            return y
    
    def add_background_noise(self, y, sr, noise_type='white', snr_db=20):
        """Add different types of background noise"""
        try:
            signal_power = np.mean(y ** 2)
            if signal_power < 1e-10:
                return y
            
            snr_linear = 10 ** (snr_db / 10)
            noise_power = signal_power / snr_linear
            
            if noise_type == 'white':
                noise = np.random.normal(0, np.sqrt(noise_power), len(y))
            elif noise_type == 'pink':
                # Pink noise (1/f noise)
                white_noise = np.random.normal(0, 1, len(y))
                # Simple pink noise approximation
                fft = np.fft.rfft(white_noise)
                freqs = np.fft.rfftfreq(len(y), 1/sr)
                # Apply 1/f filter
                fft_filtered = fft / np.sqrt(freqs + 1e-10)
                noise = np.fft.irfft(fft_filtered, len(y))
                # Normalize to desired power
                noise = noise * np.sqrt(noise_power / (np.mean(noise**2) + 1e-10))
            else:
                noise = np.random.normal(0, np.sqrt(noise_power), len(y))
            
            return y + noise
        except Exception:
            return y
    
    def volume_scale(self, y, scale):
        """Scale volume"""
        return y * scale
    
    def time_shift(self, y, shift_samples):
        """Time shift audio (circular shift)"""
        return np.roll(y, shift_samples)
    
    def speed_variation(self, y, speed_factor):
        """Vary speed (changes both pitch and duration)"""
        try:
            # Use time_stretch for speed variation
            return librosa.effects.time_stretch(y, rate=speed_factor)
        except Exception:
            return y
    
    def apply_augmentation(self, y, sr, augmentation_type, **kwargs):
        """Apply specific augmentation"""
        if augmentation_type == 'pitch_shift':
            semitones = kwargs.get('semitones', random.choice([-2, -1, 1, 2]))
            return self.pitch_shift(y, sr, semitones)
        elif augmentation_type == 'time_stretch':
            rate = kwargs.get('rate', random.uniform(0.8, 1.2))
            return self.time_stretch(y, rate)
        elif augmentation_type == 'add_noise':
            snr_db = kwargs.get('snr_db', random.choice([20, 30, 40]))
            return self.add_noise(y, snr_db)
        elif augmentation_type == 'add_background_noise':
            noise_type = kwargs.get('noise_type', 'white')
            snr_db = kwargs.get('snr_db', random.choice([20, 30, 40]))
            return self.add_background_noise(y, sr, noise_type, snr_db)
        elif augmentation_type == 'volume_scale':
            scale = kwargs.get('scale', random.uniform(0.7, 1.3))
            return self.volume_scale(y, scale)
        elif augmentation_type == 'time_shift':
            shift_ratio = kwargs.get('shift_ratio', random.uniform(0.1, 0.3))
            shift_samples = int(len(y) * shift_ratio)
            return self.time_shift(y, shift_samples)
        elif augmentation_type == 'speed_variation':
            speed_factor = kwargs.get('speed_factor', random.uniform(0.9, 1.1))
            return self.speed_variation(y, speed_factor)
        else:
            return y
    
    def augment_file(self, input_path, output_path, num_augmentations=3, 
                    augmentation_types=None):
        """
        Augment a single audio file
        
        Args:
            input_path: Path to input audio file
            output_path: Path to save augmented file
            num_augmentations: Number of augmented versions to create
            augmentation_types: List of augmentation types to use
        """
        if augmentation_types is None:
            augmentation_types = [
                'pitch_shift', 'time_stretch', 'add_noise', 
                'volume_scale', 'time_shift', 'speed_variation'
            ]
        
        try:
            # Load audio
            y, sr = librosa.load(str(input_path), sr=self.processor.sample_rate)
            if y is None or len(y) == 0:
                return 0
            
            # Preprocess
            y, sr = self.processor.preprocess_audio(y, sr)
            
            created = 0
            for i in range(num_augmentations):
                # Apply random augmentation
                aug_type = random.choice(augmentation_types)
                y_aug = self.apply_augmentation(y, sr, aug_type)
                
                # Ensure output directory exists
                output_path.parent.mkdir(parents=True, exist_ok=True)
                
                # Generate unique filename
                base_name = input_path.stem
                aug_filename = f"{base_name}_aug{i+1}_{aug_type}.wav"
                aug_path = output_path.parent / aug_filename
                
                # Save augmented audio
                try:
                    sf.write(str(aug_path), y_aug, sr)
                    created += 1
                    self.generated_count += 1
                except Exception as e:
                    print(f"  ⚠️  Error saving {aug_path}: {e}")
            
            return created
            
        except Exception as e:
            print(f"  ⚠️  Error augmenting {input_path}: {e}")
            return 0
    
    def augment_category(self, category, target_count, max_per_file=5):
        """
        Augment all files in a category to reach target count
        
        Args:
            category: Category name
            target_count: Target number of files
            max_per_file: Maximum augmentations per file
        """
        category_dir = self.data_dir / category
        output_category_dir = self.output_dir / category
        output_category_dir.mkdir(parents=True, exist_ok=True)
        
        if not category_dir.exists():
            print(f"  ⚠️  Category directory not found: {category_dir}")
            return 0
        
        # Get existing files
        existing_files = list(category_dir.glob("*.wav"))
        existing_count = len(existing_files)
        
        print(f"  📁 {category}: {existing_count} existing files")
        
        if existing_count >= target_count:
            print(f"     ✅ Already has enough files")
            return 0
        
        needed = target_count - existing_count
        print(f"     🎯 Need {needed} more files")
        
        # Calculate augmentations per file
        if existing_count > 0:
            augs_per_file = min(max_per_file, (needed // existing_count) + 1)
        else:
            print(f"     ⚠️  No files to augment")
            return 0
        
        created = 0
        for file_path in existing_files:
            if created >= needed:
                break
            
            num_aug = min(augs_per_file, needed - created)
            created += self.augment_file(
                file_path, 
                output_category_dir / file_path.name,
                num_augmentations=num_aug
            )
        
        print(f"     ✅ Created {created} augmented files")
        return created
    
    def augment_dataset(self, target_per_category=500, max_per_file=5):
        """
        Augment entire dataset
        
        Args:
            target_per_category: Target number of files per category
            max_per_file: Maximum augmentations per file
        """
        print("🔄 Starting Advanced Data Augmentation")
        print("=" * 60)
        
        if not self.data_dir.exists():
            print(f"❌ Data directory not found: {self.data_dir}")
            return
        
        # Get categories
        categories = [d for d in os.listdir(self.data_dir) 
                     if os.path.isdir(self.data_dir / d)]
        
        print(f"\n📁 Found {len(categories)} categories: {categories}")
        print(f"🎯 Target: {target_per_category} files per category")
        
        total_created = 0
        
        for category in categories:
            print(f"\n📂 Processing {category}...")
            created = self.augment_category(category, target_per_category, max_per_file)
            total_created += created
        
        print("\n" + "=" * 60)
        print(f"✅ Augmentation complete!")
        print(f"   Total files created: {total_created}")
        print(f"   Output directory: {self.output_dir}")
        print("=" * 60)
        
        return total_created


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='Advanced data augmentation')
    parser.add_argument('--data-dir', type=str, default='../data',
                       help='Path to data directory')
    parser.add_argument('--output-dir', type=str, default=None,
                       help='Path to output directory')
    parser.add_argument('--target', type=int, default=500,
                       help='Target files per category')
    parser.add_argument('--max-per-file', type=int, default=5,
                       help='Maximum augmentations per file')
    
    args = parser.parse_args()
    
    augmenter = AdvancedDataAugmenter(
        data_dir=args.data_dir,
        output_dir=args.output_dir
    )
    
    augmenter.augment_dataset(
        target_per_category=args.target,
        max_per_file=args.max_per_file
    )