#!/usr/bin/env python3
"""Frame-ablation causal test on VideoHallucer temporal_absolute.

Follows up scripts/frame_attention_probe.py. Attention showed the model "looks"
at the key region equally whether right or wrong. This asks the causal question:
does masking the key frames actually change the answer (i.e. did the model *use*
them), vs masking a matched set of non-key frames?

Per question:
  decode 32 frames -> processor -> base forward -> p_yes_base, answer_base
  then a batch of masked variants (black-out frames in a temporal bin/region):
    - 16 single-bin masks      -> per-bin causal importance profile (Δp_yes)
    - key-region mask           -> mask all bins in the prompt-referenced third
    - random-region mask        -> mask the same #bins drawn from NON-key bins
  Metrics: answer flip rate (key vs random) and causal lift = flip_key - flip_random.

p_yes = P(yes)/(P(yes)+P(no)) from last-position logits; answer = yes if p_yes>.5.
Masking = set those frames to black (matches the design's black_frames perturbation),
frame count unchanged so vision-token count is identical -> variants batch cleanly.
"""
from __future__ import annotations
import argparse, json, os, sys, time
from pathlib import Path

os.environ.setdefault("HF_HOME", "/mnt/local-fast/opd_zt/hf_cache")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
import numpy as np, torch

ROOT = Path("/mnt/local-fast/opd_zt")
DEFAULT_MODEL = str(ROOT / "hf_cache/hub/models--Qwen--Qwen2.5-VL-7B-Instruct/snapshots/"
                    "cc594898137f460bfe9f0759e9844b3ce807cfb5")
VH_TEMPORAL = ROOT / "data/benchmarks/VideoHallucer/temporal"
NFRAMES = 32
VIDEO_MAX_PIXELS = 128 * 28 * 28
VIDEO_MIN_PIXELS = 16 * 28 * 28
SUFFIX = "\nAnswer the question using 'yes' or 'no'."


def load_temporal_absolute() -> list[dict]:
    data = json.loads((VH_TEMPORAL / "temporal.json").read_text())
    vdir = VH_TEMPORAL / "videos"
    items = []
    for idx, pair in enumerate(data):
        if pair.get("type") != "temporal_absolute":
            continue
        for side in ("basic", "hallucination"):
            q = pair[side]
            items.append({"pair_id": f"temporal/{idx}", "side": side,
                          "video": str(vdir / q["video"]),
                          "question": q["question"], "answer": q["answer"].strip().lower()})
    return items


def decode_frames(path):
    from decord import VideoReader, cpu
    vr = VideoReader(path, ctx=cpu(0), num_threads=1)
    total = len(vr)
    if total < 1:
        return None
    idx = np.linspace(0, total - 1, NFRAMES).round().astype(int).clip(0, total - 1)
    return vr.get_batch(idx.tolist()).asnumpy()  # [32,H,W,C] uint8


