Trixy3/trainer/alternative_methods.py

"""
Alternative training methods for voice recognition and wakeword detection.

This module implements few-shot learning, transfer learning, and other advanced training
methodologies inspired by Porcupine and other efficient approaches that require minimal data.
"""

import os
import json
import math
from pathlib import Path
from typing import Dict, List, Any, Optional, Tuple, Union
from enum import Enum
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset, Subset
import torch.optim as optim

from .base import BaseTrainer, TrainerConfig
from .data_pipeline import AudioProcessingConfig


class TrainingMethod(Enum):
    """Enumeration of available training methods."""
    STANDARD = "standard"                    # Traditional full-dataset training
    FEW_SHOT = "few_shot"                   # Few-shot learning with 3-10 samples
    META_LEARNING = "meta_learning"         # Model-agnostic meta-learning
    TRANSFER_LEARNING = "transfer_learning" # Pre-trained backbone + fine-tuning
    SELF_SUPERVISED = "self_supervised"     # Self-supervised pre-training
    DATA_AUGMENTATION = "data_augmentation" # Heavy augmentation with minimal data
    CONTRASTIVE = "contrastive"            # Contrastive learning approach
    PROTOTYPICAL = "prototypical"          # Prototypical networks
    SIAMESE = "siamese"                    # Siamese network training


class FewShotConfig:
    """Configuration for few-shot learning approaches."""
    
    def __init__(self,
                 n_way: int = 5,              # Number of classes per episode
                 k_shot: int = 3,             # Number of samples per class
                 n_query: int = 5,            # Number of query samples per class
                 n_episodes: int = 1000,      # Number of training episodes
                 inner_lr: float = 0.01,      # Inner loop learning rate
                 outer_lr: float = 0.001,     # Outer loop learning rate
                 adaptation_steps: int = 5):   # Number of adaptation steps
        
        self.n_way = n_way
        self.k_shot = k_shot
        self.n_query = n_query
        self.n_episodes = n_episodes
        self.inner_lr = inner_lr
        self.outer_lr = outer_lr
        self.adaptation_steps = adaptation_steps


