#!/usr/bin/env python3
"""What 256 -> 16 colours actually costs, on real frames.

    python3 tools/analysis/18_text_plane_16col.py [frames_dir]

FINDINGS 46.3 opened a lead and could not price it: the X68000 text plane is
4bpp planar -- 0.5 bytes/pixel against the graphics planes' 2.0 -- so a LITERAL
uncompressed 16-colour frame is 288.0 KB/s against the shipping compressed
256-colour container's 496.7 KB/s.  42% cheaper on the wire, with no decoder.

The whole lead turns on one number nobody had computed: the quality cost of 16
colours.  This computes it, and it is deliberately generous to the 16-colour
side on every axis where the hardware allows it:

  * PER-FRAME palettes are legitimate here.  The text palette is 16 entries and
    reloading it is 16 words a frame -- nothing, against a 833,333-clock budget.
    The 256-colour path cannot do this: its palette is shared scene-wide
    (vq.scene_palette) because the codec's codebooks are indices INTO it.
  * DITHERING is free here, and only here.  The tree does not dither (vq.py:32,
    "cel art is flat") because dither destroys the inter-frame coherence SKIP
    blocks and v7 spans are built on.  A literal frame has no codec to wreck, so
    Floyd-Steinberg is available to this path at zero runtime cost.

Both are measured, so the comparison cannot be accused of hobbling the option it
is testing.  Reported against the 256-colour scene-palette ceiling (the tree's
existing "palette ceiling" figure) and against the shipping container's PSNR.
"""
import sys, os
sys.path.insert(0, "tools/encoder")
import numpy as np
from PIL import Image
import vq as VQ

FRAMES = sys.argv[1] if len(sys.argv) > 1 else "tmp/fr_singe"
SHIPPED_PSNR = 29.19          # docs/STATUS.md, --spans all, c=5, 496.7 KB/s

rgb = VQ.load_frames(FRAMES)
H, W = rgb[0].shape[:2]
n = len(rgb)
print(f"{FRAMES}: {n} frames, {W}x{H}")
print()

def recon_scene(colors, dither):
    """One palette for the whole scene -- what the 256 path is forced to do."""
    d = Image.FLOYDSTEINBERG if dither else Image.NONE
    samp = np.concatenate([r.reshape(-1, 3) for r in rgb[::3]])
    ref = Image.fromarray(samp.reshape(-1, 1, 3)).quantize(
        colors=colors, method=Image.MEDIANCUT, dither=Image.NONE)
    pal = np.array(ref.getpalette()[:colors * 3], np.uint8).reshape(-1, 3)
    return [pal[np.asarray(Image.fromarray(r).quantize(palette=ref, dither=d),
                           np.uint8)] for r in rgb]

def recon_perframe(colors, dither):
    """A fresh palette every frame -- what the text plane can afford."""
    d = Image.FLOYDSTEINBERG if dither else Image.NONE
    out = []
    for r in rgb:
        q = Image.fromarray(r).quantize(colors=colors, method=Image.MEDIANCUT,
                                        dither=d)
        pal = np.array(q.getpalette()[:colors * 3], np.uint8).reshape(-1, 3)
        out.append(pal[np.asarray(q, np.uint8)])
    return out

def report(name, recon):
    per = np.array([VQ.psnr(a, b) for a, b in zip(rgb, recon)])
    print(f"  {name:<42s} {per.mean():6.2f} dB   "
          f"(min {per.min():5.2f}  max {per.max():5.2f})")
    return per.mean()

print("PSNR vs the 24-bit source, mean over frames:")
c256 = report("256 colours, scene palette  [the tree's]", recon_scene(256, False))
report("256 colours, per-frame palette", recon_perframe(256, False))
print()
s16  = report("16 colours, scene palette", recon_scene(16, False))
p16  = report("16 colours, per-frame palette", recon_perframe(16, False))
p16d = report("16 colours, per-frame + FS dither", recon_perframe(16, True))
print()
print(f"  the 16-colour ceiling is the best of those: {max(s16, p16, p16d):.2f} dB")
print(f"  cost of 256 -> 16, at each side's best:     "
      f"{c256 - max(s16, p16, p16d):.2f} dB")
print()
print(f"  for scale, the shipping container delivers  {SHIPPED_PSNR:.2f} dB "
      f"at 496.7 KB/s")
print(f"  a 16-colour literal would deliver           "
      f"{max(s16, p16, p16d):.2f} dB at 288.0 KB/s")
delta = max(s16, p16, p16d) - SHIPPED_PSNR
print(f"  so the text-plane path is {abs(delta):.2f} dB "
      f"{'BETTER' if delta > 0 else 'WORSE'} at 58% of the bitrate")