"""
Feature Extraction for Phoneme Scoring Model
Extracts acoustic and linguistic features from phoneme segments for ML classification
"""

import numpy as np
import librosa
from typing import Dict, List, Optional, Tuple


def extract_segment_audio(audio_array: np.ndarray, 
                          start_time: float, 
                          end_time: float,
                          sr: int = 16000) -> np.ndarray:
    """
    Extract audio segment for a phoneme based on timestamps.
    
    Args:
        audio_array: Full audio samples
        start_time: Start time in seconds
        end_time: End time in seconds
        sr: Sample rate
        
    Returns:
        Audio segment as numpy array
    """
    start_sample = int(start_time * sr)
    end_sample = int(end_time * sr)
    
    # Ensure valid bounds
    start_sample = max(0, start_sample)
    end_sample = min(len(audio_array), end_sample)
    
    if start_sample >= end_sample:
        return np.array([])
    
    return audio_array[start_sample:end_sample]


def compute_mfcc_features(audio_segment: np.ndarray, 
                          sr: int = 16000,
                          n_mfcc: int = 13) -> Dict[str, float]:
    """
    Compute MFCC-based features for a phoneme segment.
    
    Args:
        audio_segment: Audio samples for the phoneme
        sr: Sample rate
        n_mfcc: Number of MFCC coefficients
        
    Returns:
        Dictionary of MFCC-based features
    """
    features = {}
    
    if len(audio_segment) < 512:  # Too short for reliable features
        # Return zero features
        for i in range(n_mfcc):
            features[f'mfcc_{i}_mean'] = 0.0
            features[f'mfcc_{i}_std'] = 0.0
        features['mfcc_delta_mean'] = 0.0
        features['mfcc_delta_std'] = 0.0
        return features
    
    # Compute MFCCs
    mfccs = librosa.feature.mfcc(y=audio_segment, sr=sr, n_mfcc=n_mfcc)
    
    # Statistical features for each MFCC coefficient
    for i in range(n_mfcc):
        features[f'mfcc_{i}_mean'] = float(np.mean(mfccs[i]))
        features[f'mfcc_{i}_std'] = float(np.std(mfccs[i]))
    
    # Delta MFCCs (first-order derivatives)
    try:
        delta_mfccs = librosa.feature.delta(mfccs)
        features['mfcc_delta_mean'] = float(np.mean(delta_mfccs))
        features['mfcc_delta_std'] = float(np.std(delta_mfccs))
    except:
        features['mfcc_delta_mean'] = 0.0
        features['mfcc_delta_std'] = 0.0
    
    return features


def compute_energy_features(audio_segment: np.ndarray, 
                            sr: int = 16000) -> Dict[str, float]:
    """
    Compute energy-based features.
    
    Args:
        audio_segment: Audio samples for the phoneme
        sr: Sample rate
        
    Returns:
        Dictionary of energy features
    """
    features = {}
    
    if len(audio_segment) == 0:
        features['energy_mean'] = 0.0
        features['energy_std'] = 0.0
        features['energy_max'] = 0.0
        features['rms_energy'] = 0.0
        return features
    
    # RMS energy
    rms = librosa.feature.rms(y=audio_segment)[0]
    features['energy_mean'] = float(np.mean(rms))
    features['energy_std'] = float(np.std(rms))
    features['energy_max'] = float(np.max(rms))
    features['rms_energy'] = float(np.sqrt(np.mean(audio_segment ** 2)))
    
    return features


def compute_spectral_features(audio_segment: np.ndarray, 
                              sr: int = 16000) -> Dict[str, float]:
    """
    Compute spectral features.
    
    Args:
        audio_segment: Audio samples for the phoneme
        sr: Sample rate
        
    Returns:
        Dictionary of spectral features
    """
    features = {}
    
    if len(audio_segment) < 512:
        features['spectral_centroid'] = 0.0
        features['spectral_rolloff'] = 0.0
        features['spectral_flatness'] = 0.0
        features['zero_crossing_rate'] = 0.0
        return features
    
    # Spectral centroid (brightness)
    centroid = librosa.feature.spectral_centroid(y=audio_segment, sr=sr)[0]
    features['spectral_centroid'] = float(np.mean(centroid))
    
    # Spectral rolloff
    rolloff = librosa.feature.spectral_rolloff(y=audio_segment, sr=sr)[0]
    features['spectral_rolloff'] = float(np.mean(rolloff))
    
    # Spectral flatness
    flatness = librosa.feature.spectral_flatness(y=audio_segment)[0]
    features['spectral_flatness'] = float(np.mean(flatness))
    
    # Zero crossing rate
    zcr = librosa.feature.zero_crossing_rate(audio_segment)[0]
    features['zero_crossing_rate'] = float(np.mean(zcr))
    
    return features


