"""Acoustic Feature Extraction Module for Bangla Speech Emotion Recognition.

Extracts:
1. MFCC Features: 40 Mel-Frequency Cepstral Coefficients + Delta + Delta-Delta,
   statistically summarized over time (Mean + Std Dev = 80-120 dimensions).
2. Mel-Spectrograms: 128 mel frequency bins, normalized for CNN and Recurrent architectures.
3. Feature Caching: Saves extracted features to disk to accelerate experimentation.
"""

import os
import numpy as np
import soundfile as sf
import librosa
from pathlib import Path
from tqdm import tqdm

from src.utils.config import (
    PROCESSED_DATA_DIR,
    RESULTS_DIR,
    TARGET_SAMPLE_RATE,
    MAX_DURATION_SECONDS
)

FEATURES_CACHE_DIR = RESULTS_DIR / "features"
FEATURES_CACHE_DIR.mkdir(parents=True, exist_ok=True)


def extract_mfcc_summary(audio_path: str, n_mfcc: int = 40) -> np.ndarray:
    """Extract 40-dimensional MFCCs plus delta and delta-deltas with temporal pooling (mean+std).
    
    Returns:
        1D numpy array of shape (n_mfcc * 2,) = 80 features.
    """
    audio, sr = sf.read(audio_path, dtype="float32")
    if len(audio) == 0:
        return np.zeros(n_mfcc * 2, dtype=np.float32)

    # Compute MFCC
    mfcc = librosa.feature.mfcc(y=audio, sr=sr, n_mfcc=n_mfcc, n_fft=1024, hop_length=512)
    mfcc_mean = np.mean(mfcc, axis=1)
    mfcc_std = np.std(mfcc, axis=1)

    return np.concatenate([mfcc_mean, mfcc_std]).astype(np.float32)


def extract_mel_spectrogram(audio_path: str, n_mels: int = 128, max_len: int = 313) -> np.ndarray:
    """Extract fixed-length 128-band log Mel-spectrogram for 2D CNN / LSTM models.
    
    Returns:
        2D numpy array of shape (n_mels, max_len) e.g. (128, 313).
    """
    audio, sr = sf.read(audio_path, dtype="float32")
    
    # Pad or truncate to max_duration (10.0s = 160,000 samples)
    target_samples = int(MAX_DURATION_SECONDS * TARGET_SAMPLE_RATE)
    if len(audio) < target_samples:
        audio = np.pad(audio, (0, target_samples - len(audio)), mode="constant")
    else:
        audio = audio[:target_samples]

    mel_spec = librosa.feature.melspectrogram(
        y=audio, sr=sr, n_fft=1024, hop_length=512, n_mels=n_mels, fmin=20, fmax=8000
    )
    log_mel_spec = librosa.power_to_db(mel_spec, ref=np.max)

    # Normalize to [0, 1] range or zero-mean unit-variance
    log_mel_spec = (log_mel_spec - np.mean(log_mel_spec)) / (np.std(log_mel_spec) + 1e-6)

    # Ensure shape is exactly (n_mels, max_len)
    if log_mel_spec.shape[1] < max_len:
        log_mel_spec = np.pad(log_mel_spec, ((0, 0), (0, max_len - log_mel_spec.shape[1])), mode="constant")
    else:
        log_mel_spec = log_mel_spec[:, :max_len]

    return log_mel_spec.astype(np.float32)


def build_and_cache_features(metadata_df, force_recompute: bool = False):
    """Precompute and cache MFCC feature matrices for all train, val, and test splits."""
    mfcc_cache_file = FEATURES_CACHE_DIR / "mfcc_summary_features.npz"

    if not force_recompute and mfcc_cache_file.exists():
        print(f"Loading cached MFCC features from {mfcc_cache_file}")
        data = np.load(mfcc_cache_file, allow_pickle=True)
        return data

    print("Extracting acoustic features across all 5,000 recordings...")
    audio_ids = metadata_df["audio_id"].values
    file_paths = metadata_df["processed_file_path"].values
    emotions = metadata_df["emotion"].values
    splits = metadata_df["split"].values

    mfcc_list = []
    for fpath in tqdm(file_paths, desc="Extracting 40-MFCCs"):
        feats = extract_mfcc_summary(fpath, n_mfcc=40)
        mfcc_list.append(feats)

    X_mfcc = np.stack(mfcc_list)

    np.savez_compressed(
        mfcc_cache_file,
        X=X_mfcc,
        audio_ids=audio_ids,
        emotions=emotions,
        splits=splits
    )
    print(f"Saved cached MFCC features (shape: {X_mfcc.shape}) to {mfcc_cache_file}")
    data = np.load(mfcc_cache_file, allow_pickle=True)
    return data


def build_and_cache_mel_spectrograms(metadata_df, force_recompute: bool = False):
    """Precompute and cache 128-band Mel-spectrogram matrices for all recordings."""
    mel_cache_file = FEATURES_CACHE_DIR / "mel_spectrogram_features.npz"

    if not force_recompute and mel_cache_file.exists():
        print(f"Loading cached Mel-spectrogram features from {mel_cache_file}")
        data = np.load(mel_cache_file, allow_pickle=True)
        return data

    print("Extracting 128-band Mel-spectrograms across all 5,000 recordings...")
    audio_ids = metadata_df["audio_id"].values
    file_paths = metadata_df["processed_file_path"].values
    emotions = metadata_df["emotion"].values
    splits = metadata_df["split"].values

    mel_list = []
    for fpath in tqdm(file_paths, desc="Extracting Mel-Spectrograms"):
        mel = extract_mel_spectrogram(fpath, n_mels=128, max_len=313)
        mel_list.append(mel)

    X_mel = np.stack(mel_list)  # (5000, 128, 313)

    np.savez_compressed(
        mel_cache_file,
        X=X_mel,
        audio_ids=audio_ids,
        emotions=emotions,
        splits=splits
    )
    print(f"Saved cached Mel-spectrogram features (shape: {X_mel.shape}) to {mel_cache_file}")
    data = np.load(mel_cache_file, allow_pickle=True)
    return data