"""
Train and compare: MLP (baseline), MLP+channel attention (prosody emphasis), and ensemble.
Reports metrics for each so you can see which gives higher accuracy/F1.
"""

import os
import sys
import random

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.utils.class_weight import compute_class_weight
from sklearn.decomposition import PCA  # --- NEW: Import PCA ---
from sklearn.metrics import (
    accuracy_score,
    f1_score,
    precision_score,
    recall_score,
    confusion_matrix,
)

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

from model import NativityClassifier, NativityClassifierWithAttention

SEED = 42
MANIFEST_PATH = "manifests/final_training_manifest.csv"
MODEL_SAVE_DIR = "models"
EPOCHS = 100
BATCH_SIZE = 32

# --- ANTI-OVERFITTING: Hyperparameter adjustments ---
LR = 1e-3  # Slowed down from 1e-3
WEIGHT_DECAY = 1e-4  # Punish extreme weights (up from 5e-4)
PATIENCE = 15  # Stop earlier if it starts memorizing
MIN_EPOCHS = 40
PCA_COMPONENTS = 256  # Compress 1000+ dims down to top 64
HIDDEN_DIMS = (256, 128)  # Shrunk from (384, 256, 128)
DROPOUT = 0.4
# ----------------------------------------------------

OUTPUT_DIM = 2
# Lower threshold for Class 1 (Non-Native) to boost recall on imbalanced data
EVAL_THRESHOLD = 0.50
# SMOTE oversamples minority (Non-Native) in embedding space; already using weighted CE
USE_SMOTE = False
SMOTE_K_NEIGHBORS = 5
# Threshold sweep on val to pick best (don't tune on test)
THRESHOLD_CANDIDATES = [0.25, 0.28, 0.30, 0.32, 0.35, 0.40]
BEST_THRESHOLD_PATH = os.path.join(MODEL_SAVE_DIR, "best_threshold.txt")

MIN_LR = 1e-6


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    else:
        torch.backends.cudnn.benchmark = True


class EmbeddingDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        df,
        embedding_path_col="embedding_path",
        label_col="label_idx",
        l2_norm=True,
    ):
        self.paths = [str(p).replace("\\", "/") for p in df[embedding_path_col].tolist()]
        self.labels = df[label_col].values
        self.l2_norm = l2_norm

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

    def __getitem__(self, idx):
        emb = np.load(self.paths[idx]).astype(np.float32)
        if self.l2_norm:
            n = np.linalg.norm(emb)
            if n > 1e-8:
                emb = emb / n
        return torch.tensor(emb, dtype=torch.float32), torch.tensor(
            self.labels[idx], dtype=torch.long
        )


def train_model(
    model,
    train_loader,
    val_loader,
    criterion,
    device,
    save_path,
    epochs=EPOCHS,
    patience=PATIENCE,
    min_epochs=MIN_EPOCHS,
):
    optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="max", factor=0.5, patience=5, min_lr=MIN_LR
    )
    best_f1 = 0.0
    patience_counter = 0

    history = {
        "train_loss": [],
        "val_loss": [],
        "train_acc": [],
        "val_acc": [],
        "val_f1": [],
    }

    for epoch in range(epochs):
        model.train()
        train_loss_sum = 0.0
        train_correct = 0
        train_total = 0
        for batch_x, batch_y in train_loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)

            # --- NEW: Inject Gaussian Noise during training to prevent memorization ---
            if model.training:
                noise = torch.randn_like(batch_x) * 0.02
                batch_x = batch_x + noise
            # ------------------------------------------------------------------------

            optimizer.zero_grad()
            logits = model(batch_x)
            loss = criterion(logits, batch_y)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

            train_loss_sum += loss.item() * batch_x.size(0)
            preds = logits.argmax(dim=1)
            train_correct += (preds == batch_y).sum().item()
            train_total += batch_x.size(0)

        train_loss = train_loss_sum / max(1, train_total)
        train_acc = train_correct / max(1, train_total)

        # Validation
        model.eval()
        val_loss_sum = 0.0
        val_correct = 0
        val_total = 0
        preds_list = []
        labels_list = []
        with torch.no_grad():
            for batch_x, batch_y in val_loader:
                batch_x, batch_y = batch_x.to(device), batch_y.to(device)
                logits = model(batch_x)
                loss = criterion(logits, batch_y)
                val_loss_sum += loss.item() * batch_x.size(0)
                preds = logits.argmax(dim=1)
                val_correct += (preds == batch_y).sum().item()
                val_total += batch_x.size(0)
                preds_list.extend(preds.cpu().numpy())
                labels_list.extend(batch_y.cpu().numpy())

        val_loss = val_loss_sum / max(1, val_total)
        val_acc = val_correct / max(1, val_total)
        val_f1 = f1_score(
            np.array(labels_list),
            np.array(preds_list),
            average="macro",
            zero_division=0,
        )

        history["train_loss"].append(train_loss)
        history["val_loss"].append(val_loss)
        history["train_acc"].append(train_acc)
        history["val_acc"].append(val_acc)
        history["val_f1"].append(val_f1)

        scheduler.step(val_f1)

        if val_f1 > best_f1:
            best_f1 = val_f1
            patience_counter = 0
            torch.save(model.state_dict(), save_path)
        else:
            patience_counter += 1

        if patience_counter >= patience and (epoch + 1) >= min_epochs:
            print(f"Early stopping triggered at epoch {epoch + 1}")
            break

    return best_f1, history