class DataAugmentationTrainer:
    """
    Heavy data augmentation approach similar to Porcupine.
    
    Uses minimal samples (3-5 per class) but applies extensive augmentation:
    - Background noise injection
    - Speed/pitch variations
    - Volume changes
    - Temporal shifts
    - Spectral augmentation
    - Mixup/CutMix
    """
    
    def __init__(self, base_trainer: BaseTrainer, augmentation_factor: int = 100):
        """
        Initialize data augmentation trainer.
        
        Args:
            base_trainer: Base trainer instance
            augmentation_factor: How many augmented samples per original sample
        """
        self.base_trainer = base_trainer
        self.augmentation_factor = augmentation_factor
        self.logger = base_trainer.logger
        
    def create_augmented_dataset(self, original_data: List[Tuple], target_dir: str) -> str:
        """
        Create heavily augmented dataset from minimal samples.
        
        Args:
            original_data: List of (audio_path, label) tuples
            target_dir: Directory to save augmented data
            
        Returns:
            Path to augmented dataset
        """
        target_path = Path(target_dir)
        target_path.mkdir(parents=True, exist_ok=True)
        
        augmented_samples = []
        
        for audio_path, label in original_data:
            # Load original audio
            audio = self._load_audio(audio_path)
            
            # Generate augmented versions
            for i in range(self.augmentation_factor):
                augmented_audio = self._apply_augmentation(audio, intensity=0.8)
                
                # Save augmented sample
                aug_filename = f"{Path(audio_path).stem}_aug_{i:03d}.wav"
                aug_path = target_path / label / aug_filename
                aug_path.parent.mkdir(parents=True, exist_ok=True)
                
                self._save_audio(augmented_audio, aug_path)
                augmented_samples.append((str(aug_path), label))
        
        # Save augmentation metadata
        metadata = {
            'original_samples': len(original_data),
            'augmented_samples': len(augmented_samples),
            'augmentation_factor': self.augmentation_factor,
            'augmentation_config': self._get_augmentation_config()
        }
        
        with open(target_path / "augmentation_metadata.json", 'w') as f:
            json.dump(metadata, f, indent=2)
        
        self.logger.info(f"Generated {len(augmented_samples)} augmented samples from {len(original_data)} originals")
        return str(target_path)
    
    def _apply_augmentation(self, audio: np.ndarray, intensity: float = 0.5) -> np.ndarray:
        """Apply comprehensive audio augmentation."""
        augmented = audio.copy()
        
        # Speed variation (0.8x to 1.2x)
        if np.random.random() < 0.7:
            speed_factor = np.random.uniform(0.8, 1.2)
            augmented = self._change_speed(augmented, speed_factor)
        
        # Pitch variation (±2 semitones)
        if np.random.random() < 0.6:
            pitch_shift = np.random.uniform(-2, 2)
            augmented = self._change_pitch(augmented, pitch_shift)
        
        # Volume variation (0.5x to 1.5x)
        if np.random.random() < 0.8:
            volume_factor = np.random.uniform(0.5, 1.5)
            augmented = augmented * volume_factor
        
        # Background noise injection
        if np.random.random() < 0.6:
            noise_level = np.random.uniform(0.01, 0.1) * intensity
            noise = np.random.normal(0, noise_level, augmented.shape)
            augmented = augmented + noise
        
        # Temporal shifting
        if np.random.random() < 0.5:
            shift_samples = int(np.random.uniform(-0.1, 0.1) * len(augmented))
            augmented = np.roll(augmented, shift_samples)
        
        # Random cropping and padding
        if np.random.random() < 0.4:
            crop_ratio = np.random.uniform(0.8, 1.0)
            crop_length = int(len(augmented) * crop_ratio)
            start_idx = np.random.randint(0, len(augmented) - crop_length)
            augmented = augmented[start_idx:start_idx + crop_length]
            
            # Pad back to original length
            pad_length = len(audio) - len(augmented)
            if pad_length > 0:
                pad_left = np.random.randint(0, pad_length + 1)
                pad_right = pad_length - pad_left
                augmented = np.pad(augmented, (pad_left, pad_right), mode='constant')
        
        # Ensure audio stays in valid range
        augmented = np.clip(augmented, -1.0, 1.0)
        
        return augmented
    
    def _change_speed(self, audio: np.ndarray, factor: float) -> np.ndarray:
        """Change audio speed using resampling."""
        # Simple linear interpolation for speed change
        indices = np.arange(0, len(audio), factor)
        indices = indices[indices < len(audio)]
        return np.interp(indices, np.arange(len(audio)), audio)
    
    def _change_pitch(self, audio: np.ndarray, semitones: float) -> np.ndarray:
        """Change pitch by shifting in frequency domain (simplified)."""
        # This is a simplified pitch shift - in practice, you'd use librosa or similar
        factor = 2 ** (semitones / 12.0)
        return self._change_speed(audio, factor)
    
    def _load_audio(self, path: str) -> np.ndarray:
        """Load audio file (simplified - implement with librosa/soundfile)."""
        # Placeholder - implement actual audio loading
        return np.random.randn(16000)  # 1 second at 16kHz
    
    def _save_audio(self, audio: np.ndarray, path: str):
        """Save audio file (simplified - implement with librosa/soundfile)."""
        # Placeholder - implement actual audio saving
        pass
    
    def _get_augmentation_config(self) -> Dict[str, Any]:
        """Get augmentation configuration for metadata."""
        return {
            'speed_variation': {'min': 0.8, 'max': 1.2, 'probability': 0.7},
            'pitch_variation': {'min': -2, 'max': 2, 'probability': 0.6},
            'volume_variation': {'min': 0.5, 'max': 1.5, 'probability': 0.8},
            'noise_injection': {'min': 0.01, 'max': 0.1, 'probability': 0.6},
            'temporal_shift': {'min': -0.1, 'max': 0.1, 'probability': 0.5},
            'random_crop': {'min_ratio': 0.8, 'max_ratio': 1.0, 'probability': 0.4}
        }


