#!/usr/bin/env python3
"""Rate control: hit a target bitrate exactly, so one encoder serves both targets.

USER DECISION (session 2): ship TWO quality modes, SASI and SCSI. The codec's
bitrate ceiling is a build parameter; the encoder is otherwise identical.

Mechanism: the hybrid encoder's lagrangian `lam` trades distortion for bytes
monotonically, so per frame we binary-search lam to land inside a byte budget.
A leaky bucket lets a quiet frame bank bytes that an action frame can spend --
without that, quiet frames waste budget and action frames stay ugly.

The ceiling is HARD: the 68000 streams at a fixed rate off the disk, and a frame
that overruns is a dropped frame, not a slow frame.
"""
import numpy as np
import vq_hybrid as H

# Profiles. Bandwidths are the sustained-read figures the player can rely on;
# see docs/FINDINGS.md 5 -- these are FOLKLORE-grade until the disk benchmark
# is unblocked, so they are deliberately conservative fractions of the quoted
# ceiling (audio, seeks and container overhead come out of the same pipe).
# Calibrated against the CORRECTED rate-distortion measurement (FINDINGS 14).
#
# k=256 with 1-byte indices beats k=1024 with 2-byte indices at every matched
# bitrate. The earlier "+2.4 dB for k=1024" was an artifact of a rate model that
# charged 1 byte for a 10-bit index. 1-byte indices also mean the 68000 decoder
# reads a plain move.b with no alignment case, and the codebook is 8 KB not 32 KB.
#
# The two profiles are the SAME codec, decoder and bitstream -- only `lam` differs.
PROFILES = {
    "sasi": dict(kbps=45,  lam=300.0, k1=256, k4=256,
                 desc="stock 10MHz ACE/EXPERT, SASI",
                 quality="34.8 dB on 00020 / 28.3 dB on 00146"),
    "scsi": dict(kbps=75,  lam=100.0, k1=256, k4=256,
                 desc="Super/XVI, or CZ-6BS1 board in a 10MHz machine",
                 quality="35.9 dB on 00020 / 29.0 dB on 00146"),
}
# Not a shipping profile, but the curve continues: lam=25 is ~185 KB/s at ~38.7 dB
# with 26% RAW blocks, and lam->0 is pixel-exact (0.00 dB loss). Entropy-coding
# the payload (NOT YET IMPLEMENTED) should shift the whole curve ~1.4x left.

AUDIO_KBPS = 7.8          # MSM6258 ADPCM 15.6kHz mono -- comes out of the same budget


def frame_budget(kbps, fps=12, audio=AUDIO_KBPS):
    """bytes per video frame after audio takes its cut"""
    return (kbps - audio) * 1024.0 / fps


def encode_rate_controlled(m, target_kbps, fps=12, bucket_frames=8,
                           lam_lo=1.0, lam_hi=2e5, steps=9, verbose=False):
    budget = frame_budget(target_kbps, fps)
    bucket = 0.0                       # banked bytes, capped at bucket_frames*budget
    cap = bucket_frames * budget
    out_recon, out_modes, out_sizes, out_lam = [], [], [], []

    # encode() is whole-sequence; drive it per-lam and pick per frame.
    # Cheaper than re-running the whole encoder per frame: precompute the ladder.
    ladder = []
    lams = np.geomspace(lam_lo, lam_hi, steps)
    for lam in lams:
        e = H.encode(m, lam=float(lam))
        ladder.append(e)
        if verbose:
            print(f"  lam={lam:9.0f}  mean {e['sizes'].mean():6.0f} B/frame")

    nf = len(m["idx"])
    for f in range(nf):
        allow = budget + bucket
        # cheapest lam (highest quality) whose size fits the allowance
        pick = len(lams) - 1
        for i in range(len(lams)):
            if ladder[i]["sizes"][f] <= allow:
                pick = i; break
        sz = ladder[pick]["sizes"][f]
        bucket = min(cap, bucket + budget - sz)
        out_recon.append(ladder[pick]["recon"][f])
        out_modes.append(ladder[pick]["modes"][f])
        out_sizes.append(sz); out_lam.append(lams[pick])

    return dict(recon=out_recon, modes=out_modes, sizes=np.array(out_sizes),
                lam=np.array(out_lam), nb=ladder[0]["nb"], budget=budget)


def summarise(m, enc, target_kbps, fps=12):
    import vq as VQ
    pal = m["pal"]
    rec = [pal[i] for i in enc["recon"]]
    src = [pal[i] for i in m["idx"]]
    p = np.mean([VQ.psnr(o, v) for o, v in zip(m["rgb"], rec)])
    pp = np.mean([VQ.psnr(o, v) for o, v in zip(m["rgb"], src)])
    sz = enc["sizes"]
    mo = np.concatenate(enc["modes"])
    return dict(target=target_kbps, psnr=p, pal=pp, loss=pp - p,
                mean_B=sz.mean(), max_B=sz.max(), budget=enc["budget"],
                kbps=sz.mean() * fps / 1024 + AUDIO_KBPS,
                over=100.0 * np.mean(sz > enc["budget"]),
                skip=100 * (mo == 0).mean(), v1=100 * (mo == 1).mean(),
                v4=100 * (mo == 2).mean())