def key_bins(question: str, grid_t: int):
    q = question.lower()
    third = max(1, grid_t // 3)
    if "beginning" in q or "start of the video" in q:
        return list(range(0, third))
    if "end" in q:
        return list(range(grid_t - third, grid_t))
    return None


def mask_frames(frames, bins, grid_t):
    """Return a copy with the 2 frames of each temporal bin set to black."""
    f = frames.copy()
    per = max(1, frames.shape[0] // grid_t)  # frames per temporal bin (=2 for 32f/gt16)
    for b in bins:
        lo = b * per
        f[lo: lo + per] = 0
    return f


@torch.no_grad()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=DEFAULT_MODEL)
    ap.add_argument("--outdir", default=str(ROOT / "outputs/frame_ablation"))
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--n_random", type=int, default=3, help="random-control draws to average")
    ap.add_argument("--device", default="cuda:0")
    args = ap.parse_args()
    outdir = Path(args.outdir); outdir.mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(0)

    from PIL import Image
    from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
    proc = AutoProcessor.from_pretrained(args.model, trust_remote_code=True,
                                         max_pixels=VIDEO_MAX_PIXELS, min_pixels=VIDEO_MIN_PIXELS)
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        args.model, torch_dtype=torch.bfloat16, attn_implementation="sdpa",
        trust_remote_code=True).to(args.device).eval()

    tok = proc.tokenizer
    def first_ids(words):
        s = set()
        for w in words:
            ids = tok.encode(w, add_special_tokens=False)
            if ids:
                s.add(ids[0])
        return torch.tensor(sorted(s), device=args.device)
    yes_ids = first_ids(["yes", "Yes", " yes", " Yes", "YES", " YES"])
    no_ids = first_ids(["no", "No", " no", " No", "NO", " NO"])
    print(f"[ids] yes={yes_ids.tolist()} no={no_ids.tolist()}")

    def p_yes_from_logits(last_logits):  # [B, V]
        sm = torch.softmax(last_logits.float(), -1)
        py = sm[:, yes_ids].sum(-1); pn = sm[:, no_ids].sum(-1)
        return (py / (py + pn + 1e-9)).cpu().numpy()

    items = load_temporal_absolute()
    if args.limit:
        items = items[: args.limit]
    print(f"[load] {len(items)} temporal_absolute questions")

    recs = []
    t0 = time.time()
    for qi, it in enumerate(items):
        frames = decode_frames(it["video"])
        if frames is None:
            continue
        # determine grid_t once
        base_pil = [Image.fromarray(f) for f in frames]
        messages = [{"role": "user", "content": [{"type": "video"},
                    {"type": "text", "text": it["question"] + SUFFIX}]}]
        text = proc.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        gthw = proc(text=[text], videos=[base_pil], return_tensors="pt")["video_grid_thw"][0]
        grid_t = int(gthw[0])
        kb = key_bins(it["question"], grid_t)
        if kb is None:
            continue

        # build variants: base, 16 single-bin, key-region, n_random controls
        variants = [("base", [])]
        for b in range(grid_t):
            variants.append((f"bin{b}", [b]))
        variants.append(("key", kb))
        non_key = [b for b in range(grid_t) if b not in kb]
        rand_masks = []
        for r in range(args.n_random):
            sel = sorted(rng.choice(non_key, size=min(len(kb), len(non_key)), replace=False).tolist())
            rand_masks.append(sel)
            variants.append((f"rand{r}", sel))

        vids = [[Image.fromarray(f) for f in mask_frames(frames, bins, grid_t)] for _, bins in variants]
        batch = proc(text=[text] * len(variants), videos=vids, return_tensors="pt")
        batch = {k: (v.to(model.device) if hasattr(v, "to") else v) for k, v in batch.items()}
        out = model(**batch, use_cache=False)
        last = out.logits[:, -1, :]
        py = p_yes_from_logits(last)  # [n_variants]

        names = [n for n, _ in variants]
        pyd = dict(zip(names, py))
        ans = {n: ("yes" if pyd[n] > 0.5 else "no") for n in names}
        base_py, base_ans = pyd["base"], ans["base"]
        correct = base_ans == it["answer"]

        per_bin_drop = [float(base_py - pyd[f"bin{b}"]) for b in range(grid_t)]
        flip_key = int(ans["key"] != base_ans)
        flip_rand = float(np.mean([ans[f"rand{r}"] != base_ans for r in range(args.n_random)]))
        dp_key = float(base_py - pyd["key"])
        dp_rand = float(np.mean([base_py - pyd[f"rand{r}"] for r in range(args.n_random)]))

        recs.append({**{k: it[k] for k in ("pair_id", "side", "question", "answer", "video")},
                     "grid_t": grid_t, "key_bins": kb,
                     "base_p_yes": float(base_py), "base_answer": base_ans, "correct": correct,
                     "flip_key": flip_key, "flip_random": flip_rand,
                     "dp_yes_key": dp_key, "dp_yes_random": dp_rand,
                     "per_bin_p_yes_drop": [round(x, 4) for x in per_bin_drop]})
        if (qi + 1) % 10 == 0:
            print(f"[run] {qi+1}/{len(items)} ({time.time()-t0:.0f}s)", flush=True)
        del out, batch

    with (outdir / "records.jsonl").open("w") as f:
        for r in recs:
            f.write(json.dumps(r) + "\n")

    # ---- aggregate, in the user's framing ----
    def grp_stats(rows):
        if not rows:
            return None
        fk = np.mean([r["flip_key"] for r in rows]); fr = np.mean([r["flip_random"] for r in rows])
        return {"n": len(rows),
                "flip_rate_key": round(float(fk), 3), "flip_rate_random": round(float(fr), 3),
                "causal_lift_flip": round(float(fk - fr), 3),
                "ratio": round(float(fk / fr), 1) if fr > 0 else None,
                "dp_yes_key_mean": round(float(np.mean([r["dp_yes_key"] for r in rows])), 3),
                "dp_yes_random_mean": round(float(np.mean([r["dp_yes_random"] for r in rows])), 3)}
    cor = [r for r in recs if r["correct"]]; wro = [r for r in recs if not r["correct"]]
    summary = {"model": args.model, "n": len(recs),
               "all": grp_stats(recs), "correct_only": grp_stats(cor), "wrong_only": grp_stats(wro)}
    (outdir / "summary.json").write_text(json.dumps(summary, indent=2))
    print(json.dumps(summary, indent=2))

    # ---- figure: causal importance vs attention ----
    try:
        import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
        gt = 16
        sub = [r for r in recs if r["grid_t"] == gt]
        caus = np.array([r["per_bin_p_yes_drop"] for r in sub])  # [n,16]
        caus_imp = np.abs(caus).mean(0)
        # attention profile from the earlier probe, if available
        att = None
        ap_path = ROOT / "outputs/frame_attention/records.jsonl"
        if ap_path.exists():
            ar = [json.loads(l) for l in open(ap_path)]
            aa = np.array([r["per_frame_attention"] for r in ar if r["grid_t"] == gt])
            att = aa.mean(0)
        fig, ax1 = plt.subplots(figsize=(8, 4.4))
        x = np.arange(gt)
        ax1.bar(x, caus_imp, color="#1f77b4", alpha=0.7, label="causal importance |Δp_yes| (ablation)")
        ax1.set_xlabel("temporal bin (0=start .. 15=end)"); ax1.set_ylabel("mean |Δp_yes| when bin masked", color="#1f77b4")
        if att is not None:
            ax2 = ax1.twinx()
            ax2.plot(x, att, "-o", color="#d62728", ms=4, label="attention (frame_attention_probe)")
            ax2.set_ylabel("mean frame-attention", color="#d62728")
        ax1.set_title("Causal frame importance (ablation) vs attention\nVideoHallucer temporal_absolute")
        fig.tight_layout(); fig.savefig(outdir / "fig_causal_vs_attention.png", dpi=130); plt.close(fig)
        print(f"[fig] wrote {outdir/'fig_causal_vs_attention.png'}")
    except Exception as e:
        print("[fig] skipped:", e)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())