# -*- coding: utf-8 -*-

import os
import json
import random
import argparse
from typing import List, Optional, Tuple

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


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 BinaryClassifier(nn.Module):
    def __init__(self, input_dim: int = 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)


class VideoDataset(Dataset):
    def __init__(self, video_ids, tokens, labels, attn_t2f, attn_f2t):
        self.video_ids = list(video_ids)
        self.tokens    = tokens
        self.labels    = labels
        self.attn_t2f  = attn_t2f
        self.attn_f2t  = attn_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.attn_t2f[idx],
            self.attn_f2t[idx],
        )


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


def compute_loss(logits, targets, attn_t2f, attn_f2t,
                 pos_weight, alpha, beta, use_entropy):
    bce = nn.BCEWithLogitsLoss(pos_weight=pos_weight)(logits, targets)
    if use_entropy:
        return bce + alpha * attn_t2f.mean() + beta * attn_f2t.mean()
    return bce


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


def load_and_merge(annotations_csv: str,
                   attention_csv: Optional[str]) -> pd.DataFrame:
    dfann = pd.read_csv(annotations_csv).drop_duplicates(
        subset="video_id", keep="first")
    if attention_csv:
        dfattn = pd.read_csv(attention_csv)
        return pd.merge(dfattn, dfann, on="video_id", how="inner")
    return dfann


def filter_and_align(video_ids:   List[str],
                     tensors:     List[torch.Tensor],
                     merged_df:   pd.DataFrame,
                     use_entropy: bool):
    keep_ids, 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_ids.append(vid)
        keep_tensors.append(ten)
        labels.append(int(row["Actionable"].iloc[0]))
        if use_entropy:
            ent_t2f.append(float(row["entropy_t2f"].iloc[0]))
            ent_f2t.append(float(row["entropy_f2t"].iloc[0]))

    if not keep_ids:
        raise RuntimeError(
            "No overlap between embeddings and metadata after merging"
        )

    X   = torch.cat(keep_tensors, dim=0)
    y   = torch.tensor(labels, dtype=torch.float32)

    if use_entropy:
        t2f = torch.tensor(ent_t2f, dtype=torch.float32)
        f2t = torch.tensor(ent_f2t, dtype=torch.float32)
    else:
        t2f = torch.zeros(len(keep_ids), dtype=torch.float32)
        f2t = torch.zeros(len(keep_ids), dtype=torch.float32)

    return keep_ids, X, y, t2f, f2t


