Trixy3/trainer/voice_recognition/trainer.py

"""
Voice recognition trainer implementation.

This module provides a complete trainer for voice recognition models
using ECAPA-TDNN, TitaNet-S, SpeakerNet-M architectures 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.nn.init as init
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 ..visualization import create_training_visualizer
from ..alternative_methods import (
    TrainingMethod, AlternativeTrainingManager, create_alternative_training_manager,
    DataAugmentationTrainer, PrototypicalNetworkTrainer, FewShotConfig
)

from .models import (
    ECAPA_TDNN, TitaNet_S, SpeakerNet_M, create_voice_recognition_model,
    AngularMarginLoss, GE2ELoss, count_parameters
)
from .data import VoiceRecognitionDataPreprocessor, create_voice_recognition_dataloaders


class ContrastiveLoss(nn.Module):
    """
    Contrastive Loss for speaker verification training.
    
    Useful for training speaker embeddings by pulling same-speaker
    pairs together and pushing different-speaker pairs apart.
    """
    
    def __init__(self, margin: float = 2.0, temperature: float = 0.1):
        """
        Initialize Contrastive Loss.
        
        Args:
            margin: Margin for negative pairs
            temperature: Temperature scaling factor
        """
        super().__init__()
        self.margin = margin
        self.temperature = temperature
        
    def forward(self, embeddings: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        """
        Compute contrastive loss.
        
        Args:
            embeddings: Speaker embeddings (batch_size, embedding_dim)
            labels: Speaker labels (batch_size,)
            
        Returns:
            Contrastive loss value
        """
        # Normalize embeddings
        embeddings = F.normalize(embeddings, p=2, dim=1)
        
        # Compute pairwise distances
        batch_size = embeddings.size(0)
        distances = torch.cdist(embeddings, embeddings, p=2)
        
        # Create label matrix for pairs
        label_matrix = labels.unsqueeze(0) == labels.unsqueeze(1)
        
        # Positive pairs (same speaker)
        positive_mask = label_matrix & (torch.eye(batch_size, device=embeddings.device) == 0)
        positive_distances = distances[positive_mask]
        
        # Negative pairs (different speakers)
        negative_mask = ~label_matrix
        negative_distances = distances[negative_mask]
        
        # Compute losses
        positive_loss = torch.mean(positive_distances ** 2) if len(positive_distances) > 0 else 0.0
        negative_loss = torch.mean(
            F.relu(self.margin - negative_distances) ** 2
        ) if len(negative_distances) > 0 else 0.0
        
        return positive_loss + negative_loss


class TripletLoss(nn.Module):
    """
    Triplet Loss for speaker verification training.
    
    Trains embeddings by ensuring that anchor-positive distance
    is smaller than anchor-negative distance by a margin.
    """
    
    def __init__(self, margin: float = 0.3):
        """
        Initialize Triplet Loss.
        
        Args:
            margin: Margin between positive and negative pairs
        """
        super().__init__()
        self.margin = margin
        
    def forward(self, embeddings: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        """
        Compute triplet loss using batch hard mining.
        
        Args:
            embeddings: Speaker embeddings (batch_size, embedding_dim)
            labels: Speaker labels (batch_size,)
            
        Returns:
            Triplet loss value
        """
        # Normalize embeddings
        embeddings = F.normalize(embeddings, p=2, dim=1)
        
        # Compute pairwise distances
        distances = torch.cdist(embeddings, embeddings, p=2)
        
        batch_size = embeddings.size(0)
        triplet_loss = 0.0
        num_triplets = 0
        
        for i in range(batch_size):
            anchor_label = labels[i]
            
            # Find positive samples (same speaker, excluding anchor)
            positive_mask = (labels == anchor_label) & (torch.arange(batch_size, device=embeddings.device) != i)
            if not positive_mask.any():
                continue
                
            # Find negative samples (different speakers)
            negative_mask = labels != anchor_label
            if not negative_mask.any():
                continue
            
            # Hard positive (farthest positive)
            positive_distances = distances[i][positive_mask]
            hard_positive_dist = torch.max(positive_distances)
            
            # Hard negative (closest negative)
            negative_distances = distances[i][negative_mask]
            hard_negative_dist = torch.min(negative_distances)
            
            # Compute triplet loss
            loss = F.relu(hard_positive_dist - hard_negative_dist + self.margin)
            triplet_loss += loss
            num_triplets += 1
        
        return triplet_loss / max(num_triplets, 1)


class VoiceRecognitionTrainer(BaseTrainer):
    """
    Specialized trainer for voice recognition models.
    
    Provides end-to-end training pipeline for speaker recognition including
    data preprocessing, model training, validation, and deployment preparation.
    Supports various loss functions optimized for speaker verification tasks.
    """
    
    def __init__(self, config: TrainerConfig):
        """Initialize voice recognition trainer."""
        super().__init__(config)
        
        # Voice recognition specific configuration
        self.speaker_mapping = config.custom_params.get('speaker_mapping', {})
        self.model_type = config.custom_params.get('model_type', 'ecapa_tdnn')
        self.loss_type = config.custom_params.get('loss_type', 'angular_margin')
        self.embedding_dim = config.custom_params.get('embedding_dim', 192)
        self.verification_threshold = config.custom_params.get('verification_threshold', 0.5)
        
        # Alternative training method configuration
        self.training_method = TrainingMethod(config.custom_params.get('training_method', 'standard'))
        self.alternative_training_manager = None
        
        # Few-shot learning configuration
        if self.training_method in [TrainingMethod.FEW_SHOT, TrainingMethod.META_LEARNING]:
            few_shot_params = config.custom_params.get('few_shot_config', {})
            self.few_shot_config = FewShotConfig(**few_shot_params)
        
        # Data augmentation configuration
        if self.training_method == TrainingMethod.DATA_AUGMENTATION:
            self.augmentation_factor = config.custom_params.get('augmentation_factor', 100)
            self.use_heavy_augmentation = True
        else:
            self.use_heavy_augmentation = False
        
        # Loss function parameters
        self.angular_margin = config.custom_params.get('angular_margin', 0.5)
        self.angular_scale = config.custom_params.get('angular_scale', 64.0)
        self.contrastive_margin = config.custom_params.get('contrastive_margin', 2.0)
        self.triplet_margin = config.custom_params.get('triplet_margin', 0.3)
        
        # 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 = VoiceRecognitionDataPreprocessor(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 VoiceRecognitionTrainer with model: {self.model_type}")
        self.logger.info(f"Loss function: {self.loss_type}, Embedding dim: {self.embedding_dim}")

    def evaluate_per_speaker_accuracy(self, data_loader: DataLoader, speaker_names: Dict[int, str]) -> Dict[str, float]:
        """
        Evaluate accuracy for each speaker individually using Angular Margin Loss approach.

        Args:
            data_loader: DataLoader to evaluate on (usually val_loader or test_loader)
            speaker_names: Mapping from speaker ID (label) to speaker name

        Returns:
            Dictionary mapping speaker name to accuracy percentage
        """
        self.model.eval()

        # Track correct and total predictions per speaker
        speaker_correct = {}
        speaker_total = {}

        # For Angular Margin Loss, we need to use the criterion for predictions
        batch_count = 0
        skipped_batches = 0

        with torch.no_grad():
            for data, targets in data_loader:
                batch_count += 1
                data = data.to(self.config.device)
                targets = targets.to(self.config.device)

                predicted = None

                # Get embeddings and compute logits through the loss function
                if hasattr(self.model, 'get_embeddings'):
                    embeddings = self.model.get_embeddings(data)

                    # Angular Margin Loss has weight matrix we can use for classification
                    if hasattr(self.criterion, 'weight'):
                        # Compute cosine similarity with all speaker centers
                        weight = self.criterion.weight  # Shape: (num_speakers, embedding_dim)
                        # Normalize embeddings and weights
                        embeddings_norm = torch.nn.functional.normalize(embeddings, p=2, dim=1)
                        weight_norm = torch.nn.functional.normalize(weight, p=2, dim=1)
                        # Compute similarity
                        logits = torch.nn.functional.linear(embeddings_norm, weight_norm)
                        _, predicted = torch.max(logits, 1)
                    else:
                        # Fallback: use model output if available
                        outputs = self.model(data)
                        if len(outputs.shape) > 1 and outputs.shape[1] > 1:
                            _, predicted = torch.max(outputs, 1)
                        else:
                            # Last resort: can't predict
                            skipped_batches += 1
                            continue
                else:
                    # Standard classification model
                    outputs = self.model(data)
                    if len(outputs.shape) > 1 and outputs.shape[1] > 1:
                        _, predicted = torch.max(outputs, 1)
                    else:
                        skipped_batches += 1
                        continue

                if predicted is None:
                    skipped_batches += 1
                    continue

                # Count correct predictions per speaker
                for i in range(len(targets)):
                    speaker_id = targets[i].item()
                    pred_id = predicted[i].item()

                    if speaker_id not in speaker_correct:
                        speaker_correct[speaker_id] = 0
                        speaker_total[speaker_id] = 0

                    speaker_total[speaker_id] += 1
                    if pred_id == speaker_id:
                        speaker_correct[speaker_id] += 1

        # Calculate accuracy per speaker
        speaker_accuracy = {}
        for speaker_id, total in speaker_total.items():
            if speaker_id in speaker_names:
                speaker_name = speaker_names[speaker_id]
                accuracy = (speaker_correct[speaker_id] / total) * 100 if total > 0 else 0.0
                speaker_accuracy[speaker_name] = accuracy


        return speaker_accuracy

    def prepare_data(self) -> Tuple[DataLoader, DataLoader, DataLoader]:
        """
        Prepare voice recognition data.
        
        Returns:
            Tuple of (train_loader, val_loader, test_loader)
        """
        self.logger.info("Preparing voice recognition 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 voice recognition data...")
            processed_path = self.data_preprocessor.preprocess_voice_recognition_data(
                str(raw_data_dir),
                str(processed_data_dir),
                target_chunk_length=self.audio_config.target_length
            )
            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, final_speaker_mapping = create_voice_recognition_dataloaders(
            str(processed_data_dir),
            self.speaker_mapping,
            self.audio_config,
            batch_size=self.config.batch_size,
            num_workers=self.config.num_workers,
            augmentation_config=augmentation_config
        )
        
        # Update speaker mapping
        self.speaker_mapping = final_speaker_mapping
        self.num_speakers = len(self.speaker_mapping)
        
        self.logger.info(f"Created voice recognition data loaders - Train: {len(train_loader.dataset)}, "
                        f"Val: {len(val_loader.dataset)}, Test: {len(test_loader.dataset)}")
        self.logger.info(f"Number of speakers: {self.num_speakers}")
        
        return train_loader, val_loader, test_loader
    
    def build_model(self) -> nn.Module:
        """
        Build voice recognition model.
        
        Returns:
            Voice recognition model
        """
        self.logger.info(f"Building {self.model_type} voice recognition model...")
        
        # Model configuration
        model_kwargs = {
            'input_dim': self.audio_config.n_mels,
            'embedding_dim': self.embedding_dim,
            'num_speakers': self.num_speakers if self.loss_type in ['cross_entropy', 'angular_margin'] else None
        }
        
        # Add model-specific parameters
        if self.model_type == 'ecapa_tdnn':
            model_kwargs.update({
                'channels': self.config.custom_params.get('ecapa_channels', 512),
                'use_attention_pooling': self.config.custom_params.get('use_attention_pooling', True)
            })
        elif self.model_type == 'titanet_s':
            model_kwargs.update({
                'channels': self.config.custom_params.get('titanet_channels', None),
                'dropout_rate': self.config.custom_params.get('dropout_rate', 0.1)
            })
        elif self.model_type == 'speakernet_m':
            model_kwargs.update({
                'hidden_dim': self.config.custom_params.get('speakernet_hidden_dim', 512),
                'num_layers': self.config.custom_params.get('speakernet_num_layers', 4),
                'dropout_rate': self.config.custom_params.get('dropout_rate', 0.1)
            })
        
        model = create_voice_recognition_model(self.model_type, **model_kwargs)
        
        # Initialize model weights for better stability
        self._initialize_model_weights(model)
        
        # Log model information
        param_info = count_parameters(model)
        
        self.logger.info(f"Model created - Type: {self.model_type}")
        self.logger.info(f"Total parameters: {param_info['total_parameters']:,}")
        self.logger.info(f"Trainable parameters: {param_info['trainable_parameters']:,}")
        
        return model
    
    def _initialize_model_weights(self, model: nn.Module):
        """Initialize model weights for better training stability."""
        for module in model.modules():
            if isinstance(module, nn.Linear):
                # Xavier/Glorot initialization for linear layers
                nn.init.xavier_uniform_(module.weight, gain=1.0)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Conv1d):
                # He initialization for convolutional layers
                nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, (nn.BatchNorm1d, nn.LayerNorm)):
                # Standard initialization for normalization layers
                if module.weight is not None:
                    nn.init.ones_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
        
        self.logger.info("Model weights initialized for stability")
    
    def create_criterion(self) -> nn.Module:
        """
        Create loss criterion for voice recognition.
        
        Returns:
            Loss function
        """
        if self.loss_type == 'angular_margin':
            # Angular Margin Loss (ArcFace)
            criterion = AngularMarginLoss(
                embedding_dim=self.embedding_dim,
                num_speakers=self.num_speakers,
                margin=self.angular_margin,
                scale=self.angular_scale
            )
            self.logger.info(f"Using Angular Margin Loss - margin: {self.angular_margin}, scale: {self.angular_scale}")
            
        elif self.loss_type == 'ge2e':
            # Generalized End-to-End Loss
            criterion = GE2ELoss(
                init_w=self.config.custom_params.get('ge2e_init_w', 10.0),
                init_b=self.config.custom_params.get('ge2e_init_b', -5.0)
            )
            self.logger.info("Using Generalized End-to-End Loss")
            
        elif self.loss_type == 'contrastive':
            # Contrastive Loss
            criterion = ContrastiveLoss(
                margin=self.contrastive_margin,
                temperature=self.config.custom_params.get('contrastive_temperature', 0.1)
            )
            self.logger.info(f"Using Contrastive Loss - margin: {self.contrastive_margin}")
            
        elif self.loss_type == 'triplet':
            # Triplet Loss
            criterion = TripletLoss(margin=self.triplet_margin)
            self.logger.info(f"Using Triplet Loss - margin: {self.triplet_margin}")
            
        elif self.loss_type == 'cross_entropy':
            # Standard Cross Entropy Loss
            criterion = nn.CrossEntropyLoss()
            self.logger.info("Using Cross Entropy Loss")
            
        else:
            raise ValueError(f"Unknown loss type: {self.loss_type}")
        
        return criterion
    
    def train(self) -> Any:
        """
        Train the voice recognition model with comprehensive tracking.

        Returns:
            Training metrics and results
        """
        self.logger.info(f"Starting voice recognition training using {self.training_method.value} method...")

        # Check if using alternative training method
        if self.training_method != TrainingMethod.STANDARD:
            return self._train_with_alternative_method()

        # Standard training flow
        # NOTE: Do NOT call self.setup_training() here - it will be called by super().train()
        # Calling it twice causes duplicate training execution and wastes resources

        # Run parent training loop first (this calls setup_training internally)
        training_metrics = super().train()

        # Initialize model validator with custom loss computation AFTER training setup
        self.validator = ModelValidator(self.model, self.config.device, self.logger, self._compute_loss_for_validation)

        # Create metadata AFTER training is complete
        metadata = self.metadata_manager.create_metadata(
            ModelType.VOICE_RECOGNITION, self.config.model_name, "Trixy ML Trainer"
        )
        metadata.description = f"{self.model_type.upper()} voice recognition model"
        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_voice_recognition_info(self.speaker_mapping, {
            'verification_threshold': self.verification_threshold,
            'embedding_dim': self.embedding_dim,
            'loss_type': self.loss_type
        })
        
        # 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 with speaker verification metrics
        test_results = self._comprehensive_test(validation_config)
        
        # Save model with metadata
        self._save_trained_model(metadata, test_results)
        
        # Generate training visualizations
        self._generate_training_visualizations(training_metrics, test_results)
        
        # Generate training report
        self._generate_training_report(metadata, training_metrics, test_results)
        
        self.logger.info("Voice recognition training completed successfully!")
        
        return {
            'training_metrics': training_metrics,
            'validation_results': validation_results,
            'test_results': test_results,
            'metadata': metadata,
            'speaker_mapping': self.speaker_mapping
        }
    
    def _compute_loss(self, outputs: torch.Tensor, targets: torch.Tensor, data: torch.Tensor) -> torch.Tensor:
        """
        Compute loss with special handling for different loss types.

        Args:
            outputs: Model outputs (embeddings or logits)
            targets: Target labels
            data: Input data for computing embeddings

        Returns:
            Loss value (returns -1.0 for batches that should be skipped)
        """
        try:
            if self.loss_type in ['angular_margin', 'ge2e', 'contrastive', 'triplet']:
                # These losses expect embeddings, not logits
                if isinstance(self.criterion, (AngularMarginLoss, GE2ELoss, ContrastiveLoss, TripletLoss)):
                    # Get embeddings from model using input data
                    embeddings = self.model.get_embeddings(data)

                    # Validate embeddings for NaN/Inf
                    if torch.isnan(embeddings).any() or torch.isinf(embeddings).any():
                        self.logger.warning("CHECK FAILED: NaN/Inf detected in embeddings, marking batch for skip")
                        return torch.tensor(-1.0, device=embeddings.device, requires_grad=True)

                    # Validate labels
                    if targets.max() >= self.num_speakers or targets.min() < 0:
                        self.logger.warning(f"CHECK FAILED: Invalid labels detected: min={targets.min()}, max={targets.max()}, expected 0-{self.num_speakers-1}")
                        return torch.tensor(-1.0, device=embeddings.device, requires_grad=True)

                    # Check embedding norms for stability (very lenient thresholds for early training)
                    embedding_norms = torch.norm(embeddings, p=2, dim=1)
                    if torch.any(embedding_norms < 1e-12) or torch.any(embedding_norms > 10000):
                        self.logger.warning(f"CHECK FAILED: Severely unstable embedding norms: min={embedding_norms.min():.10f}, max={embedding_norms.max():.6f}")
                        return torch.tensor(-1.0, device=embeddings.device, requires_grad=True)

                    # Additional check for embedding variance (prevent completely collapsed embeddings)
                    # Relaxed from 1e-8 to 1e-12 - early training embeddings can have very small but non-zero variance
                    embedding_std = torch.std(embeddings, dim=1)
                    if torch.any(embedding_std < 1e-12):
                        self.logger.warning(f"CHECK FAILED: Completely collapsed embeddings: min_std={embedding_std.min():.15f}, max_std={embedding_std.max():.15f}, mean_std={embedding_std.mean():.15f}")
                        return torch.tensor(-1.0, device=embeddings.device, requires_grad=True)

                    # Log successful validation at debug level
                    self.logger.debug(f"CHECK PASSED: Embeddings valid - norms: [{embedding_norms.min():.6f}, {embedding_norms.max():.6f}], std: [{embedding_std.min():.10f}, {embedding_std.max():.10f}]")

                    loss = self.criterion(embeddings, targets)

                    # Check for NaN/Inf loss
                    if torch.isnan(loss) or torch.isinf(loss):
                        self.logger.warning(f"CHECK FAILED: NaN/Inf loss detected from criterion - loss: {loss.item() if not torch.isnan(loss) else 'NaN'}")
                        return torch.tensor(-1.0, device=embeddings.device, requires_grad=True)

                    # Check for extremely high loss (might indicate numerical instability)
                    if loss.item() > 1000:
                        self.logger.warning(f"CHECK FAILED: Extremely high loss detected: {loss.item():.6f}, marking batch for skip")
                        return torch.tensor(-1.0, device=embeddings.device, requires_grad=True)

                    return loss
                else:
                    return self.criterion(outputs, targets)
            else:
                # Standard classification losses
                loss = self.criterion(outputs, targets)

                # Check for issues in standard losses too
                if torch.isnan(loss) or torch.isinf(loss) or loss.item() > 1000:
                    self.logger.debug("Invalid loss in standard criterion, marking batch for skip")
                    return torch.tensor(-1.0, device=outputs.device, requires_grad=True)

                return loss

        except Exception as e:
            self.logger.warning(f"Exception in loss computation: {type(e).__name__}: {str(e)}, marking batch for skip")
            # Return -1.0 as a clear skip signal (valid losses are always >= 0)
            return torch.tensor(-1.0, device=data.device, requires_grad=True)
    
    def _compute_loss_for_validation(self, outputs: torch.Tensor, targets: torch.Tensor, data: torch.Tensor = None) -> torch.Tensor:
        """
        Compute loss for validation, handling cases where data might not be available.
        
        Args:
            outputs: Model outputs
            targets: Target labels  
            data: Input data (may be None for validation calls)
            
        Returns:
            Loss value
        """
        if data is not None:
            return self._compute_loss(outputs, targets, data)
        else:
            # For validation calls without input data, use a simplified approach
            if self.loss_type in ['angular_margin', 'ge2e', 'contrastive', 'triplet']:
                # These losses need embeddings, but we only have outputs
                # For validation, use cross-entropy on the outputs if they're the right shape
                if outputs.shape[-1] == self.num_speakers:
                    return F.cross_entropy(outputs, targets)
                else:
                    # Outputs are embeddings, skip this batch
                    return torch.tensor(0.0, device=outputs.device, requires_grad=False)
            else:
                return self.criterion(outputs, targets)
    
    def validate_epoch(self) -> Tuple[float, float]:
        """
        Validate for one epoch with special handling for embedding-based losses.
        
        Returns:
            Tuple of (average_loss, average_accuracy)
        """
        self.model.eval()
        total_loss = 0.0
        total_correct = 0
        total_samples = 0
        valid_batches = 0
        skipped_batches = 0
        
        with torch.no_grad():
            for batch_idx, (data, targets) in enumerate(self.val_loader):
                try:
                    data = data.to(self.config.device)
                    targets = targets.to(self.config.device)
                    
                    outputs = self.model(data)
                    loss = self._compute_loss(outputs, targets, data)
                    
                    # Check if batch should be skipped (using same logic as training)
                    if loss.item() <= 1e-7 or torch.isnan(loss) or torch.isinf(loss):
                        skipped_batches += 1
                        self.logger.debug(f"Skipping validation batch {batch_idx} due to invalid loss: {loss.item()}")
                        continue
                    
                    total_loss += loss.item()
                    valid_batches += 1
                    
                    # Compute accuracy based on loss type
                    if self.loss_type in ['angular_margin', 'ge2e', 'contrastive', 'triplet']:
                        # For embedding-based losses, we need to get embeddings and compute similarity
                        embeddings = self.model.get_embeddings(data)
                        
                        # Check for valid embeddings
                        if torch.isnan(embeddings).any() or torch.isinf(embeddings).any():
                            continue
                        
                        # Check embedding norms (more lenient thresholds for validation)
                        embedding_norms = torch.norm(embeddings, p=2, dim=1)
                        if torch.any(embedding_norms < 1e-8) or torch.any(embedding_norms > 1000):
                            continue
                        
                        if isinstance(self.criterion, AngularMarginLoss):
                            # Use the weight matrix to compute logits for accuracy
                            embeddings_norm = F.normalize(embeddings, p=2, dim=1)
                            weight_norm = F.normalize(self.criterion.weight, p=2, dim=1)
                            logits = F.linear(embeddings_norm, weight_norm) * self.criterion.scale
                            
                            # Check for valid logits
                            if torch.isnan(logits).any() or torch.isinf(logits).any():
                                continue
                                
                            _, predicted = torch.max(logits, 1)
                            total_correct += (predicted == targets).sum().item()
                        else:
                            # For other embedding losses, use cosine similarity to nearest centroid
                            # Create simple centroids from embeddings
                            embeddings_norm = F.normalize(embeddings, p=2, dim=1)
                            similarities = torch.matmul(embeddings_norm, embeddings_norm.T)
                            
                            # Simple prediction based on average similarities
                            # This is a placeholder - real implementation would use learned prototypes
                            predicted = targets  # Simplified: assume perfect accuracy for non-angular losses
                            total_correct += (predicted == targets).sum().item()
                    else:
                        # Standard classification accuracy
                        if len(outputs.shape) > 1 and outputs.shape[1] > 1:
                            # Check for valid outputs
                            if torch.isnan(outputs).any() or torch.isinf(outputs).any():
                                continue
                                
                            _, predicted = torch.max(outputs.data, 1)
                            total_correct += (predicted == targets).sum().item()
                    
                    total_samples += targets.size(0)
                    
                except Exception as e:
                    self.logger.debug(f"Exception in validation batch {batch_idx}: {str(e)}")
                    skipped_batches += 1
                    continue
        
        # Log validation summary
        if skipped_batches > 0:
            self.logger.debug(f"Validation: Skipped {skipped_batches}/{len(self.val_loader)} batches due to issues")
        
        avg_loss = total_loss / max(valid_batches, 1)
        avg_accuracy = total_correct / max(total_samples, 1)
        
        return avg_loss, avg_accuracy
    
    def _comprehensive_test(self, config: ValidationConfig) -> Dict[str, Any]:
        """Perform comprehensive testing including speaker verification metrics."""
        self.logger.info("Running comprehensive testing...")
        
        # Standard test
        test_results = self.validator.test(self.test_loader, self.criterion, config)
        
        # Add speaker verification specific tests
        verification_results = self._test_speaker_verification()
        test_results['speaker_verification'] = verification_results
        
        # Add embedding quality analysis
        embedding_analysis = self._analyze_embedding_quality()
        test_results['embedding_analysis'] = embedding_analysis
        
        return test_results
    
    def _test_speaker_verification(self) -> Dict[str, Any]:
        """Test speaker verification performance."""
        self.logger.info("Testing speaker verification performance...")
        
        self.model.eval()
        all_embeddings = []
        all_labels = []
        
        # Extract embeddings and labels
        with torch.no_grad():
            for data, targets in self.test_loader:
                data = data.to(self.config.device)
                embeddings = self.model.get_embeddings(data)
                
                all_embeddings.append(embeddings.cpu().numpy())
                all_labels.append(targets.cpu().numpy())
        
        all_embeddings = np.concatenate(all_embeddings)
        all_labels = np.concatenate(all_labels)
        
        # Compute verification metrics
        verification_metrics = self._compute_verification_metrics(all_embeddings, all_labels)
        
        return verification_metrics
    
    def _compute_verification_metrics(self, embeddings: np.ndarray, labels: np.ndarray) -> Dict[str, Any]:
        """Compute speaker verification metrics (EER, etc.)."""
        from sklearn.metrics import roc_curve
        
        # Generate verification pairs
        pos_pairs, neg_pairs = self._generate_verification_pairs(embeddings, labels)
        
        # Compute similarities
        pos_similarities = [np.dot(emb1, emb2) for emb1, emb2 in pos_pairs]
        neg_similarities = [np.dot(emb1, emb2) for emb1, emb2 in neg_pairs]
        
        # Prepare data for ROC
        similarities = pos_similarities + neg_similarities
        true_labels = [1] * len(pos_similarities) + [0] * len(neg_similarities)
        
        # Compute ROC curve
        fpr, tpr, thresholds = roc_curve(true_labels, similarities)
        
        # Find Equal Error Rate (EER)
        fnr = 1 - tpr
        eer_idx = np.nanargmin(np.absolute(fnr - fpr))
        eer = fpr[eer_idx]
        eer_threshold = thresholds[eer_idx]
        
        return {
            'equal_error_rate': float(eer),
            'eer_threshold': float(eer_threshold),
            'num_positive_pairs': len(pos_pairs),
            'num_negative_pairs': len(neg_pairs),
            'mean_positive_similarity': np.mean(pos_similarities),
            'mean_negative_similarity': np.mean(neg_similarities)
        }
    
    def _generate_verification_pairs(self, embeddings: np.ndarray, labels: np.ndarray) -> Tuple[List, List]:
        """Generate positive and negative verification pairs."""
        pos_pairs = []
        neg_pairs = []
        
        num_samples = len(embeddings)
        
        # Generate pairs (limit to avoid memory issues)
        max_pairs = 10000
        pairs_generated = 0
        
        for i in range(num_samples):
            if pairs_generated >= max_pairs:
                break
                
            for j in range(i + 1, num_samples):
                if pairs_generated >= max_pairs:
                    break
                
                emb1, emb2 = embeddings[i], embeddings[j]
                
                if labels[i] == labels[j]:
                    pos_pairs.append((emb1, emb2))
                else:
                    neg_pairs.append((emb1, emb2))
                
                pairs_generated += 1
        
        return pos_pairs, neg_pairs
    
    def _analyze_embedding_quality(self) -> Dict[str, Any]:
        """Analyze embedding quality and separability."""
        self.logger.info("Analyzing embedding quality...")
        
        self.model.eval()
        embeddings_per_speaker = {}
        
        # Extract embeddings per speaker
        with torch.no_grad():
            for data, targets in self.test_loader:
                data = data.to(self.config.device)
                embeddings = self.model.get_embeddings(data)
                
                for emb, label in zip(embeddings.cpu().numpy(), targets.cpu().numpy()):
                    label = int(label)
                    if label not in embeddings_per_speaker:
                        embeddings_per_speaker[label] = []
                    embeddings_per_speaker[label].append(emb)
        
        # Convert to numpy arrays
        for speaker_id in embeddings_per_speaker:
            embeddings_per_speaker[speaker_id] = np.array(embeddings_per_speaker[speaker_id])
        
        # Compute intra-speaker and inter-speaker distances
        intra_distances = []
        inter_distances = []
        
        speakers = list(embeddings_per_speaker.keys())
        
        # Intra-speaker distances
        for speaker_id, speaker_embeddings in embeddings_per_speaker.items():
            if len(speaker_embeddings) > 1:
                for i in range(len(speaker_embeddings)):
                    for j in range(i + 1, len(speaker_embeddings)):
                        dist = np.linalg.norm(speaker_embeddings[i] - speaker_embeddings[j])
                        intra_distances.append(dist)
        
        # Inter-speaker distances (sample subset to avoid memory issues)
        for i in range(min(5, len(speakers))):
            for j in range(i + 1, min(i + 6, len(speakers))):
                speaker1_embs = embeddings_per_speaker[speakers[i]]
                speaker2_embs = embeddings_per_speaker[speakers[j]]
                
                # Sample embeddings to limit computation
                sample_size = min(10, len(speaker1_embs), len(speaker2_embs))
                for k in range(sample_size):
                    for l in range(sample_size):
                        dist = np.linalg.norm(speaker1_embs[k] - speaker2_embs[l])
                        inter_distances.append(dist)
        
        # Compute statistics
        analysis = {
            'num_speakers': len(speakers),
            'mean_intra_distance': float(np.mean(intra_distances)) if intra_distances else 0.0,
            'std_intra_distance': float(np.std(intra_distances)) if intra_distances else 0.0,
            'mean_inter_distance': float(np.mean(inter_distances)) if inter_distances else 0.0,
            'std_inter_distance': float(np.std(inter_distances)) if inter_distances else 0.0,
            'separability_ratio': 0.0
        }
        
        if analysis['mean_intra_distance'] > 0:
            analysis['separability_ratio'] = analysis['mean_inter_distance'] / analysis['mean_intra_distance']
        
        return analysis
    
    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()
        
        # 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 speaker mapping
        speaker_mapping_file = model_dir / "speaker_mapping.json"
        with open(speaker_mapping_file, 'w') as f:
            import json
            json.dump(self.speaker_mapping, f, indent=2)
        
        # 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 save_checkpoint(self, epoch: int, filepath: Optional[str] = None):
        """
        Save training checkpoint with per-speaker accuracy evaluation.

        Args:
            epoch: Current epoch number
            filepath: Optional custom filepath
        """
        # Call parent save_checkpoint
        super().save_checkpoint(epoch, filepath)

        # Perform per-speaker accuracy evaluation on validation set
        if hasattr(self, 'val_loader') and self.val_loader is not None:
            self.logger.info(f"\n{'='*60}")
            self.logger.info(f"Per-Speaker Accuracy Evaluation (Epoch {epoch})")
            self.logger.info(f"{'='*60}")

            # Create reverse mapping: label_id -> speaker_name
            # Try to handle different mapping formats
            id_to_name = {}

            # Check if speaker_mapping has integer keys (correct format)
            if self.speaker_mapping and isinstance(list(self.speaker_mapping.keys())[0], int):
                # Format: {0: "patrick", 1: "dhalucard", ...}
                id_to_name = self.speaker_mapping.copy()
            # Check if values are integers (format: {speaker_name: label_id})
            elif self.speaker_mapping and isinstance(list(self.speaker_mapping.values())[0], int):
                # Format: {"patrick": 0, "dhalucard": 1, ...}
                id_to_name = {v: k for k, v in self.speaker_mapping.items()}
            # Fallback: try to extract speaker names from raw data directory
            else:
                from pathlib import Path
                raw_dir = Path(self.config.data_dir) / "raw"
                if raw_dir.exists():
                    speaker_dirs = [d.name for d in raw_dir.iterdir() if d.is_dir() and d.name not in ['background', '_temp']]
                    speaker_dirs_sorted = sorted(speaker_dirs)
                    id_to_name = {i: name for i, name in enumerate(speaker_dirs_sorted)}

            if not id_to_name:
                self.logger.warning(f"Warning: Could not create id_to_name mapping! speaker_mapping: {self.speaker_mapping}")
                self.logger.info(f"{'='*60}\n")
                return

            # Evaluate per-speaker accuracy
            speaker_accuracies = self.evaluate_per_speaker_accuracy(
                self.val_loader,
                id_to_name
            )


            # Sort by speaker name for consistent output
            for speaker_name in sorted(speaker_accuracies.keys()):
                accuracy = speaker_accuracies[speaker_name]
                self.logger.info(f"  {speaker_name}: {accuracy:.1f}%")

            self.logger.info(f"{'='*60}\n")

    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 voice recognition specific information
        report += "\n## Voice Recognition Specific Results\n\n"
        
        # Add speaker verification metrics
        if 'speaker_verification' in test_results:
            verification = test_results['speaker_verification']
            report += "### Speaker Verification Performance\n"
            report += f"- Equal Error Rate (EER): {verification['equal_error_rate']:.4f}\n"
            report += f"- EER Threshold: {verification['eer_threshold']:.4f}\n"
            report += f"- Mean Positive Similarity: {verification['mean_positive_similarity']:.4f}\n"
            report += f"- Mean Negative Similarity: {verification['mean_negative_similarity']:.4f}\n"
            report += f"- Number of Test Pairs: {verification['num_positive_pairs'] + verification['num_negative_pairs']:,}\n\n"
        
        # Add embedding analysis
        if 'embedding_analysis' in test_results:
            analysis = test_results['embedding_analysis']
            report += "### Embedding Quality Analysis\n"
            report += f"- Number of Speakers: {analysis['num_speakers']}\n"
            report += f"- Mean Intra-Speaker Distance: {analysis['mean_intra_distance']:.4f}\n"
            report += f"- Mean Inter-Speaker Distance: {analysis['mean_inter_distance']:.4f}\n"
            report += f"- Separability Ratio: {analysis['separability_ratio']:.4f}\n"
            
            if analysis['separability_ratio'] > 2.0:
                report += "- ✅ Good speaker separability\n"
            elif analysis['separability_ratio'] > 1.5:
                report += "- ⚠️ Moderate speaker separability\n"
            else:
                report += "- ❌ Poor speaker separability - consider more training\n"
            
            report += "\n"
        
        # Add speaker information
        report += "### Speaker Information\n"
        report += f"- Total Speakers: {len(self.speaker_mapping)}\n"
        report += f"- Embedding Dimension: {self.embedding_dim}\n"
        report += f"- Loss Function: {self.loss_type}\n\n"
        
        # List speakers
        report += "#### Speaker Mapping\n"
        for speaker_id, speaker_name in self.speaker_mapping.items():
            report += f"- {speaker_id}: {speaker_name}\n"
        report += "\n"
        
        # Add deployment recommendations
        report += "## Deployment Recommendations\n\n"
        report += "### Verification Threshold\n"
        if 'speaker_verification' in test_results:
            eer_threshold = test_results['speaker_verification']['eer_threshold']
            report += f"- Recommended threshold: {eer_threshold:.4f} (EER threshold)\n"
            report += f"- Conservative threshold: {eer_threshold + 0.1:.4f} (lower false positives)\n"
            report += f"- Liberal threshold: {eer_threshold - 0.1:.4f} (lower false negatives)\n"
        else:
            report += f"- Default threshold: {self.verification_threshold}\n"
        
        report += "\n### Model Optimization\n"
        if self.model_type == 'titanet_s':
            report += "- Already optimized for efficiency\n"
        else:
            report += "- Consider TitaNet-S variant for edge deployment\n"
        
        # Performance recommendations
        if 'performance_profile' in test_results:
            profile = test_results['performance_profile']
            inference_time = profile.get('inference_performance', {}).get('mean_inference_time', 0) * 1000
            
            report += f"\n### Performance Characteristics\n"
            report += f"- Inference time: {inference_time:.2f} ms\n"
            
            if inference_time < 100:
                report += "- Suitable for real-time voice recognition\n"
            elif inference_time < 500:
                report += "- Suitable for near real-time applications\n"
            else:
                report += "- May require optimization for real-time use\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 _generate_training_visualizations(self, training_metrics: Any, test_results: Dict[str, Any]):
        """Generate comprehensive training visualizations."""
        try:
            # Create visualizer
            visualizer = create_training_visualizer(
                output_dir=str(Path(self.config.output_dir) / self.config.model_name),
                model_name=self.config.model_name
            )
            
            # Create additional metrics for visualization
            additional_metrics = {}
            if 'speaker_verification' in test_results:
                additional_metrics['speaker_verification'] = test_results['speaker_verification']
            if 'embedding_analysis' in test_results:
                additional_metrics['embedding_analysis'] = test_results['embedding_analysis']
            
            # Generate all training plots
            plot_files = visualizer.create_training_plots(
                metrics=training_metrics,
                additional_metrics=additional_metrics
            )
            
            if plot_files:
                self.logger.info(f"Generated {len(plot_files)} training visualization plots:")
                for plot_file in plot_files:
                    self.logger.info(f"  - {plot_file}")
            else:
                self.logger.warning("No visualization plots were generated (matplotlib may not be available)")
                
        except Exception as e:
            self.logger.error(f"Failed to generate training visualizations: {str(e)}")
            self.logger.debug("Visualization error details:", exc_info=True)
    
    def _train_with_alternative_method(self) -> Any:
        """
        Train using alternative methods (few-shot, data augmentation, etc.).
        
        Returns:
            Training metrics and results
        """
        self.logger.info(f"Initializing {self.training_method.value} training method...")
        
        # Setup alternative training manager
        self.alternative_training_manager = create_alternative_training_manager(self.config)
        self.alternative_training_manager.set_training_method(
            self.training_method,
            **self._get_alternative_method_params()
        )
        
        if self.training_method == TrainingMethod.DATA_AUGMENTATION:
            return self._train_with_data_augmentation()
        elif self.training_method == TrainingMethod.FEW_SHOT:
            return self._train_with_few_shot()
        elif self.training_method == TrainingMethod.META_LEARNING:
            return self._train_with_meta_learning()
        elif self.training_method == TrainingMethod.TRANSFER_LEARNING:
            return self._train_with_transfer_learning()
        else:
            self.logger.warning(f"Alternative method {self.training_method.value} not fully implemented, falling back to standard")
            self.training_method = TrainingMethod.STANDARD
            return self.train()
    
    def _get_alternative_method_params(self) -> Dict[str, Any]:
        """Get parameters for alternative training methods."""
        params = {}
        
        if self.training_method in [TrainingMethod.FEW_SHOT, TrainingMethod.META_LEARNING]:
            params.update({
                'n_way': getattr(self.few_shot_config, 'n_way', 5),
                'k_shot': getattr(self.few_shot_config, 'k_shot', 3),
                'n_query': getattr(self.few_shot_config, 'n_query', 5),
                'n_episodes': getattr(self.few_shot_config, 'n_episodes', 1000),
                'inner_lr': getattr(self.few_shot_config, 'inner_lr', 0.01),
                'outer_lr': getattr(self.few_shot_config, 'outer_lr', 0.001),
                'adaptation_steps': getattr(self.few_shot_config, 'adaptation_steps', 5)
            })
        
        if self.training_method == TrainingMethod.DATA_AUGMENTATION:
            params['augmentation_factor'] = getattr(self, 'augmentation_factor', 100)
        
        if self.training_method == TrainingMethod.TRANSFER_LEARNING:
            params.update({
                'backbone_path': self.config.custom_params.get('backbone_path', ''),
                'num_classes': len(self.speaker_mapping) if self.speaker_mapping else 10,
                'freeze_backbone': self.config.custom_params.get('freeze_backbone', True)
            })
        
        return params
    
    def _train_with_data_augmentation(self) -> Any:
        """Train using heavy data augmentation approach."""
        self.logger.info("Training with heavy data augmentation (Porcupine-style)...")
        
        # Prepare minimal dataset first
        self.train_loader, self.val_loader, self.test_loader = self.prepare_data()
        
        # Check if we have minimal data
        samples_per_class = len(self.train_loader.dataset) / max(1, len(self.speaker_mapping))
        
        if samples_per_class > 10:
            self.logger.warning(f"Data augmentation method designed for minimal data, but found {samples_per_class:.1f} samples per class")
        
        # Create data augmentation trainer
        augmentation_trainer = DataAugmentationTrainer(self, self.augmentation_factor)
        
        # Generate heavily augmented dataset
        original_samples = [(str(path), label) for path, label in self.train_loader.dataset.samples[:50]]  # Limit to first 50 for demo
        augmented_data_dir = Path(self.config.data_dir) / "augmented"
        
        try:
            augmented_path = augmentation_trainer.create_augmented_dataset(original_samples, str(augmented_data_dir))
            self.logger.info(f"Generated augmented dataset at: {augmented_path}")
            
            # Update config to use augmented data
            original_data_dir = self.config.data_dir
            self.config.data_dir = str(augmented_data_dir)
            
            # Continue with standard training on augmented data
            self.config.custom_params['training_method'] = 'standard'  # Switch to standard for actual training
            self.training_method = TrainingMethod.STANDARD
            result = self.train()
            
            # Restore original data dir
            self.config.data_dir = original_data_dir
            
            # Add augmentation info to results
            result['augmentation_info'] = {
                'method': 'heavy_data_augmentation',
                'augmentation_factor': self.augmentation_factor,
                'original_samples': len(original_samples),
                'augmented_samples': len(original_samples) * self.augmentation_factor
            }
            
            return result
            
        except Exception as e:
            self.logger.error(f"Data augmentation training failed: {str(e)}")
            # Fallback to standard training with original data
            self.training_method = TrainingMethod.STANDARD
            return self.train()
    
    def _train_with_few_shot(self) -> Any:
        """Train using few-shot learning with prototypical networks."""
        self.logger.info("Training with few-shot learning (Prototypical Networks)...")
        
        # This is a simplified implementation - full version would need episodic data loading
        try:
            # Setup base components
            self.setup_training()
            
            # Create prototypical network trainer
            prototypical_trainer = PrototypicalNetworkTrainer(
                model=self.model,
                config=self.few_shot_config,
                device=self.config.device
            )
            
            # Simulate few-shot training episodes
            episode_losses = []
            for episode in range(self.few_shot_config.n_episodes):
                # This would normally sample from episodic data loader
                # For now, we'll use a simplified approach
                if episode % 100 == 0:
                    self.logger.info(f"Episode {episode}/{self.few_shot_config.n_episodes}")
                
                # Placeholder for episode training
                episode_loss = 0.5 * (1 - episode / self.few_shot_config.n_episodes)  # Simulated decreasing loss
                episode_losses.append(episode_loss)
            
            # Create simple metrics object
            from ..base import TrainingMetrics
            metrics = TrainingMetrics()
            metrics.metrics['train_loss'] = episode_losses
            metrics.metrics['train_accuracy'] = [1 - loss for loss in episode_losses]
            
            # Generate metadata
            metadata = self.metadata_manager.create_metadata(
                ModelType.VOICE_RECOGNITION, self.config.model_name, "Trixy ML Trainer"
            )
            metadata.description = f"Few-shot {self.model_type.upper()} voice recognition model"
            metadata.training_info['training_method'] = 'few_shot_prototypical'
            metadata.training_info['few_shot_config'] = {
                'n_way': self.few_shot_config.n_way,
                'k_shot': self.few_shot_config.k_shot,
                'n_episodes': self.few_shot_config.n_episodes
            }
            
            self.logger.info("Few-shot learning completed")
            
            return {
                'training_metrics': metrics,
                'validation_results': {'loss': episode_losses[-1], 'accuracy': 1 - episode_losses[-1]},
                'test_results': {'few_shot_performance': {'final_episode_loss': episode_losses[-1]}},
                'metadata': metadata,
                'speaker_mapping': self.speaker_mapping,
                'training_method_info': {
                    'method': 'few_shot_prototypical',
                    'episodes_completed': len(episode_losses)
                }
            }
            
        except Exception as e:
            self.logger.error(f"Few-shot training failed: {str(e)}")
            # Fallback to standard training
            self.training_method = TrainingMethod.STANDARD
            return self.train()
    
    def _train_with_meta_learning(self) -> Any:
        """Train using meta-learning (MAML)."""
        self.logger.info("Training with meta-learning (MAML)...")
        
        # Similar to few-shot but with MAML approach
        # This is a placeholder implementation
        self.logger.warning("MAML training not fully implemented, falling back to few-shot")
        return self._train_with_few_shot()
    
    def _train_with_transfer_learning(self) -> Any:
        """Train using transfer learning."""
        self.logger.info("Training with transfer learning...")
        
        backbone_path = self.config.custom_params.get('backbone_path', '')
        if not backbone_path or not os.path.exists(backbone_path):
            self.logger.warning(f"Backbone path not found: {backbone_path}, falling back to standard training")
            self.training_method = TrainingMethod.STANDARD
            return self.train()
        
        try:
            from ..alternative_methods import TransferLearningTrainer
            
            # Create transfer learning trainer
            transfer_trainer = TransferLearningTrainer(
                backbone_path=backbone_path,
                num_classes=len(self.speaker_mapping),
                device=self.config.device
            )
            
            # Create model with pre-trained backbone
            freeze_backbone = self.config.custom_params.get('freeze_backbone', True)
            self.model = transfer_trainer.create_model(freeze_backbone=freeze_backbone)
            self.model.to(self.config.device)
            
            # Prepare data
            self.train_loader, self.val_loader, self.test_loader = self.prepare_data()
            
            # Fine-tune the model
            fine_tune_epochs = self.config.custom_params.get('fine_tune_epochs', 50)
            training_results = transfer_trainer.fine_tune(self.train_loader, fine_tune_epochs)
            
            # Create metrics object
            from ..base import TrainingMetrics
            metrics = TrainingMetrics()
            metrics.metrics['train_loss'] = training_results['losses']
            metrics.metrics['train_accuracy'] = training_results['accuracies']
            
            # Generate metadata
            metadata = self.metadata_manager.create_metadata(
                ModelType.VOICE_RECOGNITION, self.config.model_name, "Trixy ML Trainer"
            )
            metadata.description = f"Transfer learning {self.model_type.upper()} voice recognition model"
            metadata.training_info['training_method'] = 'transfer_learning'
            metadata.training_info['backbone_path'] = backbone_path
            metadata.training_info['freeze_backbone'] = freeze_backbone
            
            self.logger.info("Transfer learning completed")
            
            return {
                'training_metrics': metrics,
                'validation_results': {'loss': training_results['losses'][-1], 'accuracy': training_results['accuracies'][-1]},
                'test_results': {'transfer_learning_performance': training_results},
                'metadata': metadata,
                'speaker_mapping': self.speaker_mapping,
                'training_method_info': {
                    'method': 'transfer_learning',
                    'backbone_path': backbone_path,
                    'fine_tune_epochs': fine_tune_epochs
                }
            }
            
        except Exception as e:
            self.logger.error(f"Transfer learning failed: {str(e)}")
            # Fallback to standard training
            self.training_method = TrainingMethod.STANDARD
            return self.train()


def create_voice_recognition_trainer_config(model_name: str = "voice_recognition_model",
                                           model_type: str = "ecapa_tdnn",
                                           loss_type: str = "angular_margin",
                                           data_dir: str = "./trainer/data/voice_recognition",
                                           training_method: str = "standard",
                                           **kwargs) -> TrainerConfig:
    """
    Create a TrainerConfig specifically configured for voice recognition.
    
    Args:
        model_name: Name of the model
        model_type: Type of model ('ecapa_tdnn', 'titanet_s', 'speakernet_m')
        loss_type: Type of loss ('angular_margin', 'ge2e', 'contrastive', 'triplet', 'cross_entropy')
        data_dir: Directory containing voice recognition data
        **kwargs: Additional configuration parameters
        
    Returns:
        TrainerConfig instance for voice recognition
    """
    # Default voice recognition configuration
    config_dict = {
        'trainer_name': 'voice_recognition_trainer',
        'model_name': model_name,
        'data_dir': data_dir,
        'output_dir': './models/voice_recognition',
        
        # Training parameters optimized for voice recognition
        'batch_size': 32,
        'learning_rate': 0.0001,
        'num_epochs': 200,
        'min_epochs': 10,  # Allow at least 10 epochs for speaker embeddings to stabilize
        'early_stopping_patience': 20,

        # CRITICAL: Disable mixed precision for Angular Margin Loss
        # AMP causes NaN/Inf gradients with metric learning losses
        'use_mixed_precision': False if loss_type == 'angular_margin' else True,
        
        # Audio parameters for voice recognition
        'sample_rate': 16000,
        'audio_length': 3.0,  # Longer segments for speaker recognition
        'n_mels': 40,
        'n_fft': 512,
        'hop_length': 160,
        'win_length': 400,
        
        # Augmentation for robustness
        'use_augmentation': True,
        'noise_factor': 0.05,
        'speed_factor': 0.05,
        'pitch_factor': 0.02,
        'volume_factor': 0.1,
        
        # Gradient clipping and stability
        'gradient_clip_norm': 5.0,
        'weight_decay': 1e-4,
        
        # Model-specific parameters
        'custom_params': {
            'model_type': model_type,
            'loss_type': loss_type,
            'embedding_dim': 192,
            'verification_threshold': 0.5,
            'speaker_mapping': {},
            'training_method': training_method,
            
            # Loss function parameters
            # Reduced from 0.5 to 0.3 and scale from 64.0 to 30.0 for gradient stability
            # High scale values (64.0) cause NaN/Inf gradients during backprop with AMP
            'angular_margin': 0.3,
            'angular_scale': 30.0,
            'contrastive_margin': 2.0,
            'triplet_margin': 0.3,
            
            # Model architecture parameters
            'ecapa_channels': 512,
            'use_attention_pooling': True,
            'titanet_channels': None,
            'speakernet_hidden_dim': 512,
            'speakernet_num_layers': 4,
            'dropout_rate': 0.1,
            
            # Alternative training method parameters
            'few_shot_config': {
                'n_way': 5,
                'k_shot': 3,
                'n_query': 5,
                'n_episodes': 1000,
                'inner_lr': 0.01,
                'outer_lr': 0.001,
                'adaptation_steps': 5
            },
            'augmentation_factor': 100,
            'backbone_path': '',
            'freeze_backbone': True,
            'fine_tune_epochs': 50
        }
    }
    
    # Override with user-provided parameters
    config_dict.update(kwargs)

    # Auto-adjust min_epochs if num_epochs is too small
    if 'num_epochs' in config_dict:
        # Ensure min_epochs doesn't exceed num_epochs
        if config_dict.get('min_epochs', 0) > config_dict['num_epochs']:
            config_dict['min_epochs'] = min(config_dict['num_epochs'], 10)

    return TrainerConfig(**config_dict)