#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os, random, json
import argparse
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.metrics import (
    roc_auc_score, f1_score, precision_score, recall_score
)
from scipy.stats import ttest_rel, t


def set_seed(seed: int = 42):
    os.environ["PYTHONHASHSEED"] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.use_deterministic_algorithms(True)


class VideoDataset(Dataset):
    def __init__(self, video_ids, tokens, labels, attention_t2f, attention_f2t):
        self.video_ids = list(video_ids)
        self.tokens = tokens
        self.labels = labels
        self.attention_t2f = attention_t2f
        self.attention_f2t = attention_f2t

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

    def __getitem__(self, idx):
        return (self.video_ids[idx],
                self.tokens[idx],
                self.labels[idx],
                self.attention_t2f[idx],
                self.attention_f2t[idx])


class BinaryClassifier(nn.Module):
    def __init__(self, input_dim=768):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
        )

    def forward(self, x):
        return self.fc(x)

def get_pos_weight_from_loader(loader):
    total_pos, total_neg = 0, 0
    for batch in loader:
        labels = batch[2].view(-1)
        total_pos += (labels == 1).sum().item()
        total_neg += (labels == 0).sum().item()
    pos_weight = total_neg / max(1, total_pos)
    return torch.tensor([pos_weight], dtype=torch.float32)

def weighted_bce_with_attention_loss(logits, targets, weights_t2f, weights_f2t, pos_weight, alpha, beta):
    bce_loss = nn.BCEWithLogitsLoss(pos_weight=pos_weight)(logits, targets)
    penalty_t2f = weights_t2f.mean()
    penalty_f2t = weights_f2t.mean()
    return bce_loss + alpha * penalty_t2f + beta * penalty_f2t


def load_embeddings_from_dir(pt_dir: str):
    video_ids, tensors = [], []
    for fn in sorted(os.listdir(pt_dir)):  
        if not fn.endswith(".pt"):
            continue
        video_id = os.path.splitext(fn)[0]
        path = os.path.join(pt_dir, fn)
        video_ids.append(video_id)
        tensors.append(torch.load(path, map_location="cpu"))
    if not tensors:
        raise RuntimeError(f"No .pt files found in: {pt_dir}")
    return video_ids, tensors

def build_merged_table(attn_csv: str, annotations_csv: str):
    dfattn = pd.read_csv(attn_csv)
    dfann = pd.read_csv(annotations_csv).drop_duplicates(subset="video_id", keep="first")
    return pd.merge(dfattn, dfann, on="video_id", how="inner")

def filter_to_available_labels(video_ids, tensors, merged_df):
    keep_vids, keep_tensors = [], []
    labels, ent_t2f, ent_f2t = [], [], []

    for vid, ten in zip(video_ids, tensors):
        row = merged_df[merged_df["video_id"] == vid]
        if row.empty:
            continue
        keep_vids.append(vid)
        keep_tensors.append(ten)
        labels.append(int(row["Actionable"].iloc[0]))
        ent_t2f.append(float(row["entropy_t2f"].iloc[0]))
        ent_f2t.append(float(row["entropy_f2t"].iloc[0]))

    if not keep_vids:
        raise RuntimeError("After merging labels/entropy, no video IDs remained. Check IDs and CSVs.")

    clstokens = torch.cat(keep_tensors, dim=0)
    labelst = torch.tensor(labels, dtype=torch.float32)
    t2f = torch.tensor(ent_t2f, dtype=torch.float32)
    f2t = torch.tensor(ent_f2t, dtype=torch.float32)
    return keep_vids, clstokens, labelst, t2f, f2t

def load_test_ids(test_json: str):
    with open(test_json, "r") as f:
        return set(json.load(f))