def run_one_seed(
    video_ids:      List[str],
    X:              torch.Tensor,
    y:              torch.Tensor,
    t2f:            torch.Tensor,
    f2t:            torch.Tensor,
    fixed_test_ids: set,
    input_dim:      int,
    lr:             float,
    wd:             float,
    alpha:          float,
    beta:           float,
    use_entropy:    bool,
    n_train:        int   = None,  
    seed:           int   = 42,
    batch_size:     int   = 32,
    epochs:         int   = 50,
    patience:       int   = 5,
    device:         str   = "cpu",
) -> dict:
    set_seed(seed)

    test_idx      = [i for i, v in enumerate(video_ids) if v     in fixed_test_ids]
    train_val_idx = [i for i, v in enumerate(video_ids) if v not in fixed_test_ids]

    if not test_idx:
        raise RuntimeError(
            "No test videos found in embeddings."
        )
    if not train_val_idx:
        raise RuntimeError("No train/val videos left after removing test set.")

    tv_ids    = [video_ids[i] for i in train_val_idx]
    test_vids = [video_ids[i] for i in test_idx]

    X_tv, y_tv = X[train_val_idx], y[train_val_idx]
    X_te, y_te = X[test_idx],      y[test_idx]
    t2f_tv, f2t_tv = t2f[train_val_idx], f2t[train_val_idx]
    t2f_te, f2t_te = t2f[test_idx],      f2t[test_idx]

    idxs = np.arange(len(tv_ids))
    tr_idx, va_idx = train_test_split(
        idxs, test_size=0.2, random_state=seed,
        stratify=y_tv.numpy().astype(int)
    )

    if n_train is not None:
        default_n_train = len(tr_idx)
        diff            = n_train - default_n_train

        if diff < 0:
            rng    = np.random.default_rng(seed)
            tr_idx = rng.choice(tr_idx, size=n_train, replace=False)

        elif diff > 0:
            if diff > len(va_idx):
                raise ValueError(
                    f"Cannot increase training by {diff},"
                    f"val only has {len(va_idx)} records."
                )
            rng      = np.random.default_rng(seed)
            move_idx = rng.choice(
                np.arange(len(va_idx)), size=diff, replace=False)
            tr_idx   = np.concatenate([tr_idx, va_idx[move_idx]])
            va_idx   = np.delete(va_idx, move_idx)

    train_ds = VideoDataset(
        [tv_ids[i] for i in tr_idx],
        X_tv[tr_idx], y_tv[tr_idx],
        t2f_tv[tr_idx], f2t_tv[tr_idx],
    )
    val_ds = VideoDataset(
        [tv_ids[i] for i in va_idx],
        X_tv[va_idx], y_tv[va_idx],
        t2f_tv[va_idx], f2t_tv[va_idx],
    )
    test_ds = VideoDataset(
        test_vids, X_te, y_te, t2f_te, f2t_te
    )

    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=input_dim).to(device)
    pos_weight = get_pos_weight_from_loader(train_loader).to(device)
    optimizer  = optim.Adam(model.parameters(), lr=lr, weight_decay=wd)

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

    for _ in range(epochs):
        model.train()
        for _, tokens, labels, a_t2f, a_f2t in train_loader:
            tokens = tokens.to(device)
            labels = labels.to(device).float()
            a_t2f  = a_t2f.to(device)
            a_f2t  = a_f2t.to(device)
            optimizer.zero_grad()
            logits = model(tokens).squeeze()
            loss   = compute_loss(logits, labels, a_t2f, a_f2t,
                                  pos_weight, alpha, beta, use_entropy)
            loss.backward()
            optimizer.step()

        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for _, tokens, labels, a_t2f, a_f2t in val_loader:
                tokens = tokens.to(device)
                labels = labels.to(device).float()
                a_t2f  = a_t2f.to(device)
                a_f2t  = a_f2t.to(device)
                logits = model(tokens).squeeze()
                val_loss += compute_loss(logits, labels, a_t2f, a_f2t,
                                         pos_weight, alpha, beta,
                                         use_entropy).item()
        val_loss /= max(1, len(val_loader))

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

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

    model.eval()
    vids_out, y_true_out, y_score_out = [], [], []
    with torch.no_grad():
        for batch_vids, tokens, labels, _, _ in test_loader:
            logits = model(tokens.to(device)).squeeze()
            vids_out.extend(list(batch_vids))
            y_true_out.extend(labels.numpy().astype(int))
            y_score_out.extend(logits.cpu().numpy())

    y_true  = np.asarray(y_true_out).astype(int)
    y_score = np.asarray(y_score_out)
    y_pred  = (y_score >= 0.5).astype(int)

    return {
        "seed":           seed,
        "test_auc":       float(roc_auc_score(y_true, y_score)),
        "test_f1":        float(f1_score(y_true, y_pred, average="weighted", zero_division=0)),
        "test_precision": float(precision_score(y_true, y_pred, average="weighted", zero_division=0)),
        "test_recall":    float(recall_score(y_true, y_pred, average="weighted", zero_division=0)),
        "preds_df": pd.DataFrame({
            "seed":     seed,
            "video_id": vids_out,
            "y_true":   y_true,
            "y_score":  y_score,
        }),
    }