def compute_duration_features(duration: float, 
                              expected_duration: Optional[float] = None) -> Dict[str, float]:
    """
    Compute duration-based features.
    
    Args:
        duration: Actual phoneme duration in seconds
        expected_duration: Expected duration (if available)
        
    Returns:
        Dictionary of duration features
    """
    features = {
        'duration': duration,
        'duration_log': np.log1p(duration),  # log(1 + duration)
    }
    
    if expected_duration is not None and expected_duration > 0:
        features['duration_deviation'] = abs(duration - expected_duration)
        features['duration_ratio'] = duration / expected_duration
        features['duration_deviation_norm'] = features['duration_deviation'] / expected_duration
    else:
        # Use average phoneme duration as baseline (~0.08 seconds)
        avg_duration = 0.08
        features['duration_deviation'] = abs(duration - avg_duration)
        features['duration_ratio'] = duration / avg_duration
        features['duration_deviation_norm'] = features['duration_deviation'] / avg_duration
    
    return features


def compute_confidence_features(confidence: float, 
                                confidence_scores: Optional[np.ndarray] = None) -> Dict[str, float]:
    """
    Compute ASR confidence-based features.
    
    Args:
        confidence: Mean confidence for the phoneme
        confidence_scores: Frame-level confidence scores (if available)
        
    Returns:
        Dictionary of confidence features
    """
    features = {
        'confidence': confidence,
        'confidence_log': np.log1p(confidence),  # log(1 + conf)
    }
    
    if confidence_scores is not None and len(confidence_scores) > 0:
        features['confidence_std'] = float(np.std(confidence_scores))
        features['confidence_min'] = float(np.min(confidence_scores))
        features['confidence_max'] = float(np.max(confidence_scores))
    else:
        features['confidence_std'] = 0.0
        features['confidence_min'] = confidence
        features['confidence_max'] = confidence
    
    return features


def extract_phoneme_features(audio_array: np.ndarray,
                            phoneme_info: Dict,
                            sr: int = 16000,
                            n_mfcc: int = 13) -> Dict[str, float]:
    """
    Extract all features for a single phoneme.
    
    Args:
        audio_array: Full audio samples
        phoneme_info: Dictionary with phoneme information:
                     - timestamp: Start time in seconds
                     - duration: Duration in seconds
                     - confidence: ASR confidence score
                     - symbol: Phoneme symbol
        sr: Sample rate
        n_mfcc: Number of MFCC coefficients
        
    Returns:
        Dictionary of all extracted features
    """
    # Extract audio segment
    start_time = phoneme_info.get('timestamp', 0.0)
    duration = phoneme_info.get('duration', 0.0)
    end_time = start_time + duration
    
    audio_segment = extract_segment_audio(audio_array, start_time, end_time, sr)
    
    # Initialize features dictionary
    all_features = {}
    
    # 1. Duration features
    duration_feats = compute_duration_features(duration)
    all_features.update(duration_feats)
    
    # 2. Confidence features
    confidence = phoneme_info.get('confidence', 0.5)
    confidence_feats = compute_confidence_features(confidence)
    all_features.update(confidence_feats)
    
    # 3. MFCC features
    mfcc_feats = compute_mfcc_features(audio_segment, sr, n_mfcc)
    all_features.update(mfcc_feats)
    
    # 4. Energy features
    energy_feats = compute_energy_features(audio_segment, sr)
    all_features.update(energy_feats)
    
    # 5. Spectral features
    spectral_feats = compute_spectral_features(audio_segment, sr)
    all_features.update(spectral_feats)
    
    # 6. Metadata
    all_features['phoneme_symbol'] = phoneme_info.get('symbol', '')
    all_features['expected'] = int(phoneme_info.get('expected', True))
    
    return all_features