class PrototypicalNetworkTrainer:
    """
    Prototypical Networks for few-shot learning.
    
    Based on "Prototypical Networks for Few-shot Learning" by Snell et al.
    Creates class prototypes from support samples and classifies query samples
    based on distance to prototypes.
    """
    
    def __init__(self, model: nn.Module, config: FewShotConfig, device: str):
        """
        Initialize prototypical network trainer.
        
        Args:
            model: Embedding model (backbone)
            config: Few-shot learning configuration
            device: Training device
        """
        self.model = model
        self.config = config
        self.device = device
        
    def train_episode(self, support_data: torch.Tensor, support_labels: torch.Tensor,
                     query_data: torch.Tensor, query_labels: torch.Tensor) -> float:
        """
        Train on a single episode using prototypical networks.
        
        Args:
            support_data: Support set data (n_way * k_shot, ...)
            support_labels: Support set labels
            query_data: Query set data (n_way * n_query, ...)
            query_labels: Query set labels
            
        Returns:
            Episode loss
        """
        self.model.train()
        
        # Compute embeddings
        support_embeddings = self.model(support_data)
        query_embeddings = self.model(query_data)
        
        # Compute prototypes (class centroids)
        n_way = self.config.n_way
        k_shot = self.config.k_shot
        
        prototypes = []
        for class_idx in range(n_way):
            class_embeddings = support_embeddings[class_idx * k_shot:(class_idx + 1) * k_shot]
            prototype = torch.mean(class_embeddings, dim=0)
            prototypes.append(prototype)
        
        prototypes = torch.stack(prototypes)  # (n_way, embedding_dim)
        
        # Compute distances from query embeddings to prototypes
        distances = self._compute_distances(query_embeddings, prototypes)
        
        # Convert distances to logits (negative distances)
        logits = -distances
        
        # Compute loss
        loss = F.cross_entropy(logits, query_labels)
        
        return loss.item()
    
    def _compute_distances(self, query_embeddings: torch.Tensor, 
                          prototypes: torch.Tensor) -> torch.Tensor:
        """
        Compute Euclidean distances between query embeddings and prototypes.
        
        Args:
            query_embeddings: Query embeddings (n_query, embedding_dim)
            prototypes: Class prototypes (n_way, embedding_dim)
            
        Returns:
            Distance matrix (n_query, n_way)
        """
        n_query = query_embeddings.size(0)
        n_way = prototypes.size(0)
        
        # Expand dimensions for broadcasting
        query_expanded = query_embeddings.unsqueeze(1).expand(n_query, n_way, -1)
        prototype_expanded = prototypes.unsqueeze(0).expand(n_query, n_way, -1)
        
        # Compute Euclidean distances
        distances = torch.pow(query_expanded - prototype_expanded, 2).sum(dim=2)
        
        return distances


