Trixy3/trainer/wakeword/trainer.py

"""
Wakeword detection trainer implementation.

This module provides a complete trainer for wakeword detection models
using RepCNN architecture with comprehensive training, validation, and
deployment capabilities.
"""

import os
import time
import logging
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Any
import numpy as np

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

from ..base import BaseTrainer, TrainerConfig, TrainingState
from ..metadata import ModelMetadata, MetadataManager, ModelType
from ..model_formats import ModelFormatManager
from ..data_pipeline import AudioProcessingConfig
from ..utils import TrainerLogger, ProgressMonitor, ValidationMetrics
from ..validation import ModelValidator, ValidationConfig

from .models import RepCNN, ImprovedRepCNN, LightweightRepCNN, create_repcnn_model
from .data import WakewordDataPreprocessor, create_wakeword_dataloaders


class FocalLoss(nn.Module):
    """
    Focal Loss for addressing class imbalance in wakeword detection.
    
    Particularly useful when dealing with imbalanced datasets where
    negative samples significantly outnumber positive wakeword samples.
    """
    
    def __init__(self, alpha: float = 1.0, gamma: float = 2.0, 
                 class_weights: Optional[torch.Tensor] = None):
        """
        Initialize Focal Loss.
        
        Args:
            alpha: Weighting factor for rare class
            gamma: Focusing parameter
            class_weights: Optional class weights
        """
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.class_weights = class_weights
        
    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        """Compute focal loss."""
        ce_loss = F.cross_entropy(inputs, targets, weight=self.class_weights, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss
        return focal_loss.mean()


class WakewordTrainer(BaseTrainer):
    """
    Specialized trainer for wakeword detection models.
    
    Provides end-to-end training pipeline for wakeword detection including
    data preprocessing, model training, validation, and deployment preparation.
    """
    
    def __init__(self, config: TrainerConfig):
        """Initialize wakeword trainer."""
        super().__init__(config)
        
        # Wakeword-specific configuration
        self.wakeword_classes = config.custom_params.get('wakeword_classes', {
            0: "custom",
            1: "system_command", 
            2: "negative"
        })
        
        self.model_type = config.custom_params.get('model_type', 'standard')
        self.detection_threshold = config.custom_params.get('detection_threshold', 0.5)
        self.use_focal_loss = config.custom_params.get('use_focal_loss', True)
        self.class_weights = config.custom_params.get('class_weights', None)
        
        # Audio processing configuration
        self.audio_config = AudioProcessingConfig(
            sample_rate=config.sample_rate,
            target_length=config.audio_length,
            n_mels=config.n_mels,
            n_fft=config.n_fft,
            hop_length=config.hop_length,
            win_length=config.win_length
        )
        
        # Initialize specialized components
        self.data_preprocessor = WakewordDataPreprocessor(self.audio_config, self.logger)
        self.metadata_manager = MetadataManager(self.logger)
        self.format_manager = ModelFormatManager(self.logger)
        
        # Model validation
        self.validator = None
        
        self.logger.info(f"Initialized WakewordTrainer with classes: {self.wakeword_classes}")
    
    def prepare_data(self) -> Tuple[DataLoader, DataLoader, DataLoader]:
        """
        Prepare wakeword detection data.
        
        Returns:
            Tuple of (train_loader, val_loader, test_loader)
        """
        self.logger.info("Preparing wakeword detection data...")
        
        # Check if data is already preprocessed
        processed_data_dir = Path(self.config.data_dir) / "processed"
        
        if not processed_data_dir.exists() or not any(processed_data_dir.iterdir()):
            # Preprocess raw data
            raw_data_dir = Path(self.config.data_dir) / "raw"
            if not raw_data_dir.exists():
                raise FileNotFoundError(f"Raw data directory not found: {raw_data_dir}")
            
            self.logger.info("Preprocessing raw wakeword data...")
            processed_path = self.data_preprocessor.preprocess_wakeword_data(
                str(raw_data_dir),
                str(processed_data_dir),
                self.wakeword_classes
            )
            self.logger.info(f"Data preprocessing completed: {processed_path}")
        
        # Create data loaders
        augmentation_config = {
            'noise_factor': self.config.noise_factor,
            'speed_factor': self.config.speed_factor,
            'pitch_factor': self.config.pitch_factor,
            'volume_factor': self.config.volume_factor,
            'time_shift_factor': 0.1
        } if self.config.use_augmentation else None
        
        train_loader, val_loader, test_loader = create_wakeword_dataloaders(
            str(processed_data_dir),
            self.wakeword_classes,
            self.audio_config,
            batch_size=self.config.batch_size,
            num_workers=self.config.num_workers,
            augmentation_config=augmentation_config,
            balance_training=True
        )
        
        self.logger.info(f"Created data loaders - Train: {len(train_loader.dataset)}, "
                        f"Val: {len(val_loader.dataset)}, Test: {len(test_loader.dataset)}")
        
        return train_loader, val_loader, test_loader
    
    def build_model(self) -> nn.Module:
        """
        Build RepCNN model for wakeword detection.
        
        Returns:
            RepCNN model
        """
        self.logger.info(f"Building {self.model_type} RepCNN model...")
        
        # Calculate input dimensions
        time_frames = int(self.audio_config.target_length * self.audio_config.sample_rate 
                         // self.audio_config.hop_length) + 1
        
        model_kwargs = {
            'num_classes': len(self.wakeword_classes),
            'input_channels': 1,
            'num_mels': self.audio_config.n_mels,
            'time_frames': time_frames,
            'width_multiplier': self.config.custom_params.get('width_multiplier', 1.0),
            'use_se': self.config.custom_params.get('use_se', False),
            'dropout_rate': self.config.custom_params.get('dropout_rate', 0.2)
        }
        
        # Add model-specific parameters
        if self.model_type == 'improved':
            model_kwargs['use_temporal_attention'] = self.config.custom_params.get('use_temporal_attention', True)
        elif self.model_type == 'lightweight':
            model_kwargs['width_multiplier'] = min(model_kwargs['width_multiplier'], 0.75)
        
        model = create_repcnn_model(self.model_type, **model_kwargs)
        
        # Log model information
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        
        self.logger.info(f"Model created - Type: {self.model_type}")
        self.logger.info(f"Total parameters: {total_params:,}")
        self.logger.info(f"Trainable parameters: {trainable_params:,}")
        
        return model
    
    def create_criterion(self) -> nn.Module:
        """
        Create loss criterion for wakeword detection.
        
        Returns:
            Loss function
        """
        if self.use_focal_loss:
            # Use Focal Loss for imbalanced datasets
            class_weights = None
            if self.class_weights:
                class_weights = torch.tensor(self.class_weights, dtype=torch.float32)
            
            criterion = FocalLoss(
                alpha=self.config.custom_params.get('focal_alpha', 1.0),
                gamma=self.config.custom_params.get('focal_gamma', 2.0),
                class_weights=class_weights
            )
            self.logger.info("Using Focal Loss for training")
        else:
            # Use standard Cross Entropy Loss
            class_weights = None
            if self.class_weights:
                class_weights = torch.tensor(self.class_weights, dtype=torch.float32)
            
            criterion = nn.CrossEntropyLoss(weight=class_weights)
            self.logger.info("Using Cross Entropy Loss for training")
        
        return criterion
    
    def train(self) -> Any:
        """
        Train the wakeword detection model with comprehensive tracking.
        
        Returns:
            Training metrics and results
        """
        self.logger.info("Starting wakeword detection training...")
        
        # Setup training components
        self.setup_training()
        
        # Initialize model validator
        self.validator = ModelValidator(self.model, self.config.device, self.logger)
        
        # Create metadata
        metadata = self.metadata_manager.create_metadata(
            ModelType.WAKEWORD, self.config.model_name, "Trixy ML Trainer"
        )
        metadata.description = f"RepCNN wakeword detection model ({self.model_type})"
        metadata.update_from_training_config(self.config)
        metadata.update_from_model(self.model)
        metadata.update_from_dataset(self.train_loader, self.val_loader, self.test_loader)
        metadata.update_wakeword_info(self.wakeword_classes, {
            'detection_threshold': self.detection_threshold
        })
        
        # Run parent training loop
        training_metrics = super().train()
        
        # Update metadata with training results
        metadata.update_from_training_results(training_metrics)
        
        # Validate model
        validation_config = ValidationConfig(
            test_augmentations=True,
            robustness_tests=True,
            performance_profiling=True
        )
        
        self.logger.info("Running comprehensive model validation...")
        validation_results = self.validator.validate(self.val_loader, self.criterion, validation_config)
        metadata.add_test_results(validation_results)
        
        # Test model
        test_results = self.validator.test(self.test_loader, self.criterion, validation_config)
        
        # Save model with metadata
        self._save_trained_model(metadata, test_results)
        
        # Generate training report
        self._generate_training_report(metadata, training_metrics, test_results)
        
        self.logger.info("Wakeword training completed successfully!")
        
        return {
            'training_metrics': training_metrics,
            'validation_results': validation_results,
            'test_results': test_results,
            'metadata': metadata
        }
    
    def _save_trained_model(self, metadata: ModelMetadata, test_results: Dict[str, Any]):
        """Save the trained model with comprehensive metadata."""
        model_dir = Path(self.config.output_dir) / self.config.model_name
        model_dir.mkdir(parents=True, exist_ok=True)
        
        # Prepare model for saving
        self.model.eval()
        
        # Reparameterize RepCNN if applicable
        if hasattr(self.model, 'reparameterize'):
            self.logger.info("Reparameterizing model for inference...")
            self.model.reparameterize()
        
        # Update metadata with final model info
        metadata.set_file_info(str(model_dir / f"{self.config.model_name}.pth"))
        metadata.add_test_results(test_results)
        
        # Save model in requested format
        model_file = model_dir / f"{self.config.model_name}{self.config.model_format.value}"
        
        success = self.format_manager.save_model(
            self.model,
            str(model_file),
            metadata.to_dict(),
            password=self.config.password if self.config.use_password_protection else None
        )
        
        if success:
            self.logger.info(f"Model saved successfully: {model_file}")
        else:
            self.logger.error(f"Failed to save model: {model_file}")
        
        # Save standalone metadata file
        metadata_file = model_dir / "metadata.json"
        metadata.save_to_file(str(metadata_file))
        
        # Save training configuration
        config_file = model_dir / "training_config.json"
        with open(config_file, 'w') as f:
            import json
            json.dump(self.config.to_dict(), f, indent=2)
        
        # Save validation results
        if hasattr(self, 'validator') and self.validator:
            self.validator.save_results(str(model_dir / "validation"))
    
    def _generate_training_report(self, metadata: ModelMetadata, 
                                 training_metrics: Any, test_results: Dict[str, Any]):
        """Generate comprehensive training report."""
        report_dir = Path(self.config.output_dir) / self.config.model_name / "reports"
        report_dir.mkdir(parents=True, exist_ok=True)
        
        # Generate metadata report
        report = self.metadata_manager.create_training_report(metadata)
        
        # Add wakeword-specific information
        report += "\n## Wakeword Detection Specific Results\n\n"
        
        # Add detection metrics
        if 'standard_test' in test_results:
            std_test = test_results['standard_test']
            if 'metrics' in std_test:
                metrics = std_test['metrics']
                
                report += "### Classification Performance\n"
                for class_id, class_name in self.wakeword_classes.items():
                    if f'f1_{class_name}' in metrics:
                        precision = metrics.get(f'precision_{class_name}', 0)
                        recall = metrics.get(f'recall_{class_name}', 0)
                        f1 = metrics.get(f'f1_{class_name}', 0)
                        
                        report += f"- **{class_name.title()}**:\n"
                        report += f"  - Precision: {precision:.4f}\n"
                        report += f"  - Recall: {recall:.4f}\n"
                        report += f"  - F1-Score: {f1:.4f}\n"
                
                # Add confusion matrix info
                if 'confusion_matrix' in metrics:
                    report += "\n### Confusion Matrix\n"
                    cm = metrics['confusion_matrix']
                    report += "```\n"
                    class_names = list(self.wakeword_classes.values())
                    
                    # Header
                    report += "Actual\\Predicted  "
                    for name in class_names:
                        report += f"{name[:8]:>8} "
                    report += "\n"
                    
                    # Matrix rows
                    for i, row in enumerate(cm):
                        if i < len(class_names):
                            report += f"{class_names[i][:12]:>12}  "
                            for val in row:
                                report += f"{val:>8} "
                            report += "\n"
                    report += "```\n\n"
        
        # Add robustness test results
        if 'robustness_tests' in test_results:
            report += "### Robustness Tests\n"
            robustness = test_results['robustness_tests']
            
            if 'input_corruptions' in robustness:
                report += "#### Input Corruption Robustness\n"
                for corruption, results in robustness['input_corruptions'].items():
                    accuracy = results.get('accuracy', 0)
                    report += f"- {corruption}: {accuracy:.4f} accuracy\n"
                report += "\n"
        
        # Add performance profiling
        if 'performance_profile' in test_results:
            profile = test_results['performance_profile']
            report += "### Performance Profile\n"
            
            if 'inference_performance' in profile:
                perf = profile['inference_performance']
                report += f"- Inference time: {perf['mean_inference_time']*1000:.2f} ms\n"
                report += f"- Throughput: {perf['throughput_samples_per_second']:.1f} samples/sec\n"
            
            if 'model_size' in profile:
                size = profile['model_size']
                report += f"- Model size: {size['total_size_mb']:.2f} MB\n"
            
            report += "\n"
        
        # Add deployment recommendations
        report += "## Deployment Recommendations\n\n"
        report += "### Detection Threshold\n"
        report += f"- Recommended threshold: {self.detection_threshold}\n"
        report += "- Monitor false positive/negative rates in production\n"
        report += "- Consider threshold adjustment based on use case requirements\n\n"
        
        report += "### Hardware Requirements\n"
        if 'performance_profile' in test_results:
            profile = test_results['performance_profile']
            inference_time = profile.get('inference_performance', {}).get('mean_inference_time', 0) * 1000
            
            if inference_time < 50:
                report += "- Suitable for real-time processing on most devices\n"
            elif inference_time < 100:
                report += "- Suitable for real-time processing on modern devices\n"
            else:
                report += "- May require optimization for real-time processing\n"
        
        report += "\n### Model Optimization\n"
        if self.model_type == 'lightweight':
            report += "- Already optimized for edge deployment\n"
        else:
            report += "- Consider using lightweight variant for edge deployment\n"
        
        if hasattr(self.model, 'reparameterize'):
            report += "- Model has been reparameterized for efficient inference\n"
        
        # Save report
        report_file = report_dir / "training_report.md"
        with open(report_file, 'w', encoding='utf-8') as f:
            f.write(report)
        
        self.logger.info(f"Training report saved: {report_file}")
    
    def evaluate_detection_performance(self, test_loader: DataLoader, 
                                     thresholds: List[float] = None) -> Dict[str, Any]:
        """
        Evaluate wakeword detection performance at different thresholds.
        
        Args:
            test_loader: Test data loader
            thresholds: List of detection thresholds to evaluate
            
        Returns:
            Detection performance results
        """
        if thresholds is None:
            thresholds = [0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
        
        self.logger.info("Evaluating detection performance...")
        
        self.model.eval()
        all_predictions = []
        all_probabilities = []
        all_targets = []
        
        with torch.no_grad():
            for data, targets in test_loader:
                data = data.to(self.config.device)
                targets = targets.to(self.config.device)
                
                outputs = self.model(data)
                probabilities = F.softmax(outputs, dim=1)
                
                all_probabilities.append(probabilities.cpu().numpy())
                all_targets.append(targets.cpu().numpy())
        
        all_probabilities = np.concatenate(all_probabilities)
        all_targets = np.concatenate(all_targets)
        
        # Evaluate at different thresholds
        results = {}
        
        for threshold in thresholds:
            # Get predictions based on threshold
            max_probs = np.max(all_probabilities, axis=1)
            predictions = np.argmax(all_probabilities, axis=1)
            
            # Apply threshold (predict negative if below threshold)
            negative_class_id = max([id for id, name in self.wakeword_classes.items() 
                                   if name in ['negative', 'background']], default=2)
            
            low_confidence_mask = max_probs < threshold
            predictions[low_confidence_mask] = negative_class_id
            
            # Calculate metrics
            accuracy = np.mean(predictions == all_targets)
            
            # Calculate per-class metrics
            class_metrics = {}
            for class_id, class_name in self.wakeword_classes.items():
                mask = all_targets == class_id
                if np.sum(mask) > 0:
                    class_acc = np.mean(predictions[mask] == all_targets[mask])
                    class_metrics[class_name] = {
                        'accuracy': class_acc,
                        'samples': np.sum(mask)
                    }
            
            results[f"threshold_{threshold}"] = {
                'threshold': threshold,
                'overall_accuracy': accuracy,
                'class_metrics': class_metrics
            }
        
        self.logger.info("Detection performance evaluation completed")
        return results


def create_wakeword_trainer_config(model_name: str = "wakeword_model",
                                  model_type: str = "standard",
                                  data_dir: str = "./trainer/data/wakeword",
                                  **kwargs) -> TrainerConfig:
    """
    Create a TrainerConfig specifically configured for wakeword detection.
    
    Args:
        model_name: Name of the model
        model_type: Type of RepCNN model ('standard', 'improved', 'lightweight')
        data_dir: Directory containing wakeword data
        **kwargs: Additional configuration parameters
        
    Returns:
        TrainerConfig instance for wakeword detection
    """
    # Default wakeword configuration
    config_dict = {
        'trainer_name': 'wakeword_trainer',
        'model_name': model_name,
        'data_dir': data_dir,
        'output_dir': './models/wakeword',
        
        # Training parameters optimized for wakeword detection
        'batch_size': 64,
        'learning_rate': 0.001,
        'num_epochs': 100,
        'min_epochs': 20,  # Allow at least 20 epochs for wakeword model to learn patterns
        'early_stopping_patience': 15,
        
        # Audio parameters for wakeword detection
        'sample_rate': 16000,
        'audio_length': 1.5,  # Typical wakeword length
        'n_mels': 40,
        'n_fft': 512,
        'hop_length': 160,
        'win_length': 400,
        
        # Augmentation for robustness
        'use_augmentation': True,
        'noise_factor': 0.1,
        'speed_factor': 0.1,
        'pitch_factor': 0.05,
        'volume_factor': 0.2,
        
        # Model-specific parameters
        'custom_params': {
            'model_type': model_type,
            'wakeword_classes': {0: "custom", 1: "system_command", 2: "negative"},
            'detection_threshold': 0.5,
            'use_focal_loss': True,
            'focal_alpha': 1.0,
            'focal_gamma': 2.0,
            'width_multiplier': 1.0,
            'use_se': False,
            'dropout_rate': 0.2
        }
    }
    
    # Separate custom parameters from standard TrainerConfig parameters
    custom_param_keys = {
        'use_focal_loss', 'detection_threshold', 'focal_alpha', 'focal_gamma',
        'width_multiplier', 'use_se', 'dropout_rate', 'wakeword_classes',
        'use_temporal_attention'
    }
    
    # Extract custom parameters from kwargs
    custom_params = config_dict['custom_params'].copy()
    for key, value in kwargs.items():
        if key in custom_param_keys:
            custom_params[key] = value
        else:
            config_dict[key] = value
    
    # Handle model_type specially
    if 'model_type' in kwargs:
        custom_params['model_type'] = kwargs['model_type']
    
    # Update custom_params
    config_dict['custom_params'] = custom_params
    
    return TrainerConfig(**config_dict)