#!/usr/bin/env python3
"""Render what the codec is actually DOING, per block, per frame.

Three panels at 12 fps: the palettised source (the real quality ceiling, not
1080p -- FINDINGS 11), the decoded output, and a block-mode map.

The mode map is not decoration. FINDINGS 24.5 makes the per-frame non-SKIP
fraction the number that selects the decoder's inner loop, and a percentile
cannot show you that the non-SKIP blocks are CLUSTERED (a moving character on a
held background) rather than scattered. Clustering is what a run-length over
the mode headers would exploit.

  SKIP  left as the previous frame, costs the 68000 nothing
  V1    one 4x4 codeword, 1 byte
  V4    four 2x2 codewords, 4 bytes
  RAW   16 literal palette indices -- the escape that makes lam=0 pixel-exact

Usage: python3 tools/analysis/08_mode_map.py <frames_dir> <out.webm> [--profile p]
                                            [--scale N]
Output format follows the extension. Prefer .webm: GIF re-quantises to 256
colours, which is a poor fit for output whose subject is colour fidelity.
"""
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "encoder"))
import numpy as np
from PIL import Image
import vq as VQ, vq_hybrid as H, ratectl as RC

MODE_RGB = np.array([[ 20,  22,  30],     # SKIP - near black, costs nothing
                     [ 60, 150, 230],     # V1   - blue
                     [ 80, 200, 120],     # V4   - green
                     [235,  90,  70]],    # RAW  - red, the expensive escape
                    dtype=np.uint8)
LABEL = ["SKIP", "V1", "V4", "RAW"]

SCALE = 1
LOSSLESS = False

def main():
    global SCALE, LOSSLESS
    LOSSLESS = "--lossless" in sys.argv
    src, out = sys.argv[1], sys.argv[2]
    if "--scale" in sys.argv:
        SCALE = int(sys.argv[sys.argv.index("--scale")+1])
    prof = RC.PROFILES[sys.argv[sys.argv.index("--profile")+1]
                       if "--profile" in sys.argv else "sasi"]
    m   = H.build(src, k1=prof["k1"], k4=prof["k4"])
    enc = H.encode(m, lam=prof["lam"])
    pal, H_, W_ = m["pal"], m["H"], m["W"]
    nbx, nby = W_ // 4, H_ // 4

    frames, stats = [], []
    for f, (rec, mode) in enumerate(zip(enc["recon"], enc["modes"])):
        srcp = pal[m["idx"][f]]
        decp = pal[rec]
        mmap = MODE_RGB[mode.reshape(nby, nbx)].repeat(4, 0).repeat(4, 1)
        # tint the mode map with the decoded luma so the action stays legible
        luma = decp.mean(2, keepdims=True) / 255.0
        mmap = (mmap * (0.45 + 0.55 * luma)).astype(np.uint8)
        gap = np.full((H_, 3, 3), 60, np.uint8)
        panel = np.hstack([srcp, gap, decp, gap, mmap])
        im = Image.fromarray(panel)
        if SCALE != 1:
            im = im.resize((panel.shape[1]*SCALE, H_*SCALE), Image.NEAREST)
        frames.append(im)
        stats.append([(mode == i).mean() for i in range(4)])

    if out.endswith(".webm") or out.endswith(".mp4"):
        # Preferred. GIF would impose its own 256-colour palette on top of
        # output whose entire subject is colour fidelity, and costs ~4x the
        # bytes doing it.  -lossless keeps the panels pixel-exact.
        import subprocess
        w, h = frames[0].size
        cmd = ["ffmpeg", "-v", "error", "-y", "-f", "rawvideo", "-pix_fmt", "rgb24",
               "-s", f"{w}x{h}", "-r", "12", "-i", "-"]
        # yuv444p, not 420: the mode map is flat saturated colour on a 4-pixel
        # grid, and chroma subsampling smears exactly those edges. Lossless is
        # available but runs larger than the GIF on this content; crf 18 in 444
        # is visually clean at a quarter the size.
        if out.endswith(".webm"):
            cmd += ["-c:v", "libvpx-vp9", "-pix_fmt", "yuv444p", "-row-mt", "1"]
            cmd += ["-lossless", "1"] if LOSSLESS else ["-crf", "18", "-b:v", "0"]
        else:
            cmd += ["-c:v", "libx264", "-crf", "12", "-pix_fmt", "yuv444p"]
        p = subprocess.Popen(cmd + [out], stdin=subprocess.PIPE)
        for f in frames:
            p.stdin.write(np.asarray(f.convert("RGB")).tobytes())
        p.stdin.close(); p.wait()
    else:
        # One shared adaptive palette: per-frame palettes are what make a naive
        # GIF of this enormous, and a stable palette also stops the mode-map
        # colours shimmering between frames.
        shared = frames[0].quantize(colors=192, method=Image.MEDIANCUT)
        q = [f.quantize(palette=shared, dither=Image.NONE) for f in frames]
        q[0].save(out, save_all=True, append_images=q[1:],
                  duration=1000//12, loop=0, optimize=True)
    st = np.array(stats)
    print(f"{len(frames)} frames -> {out}  ({os.path.getsize(out)/1024:.0f} KB)")
    print("  panels: palettised source | decoded | block mode map")
    for i, n in enumerate(LABEL):
        print(f"  {n:4s} mean {100*st[:,i].mean():5.1f}%   "
              f"per-frame range {100*st[:,i].min():5.1f}% .. {100*st[:,i].max():5.1f}%")
    ns = 100 * (1 - st[:, 0])
    print(f"  non-SKIP: median {np.median(ns):.1f}%  p90 {np.percentile(ns,90):.1f}%")

if __name__ == "__main__":
    main()