class MAMLTrainer:
    """
    Model-Agnostic Meta-Learning (MAML) trainer.
    
    Based on "Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks"
    by Finn et al. Learns initialization parameters that can quickly adapt to new tasks.
    """
    
    def __init__(self, model: nn.Module, config: FewShotConfig, device: str):
        """
        Initialize MAML trainer.
        
        Args:
            model: Model to meta-learn
            config: Few-shot learning configuration
            device: Training device
        """
        self.model = model
        self.config = config
        self.device = device
        self.meta_optimizer = optim.Adam(self.model.parameters(), lr=config.outer_lr)
        
    def meta_train_step(self, tasks: List[Tuple]) -> float:
        """
        Perform one meta-training step across multiple tasks.
        
        Args:
            tasks: List of (support_data, support_labels, query_data, query_labels) tuples
            
        Returns:
            Meta-loss across all tasks
        """
        self.meta_optimizer.zero_grad()
        meta_loss = 0.0
        
        for task_data in tasks:
            support_data, support_labels, query_data, query_labels = task_data
            
            # Create a copy of the model for this task
            fast_weights = self._get_model_params()
            
            # Inner loop: adapt to support set
            for _ in range(self.config.adaptation_steps):
                support_loss = self._compute_loss(support_data, support_labels, fast_weights)
                
                # Compute gradients with respect to fast weights
                grads = torch.autograd.grad(support_loss, fast_weights, create_graph=True)
                
                # Update fast weights
                fast_weights = [w - self.config.inner_lr * g for w, g in zip(fast_weights, grads)]
            
            # Outer loop: compute meta-loss on query set
            query_loss = self._compute_loss(query_data, query_labels, fast_weights)
            meta_loss += query_loss
        
        # Average meta-loss
        meta_loss = meta_loss / len(tasks)
        
        # Backpropagation for meta-parameters
        meta_loss.backward()
        self.meta_optimizer.step()
        
        return meta_loss.item()
    
    def _get_model_params(self) -> List[torch.Tensor]:
        """Get model parameters as a list."""
        return [p.clone() for p in self.model.parameters()]
    
    def _compute_loss(self, data: torch.Tensor, labels: torch.Tensor, 
                     weights: List[torch.Tensor]) -> torch.Tensor:
        """
        Compute loss using specific weights.
        
        Args:
            data: Input data
            labels: Target labels
            weights: Model weights to use
            
        Returns:
            Loss value
        """
        # This is a simplified version - actual implementation would need to
        # substitute weights into the model forward pass
        outputs = self.model(data)  # Would use weights parameter
        return F.cross_entropy(outputs, labels)


class TransferLearningTrainer:
    """
    Transfer learning trainer using pre-trained backbones.
    
    Uses pre-trained models (e.g., from speech recognition, speaker verification)
    and fine-tunes for specific tasks with minimal data.
    """
    
    def __init__(self, backbone_path: str, num_classes: int, device: str):
        """
        Initialize transfer learning trainer.
        
        Args:
            backbone_path: Path to pre-trained backbone model
            num_classes: Number of target classes
            device: Training device
        """
        self.backbone_path = backbone_path
        self.num_classes = num_classes
        self.device = device
        self.model = None
        
    def create_model(self, freeze_backbone: bool = True) -> nn.Module:
        """
        Create transfer learning model.
        
        Args:
            freeze_backbone: Whether to freeze backbone parameters
            
        Returns:
            Transfer learning model
        """
        # Load pre-trained backbone
        backbone = self._load_backbone(self.backbone_path)
        
        # Freeze backbone if requested
        if freeze_backbone:
            for param in backbone.parameters():
                param.requires_grad = False
        
        # Add classification head
        feature_dim = self._get_backbone_feature_dim(backbone)
        classifier = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(feature_dim, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, self.num_classes)
        )
        
        # Combine backbone and classifier
        self.model = nn.Sequential(backbone, classifier)
        return self.model
    
    def fine_tune(self, data_loader: DataLoader, num_epochs: int = 50) -> Dict[str, List[float]]:
        """
        Fine-tune the model on target data.
        
        Args:
            data_loader: Target dataset loader
            num_epochs: Number of fine-tuning epochs
            
        Returns:
            Training metrics
        """
        optimizer = optim.Adam(self.model.parameters(), lr=0.0001)
        criterion = nn.CrossEntropyLoss()
        
        losses = []
        accuracies = []
        
        for epoch in range(num_epochs):
            epoch_loss = 0.0
            epoch_correct = 0
            epoch_total = 0
            
            for data, targets in data_loader:
                data, targets = data.to(self.device), targets.to(self.device)
                
                optimizer.zero_grad()
                outputs = self.model(data)
                loss = criterion(outputs, targets)
                loss.backward()
                optimizer.step()
                
                epoch_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                epoch_total += targets.size(0)
                epoch_correct += (predicted == targets).sum().item()
            
            epoch_loss /= len(data_loader)
            epoch_acc = epoch_correct / epoch_total
            
            losses.append(epoch_loss)
            accuracies.append(epoch_acc)
        
        return {'losses': losses, 'accuracies': accuracies}
    
    def _load_backbone(self, path: str) -> nn.Module:
        """Load pre-trained backbone model."""
        # Placeholder - implement actual model loading
        # This would load a pre-trained model and remove the classification head
        return nn.Sequential()
    
    def _get_backbone_feature_dim(self, backbone: nn.Module) -> int:
        """Get feature dimension of backbone output."""
        # Placeholder - determine feature dimension
        return 512


