# train_emotion_recognition.py
"""
Complete PyTorch implementation of the CNN-BiGRU-Attention model for continuous music emotion recognition
as described in the uploaded manuscript.

Requirements (suggested):
  - Python 3.8+
  - torch >= 1.8
  - torchaudio
  - librosa
  - numpy, scipy, scikit-learn, pandas
  - tqdm

Usage:
  - Prepare CSV file with columns: filepath, valence, arousal, dominance
    (one row per audio segment). Audio files can be wav/mp3/etc.
  - python train_emotion_recognition.py --data_csv path/to/labels.csv --out_dir outputs/
"""

import os
import argparse
import random
from pathlib import Path
from typing import Tuple

import numpy as np
import pandas as pd
import librosa
import soundfile as sf
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from scipy.stats import pearsonr
from tqdm import tqdm

# ---------------------------
# Utilities and preprocessing
# ---------------------------

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

def spectral_gate(y, sr, prop_decrease=0.9, n_fft=2048, hop_length=512):
    """
    Approximate spectral gating (simple noise suppression).
    Not a production-grade noise reduction; works as a light pre-processing step.
    """
    # compute STFT magnitude
    D = librosa.stft(y, n_fft=n_fft, hop_length=hop_length)
    mag, phase = np.abs(D), np.angle(D)
    # estimate noise profile as median across time of lower quantile
    noise_profile = np.median(mag, axis=1, keepdims=True)
    # gating mask
    mask = mag >= (noise_profile * 1.5)
    mag_denoised = mag * mask + mag * (1 - mask) * (1 - prop_decrease)
    D_denoised = mag_denoised * np.exp(1j * phase)
    y_denoised = librosa.istft(D_denoised, hop_length=hop_length)
    return y_denoised

def load_audio(path, sr=22050, duration=None, offset=0.0):
    y, orig_sr = sf.read(path)
    # if multi-channel, mix to mono
    if y.ndim > 1:
        y = np.mean(y, axis=1)
    if orig_sr != sr:
        y = librosa.resample(y.astype(float), orig_sr, sr)
    if duration is not None:
        max_len = int(sr * duration)
        if y.shape[0] < max_len:
            # pad
            y = np.pad(y, (0, max_len - y.shape[0]))
        else:
            y = y[:max_len]
    return y

def extract_features(y, sr=22050, n_mels=128, n_mfcc=20, hop_length=512, n_fft=2048):
    # Mel spectrogram (log)
    mel = librosa.feature.melspectrogram(y=y, sr=sr, n_fft=n_fft, hop_length=hop_length, n_mels=n_mels)
    mel_db = librosa.power_to_db(mel, ref=np.max)
    # MFCC
    mfcc = librosa.feature.mfcc(S=mel_db, n_mfcc=n_mfcc)
    return mel_db.astype(np.float32), mfcc.astype(np.float32)

# ---------------------------
# Dataset
# ---------------------------

class MusicEmotionDataset(Dataset):
    def __init__(self, csv_path, sr=22050, duration=4.0, augment=False, transform=None, noise_aug=False):
        """
        CSV must have: filepath, valence, arousal, dominance
        valence/arousal/dominance are floats (continuous)
        """
        self.df = pd.read_csv(csv_path)
        self.sr = sr
        self.duration = duration
        self.augment = augment
        self.transform = transform
        self.noise_aug = noise_aug

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

    def _augment(self, y):
        # simple augmentations: time stretch, pitch shift, add noise
        if random.random() < 0.3:
            rate = random.uniform(0.9, 1.1)
            y = librosa.effects.time_stretch(y, rate)
        if random.random() < 0.3:
            n_steps = random.uniform(-1, 1)
            y = librosa.effects.pitch_shift(y, sr=self.sr, n_steps=n_steps)
        if self.noise_aug and random.random() < 0.3:
            noise = np.random.randn(len(y)) * 0.005
            y = y + noise
        return y

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        path = row['filepath']
        y = load_audio(path, sr=self.sr, duration=self.duration)
        # spectral gating denoise (paper uses spectral gating)
        y = spectral_gate(y, sr=self.sr)
        if self.augment:
            y = self._augment(y)

        mel, mfcc = extract_features(y, sr=self.sr)
        # normalize per-sample (zero mean unit var)
        mel = (mel - mel.mean()) / (mel.std() + 1e-9)
        mfcc = (mfcc - mfcc.mean()) / (mfcc.std() + 1e-9)

        # make shapes consistent: (C, H, W) for CNN input; we treat mel as single-channel image
        mel_tensor = torch.from_numpy(mel).unsqueeze(0)  # (1, n_mels, time)
        mfcc_tensor = torch.from_numpy(mfcc).unsqueeze(0)

        # high-level label vector (Valence, Arousal, Dominance)
        label = torch.tensor([row['valence'], row['arousal'], row['dominance']], dtype=torch.float32)

        return {
            'mel': mel_tensor,
            'mfcc': mfcc_tensor,
            'label': label,
            'path': path
        }

