#!/usr/bin/env python3
"""REGRESSION TEST for the ratectl lam-ladder desync (FINDINGS 26). PASSES as
of session 6 -- keep it passing.

Exits non-zero if the encoder ever again reports a reconstruction that a
decoder would not produce. That is the acceptance criterion for any change to
rate control, and it is not a property a PSNR number can show you.

The bug it was written for: encode_rate_controlled() ran H.encode() once per lam
over the WHOLE sequence, then picked each frame from whichever rung fit the
budget. H.encode() is temporally recursive -- a frame's SKIP blocks are copied
from the PREVIOUS RECONSTRUCTION of that same rung -- so when frame f came from
rung i and frame f-1 was emitted from rung j != i, the SKIP blocks in f
referenced a frame the decoder never saw. 111 of 120 frames drifted, worst frame
43.4%. The fix was structural: the encoder is now frame-drivable and rate
control feeds back the frame it actually emitted (vq_hybrid.frame_ctx/decide/
paint), so drift is zero by construction rather than by tuning.

This replays what a real decoder does -- SKIP copies the ACTUALLY EMITTED
previous frame -- and compares it to the reconstruction ratectl recorded.

Session 8 added a SECOND controller (mu, the per-frame 68000 decode ceiling)
that also varies the mode map frame to frame, so it is exposed to exactly the
same failure and is tested here too. Both configurations must show zero drift.

Needs tmp/fr_singe (see docs/STATUS.md, reproducing the sustained-action
result). ~55 s, nearly all of it the k-means in H.build; the rate-controlled
encode of 120 frames is ~2 s.
"""
import sys, os
sys.path.insert(0, "tools/encoder")
import numpy as np
import vq as VQ, vq_hybrid as H, ratectl as RC

m = H.build("tmp/fr_singe", k1=256, k4=256, iters=16)


def check(label, cycle_budget):
    """Encode, replay as a decoder would, and return the drift in pixels."""
    print(f"\n=== {label} ===")
    m.pop("_sym", None)
    # lam_lo=1.0: let quiet frames spend the whole allowance, which is the
    # harder case for this test -- it maximises how often lam moves frame to
    # frame.
    enc = RC.encode_rate_controlled(m, target_kbps=110, lam_lo=1.0,
                                    cycle_budget=cycle_budget)
    lam = enc["lam"]
    sw = int((np.diff(lam) != 0).sum())
    print(f"frames={len(lam)}  distinct lam used={len(set(lam.tolist()))}  "
          f"lam changes frame-to-frame={sw}  "
          f"overruns={int(enc['overrun'].sum())}")

    pal, nbx = m["pal"], m["W"] // 4
    emitted = []
    drift_px, drift_db = [], []
    for f, (rec, mode) in enumerate(zip(enc["recon"], enc["modes"])):
        out = rec.copy()
        if f > 0:
            prev_true = emitted[-1]
            for b in np.flatnonzero(mode == 0):          # SKIP blocks
                by, bx = divmod(int(b), nbx)
                y, x = by*4, bx*4
                out[y:y+4, x:x+4] = prev_true[y:y+4, x:x+4]
        emitted.append(out)
        d = (out != rec).sum()
        drift_px.append(d)
        drift_db.append(VQ.psnr(pal[rec], pal[out]))

    drift_px = np.array(drift_px)
    print(f"pixels differing from what the encoder recorded:")
    print(f"  frames with ANY drift: {int((drift_px>0).sum())}/{len(drift_px)}")
    print(f"  max {drift_px.max()} px ({100*drift_px.max()/(m['H']*m['W']):.1f}% of frame)")
    print(f"  mean {drift_px.mean():.0f} px")
    fin = [d for d in drift_db if np.isfinite(d)]
    if fin:
        print(f"  encoder-vs-decoder agreement: min {min(fin):.1f} dB "
              f"(inf = identical on {len(drift_db)-len(fin)} frames)")

    r = RC.summarise(m, enc, 110)
    print(f"\nratectl reports PSNR {r['psnr']:.2f} dB, {r['kbps']:.1f} KB/s "
          f"(target 110), {r['over']:.0f}% of frames over budget")
    tp = np.mean([VQ.psnr(o, pal[e]) for o, e in zip(m["rgb"], emitted)])
    print(f"what a decoder actually reconstructs: {tp:.2f} dB "
          f"-> overstated by {r['psnr']-tp:.2f} dB")
    return drift_px


# Acceptance criterion for the fix: a decoder replaying the emitted stream must
# reconstruct exactly what the encoder recorded -- under either controller.
bad = 0
for label, cb in (("bytes only (session 6)", None),
                  ("bytes + CPU ceiling (session 8)", RC.FRAME_CYCLES)):
    d = check(label, cb)
    bad += int((d > 0).any())
sys.exit(1 if bad else 0)