"""
Train and compare: 1D-CNN (Chronological Sequence) and LightGBM 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.nn.functional as F
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.metrics import (
    accuracy_score,
    f1_score,
    precision_score,
    recall_score,
    confusion_matrix,
)

import matplotlib

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

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

# --- HYPERPARAMETERS for 1D-CNN ---
LR = 1e-3
WEIGHT_DECAY = 1e-4
PATIENCE = 15
MIN_EPOCHS = 40
DROPOUT = 0.4
INPUT_DIM = 1024
MAX_TIME_STEPS = 150  # Standardizes varying chunk lengths to exactly 3 seconds
# ----------------------------------

OUTPUT_DIM = 2
EVAL_THRESHOLD = 0.50
BEST_THRESHOLD_PATH = os.path.join(MODEL_SAVE_DIR, "best_threshold.txt")
MIN_LR = 1e-6
THRESHOLD_CANDIDATES = [0.25, 0.30, 0.35, 0.40, 0.45, 0.50, 0.55]


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


# =========================================================================
# --- NEW: 1D Convolutional Neural Network ---
# =========================================================================
class NativityCNN1D(nn.Module):
    def __init__(self, input_dim=1024, num_classes=2, dropout=0.5):
        super(NativityCNN1D, self).__init__()

        # Scans chronologically for phonetic mispronunciations
        self.conv1 = nn.Conv1d(
            in_channels=input_dim, out_channels=256, kernel_size=3, padding=1
        )
        self.bn1 = nn.BatchNorm1d(256)

        self.conv2 = nn.Conv1d(
            in_channels=256, out_channels=128, kernel_size=3, padding=1
        )
        self.bn2 = nn.BatchNorm1d(128)

        self.conv3 = nn.Conv1d(
            in_channels=128, out_channels=64, kernel_size=3, padding=1
        )
        self.bn3 = nn.BatchNorm1d(64)

        # Squeezes the sequence down to a final decision at the end
        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(64, num_classes)

    def forward(self, x):
        # Incoming shape: (Batch, Time, Features) -> Convert to (Batch, Features, Time) for PyTorch Conv1d
        x = x.permute(0, 2, 1)

        x = F.relu(self.bn1(self.conv1(x)))
        x = F.max_pool1d(x, kernel_size=2)

        x = F.relu(self.bn2(self.conv2(x)))
        x = F.max_pool1d(x, kernel_size=2)

        x = F.relu(self.bn3(self.conv3(x)))

        x = self.global_pool(x).squeeze(-1)
        x = self.dropout(x)
        logits = self.fc(x)

        return logits


# =========================================================================


class EmbeddingDataset(torch.utils.data.Dataset):
    def __init__(self, df, embedding_path_col="embedding_path", label_col="label_idx"):
        self.paths = df[embedding_path_col].tolist()
        self.labels = df[label_col].values

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

    def __getitem__(self, idx):
        emb = np.load(self.paths[idx]).astype(np.float32)

        # Handle 1D arrays if older data is mixed in
        if len(emb.shape) == 1:
            emb = np.expand_dims(emb, axis=0)

        time_len, feat_dim = emb.shape

        # --- NEW: Sequence Padding/Truncation ---
        # Ensures all batch items have the exact same time dimension
        if time_len < MAX_TIME_STEPS:
            pad_width = MAX_TIME_STEPS - time_len
            emb = np.pad(emb, ((0, pad_width), (0, 0)), mode="constant")
        elif time_len > MAX_TIME_STEPS:
            emb = emb[:MAX_TIME_STEPS, :]

        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)

            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)

        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.50):
    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():
    import argparse
    ap = argparse.ArgumentParser(description="Train 1D-CNN on sequential embeddings.")
    ap.add_argument("--manifest", default=MANIFEST_PATH, help="Manifest CSV with embedding_path (default: final_training_manifest.csv)")
    args = ap.parse_args()
    manifest_path = args.manifest

    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)

    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)
    val_set = EmbeddingDataset(val_df)

    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)

    # Removed class weights to prevent 97% class-1 guessing
    criterion = nn.CrossEntropyLoss()

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

    # ---- 1. Train 1D-CNN Sequence Classifier ----
    print("\n" + "=" * 60)
    print("1. Training 1D-CNN (Chronological Sequence)")
    print("=" * 60)
    model_cnn = NativityCNN1D(
        input_dim=INPUT_DIM,
        dropout=DROPOUT,
    ).to(device)

    best_f1_cnn, hist_cnn = train_model(
        model_cnn,
        train_loader,
        val_loader,
        criterion,
        device,
        os.path.join(MODEL_SAVE_DIR, "nativity_cnn1d_model.pth"),
    )
    try:
        plot_history(hist_cnn, "cnn1d_model", MODEL_SAVE_DIR)
        print(f"Saved training plots for cnn1d_model in {MODEL_SAVE_DIR}")
    except Exception as e:
        print(f"Plotting cnn1d_model failed: {e}")

    # ---- 2. LightGBM on flattened embeddings ----
    print("\n2. Training LightGBM on mean-pooled vectors")
    print("=" * 60)

    # LightGBM cannot read 2D Sequences, so we mathematically average the sequence just for the tree
    X_train_seq = np.stack([train_set[i][0].numpy() for i in range(len(train_set))])
    y_train_lgb = np.array([train_set[i][1].item() for i in range(len(train_set))])
    X_val_seq = np.stack([val_set[i][0].numpy() for i in range(len(val_set))])

    X_train_lgb = X_train_seq.mean(
        axis=1
    )  # Flatten (Batch, 150, 1024) -> (Batch, 1024)
    X_val_lgb = X_val_seq.mean(axis=1)

    try:
        import lightgbm as lgb

        lgb_train = lgb.Dataset(X_train_lgb, y_train_lgb)
        lgb_val = lgb.Dataset(X_val_lgb, reference=lgb_train)
        params = {
            "objective": "binary",
            "metric": "binary_logloss",
            "is_unbalance": True,
            "num_leaves": 31,
            "max_depth": 5,
            "learning_rate": 0.01,
            "feature_fraction": 0.6,
            "bagging_fraction": 0.7,
            "bagging_freq": 1,
            "min_data_in_leaf": 20,
            "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_lgb)
    except Exception as e:
        print(f"LightGBM skipped: {e}")
        prob_lgb = None

    # ---- 3. Evaluate and Ensemble ----
    model_cnn.load_state_dict(
        torch.load(
            os.path.join(MODEL_SAVE_DIR, "nativity_cnn1d_model.pth"),
            map_location=device,
        )
    )

    pred_cnn, prob_cnn = eval_model(
        model_cnn, val_loader, device, threshold=EVAL_THRESHOLD
    )

    # Convert binary LightGBM probabilities to a pseudo-multiclass output so the ensemble stack works
    if prob_lgb is not None:
        prob_lgb_multi = np.vstack([1 - prob_lgb, prob_lgb]).T
        prob_list = [prob_cnn, prob_lgb_multi]
    else:
        prob_list = [prob_cnn]

    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_final = (prob_lgb >= 0.50).astype(int) if prob_lgb is not None else None

    results = [
        metrics_row("1D-CNN (Sequence)", pred_cnn),
        metrics_row(f"Ensemble (CNN + LGB, t={best_threshold})", pred_ens),
    ]
    if pred_lgb_final is not None:
        results.append(metrics_row("LightGBM", pred_lgb_final))

    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']}")


if __name__ == "__main__":
    main()