"""Classical Machine Learning Baselines: MFCC + SVM and MFCC + Random Forest.

Reproduces:
- Baseline 1: 40-MFCC + PCA + Support Vector Machine (SVM)
- Baseline 2: 40-MFCC + Random Forest (RF)
Strictly trains on Train, validates on Val, evaluates on Test with zero leakage.
"""

import os
import joblib
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier

from src.utils.config import (
    METADATA_DIR,
    RANDOM_SEED,
    EMOTION_CLASSES
)
from src.features.extract import build_and_cache_features
from src.evaluation.evaluator import ExperimentLogger


def train_baseline_svm(n_components: int = 40):
    """Baseline 1: MFCC + PCA + SVM."""
    print("\n" + "=" * 70)
    print("RUNNING BASELINE 1: MFCC + PCA + Support Vector Machine (SVM)")
    print("=" * 70)

    # 1. Load metadata and features
    meta_df = pd.read_csv(METADATA_DIR / "metadata.csv")
    feat_data = build_and_cache_features(meta_df)

    X_all = feat_data["X"]
    y_all = feat_data["emotions"]
    splits = feat_data["splits"]
    audio_ids = feat_data["audio_ids"]

    train_idx = np.where(splits == "train")[0]
    val_idx = np.where(splits == "val")[0]
    test_idx = np.where(splits == "test")[0]

    X_train, y_train = X_all[train_idx], y_all[train_idx]
    X_val, y_val = X_all[val_idx], y_all[val_idx]
    X_test, y_test = X_all[test_idx], y_all[test_idx]

    # 2. Standard Scaling & PCA (fit strictly on train to prevent leakage)
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_val_scaled = scaler.transform(X_val)
    X_test_scaled = scaler.transform(X_test)

    pca = PCA(n_components=min(n_components, X_train.shape[1]), random_state=RANDOM_SEED)
    X_train_pca = pca.fit_transform(X_train_scaled)
    X_val_pca = pca.transform(X_val_scaled)
    X_test_pca = pca.transform(X_test_scaled)

    # 3. Train SVM
    svm = SVC(kernel="rbf", C=10.0, gamma="scale", probability=True, random_state=RANDOM_SEED)
    svm.fit(X_train_pca, y_train)

    # 4. Evaluate on Validation Set
    val_preds = svm.predict(X_val_pca)
    val_proba = svm.predict_proba(X_val_pca)

    # 5. Evaluate on Held-Out Test Set
    test_preds = svm.predict(X_test_pca)
    test_proba = svm.predict_proba(X_test_pca)

    # 6. Log Experiment
    logger = ExperimentLogger(experiment_id="EXP001_Baseline_SVM", description="MFCC + PCA + SVM Baseline")
    config = {
        "experiment_id": "EXP001_Baseline_SVM",
        "model_type": "Support Vector Machine (SVC)",
        "kernel": "rbf",
        "C": 10.0,
        "features": "40 MFCC (Mean + Std = 80 dims)",
        "pca_components": n_components,
        "random_seed": RANDOM_SEED,
        "split_version": "v1.0"
    }
    logger.log_config(config)

    # Save model and preprocessors
    joblib.dump({"scaler": scaler, "pca": pca, "model": svm}, logger.model_dir / "svm_pipeline.joblib")

    # Evaluate test split
    metrics = logger.evaluate_and_log(
        y_true=y_test.tolist(),
        y_pred=test_preds.tolist(),
        y_proba=test_proba,
        audio_ids=audio_ids[test_idx].tolist(),
        split_name="test",
        extra_metrics={
            "val_accuracy": float(round(np.mean(val_preds == y_val), 4)),
            "paper_benchmark_accuracy": 85.40,
            "architecture": "MFCC + PCA + SVM (RBF)"
        }
    )
    return metrics


def train_baseline_rf(n_estimators: int = 300):
    """Baseline 2: MFCC + Random Forest."""
    print("\n" + "=" * 70)
    print("RUNNING BASELINE 2: MFCC + Random Forest (RF)")
    print("=" * 70)

    # 1. Load metadata and features
    meta_df = pd.read_csv(METADATA_DIR / "metadata.csv")
    feat_data = build_and_cache_features(meta_df)

    X_all = feat_data["X"]
    y_all = feat_data["emotions"]
    splits = feat_data["splits"]
    audio_ids = feat_data["audio_ids"]

    train_idx = np.where(splits == "train")[0]
    val_idx = np.where(splits == "val")[0]
    test_idx = np.where(splits == "test")[0]

    X_train, y_train = X_all[train_idx], y_all[train_idx]
    X_val, y_val = X_all[val_idx], y_all[val_idx]
    X_test, y_test = X_all[test_idx], y_all[test_idx]

    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_val_scaled = scaler.transform(X_val)
    X_test_scaled = scaler.transform(X_test)

    # 2. Train Random Forest
    rf = RandomForestClassifier(n_estimators=n_estimators, max_depth=20, random_state=RANDOM_SEED, n_jobs=-1)
    rf.fit(X_train_scaled, y_train)

    # 3. Predict
    val_preds = rf.predict(X_val_scaled)
    test_preds = rf.predict(X_test_scaled)
    test_proba = rf.predict_proba(X_test_scaled)

    # 4. Log Experiment
    logger = ExperimentLogger(experiment_id="EXP002_Baseline_RF", description="MFCC + Random Forest Baseline")
    config = {
        "experiment_id": "EXP002_Baseline_RF",
        "model_type": "Random Forest Classifier",
        "n_estimators": n_estimators,
        "max_depth": 20,
        "features": "40 MFCC (Mean + Std = 80 dims)",
        "random_seed": RANDOM_SEED,
        "split_version": "v1.0"
    }
    logger.log_config(config)

    joblib.dump({"scaler": scaler, "model": rf}, logger.model_dir / "rf_pipeline.joblib")

    metrics = logger.evaluate_and_log(
        y_true=y_test.tolist(),
        y_pred=test_preds.tolist(),
        y_proba=test_proba,
        audio_ids=audio_ids[test_idx].tolist(),
        split_name="test",
        extra_metrics={
            "val_accuracy": float(round(np.mean(val_preds == y_val), 4)),
            "paper_benchmark_accuracy": 78.00,
            "architecture": "MFCC + Random Forest"
        }
    )
    return metrics


if __name__ == "__main__":
    train_baseline_svm()
    train_baseline_rf()