"""PyTorch Dataset and DataLoader with Audio Augmentation for Bangla Speech Emotion Recognition."""

import os
import random
import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
import soundfile as sf
import librosa
import pandas as pd

from src.utils.config import (
    METADATA_DIR,
    TARGET_SAMPLE_RATE,
    MAX_DURATION_SECONDS,
    EMOTION_TO_ID,
    EMOTION_CLASSES
)


class AudioAugmenter:
    """Acoustic augmentations tailored for speech emotion recognition without distorting emotion."""

    def __init__(self, p: float = 0.5):
        self.p = p

    def apply(self, y: np.ndarray, sr: int = TARGET_SAMPLE_RATE) -> np.ndarray:
        if random.random() > self.p:
            return y

        aug_type = random.choice(["noise", "gain", "shift", "pitch"])

        if aug_type == "noise":
            noise_amp = random.uniform(0.002, 0.015)
            noise = np.random.randn(len(y)) * noise_amp
            y = y + noise

        elif aug_type == "gain":
            gain_factor = random.uniform(0.8, 1.2)
            y = y * gain_factor

        elif aug_type == "shift":
            shift_samples = int(random.uniform(-0.15, 0.15) * sr)
            y = np.roll(y, shift_samples)

        elif aug_type == "pitch":
            # Mild pitch variation (-1 to +1 semitone)
            n_steps = random.choice([-1.0, -0.5, 0.5, 1.0])
            y = librosa.effects.pitch_shift(y, sr=sr, n_steps=n_steps)

        # Clip to prevent digital clipping
        y = np.clip(y, -1.0, 1.0)
        return y


class BanglaSpeechDataset(Dataset):
    """Dataset delivering log Mel-spectrograms or raw waveforms with labels."""

    def __init__(
        self,
        split: str = "train",
        feature_type: str = "melspec",
        n_mels: int = 128,
        augment: bool = False,
        cached_mel_dict: dict = None
    ):
        self.split = split
        self.feature_type = feature_type
        self.n_mels = n_mels
        self.augment = augment
        self.augmenter = AudioAugmenter(p=0.5) if augment else None
        self.target_len = int(MAX_DURATION_SECONDS * TARGET_SAMPLE_RATE)  # 160,000 samples

        meta_path = METADATA_DIR / "metadata.csv"
        df = pd.read_csv(meta_path)
        self.df = df[df["split"] == split].reset_index(drop=True)
        
        # In-memory cached mel spectrograms if provided
        self.cached_mel = cached_mel_dict.get(split) if cached_mel_dict else None

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        emotion = row["emotion"]
        label_id = EMOTION_TO_ID[emotion]
        audio_id = row["audio_id"]

        if self.feature_type == "melspec" and self.cached_mel is not None:
            # Ultra-fast indexed from RAM
            mel = self.cached_mel[idx]
            tensor_mel = torch.tensor(mel, dtype=torch.float32).unsqueeze(0)
            return tensor_mel, label_id, audio_id

        file_path = row["processed_file_path"]
        audio, sr = sf.read(file_path, dtype="float32")

        if self.augment and self.augmenter:
            audio = self.augmenter.apply(audio, sr=sr)

        # Pad or truncate to fixed length (10 seconds)
        if len(audio) < self.target_len:
            audio = np.pad(audio, (0, self.target_len - len(audio)), mode="constant")
        else:
            audio = audio[:self.target_len]

        if self.feature_type == "raw":
            return torch.tensor(audio, dtype=torch.float32), label_id, audio_id

        elif self.feature_type == "melspec":
            # Compute 128 Mel bands
            mel_spec = librosa.feature.melspectrogram(
                y=audio, sr=sr, n_fft=1024, hop_length=512, n_mels=self.n_mels, fmin=20, fmax=8000
            )
            log_mel_spec = librosa.power_to_db(mel_spec, ref=np.max)
            log_mel_spec = (log_mel_spec - np.mean(log_mel_spec)) / (np.std(log_mel_spec) + 1e-6)
            tensor_mel = torch.tensor(log_mel_spec, dtype=torch.float32).unsqueeze(0)
            return tensor_mel, label_id, audio_id


        elif self.feature_type == "mfcc":
            mfcc = librosa.feature.mfcc(y=audio, sr=sr, n_mfcc=40, n_fft=1024, hop_length=512)
            mfcc = (mfcc - np.mean(mfcc)) / (np.std(mfcc) + 1e-6)
            # Shape: (313, 40)
            tensor_mfcc = torch.tensor(mfcc.T, dtype=torch.float32)
            return tensor_mfcc, label_id, audio_id


def get_dataloader(
    split: str = "train",
    feature_type: str = "melspec",
    batch_size: int = 16,
    augment: bool = False,
    num_workers: int = 0,
    cached_mel_dict: dict = None
) -> DataLoader:
    dataset = BanglaSpeechDataset(
        split=split,
        feature_type=feature_type,
        augment=augment,
        cached_mel_dict=cached_mel_dict
    )
    shuffle = (split == "train")
    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)