#!/usr/bin/env python
"""Pick the LingBot operating points from the scans (the analogue of
build_points.py + hy_build_points_c0.py):

* ``results/scan_lb_<method>_<sched>.json``        -> results/lb_point_v2.json
  bands on measured speedup: FxxF 1.70-1.80x, Fxxx 2.70-3.00x; MotionCache must keep
  its token-level decision alive (>=25% of cacheable steps partial, >=5% tokens
  selected); else the fastest alive setting below the band, flagged.
* ``results/scan_lb_<method>_<sched>_c0FFFF.json`` -> results/lb_point_c0.json
  bands on the 8-chunk-extrapolated measured speedup: FxxF [1.6, 1.75], Fxxx
  [2.5, 2.85]; closest to the upper edge, MotionCache active_ratio >= 0.25.
"""
import glob, json, os
ROOT = os.path.dirname(os.path.abspath(__file__))
BASE = "lingbot_world_v2_1.3b_causal_fast"
BANDS = {"FxxF": (1.70, 1.80, 1.75), "Fxxx": (2.70, 3.00, 2.85)}
ACT_MIN, SEL_MIN = 0.25, 0.05
TOKENWISE = {"motioncache"}
BANDS_C0 = {"FxxF": (1.6, 1.75), "Fxxx": (2.5, 2.85)}
TARGET = {"FxxF": 1.75, "Fxxx": 2.85}

scans, scans_c0 = {}, {}
for p in glob.glob(os.path.join(ROOT, "results/scan_lb_*.json")):
    if p.endswith("_dirty.json"):
        continue
    d = json.load(open(p))
    key = (d["method"], d["schedule"])
    if d.get("first_chunk_schedule"):
        scans_c0.setdefault(key, []).extend([dict(r, _file=os.path.basename(p)) for r in d["rows"]])
    else:
        scans.setdefault(key, []).extend(d["rows"])


def choose(method, sched):
    lo, hi, _ = BANDS[sched]
    cand = scans.get((method, sched), [])
    if not cand:
        return None, "no scan"
    alive = ((lambda r: r["active_ratio"] >= ACT_MIN and r["sel_mean"] >= SEL_MIN)
             if method in TOKENWISE else (lambda r: True))
    band = [r for r in cand if lo <= r["measured_speedup"] <= hi and alive(r)]
    if band:
        return max(band, key=lambda r: r["measured_speedup"]), ""
    live = [r for r in cand if alive(r) and r["measured_speedup"] <= hi]
    if live:
        r = max(live, key=lambda r: r["measured_speedup"])
        why = ("token-wise decision goes vacuous above this point"
               if method in TOKENWISE else "reachable speedups are quantised")
        return r, f"below band ({why})"
    r = min(cand, key=lambda r: abs(r["measured_speedup"] - hi))
    return r, "above band (reachable speedups are quantised)"


rows = []
print(f"{'method':12s}{'sched':>6}{'param':>10}{'measured':>10}{'compute':>9}{'active%':>9}{'sel%':>7}  note")
for method in ("teacache", "taylorseer", "motioncache"):
    for sched in ("FxxF", "Fxxx"):
        r, note = choose(method, sched)
        if r is None:
            print(f"{method:12s}{sched:>6}  {note}"); continue
        rows.append({"base": BASE, "method": method, "param": r["param"], "value": r["value"],
                     "schedule": sched, "target_speedup": BANDS[sched][2],
                     "scan_measured_speedup": r["measured_speedup"], "scan_compute_speedup": r["compute_speedup"],
                     "scan_compute_equivalent": r["compute_equivalent"], "scan_active_step_ratio": r["active_ratio"],
                     "scan_mean_selected_fraction": r["sel_mean"], "band_note": note,
                     "note": "operating point targeted on measured wall-clock speedup"})
        print(f"{method:12s}{sched:>6}{r['value']:>10.5g}{r['measured_speedup']:>10.3f}{r['compute_speedup']:>9.3f}"
              f"{100*r['active_ratio']:>9.1f}{100*r['sel_mean']:>7.1f}  {note}")
p = os.path.join(ROOT, "results/lb_point_v2.json")
json.dump({"base": BASE, "method": "multi", "param": "-", "schedule": "-", "operating_points": [], "rows": rows},
          open(p, "w"), indent=2)
print("wrote", p, len(rows), "rows")

rows = []
for (m, sched), cands in sorted(scans_c0.items()):
    lo, hi = BANDS_C0[sched]
    inband = [r for r in cands if lo <= r["ext8_measured_speedup"] <= hi]
    if m == "motioncache" and any(r["active_ratio"] >= 0.25 for r in inband):
        inband = [r for r in inband if r["active_ratio"] >= 0.25]
    if inband:
        best = min(inband, key=lambda r: (abs(r["ext8_measured_speedup"] - hi), r["value"]))
        note = "in band"
    else:
        best = min(cands, key=lambda r: (min(abs(r["ext8_measured_speedup"] - lo), abs(r["ext8_measured_speedup"] - hi)), r["value"]))
        note = "closest to band (reachable speedups are quantised)"
    tag = "" if sched == "FxxF" else f"_{sched}"
    rows.append({"base": BASE, "method": m, "param": best["param"], "value": best["value"],
                 "schedule": sched, "first_chunk_schedule": "FFFF", "target_speedup": TARGET[sched],
                 "name": f"lb_{m}{tag}_x{TARGET[sched]:g}_c0FFFF",
                 "scan_ext8_measured_speedup": best["ext8_measured_speedup"],
                 "scan_ext8_compute_equivalent": best["ext8_compute_equivalent"],
                 "scan_measured_speedup_4chunks": best["measured_speedup"],
                 "scan_active_step_ratio": best["active_ratio"], "scan_mean_selected_fraction": best["sel_mean"],
                 "band_note": note, "source_file": best["_file"]})
    print(f"{rows[-1]['name']:36s} {best['param']}={best['value']:<8g} ext8 {best['ext8_measured_speedup']:.3f}x "
          f"(Full {best['ext8_compute_equivalent']:.1f}) act {100*best['active_ratio']:.0f}%  {note}")
p = os.path.join(ROOT, "results/lb_point_c0.json")
json.dump({"base": BASE, "source": "c0FFFF re-search, 61-frame scans extrapolated to 8 chunks", "rows": rows},
          open(p, "w"), indent=2)
print("wrote", p, len(rows), "rows")