# ---------------------------
# Model modules
# ---------------------------

class ConvBlock(nn.Module):
    def __init__(self, in_ch, out_ch, kernel=(3,3), padding='same'):
        super().__init__()
        if padding == 'same':
            pad = (kernel[0]//2, kernel[1]//2)
        else:
            pad = 0
        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size=kernel, padding=pad)
        self.bn = nn.BatchNorm2d(out_ch)
        self.pool = nn.MaxPool2d(kernel_size=(2,2))
        self.act = nn.ReLU()

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.act(x)
        x = self.pool(x)
        return x

class CNNFeatureExtractor(nn.Module):
    def __init__(self, in_ch=1, filters=[64,64,64]):
        super().__init__()
        self.blocks = nn.ModuleList()
        ch = in_ch
        for f in filters:
            self.blocks.append(ConvBlock(ch, f, kernel=(3,3)))
            ch = f
        # final projection
        self.project = nn.Conv2d(ch, ch, kernel_size=(1,1))

    def forward(self, x):
        # x: (B, 1, n_mels, T)
        for b in self.blocks:
            x = b(x)
        x = self.project(x)  # (B, C, H', W')
        # flatten spatial dims to sequence for RNN: treat time dimension as sequence
        B, C, H, W = x.shape
        # collapse frequency (H) and channel C into feature dim, sequence along W (time)
        x = x.permute(0,3,1,2).contiguous()  # (B, W, C, H)
        x = x.view(B, W, C*H)  # (B, seq_len, feat_dim)
        return x  # sequence features

class BiGRUWithAttention(nn.Module):
    def __init__(self, input_dim, hidden_dim=128, num_layers=1, attention_dim=128):
        super().__init__()
        self.bigru = nn.GRU(input_dim, hidden_dim, num_layers=num_layers, batch_first=True, bidirectional=True)
        self.att_fc = nn.Linear(hidden_dim*2, attention_dim)
        self.att_v = nn.Linear(attention_dim, 1, bias=False)

    def forward(self, x, mask=None):
        # x: (B, seq_len, feat_dim)
        out, _ = self.bigru(x)  # (B, seq_len, 2*hidden)
        # attention scoring
        score = torch.tanh(self.att_fc(out))  # (B, seq_len, att_dim)
        score = self.att_v(score).squeeze(-1)  # (B, seq_len)
        if mask is not None:
            score = score.masked_fill(~mask, -1e9)
        att_w = F.softmax(score, dim=1)  # (B, seq_len)
        context = torch.bmm(att_w.unsqueeze(1), out).squeeze(1)  # (B, 2*hidden)
        return context, att_w

