#!/usr/bin/env python3
"""4x4 vector-quantisation prototype for the X68000 Dragon's Lair port.

Pipeline mirrors what the 68000 decoder would actually do, so the measured
quality is honest:

  frames -> per-scene 256-colour palette (median cut, NO dither)
         -> 4x4 blocks of PALETTISED rgb
         -> k-means codebook (luma-weighted euclidean)
         -> each codeword's 16 pixels snapped back to a palette index

The decoder only ever copies 16 palette indices out of a table, so the codebook
entries MUST be legal palette indices -- both quantisation losses compose.

No sklearn on this box; k-means is hand-rolled (chunked, numpy).
"""
import numpy as np, glob, os, sys
from PIL import Image

BW = BH = 4                      # block size
# ITU-R BT.601 luma weights, squared -- we compare in a luma-weighted RGB space
LUMA = np.array([0.299, 0.587, 0.114], dtype=np.float32)


def load_frames(d):
    fs = sorted(glob.glob(f"{d}/f*.png"))
    return [np.asarray(Image.open(f).convert("RGB")) for f in fs]


def scene_palette(rgb, colors=256, stride=3):
    """One shared palette for the whole scene, no dithering (cel art is flat)."""
    samp = np.concatenate([r.reshape(-1, 3) for r in rgb[::stride]])
    ref = Image.fromarray(samp.reshape(-1, 1, 3)).quantize(
        colors=colors, method=Image.MEDIANCUT, dither=Image.NONE)
    pal = np.array(ref.getpalette()[:colors * 3], dtype=np.uint8).reshape(-1, 3)
    return ref, pal


def palettise(rgb, ref):
    return [np.asarray(Image.fromarray(r).quantize(palette=ref, dither=Image.NONE),
                       dtype=np.uint8) for r in rgb]


