"""
Train Native vs Non-Native classifier on Wav2Vec embeddings.
Uses early stopping, LR scheduler, class weights, and saves best model by val F1.
"""
import os
import sys
import random

# Allow importing model from same directory when run from project root
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
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.utils.class_weight import compute_class_weight
from sklearn.metrics import (
    accuracy_score,
    f1_score,
    precision_score,
    recall_score,
    confusion_matrix,
)

from model import NativityClassifier

# -----------------------------------------------------------------------------
# Config (change these if needed)
# -----------------------------------------------------------------------------
SEED = 42
MANIFEST_PATH = "manifests/final_training_manifest.csv"
MODEL_SAVE_DIR = "models"
BEST_MODEL_PATH = os.path.join(MODEL_SAVE_DIR, "nativity_model.pth")

# With more compute you can safely raise these
EPOCHS = 200        # max epochs; early stopping will usually stop earlier
BATCH_SIZE = 32
LR = 1e-3
WEIGHT_DECAY = 5e-4  # stronger L2 to reduce overfitting
PATIENCE = 15       # stricter early stopping (was 30) — stop sooner
MIN_LR = 1e-6

# Model architecture (must match model.py and embeddings dim)
HIDDEN_DIMS = (384, 256, 128)
DROPOUT = 0.5   # increased from 0.35 to reduce overfitting
INPUT_DIM = 1024  # overridden in main() from first .npy (1063 when using --mfcc)
OUTPUT_DIM = 2

# -----------------------------------------------------------------------------
# Reproducibility
# -----------------------------------------------------------------------------
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


# -----------------------------------------------------------------------------
# Dataset with optional L2 normalization
# -----------------------------------------------------------------------------
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
        )


# -----------------------------------------------------------------------------
# Training
# -----------------------------------------------------------------------------
def main():
    set_seed(SEED)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    os.makedirs(MODEL_SAVE_DIR, exist_ok=True)

    df = pd.read_csv(MANIFEST_PATH)
    input_dim = int(np.load(str(df["embedding_path"].iloc[0]).replace("\\", "/")).shape[0])
    print(f"Embedding dimension: {input_dim}")

    le = LabelEncoder()
    df["label_idx"] = le.fit_transform(df["label"])
    print(f"Classes: {le.classes_} (Native=0, Non-Native=1)")
    print(f"Total samples: {len(df)}")

    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)

    # Split by speaker (original_id) so the same speaker is never in both train and val.
    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()
    print(f"Train: {len(train_df)} chunks ({len(train_speakers)} speakers), Val: {len(val_df)} chunks ({len(val_speakers)} speakers)")

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

    train_loader = DataLoader(
        train_set,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=0,
        pin_memory=(device.type == "cuda"),
        drop_last=True,  # avoid batch size 1 with BatchNorm in training
    )
    val_loader = DataLoader(val_set, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)

    model = NativityClassifier(
        input_dim=input_dim,
        hidden_dims=HIDDEN_DIMS,
        output_dim=OUTPUT_DIM,
        dropout=DROPOUT,
    ).to(device)

    criterion = nn.CrossEntropyLoss(weight=class_weights)
    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_val_f1 = 0.0
    best_epoch = 0
    patience_counter = 0

    for epoch in range(EPOCHS):
        model.train()
        total_loss = 0.0
        for batch_x, batch_y in train_loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            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()
            total_loss += loss.item()

        train_loss = total_loss / len(train_loader)

        # Train metrics (to monitor overfitting: train_F1 >> val_F1 = overfitting)
        model.eval()
        train_preds, train_labels = [], []
        with torch.no_grad():
            for batch_x, batch_y in train_loader:
                batch_x = batch_x.to(device)
                preds = model(batch_x).argmax(dim=1).cpu().numpy()
                train_preds.extend(preds)
                train_labels.extend(batch_y.numpy())
        train_f1 = f1_score(train_labels, train_preds, average="macro", zero_division=0)

        # Validation
        all_preds = []
        all_labels = []
        val_loss = 0.0
        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)
                val_loss += criterion(logits, batch_y).item()
                preds = logits.argmax(dim=1)
                all_preds.extend(preds.cpu().numpy())
                all_labels.extend(batch_y.cpu().numpy())

        all_preds = np.array(all_preds)
        all_labels = np.array(all_labels)
        val_loss /= len(val_loader)

        acc = accuracy_score(all_labels, all_preds)
        f1 = f1_score(all_labels, all_preds, average="macro", zero_division=0)
        prec = precision_score(all_labels, all_preds, average="macro", zero_division=0)
        rec = recall_score(all_labels, all_preds, average="macro", zero_division=0)

        scheduler.step(f1)

        if f1 > best_val_f1:
            best_val_f1 = f1
            best_epoch = epoch + 1
            patience_counter = 0
            torch.save(model.state_dict(), BEST_MODEL_PATH)
            print(f"  -> New best model saved (val F1={f1:.4f})")
        else:
            patience_counter += 1

        print(
            f"Epoch {epoch + 1:3d} | loss {train_loss:.4f} | "
            f"train_F1 {train_f1:.4f} | val_F1 {f1:.4f} | val_acc {acc:.4f} | "
            f"(gap={(train_f1-f1):.3f})"
        )

        if patience_counter >= PATIENCE:
            print(f"Early stopping at epoch {epoch + 1} (no improvement for {PATIENCE} epochs).")
            break

    print(f"\nBest val F1: {best_val_f1:.4f} (epoch {best_epoch})")

    # Load best and print final val metrics + confusion matrix
    model.load_state_dict(torch.load(BEST_MODEL_PATH, map_location=device))
    model.eval()
    all_preds = []
    all_labels = []
    with torch.no_grad():
        for batch_x, batch_y in val_loader:
            batch_x = batch_x.to(device)
            logits = model(batch_x)
            all_preds.extend(logits.argmax(dim=1).cpu().numpy())
            all_labels.extend(batch_y.numpy())
    all_preds = np.array(all_preds)
    all_labels = np.array(all_labels)

    print("\nFinal validation metrics (best model):")
    print(f"  Accuracy:  {accuracy_score(all_labels, all_preds):.4f}")
    print(f"  F1 macro:  {f1_score(all_labels, all_preds, average='macro', zero_division=0):.4f}")
    print(f"  Precision: {precision_score(all_labels, all_preds, average='macro', zero_division=0):.4f}")
    print(f"  Recall:   {recall_score(all_labels, all_preds, average='macro', zero_division=0):.4f}")
    print("  Confusion matrix (rows=true, cols=pred):")
    cm = confusion_matrix(all_labels, all_preds)
    print(f"    {cm}")
    print(f"  (Native=0, Non-Native=1)")
    print(f"\nModel saved to {BEST_MODEL_PATH}")


if __name__ == "__main__":
    main()