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

import os
import json
import random
import argparse
from itertools import product
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


LRS    = [1e-4, 1e-3, 1e-2, 1e-1]
WDS    = [0.0, 1e-4, 1e-3, 1e-2, 1e-1]
ALPHAS = [0.0, 0.5, 1.0]
BETAS  = [0.0, 0.5, 1.0]

VARIANT_CONFIG = {
    "co_attn": {"use_entropy": True},
    "text":    {"use_entropy": False},
    "frames":  {"use_entropy": False},
    "t2f":     {"use_entropy": False},
    "f2t":     {"use_entropy": False},
    "mean":    {"use_entropy": False},
    "concat":  {"use_entropy": False},
}


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(labels: torch.Tensor) -> torch.Tensor:
    labels = labels.view(-1)
    pos    = (labels == 1).sum().item()
    neg    = (labels == 0).sum().item()
    return torch.tensor([neg / max(1, 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], 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, torch.cat(tensors, dim=0)


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 align_data(video_ids:   List[str],
               X:           torch.Tensor,
               df:          pd.DataFrame,
               use_entropy: bool):
    df       = df.set_index("video_id")
    keep_ids = []
    keep_row = []
    y_list, t2f_list, f2t_list = [], [], []

    for i, vid in enumerate(video_ids):
        if vid not in df.index:
            continue
        keep_ids.append(vid)
        keep_row.append(i)
        y_list.append(int(df.loc[vid, "Actionable"]))
        if use_entropy:
            t2f_list.append(float(df.loc[vid, "entropy_t2f"]))
            f2t_list.append(float(df.loc[vid, "entropy_f2t"]))

    if not keep_ids:
        raise RuntimeError(
            "No overlap between embeddings and metadata. "
            "Check video_id values in pt_dir vs annotations/attention CSVs."
        )

    X2  = X[keep_row]
    y2  = torch.tensor(y_list, dtype=torch.float32)

    if use_entropy:
        t2f = torch.tensor(t2f_list, dtype=torch.float32)
        f2t = torch.tensor(f2t_list, 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, X2, y2, t2f, f2t


def split_train_val(video_ids, X, y, t2f, f2t,
                    test_ids: set, val_size: float = 0.2,
                    seed: int = 42, n_train: int = None):

    pool_mask = [i for i, v in enumerate(video_ids) if v not in test_ids]
    pool_ids  = [video_ids[i] for i in pool_mask]
    X_pool    = X[pool_mask]
    y_pool    = y[pool_mask]
    t2f_pool  = t2f[pool_mask]
    f2t_pool  = f2t[pool_mask]

    idxs = np.arange(len(pool_ids))
    tr_idx, va_idx = train_test_split(
        idxs,
        test_size   = val_size,
        random_state= seed,
        stratify    = y_pool.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)
            print(f"  [n_train] Increased training: "
                  f"{default_n_train} -> {n_train} "
                  f"(val reduced: {len(va_idx) + diff} -> {len(va_idx)})")

    train_ds = VideoDataset(
        [pool_ids[i] for i in tr_idx],
        X_pool[tr_idx], y_pool[tr_idx],
        t2f_pool[tr_idx], f2t_pool[tr_idx],
    )
    val_ds = VideoDataset(
        [pool_ids[i] for i in va_idx],
        X_pool[va_idx], y_pool[va_idx],
        t2f_pool[va_idx], f2t_pool[va_idx],
    )
    return train_ds, val_ds


def train_one_config(
    train_ds:    VideoDataset,
    val_ds:      VideoDataset,
    input_dim:   int,
    lr:          float,
    wd:          float,
    alpha:       float,
    beta:        float,
    use_entropy: bool,
    epochs:      int,
    patience:    int,
    batch_size:  int,
    device:      str,
    seed:        int = 42,
) -> Tuple[float, float]:
    set_seed(seed)

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

    model     = BinaryClassifier(input_dim=input_dim).to(device)
    pw        = get_pos_weight(train_ds.labels).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):
        # Train
        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,
                                  pw, alpha, beta, use_entropy)
            loss.backward()
            optimizer.step()

        # Validate
        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,
                                         pw, 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 None:
        raise RuntimeError("Training produced no best state: check data.")

    model.load_state_dict(best_state)
    model.eval()
    y_true_all, y_score_all = [], []
    with torch.no_grad():
        for _, tokens, labels, _, _ in val_loader:
            logits = model(tokens.to(device)).squeeze()
            y_true_all.extend(labels.numpy().astype(int))
            y_score_all.extend(logits.cpu().numpy())

    y_true  = np.asarray(y_true_all).astype(int)
    y_score = np.asarray(y_score_all)
    val_auc = roc_auc_score(y_true, y_score)
    val_f1  = f1_score(y_true, (y_score >= 0.5).astype(int),
                       average="weighted", zero_division=0)

    return val_auc, val_f1

def run_grid_search(args):
    print("\n" + "=" * 60)
    print("  TRAIN.PY  Grid Search")
    print(f"  Variant : {args.variant}")
    print(f"  PT dir  : {args.pt_dir}")
    print(f"  Out dir : {args.out_dir}")
    print("=" * 60)

    os.makedirs(args.out_dir, exist_ok=True)
    set_seed(args.seed)

    use_entropy = VARIANT_CONFIG[args.variant]["use_entropy"]

    if use_entropy and not args.attention_csv:
        raise ValueError(
            f"variant='co_attn' requires --attention_csv (attention_scores.csv)."
        )

    video_ids, X = load_embeddings(args.pt_dir)
    df           = load_and_merge(args.annotations_csv,
                                  args.attention_csv if use_entropy else None)
    video_ids, X, y, t2f, f2t = align_data(video_ids, X, df, use_entropy)

    with open(args.test_json, "r", encoding="utf-8") as f:
        test_ids = set(json.load(f))

    print(f"\n  Total videos aligned : {len(video_ids)}")
    print(f"  Test IDs excluded    : {len([v for v in video_ids if v in test_ids])}")
    print(f"  Train/val pool       : {len([v for v in video_ids if v not in test_ids])}")
    print(f"  use_entropy          : {use_entropy}")

    train_ds, val_ds = split_train_val(
        video_ids, X, y, t2f, f2t,
        test_ids=test_ids, val_size=0.2, seed=args.seed, n_train   = args.n_train
    )
    input_dim = X.shape[1]

    if use_entropy:
        grid = list(product(LRS, WDS, ALPHAS, BETAS))
        print(f"  Grid size: {len(grid)} combinations "
              f"(lr x wd x alpha x beta)\n")
    else:
        grid = [(lr, wd, 0.0, 0.0) for lr, wd in product(LRS, WDS)]
        print(f"  Grid size: {len(grid)} combinations "
              f"(lr x wd (alpha/beta) fixed at 0 for this variant)\n")

    rows = []
    for combo_i, (lr, wd, alpha, beta) in enumerate(grid, 1):
        val_auc, val_f1 = train_one_config(
            train_ds, val_ds,
            input_dim   = input_dim,
            lr          = lr,
            wd          = wd,
            alpha       = alpha,
            beta        = beta,
            use_entropy = use_entropy,
            epochs      = args.epochs,
            patience    = args.patience,
            batch_size  = args.batch_size,
            device      = args.device,
            seed        = args.seed,
        )

        if use_entropy:
            row = {"lr": lr, "wd": wd, "alpha": alpha, "beta": beta,
                   "val_auc": val_auc, "val_f1": val_f1}
            print(f"  [{combo_i:3d}/{len(grid)}] "
                  f"lr={lr:g} wd={wd:g} alpha={alpha:g} beta={beta:g} "
                  f"-> val_auc={val_auc:.4f}  val_f1={val_f1:.4f}")
        else:
            row = {"lr": lr, "wd": wd,
                   "val_auc": val_auc, "val_f1": val_f1}
            print(f"  [{combo_i:3d}/{len(grid)}] "
                  f"lr={lr:g} wd={wd:g} "
                  f"-> val_auc={val_auc:.4f}  val_f1={val_f1:.4f}")

        rows.append(row)

    df_res   = pd.DataFrame(rows)
    grid_csv = os.path.join(args.out_dir, "grid_search_results.csv")
    df_res.to_csv(grid_csv, index=False)
    print(f"\n  Grid results saved : {grid_csv}")

    best_row = df_res.sort_values(
        by=["val_auc", "val_f1"], ascending=False
    ).iloc[0]

    best_params = {
        "variant":     args.variant,
        "use_entropy": use_entropy,
        "lr":          float(best_row["lr"]),
        "wd":          float(best_row["wd"]),
        "val_auc":      float(best_row["val_auc"]),
        "val_f1":     float(best_row["val_f1"]),
    }
    if use_entropy:
        best_params["alpha"] = float(best_row["alpha"])
        best_params["beta"]  = float(best_row["beta"])

    params_path = os.path.join(args.out_dir, "best_params.json")
    with open(params_path, "w", encoding="utf-8") as f:
        json.dump(best_params, f, indent=2)

    print(f"  Best params saved  : {params_path}")
    print("\n  BEST HYPERPARAMETERS:")
    for k, v in best_params.items():
        print(f"    {k:15s}: {v}")
    print("\n  Grid search complete.")
    print("Now run:  python scripts/evaluate.py ...")


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Grid search hyperparameters for CAAC model.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument("--variant",         required=True,
                   choices=list(VARIANT_CONFIG.keys()),
                   help="Embedding variant to search over.")
    p.add_argument("--pt_dir",          required=True,
                   help="Dir of <video_id>.pt final embeddings.")
    p.add_argument("--annotations_csv", required=True,
                   help="Path to complete_set.csv")
    p.add_argument("--test_json",       required=True,
                   help="Path to test_videoids.json (excluded from search).")
    p.add_argument("--out_dir",         required=True,
                   help="Output directory for grid_search_results.csv + best_params.json.")
    p.add_argument("--attention_csv",   default=None,
                   help="Path to attention_scores.csv (required only for co_attn).")
    p.add_argument("--batch_size",      type=int, default=32)
    p.add_argument("--epochs",          type=int, default=50)
    p.add_argument("--patience",        type=int, default=5)
    p.add_argument("--seed",            type=int, default=42,
                   help="Seed for train/val split during grid search (default: 42).")
    p.add_argument("--n_train", type=int, default=None,
                   help="Number of training records to use. Default uses all "
                   "available training records after val split. "
                   "If less than default, records are randomly dropped. "
                   "If more than default, records are taken from val.")
    p.add_argument("--device",          default="cpu",
                   choices=["cpu", "cuda"])
    return p

def main():
    args         = build_parser().parse_args()
    args.variant = args.variant.strip().lower()
    if args.variant not in VARIANT_CONFIG:
        raise ValueError(
            f"Unknown variant '{args.variant}'. "
            f"Choose from: {list(VARIANT_CONFIG.keys())}"
        )
    run_grid_search(args)


if __name__ == "__main__":
    main()