Trixy3/trixy_core/ml/voice_recognition/audio_features.py

"""
Audio feature extraction for Voice Recognition system.

This module provides comprehensive audio feature extraction capabilities including
Log-Mel spectrograms and MFCC features optimized for speaker recognition tasks.
"""

import logging
import warnings
from typing import Optional, Tuple, Union, List, Dict, Any
from dataclasses import dataclass
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
import torchaudio.transforms as T
from pathlib import Path

# Configure logging
logger = logging.getLogger(__name__)


@dataclass
class LogMelConfig:
    """Configuration for Log-Mel spectrogram extraction."""
    sample_rate: int = 16000
    n_mels: int = 40
    n_fft: int = 512
    hop_length: int = 160  # 10ms at 16kHz
    win_length: int = 400  # 25ms at 16kHz
    f_min: float = 80.0
    f_max: Optional[float] = None
    power: float = 2.0
    normalized: bool = True
    center: bool = True
    pad_mode: str = "reflect"
    
    def __post_init__(self):
        """Validate configuration parameters."""
        if self.f_max is None:
            self.f_max = self.sample_rate // 2
        
        if self.f_min >= self.f_max:
            raise ValueError("f_min must be less than f_max")
        if self.hop_length > self.win_length:
            raise ValueError("hop_length must be <= win_length")
        if self.n_mels < 1:
            raise ValueError("n_mels must be positive")


@dataclass
class MFCCConfig:
    """Configuration for MFCC feature extraction."""
    sample_rate: int = 16000
    n_mfcc: int = 40
    n_mels: int = 80
    n_fft: int = 512
    hop_length: int = 160  # 10ms at 16kHz
    win_length: int = 400  # 25ms at 16kHz
    f_min: float = 80.0
    f_max: Optional[float] = None
    dct_type: int = 2
    norm: str = "ortho"
    log_mels: bool = False
    melkwargs: Dict[str, Any] = None
    
    def __post_init__(self):
        """Validate configuration and set defaults."""
        if self.f_max is None:
            self.f_max = self.sample_rate // 2
        
        if self.melkwargs is None:
            self.melkwargs = {}
        
        if self.n_mfcc > self.n_mels:
            raise ValueError("n_mfcc must be <= n_mels")
        if self.dct_type not in [1, 2, 3, 4]:
            raise ValueError("dct_type must be 1, 2, 3, or 4")