def run_one_seed(
    videoidslist, clstokens, labelstest, t2f_sim, f2t_sim,
    fixed_test_ids,
    lr, weight_decay, alpha, beta,
    seed=42,
    batch_size=32, num_epochs=50, patience=5,
    device="cpu",
):
    set_seed(seed)

    test_idx = [i for i, vid in enumerate(videoidslist) if vid in fixed_test_ids]
    train_val_idx = [i for i, vid in enumerate(videoidslist) if vid not in fixed_test_ids]

    if not test_idx:
        raise RuntimeError("No test videos matched embeddings. Check test_videoids.json vs filenames.")
    if not train_val_idx:
        raise RuntimeError("No train/val videos left after removing test set.")

    train_val_videoids = [videoidslist[i] for i in train_val_idx]
    test_videoids = [videoidslist[i] for i in test_idx]

    train_val_inputs = clstokens[train_val_idx]
    test_inputs = clstokens[test_idx]

    train_val_labels = labelstest[train_val_idx]
    test_labels = labelstest[test_idx]

    train_val_t2f = t2f_sim[train_val_idx]
    test_t2f = t2f_sim[test_idx]

    train_val_f2t = f2t_sim[train_val_idx]
    test_f2t = f2t_sim[test_idx]

    # stratified train/val split
    y = train_val_labels.numpy().astype(int)
    idxs = np.arange(len(train_val_videoids))
    tr_idx, va_idx = train_test_split(idxs, test_size=0.2, random_state=seed, stratify=y)

    train_ds = VideoDataset(
        [train_val_videoids[i] for i in tr_idx],
        train_val_inputs[tr_idx],
        train_val_labels[tr_idx],
        train_val_t2f[tr_idx],
        train_val_f2t[tr_idx],
    )
    val_ds = VideoDataset(
        [train_val_videoids[i] for i in va_idx],
        train_val_inputs[va_idx],
        train_val_labels[va_idx],
        train_val_t2f[va_idx],
        train_val_f2t[va_idx],
    )
    test_ds = VideoDataset(test_videoids, test_inputs, test_labels, test_t2f, test_f2t)

    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False)
    test_loader  = DataLoader(test_ds,  batch_size=batch_size, shuffle=False)

    model = BinaryClassifier(input_dim=clstokens.shape[1]).to(device)
    pos_weight = get_pos_weight_from_loader(train_loader).to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)

    best_val_loss = float("inf")
    early = 0
    best_state = None

    for _ in range(num_epochs):
        model.train()
        for _, tokens, labels, attn_t2f, attn_f2t in train_loader:
            tokens   = tokens.to(device)
            labels   = labels.to(device).float()
            attn_t2f = attn_t2f.to(device)
            attn_f2t = attn_f2t.to(device)

            optimizer.zero_grad()
            logits = model(tokens).squeeze()
            loss = weighted_bce_with_attention_loss(
                logits, labels, attn_t2f, attn_f2t, pos_weight, alpha, beta
            )
            loss.backward()
            optimizer.step()

        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for _, tokens, labels, attn_t2f, attn_f2t in val_loader:
                tokens   = tokens.to(device)
                labels   = labels.to(device).float()
                attn_t2f = attn_t2f.to(device)
                attn_f2t = attn_f2t.to(device)

                logits = model(tokens).squeeze()
                val_loss += weighted_bce_with_attention_loss(
                    logits, labels, attn_t2f, attn_f2t, pos_weight, alpha, beta
                ).item()
        val_loss /= max(1, len(val_loader))

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            early = 0
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
        else:
            early += 1
            if early >= patience:
                break

    if best_state is not None:
        model.load_state_dict(best_state)

    # test predictions
    model.eval()
    vids, y_true, y_scores = [], [], []
    with torch.no_grad():
        for video_ids, tokens, labels, _, _ in test_loader:
            tokens = tokens.to(device)
            logits = model(tokens).squeeze()
            vids.extend(list(video_ids))
            y_true.extend(labels.numpy().astype(int))
            y_scores.extend(logits.cpu().numpy())

    y_true = np.asarray(y_true).astype(int)
    y_scores = np.asarray(y_scores)

    auc = roc_auc_score(y_true, y_scores)
    y_pred = (y_scores >= 0.5).astype(int)
    f1 = f1_score(y_true, y_pred, average="weighted", zero_division=0)
    pr = precision_score(y_true, y_pred, average="weighted", zero_division=0)
    re = recall_score(y_true, y_pred, average="weighted", zero_division=0)

    return {
        "seed": seed,
        "test_auc": float(auc),
        "test_f1": float(f1),
        "test_precision": float(pr),
        "test_recall": float(re),
        "preds_df": pd.DataFrame({
            "seed": seed,
            "video_id": vids,
            "y_true": y_true,
            "y_score": y_scores,
        })
    }