class AlternativeTrainingManager:
    """
    Manager for alternative training methods.
    
    Provides a unified interface for selecting and configuring different
    training approaches based on available data and requirements.
    """
    
    def __init__(self, base_config: TrainerConfig):
        """
        Initialize alternative training manager.
        
        Args:
            base_config: Base trainer configuration
        """
        self.base_config = base_config
        self.training_method = TrainingMethod.STANDARD
        self.few_shot_config = FewShotConfig()
        
    def set_training_method(self, method: TrainingMethod, **kwargs):
        """
        Set the training method and its configuration.
        
        Args:
            method: Training method to use
            **kwargs: Method-specific configuration
        """
        self.training_method = method
        
        if method == TrainingMethod.FEW_SHOT:
            self.few_shot_config = FewShotConfig(**kwargs)
        elif method == TrainingMethod.DATA_AUGMENTATION:
            self.augmentation_factor = kwargs.get('augmentation_factor', 100)
        # Add other method configurations as needed
    
    def get_recommended_method(self, dataset_size: int, samples_per_class: int) -> TrainingMethod:
        """
        Recommend training method based on available data.
        
        Args:
            dataset_size: Total dataset size
            samples_per_class: Average samples per class
            
        Returns:
            Recommended training method
        """
        if samples_per_class <= 5:
            if dataset_size < 100:
                return TrainingMethod.DATA_AUGMENTATION
            else:
                return TrainingMethod.FEW_SHOT
        elif samples_per_class <= 20:
            return TrainingMethod.TRANSFER_LEARNING
        elif samples_per_class <= 50:
            return TrainingMethod.SELF_SUPERVISED
        else:
            return TrainingMethod.STANDARD
    
    def create_trainer(self, model: nn.Module, device: str) -> Any:
        """
        Create appropriate trainer based on selected method.
        
        Args:
            model: Model to train
            device: Training device
            
        Returns:
            Trainer instance for the selected method
        """
        if self.training_method == TrainingMethod.FEW_SHOT:
            return PrototypicalNetworkTrainer(model, self.few_shot_config, device)
        elif self.training_method == TrainingMethod.META_LEARNING:
            return MAMLTrainer(model, self.few_shot_config, device)
        elif self.training_method == TrainingMethod.TRANSFER_LEARNING:
            backbone_path = self.base_config.custom_params.get('backbone_path', '')
            num_classes = self.base_config.custom_params.get('num_classes', 10)
            return TransferLearningTrainer(backbone_path, num_classes, device)
        elif self.training_method == TrainingMethod.DATA_AUGMENTATION:
            base_trainer = BaseTrainer(self.base_config)  # This would be the actual trainer
            return DataAugmentationTrainer(base_trainer, self.augmentation_factor)
        else:
            # Return standard trainer
            return BaseTrainer(self.base_config)
    
    def get_training_config(self) -> Dict[str, Any]:
        """Get configuration for the selected training method."""
        config = {
            'method': self.training_method.value,
            'base_config': self.base_config.to_dict()
        }
        
        if self.training_method in [TrainingMethod.FEW_SHOT, TrainingMethod.META_LEARNING]:
            config['few_shot_config'] = {
                'n_way': self.few_shot_config.n_way,
                'k_shot': self.few_shot_config.k_shot,
                'n_query': self.few_shot_config.n_query,
                'n_episodes': self.few_shot_config.n_episodes,
                'inner_lr': self.few_shot_config.inner_lr,
                'outer_lr': self.few_shot_config.outer_lr,
                'adaptation_steps': self.few_shot_config.adaptation_steps
            }
        
        return config


def create_alternative_training_manager(base_config: TrainerConfig) -> AlternativeTrainingManager:
    """
    Factory function to create alternative training manager.
    
    Args:
        base_config: Base trainer configuration
        
    Returns:
        AlternativeTrainingManager instance
    """
    return AlternativeTrainingManager(base_config)