def extract_features_batch(audio_array: np.ndarray,
                          phoneme_results: List[Dict],
                          sr: int = 16000,
                          n_mfcc: int = 13) -> List[Dict[str, float]]:
    """
    Extract features for multiple phonemes.
    
    Args:
        audio_array: Full audio samples
        phoneme_results: List of phoneme dictionaries from alignment
        sr: Sample rate
        n_mfcc: Number of MFCC coefficients
        
    Returns:
        List of feature dictionaries
    """
    features_list = []
    
    for phoneme_info in phoneme_results:
        features = extract_phoneme_features(audio_array, phoneme_info, sr, n_mfcc)
        features_list.append(features)
    
    return features_list


def features_to_array(features: Dict[str, float], 
                      feature_names: List[str]) -> np.ndarray:
    """
    Convert feature dictionary to numpy array for ML model input.
    
    Args:
        features: Dictionary of features
        feature_names: Ordered list of feature names
        
    Returns:
        Numpy array of feature values
    """
    return np.array([features.get(name, 0.0) for name in feature_names])


def get_feature_names(n_mfcc: int = 13) -> List[str]:
    """
    Get ordered list of all feature names for consistent ML input.
    
    Args:
        n_mfcc: Number of MFCC coefficients
        
    Returns:
        List of feature names in order
    """
    feature_names = []
    
    # Duration features
    feature_names.extend([
        'duration',
        'duration_log',
        'duration_deviation',
        'duration_ratio',
        'duration_deviation_norm'
    ])
    
    # Confidence features
    feature_names.extend([
        'confidence',
        'confidence_log',
        'confidence_std',
        'confidence_min',
        'confidence_max'
    ])
    
    # MFCC features
    for i in range(n_mfcc):
        feature_names.append(f'mfcc_{i}_mean')
        feature_names.append(f'mfcc_{i}_std')
    feature_names.extend(['mfcc_delta_mean', 'mfcc_delta_std'])
    
    # Energy features
    feature_names.extend([
        'energy_mean',
        'energy_std',
        'energy_max',
        'rms_energy'
    ])
    
    # Spectral features
    feature_names.extend([
        'spectral_centroid',
        'spectral_rolloff',
        'spectral_flatness',
        'zero_crossing_rate'
    ])
    
    return feature_names


# ============================================================================
# TESTING
# ============================================================================

if __name__ == "__main__":
    print("Phoneme Feature Extraction - Test")
    print("=" * 50)
    
    # Generate test audio (1 second of sine wave)
    sr = 16000
    duration = 1.0
    frequency = 440  # A4 note
    t = np.linspace(0, duration, int(sr * duration))
    test_audio = 0.3 * np.sin(2 * np.pi * frequency * t)
    
    # Test phoneme info
    test_phoneme = {
        'symbol': 'a',
        'timestamp': 0.0,
        'duration': 0.1,
        'confidence': 0.85,
        'expected': True
    }
    
    print("\nExtracting features for test phoneme...")
    features = extract_phoneme_features(test_audio, test_phoneme, sr=sr, n_mfcc=13)
    
    print(f"\nTotal features extracted: {len(features)}")
    print("\nSample features:")
    sample_keys = ['duration', 'confidence', 'mfcc_0_mean', 'energy_mean', 'spectral_centroid']
    for key in sample_keys:
        if key in features:
            print(f"  {key:20s}: {features[key]:.4f}")
    
    # Get feature names
    feature_names = get_feature_names(n_mfcc=13)
    print(f"\nTotal feature dimensions: {len(feature_names)}")
    print(f"Feature names (first 10): {feature_names[:10]}")
    
    # Convert to array
    feature_array = features_to_array(features, feature_names)
    print(f"\nFeature array shape: {feature_array.shape}")
    print(f"Feature array stats: mean={feature_array.mean():.4f}, std={feature_array.std():.4f}")
    
    print("\n✓ Feature extraction working correctly")