def eval_5seeds_and_save(args):
    os.makedirs(args.out_dir, exist_ok=True)

    merged = build_merged_table(args.attention_csv, args.annotations_csv)
    fixed_test_ids = load_test_ids(args.test_json)

    video_ids, tensors = load_embeddings_from_dir(args.pt_dir)
    video_ids, clstokens, labelstest, t2f_sim, f2t_sim = filter_to_available_labels(video_ids, tensors, merged)

    present = set(video_ids)
    missing = sorted(list(set(fixed_test_ids) - present))
    if missing:
        print(f"[WARN] Missing {len(missing)} test IDs in {args.pt_dir}. Example: {missing[:10]}")

    seeds = [int(s.strip()) for s in args.seeds.split(",") if s.strip()]
    if len(seeds) != 5:
        print(f"[WARN] You provided {len(seeds)} seeds; Table-style stability is best with 5.")

    metrics_rows = []
    preds_all = []

    for sd in seeds:
        res = run_one_seed(
            video_ids, clstokens, labelstest, t2f_sim, f2t_sim,
            fixed_test_ids=fixed_test_ids,
            lr=args.lr, weight_decay=args.wd, alpha=args.alpha, beta=args.beta,
            seed=sd,
            batch_size=args.batch_size, num_epochs=args.epochs, patience=args.patience,
            device=args.device
        )
        metrics_rows.append({
            "seed": sd,
            "lr": args.lr, "wd": args.wd, "alpha": args.alpha, "beta": args.beta,
            "test_auc": res["test_auc"],
            "test_f1": res["test_f1"],
            "test_precision": res["test_precision"],
            "test_recall": res["test_recall"],
        })
        preds_all.append(res["preds_df"])
        print(f"[seed={sd}] AUC={res['test_auc']:.4f} F1={res['test_f1']:.4f}")

    df_metrics = pd.DataFrame(metrics_rows)
    df_preds = pd.concat(preds_all, ignore_index=True)

    summary = {
        "lr": args.lr, "wd": args.wd, "alpha": args.alpha, "beta": args.beta,
        "n_seeds": len(seeds),
        "auc_mean": df_metrics["test_auc"].mean(),
        "auc_std": df_metrics["test_auc"].std(ddof=1) if len(seeds) > 1 else 0.0,
        "f1_mean": df_metrics["test_f1"].mean(),
        "f1_std": df_metrics["test_f1"].std(ddof=1) if len(seeds) > 1 else 0.0,
        "precision_mean": df_metrics["test_precision"].mean(),
        "precision_std": df_metrics["test_precision"].std(ddof=1) if len(seeds) > 1 else 0.0,
        "recall_mean": df_metrics["test_recall"].mean(),
        "recall_std": df_metrics["test_recall"].std(ddof=1) if len(seeds) > 1 else 0.0,
    }
    df_summary = pd.DataFrame([summary])

    df_metrics.to_csv(os.path.join(args.out_dir, "metrics_by_seed.csv"), index=False)
    df_summary.to_csv(os.path.join(args.out_dir, "metrics_summary.csv"), index=False)
    df_preds.to_csv(os.path.join(args.out_dir, "test_predictions_by_seed.csv"), index=False)

    print("\n Summary (mean and std over seeds)")
    print(df_summary.to_string(index=False))


def stars(p):
    if p < 0.01: return "***"
    if p < 0.05: return "**"
    if p < 0.10: return "*"
    return ""

def compare_methods_by_videoids_ttest(
    y_true,
    scores1,
    scores2,
    video_ids,
    n_ids=100,
    n_iters=100,
    threshold=0.5,
    random_state=42
):
    rng = np.random.default_rng(random_state)

    y_true = np.asarray(y_true)
    s1 = np.asarray(scores1)
    s2 = np.asarray(scores2)
    vids = np.asarray(video_ids)

    unique_vids = np.unique(vids)
    n_pick = min(n_ids, unique_vids.shape[0])

    diffs = {m: [] for m in ("auc", "f1", "precision", "recall")}
    m1_hist = {m: [] for m in ("auc", "f1", "precision", "recall")}
    m2_hist = {m: [] for m in ("auc", "f1", "precision", "recall")}

    def compute_metrics(y, p1, p2):
        auc1 = auc2 = np.nan
        if np.unique(y).size >= 2:
            auc1 = roc_auc_score(y, p1)
            auc2 = roc_auc_score(y, p2)
        yhat1 = (p1 >= threshold).astype(int)
        yhat2 = (p2 >= threshold).astype(int)
        f11 = f1_score(y, yhat1, average="weighted", zero_division=0)
        f12 = f1_score(y, yhat2, average="weighted", zero_division=0)
        pr1 = precision_score(y, yhat1, average="weighted", zero_division=0)
        pr2 = precision_score(y, yhat2, average="weighted", zero_division=0)
        re1 = recall_score(y, yhat1, average="weighted", zero_division=0)
        re2 = recall_score(y, yhat2, average="weighted", zero_division=0)
        return (auc1, auc2), (f11, f12), (pr1, pr2), (re1, re2)

    for _ in range(n_iters):
        picked = rng.choice(unique_vids, size=n_pick, replace=False)
        mask = np.isin(vids, picked)
        y = y_true[mask]; p1 = s1[mask]; p2 = s2[mask]
        if y.size == 0:
            continue

        (auc1, auc2), (f11, f12), (pr1, pr2), (re1, re2) = compute_metrics(y, p1, p2)

        if not np.isnan(auc1) and not np.isnan(auc2):
            diffs["auc"].append(auc1 - auc2); m1_hist["auc"].append(auc1); m2_hist["auc"].append(auc2)

        diffs["f1"].append(f11 - f12); m1_hist["f1"].append(f11); m2_hist["f1"].append(f12)
        diffs["precision"].append(pr1 - pr2); m1_hist["precision"].append(pr1); m2_hist["precision"].append(pr2)
        diffs["recall"].append(re1 - re2); m1_hist["recall"].append(re1); m2_hist["recall"].append(re2)

    results = {}
    for m in ("auc", "f1", "precision", "recall"):
        m1a = np.array(m1_hist[m], dtype=float)
        m2a = np.array(m2_hist[m], dtype=float)
        d = np.array(diffs[m], dtype=float)
        if d.size == 0:
            results[m] = {"p": np.nan, "star": "", "mean_diff": np.nan}
            continue
        t_stat, p_val = ttest_rel(m1a, m2a)
        results[m] = {
            "n": int(d.size),
            "mean_diff": float(d.mean()),
            "p": float(p_val),
            "star": stars(p_val),
            "method1_mean": float(m1a.mean()),
            "method2_mean": float(m2a.mean()),
        }
    return results

