#!/usr/bin/env python3
"""REGRESSION TEST for the ratectl lam-ladder desync (FINDINGS 26).

Exits non-zero while the bug is present. After the fix it must report ZERO
drifting frames -- that is the acceptance criterion for wiring rate control
into encode.py.

encode_rate_controlled() runs H.encode() once per lam over the WHOLE sequence,
then picks each frame from whichever rung fits the budget. But H.encode() is
temporally recursive: a frame's SKIP blocks are copied from the PREVIOUS
RECONSTRUCTION of that same rung. If frame f is taken from rung i while frame
f-1 was emitted from rung j != i, the SKIP blocks in f reference a frame the
decoder never saw.

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). Takes a few minutes: it runs `steps` full-sequence encodes and
_paint is still a Python per-block loop.
"""
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)
enc = RC.encode_rate_controlled(m, target_kbps=110, steps=5, verbose=True)

lam = enc["lam"]
sw = int((np.diff(lam) != 0).sum())
print(f"\nframes={len(lam)}  distinct lam used={len(set(lam.tolist()))}  "
      f"rung switches={sw}")

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)