def blockify(idx, pal, bw=BW, bh=BH):
    """(H,W) palette indices -> (nblocks, bh*bw*3) float32 luma-weighted RGB."""
    H, W = idx.shape
    rgb = pal[idx].astype(np.float32) * LUMA          # weight once, up front
    b = rgb.reshape(H // bh, bh, W // bw, bw, 3).transpose(0, 2, 1, 3, 4)
    return b.reshape(-1, bh * bw * 3)


def kmeans(X, k, iters=24, seed=0):
    """Chunked Lloyd's algorithm. k-means++ style seeding, deterministic."""
    rng = np.random.default_rng(seed)
    n = X.shape[0]
    if n <= k:
        return X.copy(), np.arange(n)
    # seed: farthest-point sampling on a random subsample (cheap k-means++)
    sub = X[rng.choice(n, min(n, 20000), replace=False)]
    C = np.empty((k, X.shape[1]), dtype=np.float32)
    C[0] = sub[rng.integers(len(sub))]
    d2 = ((sub - C[0]) ** 2).sum(1)
    for i in range(1, k):
        C[i] = sub[np.argmax(d2)]
        d2 = np.minimum(d2, ((sub - C[i]) ** 2).sum(1))
    lab = None
    for _ in range(iters):
        lab = assign(X, C)
        newC = C.copy()
        cnt = np.bincount(lab, minlength=k)
        s = np.zeros_like(C)
        np.add.at(s, lab, X)
        nz = cnt > 0
        newC[nz] = s[nz] / cnt[nz, None]
        # revive dead codewords on the worst-fit blocks
        if (~nz).any():
            err = ((X - newC[lab]) ** 2).sum(1)
            worst = np.argsort(err)[-int((~nz).sum()):]
            newC[~nz] = X[worst]
        if np.allclose(newC, C):
            C = newC; break
        C = newC
    return C, assign(X, C)


def assign(X, C, chunk=8192):
    """Nearest centroid, chunked to bound memory."""
    Cn = (C ** 2).sum(1)
    out = np.empty(X.shape[0], dtype=np.int32)
    for i in range(0, X.shape[0], chunk):
        x = X[i:i + chunk]
        d = Cn[None, :] - 2.0 * (x @ C.T)          # + |x|^2, constant per row
        out[i:i + chunk] = np.argmin(d, axis=1)
    return out


def snap_codebook(C, pal, bw=BW, bh=BH):
    """Centroids (luma-weighted RGB) -> legal palette indices, as the ROM stores them."""
    cb_rgb = C.reshape(-1, bh * bw, 3) / LUMA        # undo the weighting
    palw = pal.astype(np.float32) * LUMA
    flat = (cb_rgb * LUMA).reshape(-1, 3)
    d = (flat ** 2).sum(1)[:, None] - 2 * (flat @ palw.T) + (palw ** 2).sum(1)[None, :]
    return np.argmin(d, axis=1).astype(np.uint8).reshape(-1, bh * bw)


def unblockify(labels, cb_idx, H, W, bw=BW, bh=BH):
    blocks = cb_idx[labels].reshape(H // bh, W // bw, bh, bw)
    return blocks.transpose(0, 2, 1, 3).reshape(H, W)


def psnr(a, b):
    mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)
    return 99.0 if mse == 0 else 10 * np.log10(255.0 ** 2 / mse)


def encode_scene(frames_dir, k=256, bw=BW, bh=BH, iters=24):
    rgb = load_frames(frames_dir)
    H, W = rgb[0].shape[:2]
    ref, pal = scene_palette(rgb)
    idx = palettise(rgb, ref)

    X = np.concatenate([blockify(i, pal, bw, bh) for i in idx])
    C, _ = kmeans(X, k, iters)
    cb_idx = snap_codebook(C, pal, bw, bh)

    # re-assign against the SNAPPED codebook: that's what the decoder can produce
    Csnap = (pal[cb_idx].astype(np.float32) * LUMA).reshape(k, -1)
    recon, labels = [], []
    for i in idx:
        lab = assign(blockify(i, pal, bw, bh), Csnap)
        labels.append(lab)
        recon.append(unblockify(lab, cb_idx, H, W, bw, bh))
    return dict(rgb=rgb, pal=pal, idx=idx, cb_idx=cb_idx, labels=labels,
                recon=recon, H=H, W=W, k=k, bw=bw, bh=bh)


def report(r, name=""):
    pal, idx, recon = r["pal"], r["idx"], r["recon"]
    src8 = [pal[i] for i in idx]
    vq8 = [pal[i] for i in recon]
    orig = r["rgb"]
    p_pal = np.mean([psnr(o, s) for o, s in zip(orig, src8)])
    p_vq = np.mean([psnr(o, v) for o, v in zip(orig, vq8)])
    p_vq_only = np.mean([psnr(s, v) for s, v in zip(src8, vq8)])
    nb = (r["H"] // r["bh"]) * (r["W"] // r["bw"])
    idxbits = int(np.ceil(np.log2(r["k"])))
    keyf = nb * idxbits / 8
    # block-delta cost: how many block indices change frame to frame
    ch = [np.count_nonzero(r["labels"][i] != r["labels"][i - 1]) / nb
          for i in range(1, len(r["labels"]))]
    print(f"--- {name} k={r['k']} block={r['bw']}x{r['bh']} ---")
    print(f"  palette-only PSNR : {p_pal:5.2f} dB   (floor: 256c is the best we can do)")
    print(f"  after VQ     PSNR : {p_vq:5.2f} dB   (loss from VQ alone: {p_pal-p_vq:.2f} dB)")
    print(f"  VQ vs palettised  : {p_vq_only:5.2f} dB")
    print(f"  blocks/frame      : {nb}  keyframe {keyf:.0f} B  codebook {r['k']*r['bw']*r['bh']} B")
    if ch:
        print(f"  blocks changed/frm: mean {100*np.mean(ch):5.1f}%  p90 {100*np.percentile(ch,90):5.1f}%")
    return dict(p_pal=p_pal, p_vq=p_vq, nb=nb, keyf=keyf,
                chg=np.mean(ch) if ch else 0, chg90=np.percentile(ch,90) if ch else 0)


def zoom(a, f=3):
    return np.repeat(np.repeat(a, f, axis=0), f, axis=1)


def compare_png(r, frame, out, f=3):
    pal = r["pal"]
    src = pal[r["idx"][frame]]
    vq = pal[r["recon"][frame]]
    orig = r["rgb"][frame]
    gap = np.full((src.shape[0] * f, 4, 3), 40, dtype=np.uint8)
    strip = np.concatenate([zoom(orig, f), gap, zoom(src, f), gap, zoom(vq, f)], axis=1)
    Image.fromarray(strip).save(out)
    return out