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

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

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.metrics import roc_auc_score, f1_score
from sklearn.model_selection import train_test_split


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)



VARIANT_CONFIG = {
    "co_attn": {"use_entropy": True,  "default_alpha": 0.5, "default_beta": 1.0},
    "text":    {"use_entropy": False, "default_alpha": 0.0, "default_beta": 0.0},
    "frames":  {"use_entropy": False, "default_alpha": 0.0, "default_beta": 0.0},
    "t2f":     {"use_entropy": False, "default_alpha": 0.0, "default_beta": 0.0},
    "f2t":     {"use_entropy": False, "default_alpha": 0.0, "default_beta": 0.0},
}


def resolve_variant(variant: str, alpha=None, beta=None):
    v = variant.strip().lower()
    if v not in VARIANT_CONFIG:
        raise ValueError(f"Unknown variant '{variant}'. Choose one of: {list(VARIANT_CONFIG.keys())}")

    cfg = VARIANT_CONFIG[v]
    use_entropy = cfg["use_entropy"]

    if use_entropy:
        alpha = cfg["default_alpha"] if alpha is None else float(alpha)
        beta  = cfg["default_beta"]  if beta  is None else float(beta)
    else:
        alpha, beta = 0.0, 0.0

    return use_entropy, alpha, beta


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)


def get_pos_weight(labels_tensor: torch.Tensor) -> torch.Tensor:
    labels = labels_tensor.view(-1)
    pos = (labels == 1).sum().item()
    neg = (labels == 0).sum().item()
    w = neg / max(1, pos)
    return torch.tensor([w], dtype=torch.float32)


def loss_no_entropy(logits, targets, pos_weight):
    return nn.BCEWithLogitsLoss(pos_weight=pos_weight)(logits, targets)


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



def load_embeddings(pt_dir: str) -> Tuple[List[str], torch.Tensor]:
    videoids, tensors = [], []
    for fn in sorted(os.listdir(pt_dir)):  
        if not fn.endswith(".pt"):
            continue
        videoids.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}")

    X = torch.cat(tensors, dim=0)
    return videoids, X


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


def align_labels(videoids, X, df, use_entropy=False):
    df = df.set_index("video_id")
    keep_ids, keep_rows = [], []
    y, t2f, f2t = [], [], []

    for i, vid in enumerate(videoids):
        if vid not in df.index:
            continue
        keep_ids.append(vid)
        keep_rows.append(i)
        y.append(int(df.loc[vid, "Actionable"]))
        if use_entropy:
            t2f.append(float(df.loc[vid, "entropy_t2f"]))
            f2t.append(float(df.loc[vid, "entropy_f2t"]))

    if not keep_ids:
        raise RuntimeError("No overlap between embeddings and metadata video_id.")

    X2 = X[keep_rows]
    y2 = torch.tensor(y, dtype=torch.float32)

    if use_entropy:
        return keep_ids, X2, y2, torch.tensor(t2f, dtype=torch.float32), torch.tensor(f2t, dtype=torch.float32)
    return keep_ids, X2, y2, None, None


def train_eval_once(
    X_train, y_train, X_val, y_val,
    t2f_train=None, f2t_train=None, t2f_val=None, f2t_val=None,
    use_entropy=False, alpha=0.0, beta=0.0,
    lr=1e-4, wd=0.0, device="cpu",
    epochs=50, patience=5
):
    model = BinaryClassifier(input_dim=X_train.shape[1]).to(device)
    pos_weight = get_pos_weight(y_train).to(device)
    opt = optim.Adam(model.parameters(), lr=lr, weight_decay=wd)

    best_val = float("inf")
    best_state = None
    bad = 0

    for _ in range(epochs):
        model.train()
        opt.zero_grad()
        logits = model(X_train.to(device)).squeeze()

        if use_entropy:
            loss = loss_with_entropy(logits, y_train.to(device), t2f_train.to(device), f2t_train.to(device),
                                     pos_weight, alpha, beta)
        else:
            loss = loss_no_entropy(logits, y_train.to(device), pos_weight)

        loss.backward()
        opt.step()

        model.eval()
        with torch.no_grad():
            v_logits = model(X_val.to(device)).squeeze()
            if use_entropy:
                v_loss = loss_with_entropy(v_logits, y_val.to(device), t2f_val.to(device), f2t_val.to(device),
                                           pos_weight, alpha, beta).item()
            else:
                v_loss = loss_no_entropy(v_logits, y_val.to(device), pos_weight).item()

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

    if best_state is None:
        raise RuntimeError("Training did not produce a best_state (unexpected).")

    model.load_state_dict(best_state)
    model.eval()
    with torch.no_grad():
        v_scores = model(X_val.to(device)).squeeze().cpu().numpy()

    y_true = y_val.cpu().numpy().astype(int)
    val_auc = roc_auc_score(y_true, v_scores)
    val_pred = (v_scores >= 0.5).astype(int)
    val_f1 = f1_score(y_true, val_pred, average="weighted", zero_division=0)
    return val_auc, val_f1


def parse_floats_csv(s: str) -> List[float]:
    return [float(x.strip()) for x in s.split(",") if x.strip()]


