#!/usr/bin/env python
"""DAUS on benchmark/v1/val per-tick PT files, FILTERED to v5_sft_val_v6.jsonl.

Drops the 71 v6-discarded ticks before aggregation. Categories and TTAs
come from the original PT files. Joins on (video_id, frame_indices[-1]).
"""
from __future__ import annotations
import argparse, json
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import torch

ROOT = Path(__file__).resolve().parents[1]
PT_DIR = ROOT / "eval_results/benchmark_v1_val/per_tick"
V6_JSONL = ROOT / "data/cot_corpus_v3/v5_sft_val_v6.jsonl"
OUT_DIR = ROOT / "eval_results/benchmark_v1_val_v6"


@dataclass
class DausConfig:
    alpha: float = 0.60; w_R: float = 0.65; w_L: float = 0.35
    w_n: float = 1/3;    w_p: float = 1/3;  w_d: float = 1/3
    tau_star: float = 1.5; tau_starstar: float = 3.0
    L_alert: float = 5.0;  u_floor: float = 0.5
    t_recover: float = 5.0; AEPH_cap: float = 30.0


def u_lead_star(tau_lead, cfg):
    if tau_lead <= 0: return 0.0
    if tau_lead > cfg.L_alert: return 0.0
    if tau_lead <= cfg.tau_star: return tau_lead / cfg.tau_star
    if tau_lead <= cfg.tau_starstar: return 1.0
    span = cfg.L_alert - cfg.tau_starstar
    frac = (tau_lead - cfg.tau_starstar) / span
    return 1.0 - frac * (1.0 - cfg.u_floor)


def per_clip(scores, tta, category, tau, cfg):
    if category in ("safe_neg", "negative"):
        F_neg = float(np.any(scores > tau))
        return {"R_alert": np.nan, "U_lead_star": np.nan,
                "F_neg": F_neg, "F_post": np.nan,
                "post_ticks_available": False}
    pre_mask = (tta > 0) & (tta <= cfg.L_alert)
    post_mask = (tta <= 0) & (tta > -cfg.t_recover)
    pre_fires = (scores > tau) & pre_mask
    R_alert = float(pre_fires.any())
    if pre_fires.any():
        first_fire_tta = float(tta[pre_fires].max())
        Ul = u_lead_star(first_fire_tta, cfg)
    else:
        Ul = 0.0
    has_post = bool(post_mask.any())
    F_post = float(((scores > tau) & post_mask).any()) if has_post else np.nan
    return {"R_alert": R_alert, "U_lead_star": Ul,
            "F_neg": np.nan, "F_post": F_post,
            "post_ticks_available": has_post}


def build_v6_keep(jsonl_path):
    keep = set()
    for ln in open(jsonl_path):
        d = json.loads(ln)
        keep.add((d["video_id"], int(d["frame_indices"][-1])))
    return keep


def load_method(pt_path, v6_keep):
    d = torch.load(pt_path, weights_only=False, map_location="cpu")
    if "scores_binary" not in d or "tta_raw" not in d:
        return None, 0, 0
    ids = list(d["ids"])
    cat = list(d["category"])
    src = list(d["source"])
    tta = d["tta_raw"].numpy().astype(np.float64)
    sc = d["scores_binary"].numpy().astype(np.float64)
    frame_last = d["frame_indices"][:, -1].numpy().astype(np.int64)
    tick_idx = d["tick_idx"].numpy().astype(np.int64)
    N = len(ids)
    keep_mask = np.array([(ids[i], int(frame_last[i])) in v6_keep
                          for i in range(N)], dtype=bool)
    n_orig, n_kept = N, int(keep_mask.sum())
    if n_kept == 0:
        return None, n_orig, n_kept
    return {
        "ids":      [ids[i] for i in range(N) if keep_mask[i]],
        "category": [cat[i] for i in range(N) if keep_mask[i]],
        "source":   [src[i] for i in range(N) if keep_mask[i]],
        "tta":      tta[keep_mask],
        "scores":   sc[keep_mask],
        "tick_idx": tick_idx[keep_mask],
    }, n_orig, n_kept