def compare_from_saved_preds(args):
    df1 = pd.read_csv(args.preds1)
    df2 = pd.read_csv(args.preds2)

    # Aggregate across seeds per video (mean is stable)
    g1 = df1.groupby("video_id", as_index=False).agg(y_true=("y_true", "first"), y_score=("y_score", "mean"))
    g2 = df2.groupby("video_id", as_index=False).agg(y_true=("y_true", "first"), y_score=("y_score", "mean"))

    dfm = pd.merge(g1, g2, on="video_id", suffixes=("_m1", "_m2"))
    if dfm.empty:
        raise RuntimeError("No overlapping video_ids between methods.")

    if not (dfm["y_true_m1"].values == dfm["y_true_m2"].values).all():
        raise RuntimeError("y_true mismatch between methods for same video_id (labeling/split issue).")

    res = compare_methods_by_videoids_ttest(
        y_true=dfm["y_true_m1"].values,
        scores1=dfm["y_score_m1"].values,
        scores2=dfm["y_score_m2"].values,
        video_ids=dfm["video_id"].values,
        n_ids=args.n_ids,
        n_iters=args.n_iters,
        threshold=args.threshold,
        random_state=args.random_state,
    )

    print("\n Paired significance (method1 - method2)")
    for k in ("auc", "f1", "precision", "recall"):
        r = res[k]
        print(f"{k.upper():9s}  mean_diff={r['mean_diff']:+.4f}  p={r['p']:.4g} {r['star']}  "
              f"(m1={r['method1_mean']:.4f}, m2={r['method2_mean']:.4f})")


def main():
    parser = argparse.ArgumentParser(description="5-seed evaluation + paired significance testing between methods.")
    sub = parser.add_subparsers(dest="cmd", required=True)

    # EVAL one method
    p_eval = sub.add_parser("eval", help="Run 5 seeds for one method and save predictions/summary.")
    p_eval.add_argument("--annotations_csv", type=str, required=True)
    p_eval.add_argument("--attention_csv", type=str, required=True)
    p_eval.add_argument("--pt_dir", type=str, required=True)
    p_eval.add_argument("--test_json", type=str, required=True)
    p_eval.add_argument("--out_dir", type=str, required=True)

    p_eval.add_argument("--lr", type=float, required=True)
    p_eval.add_argument("--wd", type=float, required=True)
    p_eval.add_argument("--alpha", type=float, required=True)
    p_eval.add_argument("--beta", type=float, required=True)

    p_eval.add_argument("--seeds", type=str, default="42,43,44,45,46")
    p_eval.add_argument("--batch_size", type=int, default=32)
    p_eval.add_argument("--epochs", type=int, default=50)
    p_eval.add_argument("--patience", type=int, default=5)
    p_eval.add_argument("--device", type=str, default="cpu", choices=["cpu", "cuda"])

    # COMPARE two methods
    p_cmp = sub.add_parser("compare", help="Compare two methods using their saved prediction CSVs.")
    p_cmp.add_argument("--preds1", type=str, required=True, help="method1 test_predictions_by_seed.csv")
    p_cmp.add_argument("--preds2", type=str, required=True, help="method2 test_predictions_by_seed.csv")
    p_cmp.add_argument("--n_ids", type=int, default=100)
    p_cmp.add_argument("--n_iters", type=int, default=100)
    p_cmp.add_argument("--threshold", type=float, default=0.5)
    p_cmp.add_argument("--random_state", type=int, default=42)

    args = parser.parse_args()

    if args.cmd == "eval":
        eval_5seeds_and_save(args)
    else:
        compare_from_saved_preds(args)

if __name__ == "__main__":
    main()