def grid_search_lr_wd(
    variant: str,
    pt_dir: str,
    complete_csv: str,
    test_json: str,
    out_csv: str,
    attn_csv: Optional[str],
    alpha: Optional[float],
    beta: Optional[float],
    lrs: List[float],
    wds: List[float],
    seed: int,
    device: str,
    epochs: int,
    patience: int,
):
    set_seed(seed)
    use_entropy, alpha, beta = resolve_variant(variant, alpha, beta)

    if use_entropy and not attn_csv:
        raise ValueError("variant=Co_Attn requires --attn_csv (attention_scores.csv).")

    # load and align
    vids, X = load_embeddings(pt_dir)
    df = load_metadata(complete_csv, attn_csv if use_entropy else None)
    vids, X, y, t2f, f2t = align_labels(vids, X, df, use_entropy=use_entropy)

    # exclude fixed test
    with open(test_json, "r") as f:
        test_ids = set(json.load(f))
    pool_idx = [i for i, v in enumerate(vids) if v not in test_ids]
    if len(pool_idx) < 10:
        raise RuntimeError("Train/val pool too small after excluding test IDs. Check test_json and pt_dir.")

    X_pool = X[pool_idx]
    y_pool = y[pool_idx]
    if use_entropy:
        t2f_pool = t2f[pool_idx]
        f2t_pool = f2t[pool_idx]
    else:
        t2f_pool = f2t_pool = None

    # single stratified split for grid search
    idxs = np.arange(len(pool_idx))
    tr_idx, va_idx = train_test_split(
        idxs, test_size=0.25, random_state=seed, stratify=y_pool.numpy().astype(int)
    )

    X_train, y_train = X_pool[tr_idx], y_pool[tr_idx]
    X_val, y_val = X_pool[va_idx], y_pool[va_idx]

    if use_entropy:
        t2f_train, f2t_train = t2f_pool[tr_idx], f2t_pool[tr_idx]
        t2f_val, f2t_val = t2f_pool[va_idx], f2t_pool[va_idx]
    else:
        t2f_train = f2t_train = t2f_val = f2t_val = None

    rows = []
    for lr in lrs:
        for wd in wds:
            val_auc, val_f1 = train_eval_once(
                X_train, y_train, X_val, y_val,
                t2f_train, f2t_train, t2f_val, f2t_val,
                use_entropy=use_entropy, alpha=alpha, beta=beta,
                lr=lr, wd=wd, device=device,
                epochs=epochs, patience=patience
            )
            rows.append({"lr": lr, "wd": wd, "val_auc": val_auc, "val_f1": val_f1})
            print(f"lr={lr:g} wd={wd:g} -> val_auc={val_auc:.4f}, val_f1={val_f1:.4f}")

    dfres = pd.DataFrame(rows)
    os.makedirs(os.path.dirname(out_csv) or ".", exist_ok=True)
    dfres.to_csv(out_csv, index=False)

    best = dfres.sort_values(by=["val_f1", "val_auc"], ascending=False).iloc[0]
    print("\nBEST:")
    print(best.to_string())

    return best, dfres


def main():
    parser = argparse.ArgumentParser(description="Grid search lr/wd for a given embedding variant.")
    parser.add_argument("--variant", required=True, choices=list(VARIANT_CONFIG.keys()))
    parser.add_argument("--pt_dir", type=str, required=True, help="Directory of <video_id>.pt embeddings.")
    parser.add_argument("--complete_csv", type=str, required=True, help="Path to complete_set.csv")
    parser.add_argument("--test_json", type=str, required=True, help="Path to test_videoids.json")
    parser.add_argument("--out_csv", type=str, required=True, help="Where to write grid results CSV")

    # Only required for co_attn
    parser.add_argument("--attn_csv", type=str, default=None, help="Path to attention_scores.csv (required for co_attn)")
    parser.add_argument("--alpha", type=float, default=None, help="Entropy penalty alpha (co_attn only)")
    parser.add_argument("--beta", type=float, default=None, help="Entropy penalty beta (co_attn only)")

    parser.add_argument("--lrs", type=str, default="1e-4,1e-3,1e-2,1e-1", help="Comma-separated lrs")
    parser.add_argument("--wds", type=str, default="0,1e-4,1e-3,1e-2,1e-1", help="Comma-separated wds")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", type=str, default="cpu", choices=["cpu", "cuda"])
    parser.add_argument("--epochs", type=int, default=50)
    parser.add_argument("--patience", type=int, default=5)

    args = parser.parse_args()

    lrs = parse_floats_csv(args.lrs)
    wds = parse_floats_csv(args.wds)

    grid_search_lr_wd(
        variant=args.variant,
        pt_dir=args.pt_dir,
        complete_csv=args.complete_csv,
        test_json=args.test_json,
        out_csv=args.out_csv,
        attn_csv=args.attn_csv,
        alpha=args.alpha,
        beta=args.beta,
        lrs=lrs,
        wds=wds,
        seed=args.seed,
        device=args.device,
        epochs=args.epochs,
        patience=args.patience,
    )


if __name__ == "__main__":
    main()