#!/usr/bin/env python3
"""What does fitting the CPU budget cost in quality? (session 8, lever B)

    python3 tools/analysis/13_cpu_ratectl.py [frames_dir] [--profiles scsi]

Session 6 made the BYTE budget a ceiling by bisecting `lam` per frame. FINDINGS
28 then showed the binding budget is CYCLES, not bytes, and that the mode
decision cannot see them: it minimises `D + lam*R` on a machine that charges V4
1.49x a V1 block while the lagrangian charges it 4x.

`ratectl.encode_rate_controlled(cycle_budget=...)` adds the second controller --
`mu` bisected per frame against 833,333 cycles, with the lam bisection nested
inside it. This measures what that costs: PSNR, bitrate, and how many frames
still miss, against the same encode with the ceiling off.

The cycle budget is HARD, not a bucket. Bytes bank in the player's ring buffer;
there is no double buffer to decode ahead into, so a frame that misses its
decode deadline is simply late (FINDINGS 28).

Both controllers score frames with the exact clustered cost `vq_hybrid.cycles`,
validated to 1 point against the 68000 (FINDINGS 28.2) -- not with the per-block
ranking constant the mode decision uses. See vq_hybrid's note on SKIP.
"""
import argparse, os, pickle, sys, time
sys.path.insert(0, "tools/encoder")
import numpy as np
import vq as VQ, vq_hybrid as H, ratectl as RC

ap = argparse.ArgumentParser()
ap.add_argument("frames_dir", nargs="?", default="tmp/fr_singe")
ap.add_argument("--profiles", default="scsi")
ap.add_argument("--fps", type=int, default=12)
ap.add_argument("--cache", default=None, help="pickle of H.build (auto by dir)")
a = ap.parse_args()
if not os.path.isdir(a.frames_dir):
    sys.exit(f"missing {a.frames_dir} -- see tools/bench/check.sh for extraction")

BUDGET = RC.FRAME_CYCLES

# H.build is ~55 s, nearly all k-means, and it does not depend on the profile:
# both ship k1=k4=256. One build, cached, serves every row of the table.
cache = a.cache or f"tmp/model_{os.path.basename(a.frames_dir.rstrip('/'))}.pkl"
if os.path.exists(cache):
    m = pickle.load(open(cache, "rb"))
    print(f"model from {cache}")
else:
    t = time.time()
    m = H.build(a.frames_dir, k1=256, k4=256, iters=16)
    pickle.dump(m, open(cache, "wb"))
    print(f"built model in {time.time()-t:.0f} s -> {cache}")
print(f"{a.frames_dir}: {len(m['idx'])} frames, {m['nb']} blocks, "
      f"budget {BUDGET:,.0f} cycles/frame at {a.fps}fps\n")


def run(prof_name, cycle_budget):
    p = RC.PROFILES[prof_name]
    m.pop("_sym", None)              # the frame-symbol cache holds one frame
    t = time.time()
    enc = RC.encode_rate_controlled(m, p["kbps"], fps=a.fps, lam_lo=p["lam"],
                                    cycle_budget=cycle_budget)
    s = RC.summarise(m, enc, p["kbps"], fps=a.fps)
    s["secs"] = time.time() - t
    s["ns"] = float(np.mean([100*(mm != 0).mean() for mm in enc["modes"]]))
    return s, enc


rows = []
for name in a.profiles.split(","):
    for label, cb in (("bytes only", None), ("bytes + cycles", BUDGET)):
        s, enc = run(name, cb)
        rows.append((name, label, s))
        print(f"{name:5s} {label:<15s} {s['secs']:5.1f} s  "
              f"PSNR {s['psnr']:.2f} dB  {s['kbps']:6.1f} KB/s  "
              f"CPU med {100*s['cyc_med']/BUDGET:5.1f}% p90 "
              f"{100*s['cyc_p90']/BUDGET:5.1f}% max {100*s['cyc_max']/BUDGET:5.1f}%  "
              f"miss {s['cpu_miss']:3d}  late {s['late']:2d}  "
              f"mu med {s['mu_med']:.4f} max {s['mu_max']:.3f}")

print()
hdr = f"{'':<22}{'PSNR':>8}{'KB/s':>9}{'CPU med':>10}{'CPU max':>10}{'miss':>7}"
for name in a.profiles.split(","):
    r = {lab: s for n, lab, s in rows if n == name}
    b, c = r["bytes only"], r["bytes + cycles"]
    print(f"--- {name} (target {RC.PROFILES[name]['kbps']} KB/s) ---")
    print(hdr)
    for lab, s in (("bytes only", b), ("bytes + cycles", c)):
        print(f"  {lab:<20}{s['psnr']:>7.2f} {s['kbps']:>8.1f} "
              f"{100*s['cyc_med']/BUDGET:>9.1f}%{100*s['cyc_max']/BUDGET:>9.1f}%"
              f"{s['cpu_miss']:>6d}")
    print(f"  {'cost of fitting':<20}{c['psnr']-b['psnr']:>+7.2f} dB, "
          f"{c['kbps']-b['kbps']:+.1f} KB/s, "
          f"{b['cpu_miss']-c['cpu_miss']} fewer misses, "
          f"{c['late']} frames unfixable at mu={RC.MU_CLIFF:g}")
    print(f"  {'modes % (b/c)':<20}SKIP {b['skip']:.1f}/{c['skip']:.1f}  "
          f"V1 {b['v1']:.1f}/{c['v1']:.1f}  V4 {b['v4']:.1f}/{c['v4']:.1f}  "
          f"RAW {b['raw']:.1f}/{c['raw']:.1f}")
    print()

print("FINDINGS 28.7: re-coding every non-SKIP block as V1 is the floor the "
      "CURRENT mode set\nallows, and it still missed 11 frames at the retired "
      "110 KB/s profile / 12 at scsi.\nMisses above that floor are spans, not "
      "the mode decision -- and 31.3 showed the\nfloor itself was too "
      "pessimistic, because the real decision can move a block to SKIP.")