#!/usr/bin/env python
"""MotionCache threshold scan on the HY-WorldPlay base -- same diagnostic as
mc_scan.py: what the compute fraction hides, i.e. whether the cacheable steps are
making *partial* token decisions or degenerating into whole-step skips.

    CUDA_VISIBLE_DEVICES=4 python hy_mc_scan.py --schedule FxxF \
        --thresholds 0.05,0.1,0.2,0.4,0.7,1.0 --out results/hy_mc_scan_FxxF.json
"""
import argparse, json, os, statistics as st, sys
ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, ROOT)
from hycache import harness as H
H.enter_repo()
import torch
from hycache.measure import make_ctrl

ap = argparse.ArgumentParser()
ap.add_argument("--schedule", default="FxxF")
ap.add_argument("--thresholds", required=True)
ap.add_argument("--frames", type=int, default=61)
ap.add_argument("--num-cases", type=int, default=2)
ap.add_argument("--out", default=None)
args = ap.parse_args()

pipe = H.load_pipeline()
cases = H.load_cases()
jobs = [(cases[i], H.ACTIONS[0]) for i in range(args.num_cases + 1)]   # first = warm-up
forced = set(range(H.NUM_STEPS)) if False else None


def run(method, **kw):
    ctrl = make_ctrl(pipe, method, schedule=args.schedule, **kw)
    ms, recs = [], []
    for i, (case, action) in enumerate(jobs):
        _, _, t = H.generate(pipe, case, action, controller=ctrl, num_frames=args.frames, decode=False)
        if i == 0:
            continue
        ms.append(t["denoise_dit_ms"])
        recs += [r for r in ctrl.records if r["step"] not in ctrl.forced]
    return min(ms), recs, ctrl.summary()


base_ms, _, _ = run("none")
print(f"baseline FFFF min denoise {base_ms:.0f} ms ({args.frames} frames, schedule {args.schedule})", flush=True)
print(f"{'thresh':>8}{'Full':>8}{'compute_x':>10}{'measured_x':>11}{'active%':>9}{'empty%':>8}{'full%':>7}{'sel_mean%':>10}{'sel_p90%':>9}")
rows = []
for t in [float(v) for v in args.thresholds.split(",")]:
    ms, recs, summ = run("motioncache", thresh=t)
    fr = [r["compute_fraction"] for r in recs]
    act = [f for f in fr if 1e-9 < f < 1 - 1e-9]
    q = lambda v, p: sorted(v)[min(int(p * (len(v) - 1)), len(v) - 1)] if v else float("nan")
    row = {"thresh": t, "compute_equivalent": summ["compute_equivalent_forwards"],
           "compute_speedup": summ["flops_speedup_estimate"], "measured_speedup": base_ms / ms,
           "active_ratio": len(act) / len(fr), "empty_ratio": sum(f <= 1e-9 for f in fr) / len(fr),
           "full_ratio": sum(f >= 1 - 1e-9 for f in fr) / len(fr),
           "sel_mean": st.fmean(act) if act else 0.0, "sel_p90": q(act, .9), "min_denoise_ms": ms}
    rows.append(row)
    print(f"{t:>8.4g}{row['compute_equivalent']:>8.2f}{row['compute_speedup']:>10.3f}{row['measured_speedup']:>11.3f}"
          f"{100*row['active_ratio']:>9.1f}{100*row['empty_ratio']:>8.1f}{100*row['full_ratio']:>7.1f}"
          f"{100*row['sel_mean']:>10.1f}{100*row['sel_p90']:>9.1f}", flush=True)
if args.out:
    out = args.out if os.path.isabs(args.out) else os.path.join(ROOT, args.out)
    os.makedirs(os.path.dirname(out), exist_ok=True)
    json.dump({"base": "hy_worldplay", "schedule": args.schedule, "frames": args.frames,
               "baseline_ms": base_ms, "rows": rows}, open(out, "w"), indent=2)
    print("wrote", out)