class FusionRegressor(nn.Module):
    def __init__(self, low_feat_dim, high_feat_dim, gru_hidden=128, fusion_att_dim=128, out_dim=3):
        super().__init__()
        # high-level encoder (CNN->BiGRU->Att)
        self.cnn = CNNFeatureExtractor(in_ch=1, filters=[64,64,64])  # as paper: 3 conv layers, filters=64
        # after cnn, we need to know seq_len and feat_dim dynamically -> we infer at runtime
        # We'll add a linear projection to consistent dimension for GRU
        self.proj = nn.Linear(64* ( (128 // (2**3)) ), 256)  # approximate: depends on n_mels pooling; ensure large enough
        self.bigru_att = BiGRUWithAttention(input_dim=256, hidden_dim=gru_hidden, num_layers=1, attention_dim=fusion_att_dim)

        # learnable fusion weight alpha (we enforce alpha in (0,1) via sigmoid)
        self.alpha_param = nn.Parameter(torch.tensor(0.0))  # sigmoid(alpha_param) -> alpha

        # small MLP for final mapping
        fusion_dim = (gru_hidden*2) + low_feat_dim  # context vec + low-level pooled features
        self.fc1 = nn.Linear(fusion_dim, 256)
        self.dropout = nn.Dropout(0.3)
        self.fc2 = nn.Linear(256, out_dim)

    def forward(self, mel, mfcc=None, use_mfcc=False, disable_cnn=False, disable_bigru=False, disable_att=False):
        """
        mel: (B,1,n_mels, T)
        mfcc: (B,1,n_mfcc, T) (optional)
        """
        # Low-level feature: global pooled MFCC or Mel (we will pool mfcc if provided, else global avg of mel)
        if mfcc is not None and use_mfcc:
            # pool mfcc along time and freq -> vector
            low = torch.mean(mfcc, dim=[2,3])  # (B,1) -> ensure flatten
            low = low.view(low.size(0), -1)  # shape (B, feat)
            low_dim = low.shape[1]
        else:
            # use mel global pooling
            low = torch.mean(mel, dim=[2,3])  # (B,1)
            low = low.view(low.size(0), -1)
            low_dim = low.shape[1]

        # High-level semantic features via CNN->BiGRU->Attention
        if disable_cnn:
            # bypass CNN: use avg-pooled mel as proxy
            high_context = torch.zeros(mel.size(0), self.bigru_att.bigru.hidden_size*2, device=mel.device)
        else:
            seq = self.cnn(mel)  # (B, seq_len, feat)
            # ensure projection dimension matches GRU input
            if seq.size(-1) != self.bigru_att.bigru.input_size:
                # linear project
                seq = self.proj(seq)
            if disable_bigru:
                # simple temporal average
                high_context = torch.mean(seq, dim=1)
            else:
                # use BiGRU + Attention
                context, att_w = self.bigru_att(seq)
                if disable_att:
                    high_context = torch.mean(seq, dim=1)
                else:
                    high_context = context

        # fusion weight
        alpha = torch.sigmoid(self.alpha_param)  # in (0,1)
        # ensure low and high vectors have matching dims for concat; if low_dim smaller, project
        if low.shape[1] != self.bigru_att.bigru.hidden_size*2:
            # project low to that dim
            low_proj = nn.Linear(low.shape[1], self.bigru_att.bigru.hidden_size*2).to(low.device)
            low_vec = low_proj(low)
        else:
            low_vec = low

        fused = torch.cat([alpha * low_vec + (1 - alpha) * high_context, low_vec], dim=1)
        x = F.relu(self.fc1(fused))
        x = self.dropout(x)
        out = self.fc2(x)  # (B, 3) continuous outputs
        return out

# ---------------------------
# Metrics
# ---------------------------

def compute_metrics(y_true: np.ndarray, y_pred: np.ndarray):
    # y_true, y_pred: (N, 3)
    metrics = {}
    metrics['MSE'] = mean_squared_error(y_true, y_pred, multioutput='raw_values')  # per-dim
    metrics['MAE'] = mean_absolute_error(y_true, y_pred, multioutput='raw_values')
    # Pearson r per-dim (handle constant arrays)
    r_vals = []
    for i in range(y_true.shape[1]):
        try:
            r, _ = pearsonr(y_true[:, i], y_pred[:, i])
        except Exception:
            r = 0.0
        r_vals.append(r)
    metrics['Pearson'] = np.array(r_vals)
    metrics['R2'] = []
    for i in range(y_true.shape[1]):
        metrics['R2'].append(r2_score(y_true[:, i], y_pred[:, i]))
    metrics['R2'] = np.array(metrics['R2'])
    return metrics

# ---------------------------
# Training loop
# ---------------------------

def train_epoch(model, dataloader, optimizer, device):
    model.train()
    total_loss = 0.0
    for batch in tqdm(dataloader, desc='Train', leave=False):
        mel = batch['mel'].to(device)  # (B,1,n_mels,T)
        mfcc = batch['mfcc'].to(device) if 'mfcc' in batch else None
        label = batch['label'].to(device)
        optimizer.zero_grad()
        preds = model(mel, mfcc=mfcc, use_mfcc=False)
        loss = F.mse_loss(preds, label)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * mel.size(0)
    return total_loss / len(dataloader.dataset)

def eval_epoch(model, dataloader, device):
    model.eval()
    y_trues = []
    y_preds = []
    with torch.no_grad():
        for batch in tqdm(dataloader, desc='Eval', leave=False):
            mel = batch['mel'].to(device)
            mfcc = batch['mfcc'].to(device) if 'mfcc' in batch else None
            label = batch['label'].cpu().numpy()
            preds = model(mel, mfcc=mfcc, use_mfcc=False)
            preds = preds.cpu().numpy()
            y_trues.append(label)
            y_preds.append(preds)
    y_trues = np.concatenate(y_trues, axis=0)
    y_preds = np.concatenate(y_preds, axis=0)
    return compute_metrics(y_trues, y_preds), y_trues, y_preds

# ---------------------------
# Main: argument parsing and run
# ---------------------------

def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument('--data_csv', type=str, required=True, help='CSV with columns filepath,valence,arousal,dominance')
    p.add_argument('--out_dir', type=str, default='outputs', help='output dir')
    p.add_argument('--epochs', type=int, default=50)
    p.add_argument('--batch_size', type=int, default=32)
    p.add_argument('--lr', type=float, default=0.001)
    p.add_argument('--seed', type=int, default=42)
    p.add_argument('--gpu', type=int, default=0)
    p.add_argument('--augment', action='store_true')
    p.add_argument('--no_cuda', action='store_true')
    return p.parse_args()

def main():
    args = parse_args()
    set_seed(args.seed)
    device = torch.device(f'cuda:{args.gpu}' if torch.cuda.is_available() and not args.no_cuda else 'cpu')
    os.makedirs(args.out_dir, exist_ok=True)

    # Load dataset CSV and split
    df = pd.read_csv(args.data_csv)
    # simple split: 70/15/15
    df = df.sample(frac=1.0, random_state=args.seed).reset_index(drop=True)
    n = len(df)
    n_train = int(0.7 * n)
    n_val = int(0.15 * n)
    train_df = df.iloc[:n_train].reset_index(drop=True)
    val_df = df.iloc[n_train:n_train+n_val].reset_index(drop=True)
    test_df = df.iloc[n_train+n_val:].reset_index(drop=True)

    # save splits for reproducibility
    train_df.to_csv(os.path.join(args.out_dir, 'train_split.csv'), index=False)
    val_df.to_csv(os.path.join(args.out_dir, 'val_split.csv'), index=False)
    test_df.to_csv(os.path.join(args.out_dir, 'test_split.csv'), index=False)

    train_csv = os.path.join(args.out_dir, 'train_split.csv')
    val_csv = os.path.join(args.out_dir, 'val_split.csv')
    test_csv = os.path.join(args.out_dir, 'test_split.csv')

    train_ds = MusicEmotionDataset(train_csv, augment=args.augment, noise_aug=True)
    val_ds = MusicEmotionDataset(val_csv, augment=False, noise_aug=False)
    test_ds = MusicEmotionDataset(test_csv, augment=False, noise_aug=False)

    train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True)
    val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, num_workers=2, pin_memory=True)
    test_loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False, num_workers=2, pin_memory=True)

    # instantiate model with approximate dims
    # low-level feature dimension = mel global pooling dim = 1 (since we collapse channel) -> we use projection inside model
    model = FusionRegressor(low_feat_dim=1, high_feat_dim=256, gru_hidden=128).to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    best_val_loss = float('inf')

    # training loop
    for epoch in range(1, args.epochs + 1):
        print(f"Epoch {epoch}/{args.epochs}")
        train_loss = train_epoch(model, train_loader, optimizer, device)
        val_metrics, _, _ = eval_epoch(model, val_loader, device)
        val_mse_mean = np.mean(val_metrics['MSE'])
        print(f"Train Loss: {train_loss:.6f} | Val MSE per-dim: {val_metrics['MSE']} | Pearson: {val_metrics['Pearson']}")
        # save best
        if val_mse_mean < best_val_loss:
            best_val_loss = val_mse_mean
            torch.save(model.state_dict(), os.path.join(args.out_dir, 'best_model.pth'))
            print("Saved best model")

    # final test
    print("Loading best model for testing...")
    model.load_state_dict(torch.load(os.path.join(args.out_dir, 'best_model.pth'), map_location=device))
    test_metrics, y_true, y_pred = eval_epoch(model, test_loader, device)
    print("Test results:")
    print("MSE per-dim:", test_metrics['MSE'])
    print("MAE per-dim:", test_metrics['MAE'])
    print("Pearson r per-dim:", test_metrics['Pearson'])
    print("R2 per-dim:", test_metrics['R2'])

    # save predictions
    np.savez(os.path.join(args.out_dir, 'test_preds.npz'), y_true=y_true, y_pred=y_pred)

if __name__ == '__main__':
    main()