Trixy3/trainer/wakeword/data.py

"""
Wakeword-specific data handling and preprocessing.

This module provides specialized data handling for wakeword detection tasks,
including dataset creation, preprocessing, and augmentation specific to
wakeword detection requirements.
"""

import os
import logging
import random
import numpy as np
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Any, Union
from collections import defaultdict

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import torchaudio
import torchaudio.transforms as T

from ..data_pipeline import AudioDataset, AudioProcessor, AudioAugmentation, AudioProcessingConfig


class WakewordDataset(AudioDataset):
    """
    Specialized dataset for wakeword detection.
    
    Extends the base AudioDataset with wakeword-specific functionality
    including class balancing, negative sample generation, and specialized
    augmentation strategies.
    """
    
    def __init__(self, audio_files: List[str], labels: List[int],
                 class_names: Dict[int, str],
                 audio_processor: AudioProcessor,
                 augmentation: Optional[AudioAugmentation] = None,
                 augmentation_config: Optional[Dict[str, float]] = None,
                 apply_augmentation: bool = True,
                 balance_classes: bool = True,
                 negative_augmentation_factor: float = 2.0):
        """
        Initialize wakeword dataset.
        
        Args:
            audio_files: List of audio file paths
            labels: List of corresponding labels
            class_names: Mapping of class IDs to names
            audio_processor: Audio processor instance
            augmentation: Audio augmentation instance
            augmentation_config: Augmentation configuration
            apply_augmentation: Whether to apply augmentation
            balance_classes: Whether to balance classes using sampling
            negative_augmentation_factor: Extra augmentation factor for negative samples
        """
        super().__init__(
            audio_files, labels, audio_processor, augmentation,
            augmentation_config, apply_augmentation
        )
        
        self.class_names = class_names
        self.balance_classes = balance_classes
        self.negative_augmentation_factor = negative_augmentation_factor
        
        # Analyze class distribution
        self.class_distribution = self._analyze_class_distribution()
        
        # Create sampling weights for balanced training
        if balance_classes:
            self.sample_weights = self._calculate_sample_weights()
        else:
            self.sample_weights = None
    
    def _analyze_class_distribution(self) -> Dict[int, int]:
        """Analyze the distribution of classes in the dataset."""
        distribution = defaultdict(int)
        for label in self.labels:
            distribution[label] += 1
        return dict(distribution)
    
    def _calculate_sample_weights(self) -> torch.Tensor:
        """Calculate sample weights for balanced training."""
        class_counts = [self.class_distribution.get(i, 1) for i in range(len(self.class_names))]
        total_samples = sum(class_counts)
        
        # Calculate weights inversely proportional to class frequency
        class_weights = [total_samples / (len(self.class_names) * count) for count in class_counts]
        
        # Assign weights to each sample
        sample_weights = torch.zeros(len(self.labels))
        for i, label in enumerate(self.labels):
            sample_weights[i] = class_weights[label]
        
        return sample_weights
    
    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
        """Get item with wakeword-specific processing."""
        audio_file = self.audio_files[idx]
        label = self.labels[idx]
        class_name = self.class_names.get(label, "unknown")
        
        try:
            # Load and process audio
            waveform, sr = self.audio_processor.load_audio(audio_file)
            
            # Convert to mono
            if waveform.shape[0] > 1:
                waveform = torch.mean(waveform, dim=0, keepdim=True)
            
            # Resample
            waveform = self.audio_processor.resample(waveform, sr)
            
            # Apply augmentation if enabled
            if (self.apply_augmentation and self.augmentation is not None 
                and self.augmentation_config):
                
                # Apply extra augmentation for negative samples
                if class_name in ["negative", "background"]:
                    # More aggressive augmentation for negative samples
                    enhanced_config = self.augmentation_config.copy()
                    for key in enhanced_config:
                        if key.endswith('_factor'):
                            enhanced_config[key] *= self.negative_augmentation_factor
                    
                    waveform = self.augmentation.apply_augmentations(waveform, enhanced_config)
                else:
                    waveform = self.augmentation.apply_augmentations(waveform, self.augmentation_config)
            
            # Normalize length and extract features
            waveform = self.audio_processor.normalize_length(waveform)
            features = self.audio_processor.extract_features(waveform)
            
            # Ensure consistent shape
            if features.dim() == 3:  # (1, freq, time)
                features = features.squeeze(0)  # (freq, time)
            
            # Apply wakeword-specific feature normalization
            features = self._normalize_features(features)
            
            return features, label
            
        except Exception as e:
            logging.warning(f"Error processing {audio_file}: {str(e)}")
            # Return zero tensor and label on error
            dummy_features = torch.zeros(
                self.audio_processor.config.n_mels,
                self.audio_processor.config.target_samples // self.audio_processor.config.hop_length + 1
            )
            return self._normalize_features(dummy_features), label
    
    def _normalize_features(self, features: torch.Tensor) -> torch.Tensor:
        """Apply wakeword-specific feature normalization."""
        # Apply instance normalization to make model more robust
        # to volume variations and recording conditions
        
        # Calculate mean and std along time dimension
        mean = torch.mean(features, dim=-1, keepdim=True)
        std = torch.std(features, dim=-1, keepdim=True) + 1e-8
        
        # Normalize
        features = (features - mean) / std
        
        # Clip extreme values
        features = torch.clamp(features, -3.0, 3.0)
        
        return features
    
    def get_class_distribution(self) -> Dict[str, int]:
        """Get class distribution with class names."""
        return {
            self.class_names[class_id]: count 
            for class_id, count in self.class_distribution.items()
        }