def eval_model(model, val_loader, device, threshold=0.35):
    model.eval()
    preds, probs_list = [], []
    with torch.no_grad():
        for batch_x, batch_y in val_loader:
            batch_x = batch_x.to(device)
            logits = model(batch_x)
            probs = torch.softmax(logits, dim=1).cpu().numpy()
            probs_list.append(probs)
            batch_preds = (probs[:, 1] >= threshold).astype(int)
            preds.extend(batch_preds)
    probs = np.vstack(probs_list)
    return np.array(preds), probs


def plot_history(history, model_name, out_dir=MODEL_SAVE_DIR):
    os.makedirs(out_dir, exist_ok=True)
    epochs = list(range(1, len(history["train_loss"]) + 1))

    plt.figure(figsize=(8, 5))
    plt.plot(epochs, history["train_loss"], label="train loss", marker="o")
    plt.plot(epochs, history["val_loss"], label="val loss", marker="o")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.title(f"{model_name} - Loss per epoch")
    plt.legend()
    plt.grid(True)
    out_path = os.path.join(out_dir, f"{model_name}_loss.png")
    plt.tight_layout()
    plt.savefig(out_path)
    plt.close()

    plt.figure(figsize=(8, 5))
    plt.plot(epochs, history["train_acc"], label="train acc", marker="o")
    plt.plot(epochs, history["val_acc"], label="val acc", marker="o")
    plt.xlabel("Epoch")
    plt.ylabel("Accuracy")
    plt.title(f"{model_name} - Accuracy per epoch")
    plt.legend()
    plt.grid(True)
    out_path2 = os.path.join(out_dir, f"{model_name}_acc.png")
    plt.tight_layout()
    plt.savefig(out_path2)
    plt.close()