class VRAudioFeatureExtractor(nn.Module):
    """
    Audio feature extractor for voice recognition.
    
    Supports both Log-Mel spectrograms and MFCC features with optional
    preprocessing and normalization.
    """
    
    def __init__(self, 
                 feature_type: str = "log_mel",
                 log_mel_config: Optional[LogMelConfig] = None,
                 mfcc_config: Optional[MFCCConfig] = None,
                 normalize_features: bool = True,
                 apply_cmvn: bool = True,
                 add_deltas: bool = False,
                 add_delta_deltas: bool = False):
        """
        Initialize audio feature extractor.
        
        Args:
            feature_type: Type of features to extract ("log_mel" or "mfcc")
            log_mel_config: Configuration for Log-Mel spectrogram
            mfcc_config: Configuration for MFCC features
            normalize_features: Whether to normalize features
            apply_cmvn: Whether to apply Cepstral Mean and Variance Normalization
            add_deltas: Whether to add delta features
            add_delta_deltas: Whether to add delta-delta features
        """
        super().__init__()
        
        self.feature_type = feature_type.lower()
        self.normalize_features = normalize_features
        self.apply_cmvn = apply_cmvn
        self.add_deltas = add_deltas
        self.add_delta_deltas = add_delta_deltas
        
        if self.feature_type == "log_mel":
            self.config = log_mel_config or LogMelConfig()
            self._setup_log_mel_transforms()
        elif self.feature_type == "mfcc":
            self.config = mfcc_config or MFCCConfig()
            self._setup_mfcc_transforms()
        else:
            raise ValueError(f"Unsupported feature type: {feature_type}")
        
        # Statistics for normalization
        self.register_buffer('feature_mean', torch.zeros(1))
        self.register_buffer('feature_std', torch.ones(1))
        self.register_buffer('num_updates', torch.zeros(1))
        
        # Delta computation
        if self.add_deltas or self.add_delta_deltas:
            self.delta_transform = T.ComputeDeltas()
    
    def _setup_log_mel_transforms(self):
        """Setup transforms for Log-Mel spectrogram extraction."""
        self.spectrogram = T.Spectrogram(
            n_fft=self.config.n_fft,
            hop_length=self.config.hop_length,
            win_length=self.config.win_length,
            power=self.config.power,
            normalized=self.config.normalized,
            center=self.config.center,
            pad_mode=self.config.pad_mode
        )
        
        self.mel_scale = T.MelScale(
            n_mels=self.config.n_mels,
            sample_rate=self.config.sample_rate,
            f_min=self.config.f_min,
            f_max=self.config.f_max,
            n_stft=self.config.n_fft // 2 + 1
        )
        
        self.amplitude_to_db = T.AmplitudeToDB()
    
    def _setup_mfcc_transforms(self):
        """Setup transforms for MFCC feature extraction."""
        self.mfcc_transform = T.MFCC(
            sample_rate=self.config.sample_rate,
            n_mfcc=self.config.n_mfcc,
            melkwargs={
                "n_fft": self.config.n_fft,
                "hop_length": self.config.hop_length,
                "win_length": self.config.win_length,
                "f_min": self.config.f_min,
                "f_max": self.config.f_max,
                "n_mels": self.config.n_mels,
                **self.config.melkwargs
            },
            dct_type=self.config.dct_type,
            norm=self.config.norm,
            log_mels=self.config.log_mels
        )
    
    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        """
        Extract features from audio waveform.
        
        Args:
            waveform: Input audio tensor of shape (batch_size, num_samples) or (num_samples,)
            
        Returns:
            features: Extracted features of shape (batch_size, n_features, n_frames)
        """
        # Ensure batch dimension
        if waveform.dim() == 1:
            waveform = waveform.unsqueeze(0)
        
        # Validate input
        if waveform.dim() != 2:
            raise ValueError(f"Expected 1D or 2D tensor, got {waveform.dim()}D")
        
        # Extract features based on type
        if self.feature_type == "log_mel":
            features = self._extract_log_mel(waveform)
        else:  # mfcc
            features = self._extract_mfcc(waveform)
        
        # Add delta features if requested
        if self.add_deltas:
            deltas = self.delta_transform(features)
            features = torch.cat([features, deltas], dim=1)
            
            if self.add_delta_deltas:
                delta_deltas = self.delta_transform(deltas)
                features = torch.cat([features, delta_deltas], dim=1)
        
        # Apply CMVN if requested
        if self.apply_cmvn:
            features = self._apply_cmvn(features)
        
        # Normalize features if requested
        if self.normalize_features:
            features = self._normalize_features(features)
        
        return features
    
    def _extract_log_mel(self, waveform: torch.Tensor) -> torch.Tensor:
        """Extract Log-Mel spectrogram features."""
        # Compute spectrogram
        specgram = self.spectrogram(waveform)
        
        # Apply mel scale
        mel_specgram = self.mel_scale(specgram)
        
        # Convert to dB
        log_mel = self.amplitude_to_db(mel_specgram)
        
        return log_mel
    
    def _extract_mfcc(self, waveform: torch.Tensor) -> torch.Tensor:
        """Extract MFCC features."""
        return self.mfcc_transform(waveform)
    
    def _apply_cmvn(self, features: torch.Tensor) -> torch.Tensor:
        """Apply Cepstral Mean and Variance Normalization."""
        # Compute mean and std across time dimension
        mean = features.mean(dim=-1, keepdim=True)
        std = features.std(dim=-1, keepdim=True)
        
        # Avoid division by zero
        std = torch.clamp(std, min=1e-8)
        
        # Normalize
        features = (features - mean) / std
        
        return features
    
    def _normalize_features(self, features: torch.Tensor) -> torch.Tensor:
        """Normalize features using running statistics."""
        if self.training:
            # Update running statistics
            batch_mean = features.mean(dim=(0, 2), keepdim=True)
            batch_var = features.var(dim=(0, 2), keepdim=True, unbiased=False)
            
            # Update running mean and std
            momentum = 0.1
            self.feature_mean = (1 - momentum) * self.feature_mean + momentum * batch_mean.squeeze()
            
            # For std, we need to be careful with the square root
            batch_std = torch.sqrt(batch_var + 1e-8)
            if self.feature_std.numel() == 1:  # First update
                self.feature_std = batch_std.squeeze()
            else:
                self.feature_std = (1 - momentum) * self.feature_std + momentum * batch_std.squeeze()
            
            self.num_updates += 1
        
        # Apply normalization
        if self.num_updates > 0:
            mean = self.feature_mean.view(1, -1, 1)
            std = torch.clamp(self.feature_std.view(1, -1, 1), min=1e-8)
            features = (features - mean) / std
        
        return features
    
    def get_feature_dim(self) -> int:
        """Get the output feature dimension."""
        if self.feature_type == "log_mel":
            base_dim = self.config.n_mels
        else:  # mfcc
            base_dim = self.config.n_mfcc
        
        # Account for delta features
        multiplier = 1
        if self.add_deltas:
            multiplier += 1
        if self.add_delta_deltas:
            multiplier += 1
        
        return base_dim * multiplier
    
    def get_output_length(self, input_length: int) -> int:
        """Calculate output sequence length given input length."""
        if self.feature_type == "log_mel":
            hop_length = self.config.hop_length
        else:  # mfcc
            hop_length = self.config.hop_length
        
        return (input_length + hop_length // 2) // hop_length
    
    def reset_stats(self):
        """Reset normalization statistics."""
        self.feature_mean.zero_()
        self.feature_std.fill_(1.0)
        self.num_updates.zero_()


class BatchVRAudioFeatureExtractor(nn.Module):
    """
    Batch-optimized version of the audio feature extractor.
    
    Designed for efficient processing of multiple audio samples simultaneously.
    """
    
    def __init__(self, 
                 feature_extractor: VRAudioFeatureExtractor,
                 fixed_length: Optional[int] = None,
                 padding_mode: str = "constant"):
        """
        Initialize batch feature extractor.
        
        Args:
            feature_extractor: Base feature extractor to use
            fixed_length: Fixed output length (for padding/truncation)
            padding_mode: Padding mode for sequences ("constant", "reflect", "replicate")
        """
        super().__init__()
        
        self.feature_extractor = feature_extractor
        self.fixed_length = fixed_length
        self.padding_mode = padding_mode
    
    def forward(self, waveforms: torch.Tensor, lengths: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Extract features from batch of waveforms.
        
        Args:
            waveforms: Batch of waveforms [batch_size, max_length]
            lengths: Actual lengths of each waveform [batch_size]
            
        Returns:
            features: Extracted features [batch_size, n_features, n_frames]
            feature_lengths: Lengths of feature sequences [batch_size]
        """
        batch_size = waveforms.size(0)
        
        # Process each waveform
        features_list = []
        feature_lengths = []
        
        for i in range(batch_size):
            waveform = waveforms[i]
            
            # Trim to actual length if provided
            if lengths is not None:
                waveform = waveform[:lengths[i]]
            
            # Extract features
            features = self.feature_extractor(waveform)
            features_list.append(features.squeeze(0))
            feature_lengths.append(features.size(-1))
        
        # Convert to tensor
        feature_lengths = torch.tensor(feature_lengths, device=waveforms.device)
        
        # Pad/truncate to fixed length if specified
        if self.fixed_length is not None:
            max_length = self.fixed_length
        else:
            max_length = max(feature_lengths.tolist())
        
        # Pad features to same length
        padded_features = []
        for features in features_list:
            if features.size(-1) < max_length:
                # Pad
                pad_amount = max_length - features.size(-1)
                if self.padding_mode == "constant":
                    padded = F.pad(features, (0, pad_amount), mode='constant', value=0)
                else:
                    padded = F.pad(features, (0, pad_amount), mode=self.padding_mode)
            elif features.size(-1) > max_length:
                # Truncate
                padded = features[:, :max_length]
            else:
                padded = features
            
            padded_features.append(padded)
        
        # Stack into batch
        batch_features = torch.stack(padded_features, dim=0)
        
        return batch_features, feature_lengths


def create_feature_extractor(feature_type: str,
                           sample_rate: int = 16000,
                           **kwargs) -> VRAudioFeatureExtractor:
    """
    Factory function to create feature extractor with common configurations.
    
    Args:
        feature_type: Type of features ("log_mel" or "mfcc")
        sample_rate: Audio sample rate
        **kwargs: Additional configuration parameters
        
    Returns:
        Configured feature extractor
    """
    if feature_type.lower() == "log_mel":
        config = LogMelConfig(sample_rate=sample_rate, **kwargs)
        return VRAudioFeatureExtractor(
            feature_type="log_mel",
            log_mel_config=config
        )
    elif feature_type.lower() == "mfcc":
        config = MFCCConfig(sample_rate=sample_rate, **kwargs)
        return VRAudioFeatureExtractor(
            feature_type="mfcc",
            mfcc_config=config
        )
    else:
        raise ValueError(f"Unsupported feature type: {feature_type}")


def validate_feature_config(config: Union[LogMelConfig, MFCCConfig]) -> List[str]:
    """
    Validate feature extraction configuration.
    
    Args:
        config: Feature extraction configuration
        
    Returns:
        List of validation warnings
    """
    warnings = []
    
    if isinstance(config, LogMelConfig):
        # Check for common issues with Log-Mel configuration
        if config.n_mels < 20:
            warnings.append("Low number of mel bins may reduce feature quality")
        if config.n_mels > 128:
            warnings.append("High number of mel bins may be redundant")
        
        # Check frequency range
        nyquist = config.sample_rate // 2
        if config.f_max > nyquist:
            warnings.append(f"f_max ({config.f_max}) exceeds Nyquist frequency ({nyquist})")
        
        # Check hop length
        hop_ms = (config.hop_length / config.sample_rate) * 1000
        if hop_ms > 20:
            warnings.append(f"Large hop length ({hop_ms:.1f}ms) may reduce temporal resolution")
        
    elif isinstance(config, MFCCConfig):
        # Check for common issues with MFCC configuration
        if config.n_mfcc < 12:
            warnings.append("Low number of MFCC coefficients may reduce feature quality")
        if config.n_mfcc > config.n_mels:
            warnings.append("n_mfcc should not exceed n_mels")
        
        # Check mel filterbank configuration
        if config.n_mels < config.n_mfcc * 2:
            warnings.append("n_mels should be at least 2x n_mfcc for good frequency resolution")
    
    return warnings


def compute_feature_statistics(features: torch.Tensor) -> Dict[str, float]:
    """
    Compute statistics for feature analysis.
    
    Args:
        features: Feature tensor [batch_size, n_features, n_frames]
        
    Returns:
        Dictionary of feature statistics
    """
    with torch.no_grad():
        # Flatten spatial dimensions
        flat_features = features.view(features.size(0), -1)
        
        stats = {
            'mean': float(flat_features.mean()),
            'std': float(flat_features.std()),
            'min': float(flat_features.min()),
            'max': float(flat_features.max()),
            'dynamic_range': float(flat_features.max() - flat_features.min()),
            'zero_crossing_rate': float((torch.diff(torch.sign(flat_features)) != 0).float().mean()),
            'sparsity': float((flat_features.abs() < 1e-6).float().mean())
        }
        
        # Compute per-feature statistics
        feature_means = features.mean(dim=(0, 2))
        feature_stds = features.std(dim=(0, 2))
        
        stats['feature_mean_range'] = float(feature_means.max() - feature_means.min())
        stats['feature_std_range'] = float(feature_stds.max() - feature_stds.min())
        
        return stats


# Example usage and testing functions
def test_feature_extraction():
    """Test feature extraction with sample audio."""
    # Create sample audio
    sample_rate = 16000
    duration = 3.0
    num_samples = int(sample_rate * duration)
    
    # Generate test signal (sine wave + noise)
    t = torch.linspace(0, duration, num_samples)
    signal = torch.sin(2 * torch.pi * 440 * t) + 0.1 * torch.randn(num_samples)
    
    # Test Log-Mel extraction
    log_mel_extractor = create_feature_extractor("log_mel", sample_rate=sample_rate)
    log_mel_features = log_mel_extractor(signal)
    
    print(f"Log-Mel features shape: {log_mel_features.shape}")
    print(f"Log-Mel stats: {compute_feature_statistics(log_mel_features.unsqueeze(0))}")
    
    # Test MFCC extraction
    mfcc_extractor = create_feature_extractor("mfcc", sample_rate=sample_rate)
    mfcc_features = mfcc_extractor(signal)
    
    print(f"MFCC features shape: {mfcc_features.shape}")
    print(f"MFCC stats: {compute_feature_statistics(mfcc_features.unsqueeze(0))}")
    
    # Test batch processing
    batch_waveforms = signal.unsqueeze(0).repeat(4, 1)
    batch_extractor = BatchVRAudioFeatureExtractor(log_mel_extractor, fixed_length=100)
    batch_features, batch_lengths = batch_extractor(batch_waveforms)
    
    print(f"Batch features shape: {batch_features.shape}")
    print(f"Batch lengths: {batch_lengths}")


if __name__ == "__main__":
    test_feature_extraction()