class WakewordDataPreprocessor:
    """
    Specialized preprocessor for wakeword detection data.
    
    Handles the specific requirements mentioned in the project description:
    - Trimming wakeword files to consistent length
    - Splitting negative/background files into chunks
    - Creating balanced train/val/test splits
    """
    
    def __init__(self, audio_config: AudioProcessingConfig,
                 logger: Optional[logging.Logger] = None):
        """Initialize wakeword data preprocessor."""
        self.audio_config = audio_config
        self.logger = logger or logging.getLogger(__name__)
        self.audio_processor = AudioProcessor(audio_config)
    
    def preprocess_wakeword_data(self, raw_data_dir: str, output_dir: str,
                                wakeword_classes: Dict[int, str],
                                max_wakeword_length: Optional[float] = None) -> str:
        """
        Preprocess raw wakeword data according to project requirements.
        
        This method implements the preprocessing pipeline described in the
        project description:
        1. Find max wakeword length
        2. Trim all files to consistent length with silence padding
        3. Split negative/background files into chunks
        4. Create balanced train/val/test splits
        
        Args:
            raw_data_dir: Directory containing raw wakeword data
            output_dir: Output directory for processed data
            wakeword_classes: Mapping of class IDs to names
            max_wakeword_length: Maximum wakeword length (auto-detect if None)
            
        Returns:
            Path to processed data directory
        """
        self.logger.info(f"Preprocessing wakeword data from {raw_data_dir}")
        
        raw_path = Path(raw_data_dir)
        output_path = Path(output_dir)
        
        # Step 1: Analyze wakeword lengths to determine target length
        if max_wakeword_length is None:
            max_wakeword_length = self._find_max_wakeword_length(raw_path, wakeword_classes)
            self.logger.info(f"Auto-detected max wakeword length: {max_wakeword_length:.2f}s")
        
        # Update audio config with detected length
        self.audio_config.target_length = max_wakeword_length
        target_samples = int(self.audio_config.sample_rate * max_wakeword_length)
        
        # Step 2: Process each class
        processed_files_by_class = {}
        
        for class_id, class_name in wakeword_classes.items():
            self.logger.info(f"Processing class: {class_name}")
            
            class_files = self._process_class_files(
                raw_path / class_name, class_name, target_samples
            )
            
            if class_files:
                processed_files_by_class[class_id] = class_files
                self.logger.info(f"Processed {len(class_files)} files for {class_name}")
            else:
                self.logger.warning(f"No files processed for class: {class_name}")
        
        # Step 3: Save processed files
        chunked_dir = output_path / "chunked"
        self._save_processed_files(processed_files_by_class, wakeword_classes, chunked_dir)
        
        # Step 4: Create train/val/test splits
        self._create_balanced_splits(chunked_dir, output_path, wakeword_classes)
        
        self.logger.info(f"Wakeword preprocessing completed: {output_path}")
        return str(output_path)
    
    def _find_max_wakeword_length(self, raw_path: Path, 
                                 wakeword_classes: Dict[int, str]) -> float:
        """Find the maximum wakeword length after trimming silence."""
        max_length = 0.0
        
        # Only analyze actual wakeword classes (not negative/background)
        wakeword_only_classes = [name for name in wakeword_classes.values() 
                               if name not in ["negative", "background"]]
        
        for class_name in wakeword_only_classes:
            class_dir = raw_path / class_name
            if not class_dir.exists():
                continue
            
            # Get audio files
            audio_files = []
            for ext in ['.wav', '.mp3', '.flac']:
                audio_files.extend(list(class_dir.glob(f"**/*{ext}")))
            
            for audio_file in audio_files[:20]:  # Sample first 20 files
                try:
                    waveform, sr = self.audio_processor.load_audio(str(audio_file))
                    
                    # Convert to mono and resample
                    if waveform.shape[0] > 1:
                        waveform = torch.mean(waveform, dim=0, keepdim=True)
                    waveform = self.audio_processor.resample(waveform, sr)
                    
                    # Trim silence
                    trimmed = self.audio_processor.trim_silence(waveform, threshold=0.01)
                    
                    # Calculate length in seconds
                    length_seconds = trimmed.shape[-1] / self.audio_config.sample_rate
                    max_length = max(max_length, length_seconds)
                
                except Exception as e:
                    self.logger.warning(f"Error analyzing {audio_file}: {str(e)}")
                    continue
        
        # Add some padding and ensure minimum length
        max_length = max(max_length * 1.2, 1.0)  # 20% padding, minimum 1 second
        return min(max_length, 3.0)  # Cap at 3 seconds for practical reasons
    
    def _process_class_files(self, class_dir: Path, class_name: str, 
                           target_samples: int) -> List[Tuple[torch.Tensor, str]]:
        """Process all files for a specific class."""
        if not class_dir.exists():
            self.logger.warning(f"Class directory not found: {class_dir}")
            return []
        
        # Get audio files
        audio_files = []
        for ext in ['.wav', '.mp3', '.flac']:
            audio_files.extend(list(class_dir.glob(f"**/*{ext}")))
        
        processed_files = []
        
        for audio_file in audio_files:
            try:
                # Load and process audio
                waveform, sr = self.audio_processor.load_audio(str(audio_file))
                
                # Convert to mono and resample
                if waveform.shape[0] > 1:
                    waveform = torch.mean(waveform, dim=0, keepdim=True)
                waveform = self.audio_processor.resample(waveform, sr)
                
                if class_name in ["negative", "background"]:
                    # Split long files into chunks
                    chunks = self._split_into_chunks(waveform, target_samples)
                    for i, chunk in enumerate(chunks):
                        filename = f"{audio_file.stem}_chunk_{i:03d}.wav"
                        processed_files.append((chunk, filename))
                else:
                    # For wakeword files: trim silence and normalize length
                    # Trim silence first
                    trimmed = self.audio_processor.trim_silence(waveform, threshold=0.01)
                    
                    # Normalize to target length with smart padding
                    normalized = self._smart_normalize_length(trimmed, target_samples)
                    processed_files.append((normalized, audio_file.name))
            
            except Exception as e:
                self.logger.error(f"Error processing {audio_file}: {str(e)}")
                continue
        
        return processed_files
    
    def _smart_normalize_length(self, waveform: torch.Tensor, 
                               target_samples: int) -> torch.Tensor:
        """
        Smart length normalization that preserves wakeword positioning.
        
        For wakewords, we want to:
        1. Center the wakeword in the target duration
        2. Add silence padding around it
        3. Ensure the full wakeword is preserved
        """
        current_samples = waveform.shape[-1]
        
        if current_samples >= target_samples:
            # If too long, crop from center
            start = (current_samples - target_samples) // 2
            return waveform[..., start:start + target_samples]
        else:
            # If too short, pad with silence
            # Add random positioning to create variety
            max_pad_start = target_samples - current_samples
            pad_start = random.randint(0, max_pad_start // 2)  # Bias towards center
            pad_end = target_samples - current_samples - pad_start
            
            return F.pad(waveform, (pad_start, pad_end))
    
    def _split_into_chunks(self, waveform: torch.Tensor, 
                          chunk_size: int, overlap_ratio: float = 0.1) -> List[torch.Tensor]:
        """Split long audio into overlapping chunks."""
        chunks = []
        total_samples = waveform.shape[-1]
        
        if total_samples <= chunk_size:
            # If shorter than chunk size, pad to chunk size
            padded = F.pad(waveform, (0, chunk_size - total_samples))
            chunks.append(padded)
        else:
            # Split into overlapping chunks
            step_size = int(chunk_size * (1 - overlap_ratio))
            start = 0
            
            while start + chunk_size <= total_samples:
                chunk = waveform[..., start:start + chunk_size]
                chunks.append(chunk)
                start += step_size
            
            # Add final chunk if there's remaining audio
            if start < total_samples and total_samples - start > chunk_size // 2:
                final_chunk = waveform[..., -chunk_size:]
                chunks.append(final_chunk)
        
        return chunks
    
    def _save_processed_files(self, processed_files_by_class: Dict[int, List[Tuple[torch.Tensor, str]]],
                             wakeword_classes: Dict[int, str],
                             output_dir: Path):
        """Save processed files to disk."""
        for class_id, files in processed_files_by_class.items():
            class_name = wakeword_classes[class_id]
            class_dir = output_dir / class_name
            class_dir.mkdir(parents=True, exist_ok=True)
            
            for waveform, filename in files:
                output_file = class_dir / filename
                torchaudio.save(
                    str(output_file), 
                    waveform, 
                    self.audio_config.sample_rate
                )
    
    def _create_balanced_splits(self, chunked_dir: Path, output_dir: Path,
                               wakeword_classes: Dict[int, str]):
        """Create balanced train/validation/test splits."""
        # Define split ratios
        train_ratio = 0.7
        val_ratio = 0.2
        test_ratio = 0.1
        
        # Create output directories
        for split in ['train', 'val', 'test']:
            for class_name in wakeword_classes.values():
                (output_dir / split / class_name).mkdir(parents=True, exist_ok=True)
        
        # Process each class
        for class_id, class_name in wakeword_classes.items():
            class_dir = chunked_dir / class_name
            if not class_dir.exists():
                continue
            
            # Get all files for this class
            files = list(class_dir.glob("*.wav"))
            if len(files) < 3:
                self.logger.warning(f"Too few files for {class_name} - copying to train only")
                # Copy all to train
                for file in files:
                    import shutil
                    shutil.copy2(file, output_dir / "train" / class_name / file.name)
                continue
            
            # Shuffle files
            random.shuffle(files)
            
            # Calculate split indices
            total_files = len(files)
            train_end = int(total_files * train_ratio)
            val_end = int(total_files * (train_ratio + val_ratio))
            
            # Split files
            train_files = files[:train_end]
            val_files = files[train_end:val_end]
            test_files = files[val_end:]
            
            # Ensure minimum samples in each split
            if len(val_files) == 0 and len(files) > 1:
                val_files = [train_files.pop()]
            if len(test_files) == 0 and len(files) > 2:
                test_files = [train_files.pop()]
            
            # Copy files to respective directories
            splits = [
                ("train", train_files),
                ("val", val_files),
                ("test", test_files)
            ]
            
            for split_name, split_files in splits:
                for file in split_files:
                    import shutil
                    dest_dir = output_dir / split_name / class_name
                    shutil.copy2(file, dest_dir / file.name)
            
            self.logger.info(f"Split {class_name}: {len(train_files)} train, "
                           f"{len(val_files)} val, {len(test_files)} test")
    
    def create_dataset_from_directory(self, data_dir: str, split: str,
                                    wakeword_classes: Dict[int, str],
                                    augmentation_config: Dict[str, float] = None,
                                    balance_classes: bool = True) -> WakewordDataset:
        """
        Create wakeword dataset from processed directory.
        
        Args:
            data_dir: Directory containing processed data
            split: Data split ('train', 'val', 'test')
            wakeword_classes: Mapping of class IDs to names
            augmentation_config: Augmentation configuration
            balance_classes: Whether to balance classes
            
        Returns:
            WakewordDataset instance
        """
        data_path = Path(data_dir) / split
        
        # Collect all files and labels
        audio_files = []
        labels = []
        
        class_name_to_id = {name: class_id for class_id, name in wakeword_classes.items()}
        
        for class_name, class_id in class_name_to_id.items():
            class_dir = data_path / class_name
            if not class_dir.exists():
                self.logger.warning(f"Class directory not found: {class_dir}")
                continue
            
            # Get all audio files
            for audio_file in class_dir.glob("*.wav"):
                audio_files.append(str(audio_file))
                labels.append(class_id)
        
        if not audio_files:
            raise ValueError(f"No audio files found in {data_path}")
        
        # Create augmentation
        augmentation = None
        if split == "train" and augmentation_config:
            augmentation = AudioAugmentation(self.audio_config.sample_rate)
        
        # Create dataset
        dataset = WakewordDataset(
            audio_files=audio_files,
            labels=labels,
            class_names=wakeword_classes,
            audio_processor=self.audio_processor,
            augmentation=augmentation,
            augmentation_config=augmentation_config,
            apply_augmentation=(split == "train"),
            balance_classes=balance_classes
        )
        
        self.logger.info(f"Created {split} dataset with {len(dataset)} samples")
        self.logger.info(f"Class distribution: {dataset.get_class_distribution()}")
        
        return dataset


def create_wakeword_dataloaders(data_dir: str, wakeword_classes: Dict[int, str],
                               audio_config: AudioProcessingConfig,
                               batch_size: int = 32,
                               num_workers: int = 4,
                               augmentation_config: Dict[str, float] = None,
                               balance_training: bool = True) -> Tuple[DataLoader, DataLoader, DataLoader]:
    """
    Create train, validation, and test data loaders for wakeword detection.
    
    Args:
        data_dir: Directory containing processed wakeword data
        wakeword_classes: Mapping of class IDs to names
        audio_config: Audio processing configuration
        batch_size: Batch size for data loaders
        num_workers: Number of worker processes
        augmentation_config: Augmentation configuration
        balance_training: Whether to balance training data
        
    Returns:
        Tuple of (train_loader, val_loader, test_loader)
    """
    preprocessor = WakewordDataPreprocessor(audio_config)
    
    # Create datasets
    train_dataset = preprocessor.create_dataset_from_directory(
        data_dir, "train", wakeword_classes, augmentation_config, balance_training
    )
    
    val_dataset = preprocessor.create_dataset_from_directory(
        data_dir, "val", wakeword_classes, None, False
    )
    
    test_dataset = preprocessor.create_dataset_from_directory(
        data_dir, "test", wakeword_classes, None, False
    )
    
    # Create data loaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True,
        drop_last=True
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True
    )
    
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True
    )
    
    return train_loader, val_loader, test_loader