#!/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("--method", default="motioncache",
                choices=["teacache", "flowcache", "taylorseer", "motioncache"])
ap.add_argument("--values", required=True)
ap.add_argument("--coefficients", default="results/hy_coeff.json")
ap.add_argument("--frames", type=int, default=61)
ap.add_argument("--num-cases", type=int, default=2)
ap.add_argument("--out", default=None)
ap.add_argument("--first-chunk", default=None, help="schedule for chunk 0, e.g. FFFF")
args = ap.parse_args()

COEFF = None
if args.coefficients and args.method in ("teacache", "flowcache"):
    p = args.coefficients if os.path.isabs(args.coefficients) else os.path.join(ROOT, args.coefficients)
    COEFF = json.load(open(p))["coefficients"]

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, coefficients=COEFF,
                     first_chunk_schedule=args.first_chunk, **kw)
    ms, recs, chunk_ms = [], [], []
    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"])
        per = t["per_step_denoise_ms"]
        n = H.NUM_STEPS
        chunk_ms.append([sum(per[b * n:(b + 1) * n]) for b in range(len(per) // n)])
        # activity stats over the cacheable steps only (chunk 0 excluded when it is forced full)
        recs += [r for r in ctrl.records if r["step"] not in ctrl.forced
                 and not (args.first_chunk and r["block"] == 0)]
    # per-chunk denoise ms of the fastest (least contended) run, for extrapolation
    best = chunk_ms[ms.index(min(ms))]
    return min(ms), recs, ctrl.summary(), best


def extrapolate(chunks, n_target=8):
    """Denoise ms for n_target chunks from a shorter run: chunks beyond the measured
    ones cost what the last measured chunk cost (per-chunk cost is flat once the
    context window is full)."""
    return sum(chunks) + (n_target - len(chunks)) * chunks[-1]


PARAM = {"teacache": "thresh", "flowcache": "thresh", "motioncache": "thresh",
         "taylorseer": "interval"}
pname = PARAM[args.method]

base_ms, _, _, base_chunks = run("none")   # warm-up + reference point
print(f"baseline FFFF min denoise {base_ms:.0f} ms ({args.frames} frames, schedule {args.schedule})", flush=True)
print(f"{pname[:8]:>8}{'Full':>8}{'compute_x':>10}{'measured_x':>11}{'ext8_Full':>10}{'ext8_x':>8}{'active%':>9}{'empty%':>8}{'full%':>7}{'sel_mean%':>10}{'sel_p90%':>9}")
rows = []
for t in [float(v) for v in args.values.split(",")]:
    b_ms, _, _, b_chunks = run("none")
    ms, recs, summ, s_chunks = run(args.method, **{pname: t})
    n_chunks = len(b_chunks)
    # extrapolated 8-chunk (125-frame) figures: chunk 0 (forced full) is measured, the
    # cacheable chunks' per-chunk cost is taken from the last measured chunk
    ext_ms = extrapolate(s_chunks); ext_b = extrapolate(b_chunks)
    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 = {"value": t, "param": pname, "method": args.method, "compute_equivalent": summ["compute_equivalent_forwards"],
           "compute_speedup": summ["flops_speedup_estimate"], "measured_speedup": b_ms / ms, "paired_baseline_ms": b_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,
           "measured_chunks": n_chunks, "ext8_measured_speedup": ext_b / ext_ms,
           "ext8_compute_equivalent": summ["compute_equivalent_forwards"]
               + (8 - n_chunks) * (summ["compute_equivalent_forwards"] - 4.0) / max(n_chunks - 1, 1)}
    rows.append(row)
    print(f"{t:>8.4g}{row['compute_equivalent']:>8.2f}{row['compute_speedup']:>10.3f}{row['measured_speedup']:>11.3f}"
          f"{row['ext8_compute_equivalent']:>10.2f}{row['ext8_measured_speedup']:>8.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", "method": args.method, "param": pname,
               "schedule": args.schedule, "first_chunk_schedule": args.first_chunk, "frames": args.frames,
               "baseline_ms": base_ms, "rows": rows}, open(out, "w"), indent=2)
    print("wrote", out)