def main():
    set_seed(SEED)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    os.makedirs(MODEL_SAVE_DIR, exist_ok=True)

    df = pd.read_csv(MANIFEST_PATH)
    # Normalize paths and keep only rows whose embedding file exists
    df["embedding_path"] = df["embedding_path"].str.replace("\\\\", "/", regex=False)
    before = len(df)
    df = df[df["embedding_path"].apply(os.path.isfile)].copy()
    if len(df) < before:
        print(f"Dropped {before - len(df)} rows with missing embedding files (using {len(df)}).")
    if len(df) == 0:
        raise FileNotFoundError("No embedding files found. Run extract_embeddings.py first.")

    le = LabelEncoder()
    df["label_idx"] = le.fit_transform(df["label"])
    speaker_df = df.groupby("original_id")["label_idx"].first().reset_index()
    train_speakers, val_speakers = train_test_split(
        speaker_df["original_id"],
        test_size=0.2,
        stratify=speaker_df["label_idx"],
        random_state=SEED,
    )
    train_df = df[df["original_id"].isin(train_speakers)].copy()
    val_df = df[df["original_id"].isin(val_speakers)].copy()

    train_set = EmbeddingDataset(train_df, l2_norm=True)
    val_set = EmbeddingDataset(val_df, l2_norm=True)

    # =========================================================================
    # --- NEW: PCA Dimensionality Reduction (Fitted ONLY on training data) ---
    # =========================================================================
    print(f"\nApplying PCA to crush dimensions down to {PCA_COMPONENTS}...")

    # 1. Extract raw numpy arrays
    X_train_raw = np.stack([train_set[i][0].numpy() for i in range(len(train_set))])
    y_train_raw = np.array([train_set[i][1].item() for i in range(len(train_set))])

    X_val_raw = np.stack([val_set[i][0].numpy() for i in range(len(val_set))])
    y_val_raw = np.array([val_set[i][1].item() for i in range(len(val_set))])

    # 2. Fit PCA strictly on training data
    pca = PCA(n_components=PCA_COMPONENTS, random_state=SEED)
    X_train_pca = pca.fit_transform(X_train_raw)
    X_val_pca = pca.transform(X_val_raw)

    # 3. Overwrite Datasets with compressed tensors
    train_set = TensorDataset(
        torch.tensor(X_train_pca, dtype=torch.float32),
        torch.tensor(y_train_raw, dtype=torch.long),
    )
    val_set = TensorDataset(
        torch.tensor(X_val_pca, dtype=torch.float32),
        torch.tensor(y_val_raw, dtype=torch.long),
    )

    # 4. Override input_dim for the PyTorch models
    input_dim = PCA_COMPONENTS
    print(f"New Model Input Dimension: {input_dim}")
    # =========================================================================

    # Optional SMOTE (Now operating safely on the 64-dimensional PCA space)
    if USE_SMOTE:
        try:
            from imblearn.over_sampling import SMOTE

            X_train = np.stack([train_set[i][0].numpy() for i in range(len(train_set))])
            y_train = np.array([train_set[i][1].item() for i in range(len(train_set))])
            n_minority = int((y_train == 1).sum())
            k = min(SMOTE_K_NEIGHBORS, n_minority - 1) if n_minority > 1 else 1
            if k >= 1 and n_minority < (y_train == 0).sum():
                smote = SMOTE(k_neighbors=k, random_state=SEED)
                X_res, y_res = smote.fit_resample(X_train, y_train)
                train_set = TensorDataset(
                    torch.tensor(X_res, dtype=torch.float32),
                    torch.tensor(y_res, dtype=torch.long),
                )
                print(
                    f"SMOTE: {len(y_train)} -> {len(y_res)} samples (Non-Native oversampled)"
                )
            else:
                print("SMOTE: skipped (minority too small or already balanced)")
        except Exception as e:
            print(f"SMOTE skipped: {e}")

    train_loader = DataLoader(
        train_set, batch_size=BATCH_SIZE, shuffle=True, drop_last=True, num_workers=0
    )
    val_loader = DataLoader(
        val_set, batch_size=BATCH_SIZE, shuffle=False, num_workers=0
    )

    weights = compute_class_weight(
        "balanced", classes=np.unique(df["label_idx"]), y=df["label_idx"]
    )
    class_weights = torch.tensor(weights, dtype=torch.float32).to(device)

    # --- Label Smoothing included in the loss function ---
    criterion = nn.CrossEntropyLoss(weight=class_weights, label_smoothing=0.15)

    y_true = np.concatenate(
        [val_set[i][1].numpy().reshape(1) for i in range(len(val_set))]
    )

    # ---- 1. Train MLP (baseline) ----
    print("\n" + "=" * 60)
    print("1. Training MLP (baseline)")
    print("=" * 60)
    model_mlp = NativityClassifier(
        input_dim=input_dim,
        hidden_dims=HIDDEN_DIMS,
        output_dim=OUTPUT_DIM,
        dropout=DROPOUT,
    ).to(device)
    best_f1_mlp, hist_mlp = train_model(
        model_mlp,
        train_loader,
        val_loader,
        criterion,
        device,
        os.path.join(MODEL_SAVE_DIR, "nativity_model.pth"),
    )
    try:
        plot_history(hist_mlp, "mlp_baseline", MODEL_SAVE_DIR)
        print(f"Saved training plots for mlp_baseline in {MODEL_SAVE_DIR}")
    except Exception as e:
        print(f"Plotting mlp_baseline failed: {e}")

    # ---- 2. Train MLP + channel attention (prosody emphasis) ----
    print("\n2. Training MLP + channel attention (prosody emphasis)")
    print("=" * 60)
    model_att = NativityClassifierWithAttention(
        input_dim=input_dim,
        hidden_dims=HIDDEN_DIMS,
        output_dim=OUTPUT_DIM,
        dropout=DROPOUT,
        final_dropout=0.75,  # <--- Increased to 0.75 for safety
        reduction_ratio=8,
    ).to(device)
    best_f1_att, hist_att = train_model(
        model_att,
        train_loader,
        val_loader,
        criterion,
        device,
        os.path.join(MODEL_SAVE_DIR, "nativity_model_attention.pth"),
    )
    try:
        plot_history(hist_att, "mlp_attention", MODEL_SAVE_DIR)
        print(f"Saved training plots for mlp_attention in {MODEL_SAVE_DIR}")
    except Exception as e:
        print(f"Plotting mlp_attention failed: {e}")

    # ---- 3. LightGBM on same train/val embeddings ----
    print("\n3. Training LightGBM on embedding vectors")
    print("=" * 60)
    X_train = np.stack([train_set[i][0].numpy() for i in range(len(train_set))])
    y_train = np.array([train_set[i][1].item() for i in range(len(train_set))])
    X_val = np.stack([val_set[i][0].numpy() for i in range(len(val_set))])
    try:
        import lightgbm as lgb

        lgb_train = lgb.Dataset(X_train, y_train)
        lgb_val = lgb.Dataset(X_val, reference=lgb_train)
        params = {
            "objective": "multiclass",
            "num_class": 2,
            "metric": "multi_logloss",
            "num_leaves": 31,  # Reduced leaves slightly since we only have 64 dimensions now
            "learning_rate": 0.05,
            "feature_fraction": 0.8,
            "bagging_fraction": 0.8,
            "bagging_freq": 5,
            "verbose": -1,
            "seed": SEED,
        }
        callbacks = [lgb.early_stopping(stopping_rounds=25, verbose=False)]
        lgb_model = lgb.train(
            params,
            lgb_train,
            num_boost_round=500,
            valid_sets=[lgb_val],
            callbacks=callbacks,
        )
        lgb_model.save_model(os.path.join(MODEL_SAVE_DIR, "lgbm_nativity.txt"))
        prob_lgb = lgb_model.predict(X_val)
    except Exception as e:
        print(f"LightGBM skipped: {e}")
        prob_lgb = None

    # ---- 4. Load best and evaluate all + ensemble ----
    model_mlp.load_state_dict(
        torch.load(
            os.path.join(MODEL_SAVE_DIR, "nativity_model.pth"), map_location=device
        )
    )
    model_att.load_state_dict(
        torch.load(
            os.path.join(MODEL_SAVE_DIR, "nativity_model_attention.pth"),
            map_location=device,
        )
    )

    pred_mlp, prob_mlp = eval_model(
        model_mlp, val_loader, device, threshold=EVAL_THRESHOLD
    )
    pred_att, prob_att = eval_model(
        model_att, val_loader, device, threshold=EVAL_THRESHOLD
    )

    # Ensemble: average probs from MLP, attention, and LightGBM (if present)
    prob_list = [prob_mlp, prob_att]
    if prob_lgb is not None:
        prob_list.append(prob_lgb)
    prob_ens = np.stack(prob_list, axis=0).mean(axis=0)

    # Threshold sweep on validation
    best_threshold = EVAL_THRESHOLD
    best_f1_ens = 0.0
    print("\nThreshold sweep on validation (ensemble):")
    for t in THRESHOLD_CANDIDATES:
        pred_t = (prob_ens[:, 1] >= t).astype(int)
        f1_t = f1_score(y_true, pred_t, average="macro", zero_division=0)
        print(f"  t={t:.2f}  F1={f1_t:.4f}")
        if f1_t > best_f1_ens:
            best_f1_ens = f1_t
            best_threshold = t
    with open(BEST_THRESHOLD_PATH, "w") as f:
        f.write(str(best_threshold))
    print(
        f"Best threshold (val F1): {best_threshold} -> F1={best_f1_ens:.4f} (saved to {BEST_THRESHOLD_PATH})"
    )

    pred_ens = (prob_ens[:, 1] >= best_threshold).astype(int)

    def metrics_row(name, preds):
        acc = accuracy_score(y_true, preds)
        f1 = f1_score(y_true, preds, average="macro", zero_division=0)
        prec = precision_score(y_true, preds, average="macro", zero_division=0)
        rec = recall_score(y_true, preds, average="macro", zero_division=0)
        cm = confusion_matrix(y_true, preds)
        return {"name": name, "acc": acc, "f1": f1, "prec": prec, "rec": rec, "cm": cm}

    pred_lgb = np.argmax(prob_lgb, axis=1) if prob_lgb is not None else None
    n_ens = len(prob_list)
    results = [
        metrics_row("MLP (baseline)", pred_mlp),
        metrics_row("MLP + attention (prosody)", pred_att),
        metrics_row(f"Ensemble ({n_ens} models, t={best_threshold})", pred_ens),
    ]
    if pred_lgb is not None:
        results.append(metrics_row("LightGBM", pred_lgb))

    print("\n" + "=" * 60)
    print("COMPARISON (validation set)")
    print("=" * 60)
    print(
        f"{'Model':<30} {'Accuracy':>10} {'F1 macro':>10} {'Precision':>10} {'Recall':>10}"
    )
    print("-" * 72)
    for r in results:
        print(
            f"{r['name']:<30} {r['acc']:>10.4f} {r['f1']:>10.4f} {r['prec']:>10.4f} {r['rec']:>10.4f}"
        )
    print("\nConfusion matrices (rows=true, cols=pred), Native=0, Non-Native=1:")
    for r in results:
        print(f"  {r['name']}:")
        print(f"    {r['cm']}")

    best = max(results, key=lambda x: x["f1"])
    print(f"\nBest by F1: {best['name']} (F1={best['f1']:.4f})")
    print(
        "\nSaved: models/nativity_model.pth, models/nativity_model_attention.pth, models/best_threshold.txt"
    )
    print("       models/lgbm_nativity.txt (if LightGBM)")


if __name__ == "__main__":
    main()