#!/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.

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)
# 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)

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")

# Acceptance criterion for the fix: a decoder replaying the emitted stream must
# reconstruct exactly what the encoder recorded.
sys.exit(1 if (drift_px > 0).any() else 0)