def run_eval(args):
    print("\n" + "=" * 60)
    print("  EVALUATE.PY eval")
    print(f"  PT dir    : {args.pt_dir}")
    print(f"  Model dir : {args.model_dir}")
    print("=" * 60)

    os.makedirs(args.model_dir, exist_ok=True)

    params_path = os.path.join(args.model_dir, "best_params.json")
    if not os.path.exists(params_path):
        raise FileNotFoundError(
            f"best_params.json not found at: {params_path}\n"
            f"Run train.py first."
        )
    with open(params_path, "r", encoding="utf-8") as f:
        best_params = json.load(f)

    lr          = best_params["lr"]
    wd          = best_params["wd"]
    alpha       = best_params.get("alpha", 0.0)
    beta        = best_params.get("beta",  0.0)
    use_entropy = best_params["use_entropy"]
    variant     = best_params["variant"]

    print(f"\n  Loaded best_params.json:")
    print(f"    variant={variant}  lr={lr:g}  wd={wd:g}  "
          f"alpha={alpha:g}  beta={beta:g}  use_entropy={use_entropy}")

    if args.n_train is not None:
        print(f"    n_train={args.n_train}")

    video_ids, tensors = load_embeddings(args.pt_dir)
    merged             = load_and_merge(
        args.annotations_csv,
        args.attention_csv if use_entropy else None
    )
    video_ids, X, y, t2f, f2t = filter_and_align(
        video_ids, tensors, merged, use_entropy
    )
    
    if args.pt_dir_extra:
        extra_ids, extra_tensors = load_embeddings(args.pt_dir_extra)
        extra_merged = load_and_merge(
            args.annotations_csv,
            args.attention_csv_extra if use_entropy else None
        )
        extra_ids, extra_X, extra_y, extra_t2f, extra_f2t = filter_and_align(
            extra_ids, extra_tensors, extra_merged, use_entropy
        )
        video_ids = video_ids + extra_ids
        X = torch.cat([X,   extra_X],   dim=0)
        y = torch.cat([y,   extra_y],   dim=0)
        t2f = torch.cat([t2f, extra_t2f], dim=0)
        f2t = torch.cat([f2t, extra_f2t], dim=0)
    
    with open(args.test_json, "r", encoding="utf-8") as f:
        fixed_test_ids = set(json.load(f))

    missing = sorted(fixed_test_ids - set(video_ids))
    if missing:
        print(f"\n  [WARN] {len(missing)} test IDs missing from pt_dir. "
              f"Example: {missing[:5]}")

    print(f"\n  Total videos : {len(video_ids)}")
    print(f"  Test videos  : {len([v for v in video_ids if v in fixed_test_ids])}")
    print(f"  Train/val    : {len([v for v in video_ids if v not in fixed_test_ids])}")

    seeds = [int(s.strip()) for s in args.seeds.split(",") if s.strip()]
    print(f"  Seeds        : {seeds}\n")

    input_dim    = X.shape[1]
    metrics_rows = []
    preds_all    = []

    for sd in seeds:
        res = run_one_seed(
            video_ids, X, y, t2f, f2t,
            fixed_test_ids = fixed_test_ids,
            input_dim      = input_dim,
            lr             = lr,
            wd             = wd,
            alpha          = alpha,
            beta           = beta,
            use_entropy    = use_entropy,
            n_train        = args.n_train,   
            seed           = sd,
            batch_size     = args.batch_size,
            epochs         = args.epochs,
            patience       = args.patience,
            device         = args.device,
        )
        metrics_rows.append({
            "seed":           sd,
            "lr": lr, "wd": wd, "alpha": alpha, "beta": beta,
            "n_train":        args.n_train,  
            "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}]  "
              f"AUC={res['test_auc']:.4f}  "
              f"F1={res['test_f1']:.4f}  "
              f"Precision={res['test_precision']:.4f}  "
              f"Recall={res['test_recall']:.4f}")

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

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

    df_metrics.to_csv(os.path.join(args.model_dir, "metrics_by_seed.csv"),          index=False)
    df_preds.to_csv(  os.path.join(args.model_dir, "test_predictions_by_seed.csv"), index=False)
    pd.DataFrame([summary]).to_csv(
        os.path.join(args.model_dir, "metrics_summary.csv"), index=False)

    print(f"\n  Test results (mean +/- std over {n} seeds):")
    print(f"    AUC       {summary['auc_mean']:.4f}  +/-  {summary['auc_std']:.4f}")
    print(f"    F1        {summary['f1_mean']:.4f}  +/-  {summary['f1_std']:.4f}")
    print(f"    Precision {summary['precision_mean']:.4f}  +/-  {summary['precision_std']:.4f}")
    print(f"    Recall    {summary['recall_mean']:.4f}  +/-  {summary['recall_std']:.4f}")
    print(f"\n  Saved to: {args.model_dir}")


def _stars(p: float) -> str:
    if p < 0.01: return "***"
    if p < 0.05: return "**"
    if p < 0.10: return "*"
    return ""


def run_compare(args):
    print("\n" + "=" * 60)
    print("  EVALUATE.PY  --  compare")
    print(f"  Method 1 : {args.preds1}")
    print(f"  Method 2 : {args.preds2}")
    print(f"  n_iters  : {args.n_iters}   n_ids: {args.n_ids}")
    print("=" * 60)

    df1 = pd.read_csv(args.preds1)
    df2 = pd.read_csv(args.preds2)

    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. "
            "Check that both prediction files cover the same test set."
        )
    if not (dfm["y_true_m1"].values == dfm["y_true_m2"].values).all():
        raise RuntimeError(
            "y_true mismatch between methods for the same video_id. "
            "Check that both methods used the same test_videoids.json."
        )

    rng         = np.random.default_rng(args.random_state)
    y_true      = dfm["y_true_m1"].values
    s1          = dfm["y_score_m1"].values
    s2          = dfm["y_score_m2"].values
    vids        = dfm["video_id"].values
    unique_vids = np.unique(vids)
    n_pick      = min(args.n_ids, unique_vids.shape[0])

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

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

        yhat1 = (p1 >= args.threshold).astype(int)
        yhat2 = (p2 >= args.threshold).astype(int)

        m1_hist["auc"].append(roc_auc_score(yt, p1))
        m2_hist["auc"].append(roc_auc_score(yt, p2))
        m1_hist["f1"].append(f1_score(yt, yhat1, average="weighted", zero_division=0))
        m2_hist["f1"].append(f1_score(yt, yhat2, average="weighted", zero_division=0))
        m1_hist["precision"].append(precision_score(yt, yhat1, average="weighted", zero_division=0))
        m2_hist["precision"].append(precision_score(yt, yhat2, average="weighted", zero_division=0))
        m1_hist["recall"].append(recall_score(yt, yhat1, average="weighted", zero_division=0))
        m2_hist["recall"].append(recall_score(yt, yhat2, average="weighted", zero_division=0))

    print(f"\n  Paired t-test results  (method1 - method2)  "
          f"n_iters={args.n_iters}\n")
    print(f"  {'Metric':<12} {'M1 mean':>8} {'M2 mean':>8} "
          f"{'Diff':>8} {'p-value':>10}  sig")
    print(f"  {'-'*58}")

    results = {}
    for m in metrics:
        a = np.array(m1_hist[m])
        b = np.array(m2_hist[m])
        if a.size == 0:
            print(f"  {m.upper():<12}  insufficient data")
            continue
        _, p = ttest_rel(a, b)
        diff = float((a - b).mean())
        results[m] = {
            "method1_mean": float(a.mean()),
            "method2_mean": float(b.mean()),
            "mean_diff":    diff,
            "p":            float(p),
            "star":         _stars(p),
            "n":            int(a.size),
        }
        print(f"  {m.upper():<12} {a.mean():>8.4f} {b.mean():>8.4f} "
              f"{diff:>+8.4f} {p:>10.4g}  {_stars(p)}")

    print(f"\n  Significance: * p<0.10   ** p<0.05   *** p<0.01")

    if args.out_csv:
        rows = []
        for m, r in results.items():
            rows.append({"metric": m, **r})
        pd.DataFrame(rows).to_csv(args.out_csv, index=False)
        print(f"\n  Compare results saved: {args.out_csv}")


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Evaluate CAAC model: test evaluation + method comparison.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    sub = p.add_subparsers(dest="cmd", required=True)

    p_eval = sub.add_parser(
        "eval",
        help="Run test evaluation for one variant."
    )
    p_eval.add_argument("--pt_dir",          required=True,
                        help="Dir of <video_id>.pt final embeddings.")
    p_eval.add_argument("--annotations_csv", required=True,
                        help="Path to complete_set.csv")
    p_eval.add_argument("--test_json",       required=True,
                        help="Path to test_videoids.json")
    p_eval.add_argument("--model_dir",       required=True,
                        help="Dir containing best_params.json (output of train.py)")
    p_eval.add_argument("--attention_csv",   default=None,
                        help="Path to attention_scores.csv (auto used if use_entropy=True in best_params.json)")
    p_eval.add_argument("--n_train",         type=int, default=None,
                        help="Number of training records to use. Must match "
                             "the value used in train.py for this variant. "
                             "If less than default, records are randomly dropped. "
                             "If more than default, records are taken from val.")
    p_eval.add_argument("--pt_dir_extra",       default=None,
                        help="Extra embedding dir to merge with pt_dir "
                        "(e.g. COVID embeddings kept separate). "
                        "These videos are used as test set only.")
    p_eval.add_argument("--attention_csv_extra", default=None,
                        help="attention_scores.csv for extra embeddings "
                        "(e.g. data/Attention_Scores/COVID/BERT_ViT/attention_scores.csv)")
    p_eval.add_argument("--seeds",           default="42")
    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",          default="cpu",
                        choices=["cpu", "cuda"])

    p_cmp = sub.add_parser(
        "compare",
        help="Paired t-test between two methods using saved prediction CSVs."
    )
    p_cmp.add_argument("--preds1",       required=True,
                       help="method1 test_predictions_by_seed.csv")
    p_cmp.add_argument("--preds2",       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)
    p_cmp.add_argument("--out_csv",      default=None,
                       help="Optional path to save compare results as CSV.")

    return p


def main():
    args = build_parser().parse_args()
    if args.cmd == "eval":
        run_eval(args)
    else:
        run_compare(args)


if __name__ == "__main__":
    main()