def regroup(m):
    groups = defaultdict(list)
    for i, vid in enumerate(m["ids"]):
        groups[vid].append(i)
    clips = []
    for vid, idxs in groups.items():
        order = sorted(idxs, key=lambda j: int(m["tick_idx"][j]))
        cat = m["category"][order[0]]; src = m["source"][order[0]]
        tta = np.array([m["tta"][j] for j in order])
        sc  = np.array([m["scores"][j] for j in order])
        mask = np.isfinite(sc)
        tta, sc = tta[mask], sc[mask]
        if len(sc) == 0: continue
        clips.append({"vid": vid, "category": cat, "source": src,
                       "tta": tta, "scores": sc})
    return clips


def calibrate_tau(clips, q, cfg):
    pos_max = []
    for c in clips:
        if c["category"] not in ("ego_positive", "positive"): continue
        win = (c["tta"] > 0) & (c["tta"] <= cfg.L_alert)
        if not win.any(): continue
        pos_max.append(float(c["scores"][win].max()))
    if not pos_max: return 0.5
    pos_max = np.sort(np.array(pos_max))
    qi = int(np.floor((1 - q) * len(pos_max)))
    qi = min(max(qi, 0), len(pos_max) - 1)
    return float(pos_max[qi])


def aggregate(clips, tau, cfg):
    R_l, U_l, Fn_l, Fp_l = [], [], [], []
    n_pos = n_neg = n_post = 0
    for c in clips:
        m = per_clip(c["scores"], c["tta"], c["category"], tau, cfg)
        if c["category"] in ("ego_positive", "positive"):
            n_pos += 1
            R_l.append(m["R_alert"]); U_l.append(m["U_lead_star"])
            if m["post_ticks_available"]:
                Fp_l.append(m["F_post"]); n_post += 1
        elif c["category"] in ("safe_neg", "negative"):
            n_neg += 1
            Fn_l.append(m["F_neg"])

    def _mean(xs):
        a = np.array(xs, float); a = a[~np.isnan(a)]
        return float(a.mean()) if a.size else float("nan")

    R = _mean(R_l); U = _mean(U_l); Fn = _mean(Fn_l); Fp = _mean(Fp_l)
    nu = {"F_neg": Fn, "F_post": Fp, "F_drive": float("nan")}
    weights = {"F_neg": cfg.w_n, "F_post": cfg.w_p, "F_drive": cfg.w_d}
    avail = {k: v for k, v in nu.items() if not np.isnan(v)}
    if avail:
        w_total = sum(weights[k] for k in avail)
        U_minus = sum((weights[k] / w_total) * avail[k] for k in avail)
    else:
        U_minus = float("nan")
    U_plus = cfg.w_R * (R if not np.isnan(R) else 0.0) + \
             cfg.w_L * (U if not np.isnan(U) else 0.0)
    DAUS = cfg.alpha * U_plus + (1 - cfg.alpha) * (1 - U_minus
                                                    if not np.isnan(U_minus) else 1.0)
    return {"n_pos": n_pos, "n_neg": n_neg, "n_post_clips": n_post,
            "R_alert": R, "U_lead_star": U, "F_neg": Fn, "F_post": Fp,
            "U_plus": U_plus, "U_minus": U_minus, "DAUS": DAUS, "tau": tau}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pt_dir", type=Path, default=PT_DIR)
    ap.add_argument("--hit_rate", type=float, default=0.30)
    ap.add_argument("--out_json", type=Path, default=OUT_DIR / "daus_v6.json")
    ap.add_argument("--out_md", type=Path, default=OUT_DIR / "daus_v6.md")
    args = ap.parse_args()

    cfg = DausConfig()
    v6_keep = build_v6_keep(V6_JSONL)
    print(f"[v6] keep {len(v6_keep):,} (vid, last_frame) keys")

    pts = sorted(args.pt_dir.glob("*.pt"))
    print(f"[load] {len(pts)} PT files")
    rows = {}
    for p in pts:
        m, n_orig, n_kept = load_method(p, v6_keep)
        if m is None:
            print(f"  [skip] {p.name} (orig={n_orig}, kept={n_kept})")
            continue
        clips = regroup(m)
        if not clips:
            print(f"  [skip] {p.name}: no clips after regroup")
            continue
        tau = calibrate_tau(clips, args.hit_rate, cfg)
        r = aggregate(clips, tau, cfg)
        r["n_orig_ticks"] = n_orig; r["n_kept_ticks"] = n_kept
        rows[p.stem] = r
        print(f"  {p.stem:35s} kept {n_kept:5d}/{n_orig:5d}  "
              f"n+={r['n_pos']:4d} n-={r['n_neg']:4d}  tau={tau:.3f}  "
              f"R={r['R_alert']:.3f}  U*={r['U_lead_star']:.3f}  "
              f"DAUS={r['DAUS']:.4f}")

    payload = {"hit_rate": args.hit_rate, "cfg": cfg.__dict__,
                "v6_keep": len(v6_keep), "results": rows}
    args.out_json.parent.mkdir(parents=True, exist_ok=True)
    args.out_json.write_text(json.dumps(payload, indent=2,
                                         default=lambda x: None if (isinstance(x, float) and not np.isfinite(x)) else x))
    print(f"\n[save] {args.out_json}")

    # Markdown
    def f(v, p=3):
        if v is None or (isinstance(v, float) and not np.isfinite(v)): return "—"
        return f"{v:.{p}f}"
    is_vla = lambda n: "vlalert" in n.lower()
    sorted_rows = sorted(rows.items(),
                         key=lambda x: -(x[1]['DAUS']
                                          if np.isfinite(x[1]['DAUS']) else -1))
    lines = ["# DAUS — v6 labels (v5_sft_val_v6.jsonl)",
             "",
             f"Hit-rate calibration q = {args.hit_rate:.2f}.  "
             f"Config B' (alpha={cfg.alpha}, w_R={cfg.w_R}, w_L={cfg.w_L}).",
             f"v6 keep: {len(v6_keep):,} ticks (71 ticks discarded from v5).",
             "",
             "| Rank | Method | kept | n+ | n- | tau | R_alert↑ | U_lead*↑ | F_neg↓ | F_post↓ | U+↑ | U-↓ | DAUS↑ |",
             "| ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |"]
    for i, (name, r) in enumerate(sorted_rows, 1):
        marker = "**" if is_vla(name) and i == min(
            (j for j, (n, _) in enumerate(sorted_rows, 1) if is_vla(n)), default=0) else ""
        lines.append("| " + " | ".join([
            str(i), f"{marker}{name}{marker}", str(r["n_kept_ticks"]),
            str(r["n_pos"]), str(r["n_neg"]), f(r["tau"]),
            f(r["R_alert"]), f(r["U_lead_star"]),
            f(r["F_neg"]), f(r["F_post"]),
            f(r["U_plus"]), f(r["U_minus"]),
            f(r["DAUS"], 4),
        ]) + " |")

    # Highlight VLAlert winner
    vla_rows = [(n, r) for n, r in sorted_rows if is_vla(n)]
    if vla_rows:
        best_n, best_r = vla_rows[0]
        lines += ["", "## Best VLAlert variant",
                  f"**{best_n}** → DAUS = **{best_r['DAUS']:.4f}** "
                  f"(R_alert={best_r['R_alert']:.3f}, U_lead*={best_r['U_lead_star']:.3f}, "
                  f"F_neg={best_r['F_neg']:.3f}, F_post={best_r['F_post']:.3f}, "
                  f"tau={best_r['tau']:.3f})"]

    args.out_md.write_text("\n".join(lines) + "\n")
    print(f"[save] {args.out_md}")
    if vla_rows:
        print(f"\n=== BEST VLAlert (v6) ===  {vla_rows[0][0]}  DAUS={vla_rows[0][1]['DAUS']:.4f}")


if __name__ == "__main__":
    main()