#!/usr/bin/env python3
"""pack.py - turn lbrender's card output into the atlas play.cgi composites from.

  pack.py lb.idx lb.bin OUT.lba [preview-dir]

Claude Opus 5.5, October 2026.

Input: two RGB565 frames per job (cleared to magenta, then green); a pixel
equal in both is the object. Everything is 3x supersampled; we average 3x3
blocks into colour (premultiplied) and coverage.

Output (big-endian), see play.c's atlas reader:
  "LBA1" W:u16 H:u16 NSPR:u16 0:u16
  palette 256*3, LUT 32768 (RGB555 -> palette index)
  grid 81 * (x, y) i16 in 1/8 pixel: corner (f, r) of the board, white's view
  background, white's view: W*H*3; background, black's view: W*H*3
  sprites table NSPR * 24 bytes:
     colour type square 0 : u8*4, x0 y0 : i16, w h : u16, data offset : u32
     (from the file start), then the click outline: NLEV * (y, xleft, xright) i16
  The outline is the piece itself (coverage >= 1/2) at NLEV heights; play.cgi
  turns it into an image-map polygon. Table entries are 16 + 6*NLEV bytes.
  sprite data: per row  start:u16 len:u16  then len*4 bytes R G B A
Knights are rendered facing both ways; only facing 0 is used (a 180-degree
view of a left-facing knight on the mirrored square is a left-facing knight).
"""
import struct, sys, os
import numpy as np
from PIL import Image, ImageDraw, ImageFont

IDX, BIN, OUT = sys.argv[1:4]
PREV = sys.argv[4] if len(sys.argv) > 4 else None
CA, CB = 0xF81F, 0x07E0
BGC = np.array([16.0, 24.0, 40.0])     # page colour #101828: multiples of 8, so the
                                       # dither never moves it out of one 5-bit cell

lines = open(IDX).read().split("\n")
cam = [float(x) for x in lines[0].split()[1:16]]
SS, W, H = [int(x) for x in lines[0].split()[16:19]]
eye = np.array(cam[0:3]); cr = np.array(cam[3:6]); cu = np.array(cam[6:9]); cf = np.array(cam[9:12])
focal, ccx, ccy = cam[12:15]
raw = np.memmap(BIN, dtype=">u2", mode="r")

def project(p):
    q = np.asarray(p, float) - eye
    zc = q @ cf
    return ccx + focal * (q @ cr) / zc, ccy - focal * (q @ cu) / zc

def rgb565(a):
    a = a.astype(np.uint32)
    return np.dstack([((a >> 11) & 31) * 255 // 31, ((a >> 5) & 63) * 255 // 63, (a & 31) * 255 // 31]).astype(np.float64)

def frames(off, w, h):
    n = w * h
    a = np.asarray(raw[off // 2: off // 2 + n]).reshape(h, w)
    b = np.asarray(raw[off // 2 + n: off // 2 + 2 * n]).reshape(h, w)
    return a, b

def downsample(a, b, x0, y0):
    """Supersampled frames at (x0, y0) -> final-res premultiplied rgb, alpha, origin."""
    h, w = a.shape
    X0, Y0 = (x0 // SS) * SS, (y0 // SS) * SS
    X1, Y1 = -(-(x0 + w) // SS) * SS, -(-(y0 + h) // SS) * SS
    cov = np.zeros((Y1 - Y0, X1 - X0)); col = np.zeros((Y1 - Y0, X1 - X0, 3))
    m = (a == b)
    cov[y0 - Y0:y0 - Y0 + h, x0 - X0:x0 - X0 + w] = m
    col[y0 - Y0:y0 - Y0 + h, x0 - X0:x0 - X0 + w] = rgb565(a) * m[..., None]
    hh, ww = (Y1 - Y0) // SS, (X1 - X0) // SS
    cov = cov.reshape(hh, SS, ww, SS).mean(axis=(1, 3))
    col = col.reshape(hh, SS, ww, SS, 3).mean(axis=(1, 3))
    return col, cov, X0 // SS, Y0 // SS

# ---------------------------------------------------------------- background
bigc = np.zeros((H * SS, W * SS)); bigrgb = np.zeros((H * SS, W * SS, 3))
sprites = []
for ln in lines[1:]:
    f = ln.split()
    if not f:
        continue
    if f[0] == "B":
        tx, ty, x0, y0, w, h, off = map(int, f[1:])
        a, b = frames(off, w, h)
        m = a == b
        bigc[y0:y0 + h, x0:x0 + w] = m
        bigrgb[y0:y0 + h, x0:x0 + w] = rgb565(a) * m[..., None]
    elif f[0] == "S":
        sprites.append(tuple(map(int, f[1:])))
cov = bigc.reshape(H, SS, W, SS).mean(axis=(1, 3))
col = bigrgb.reshape(H, SS, W, SS, 3).mean(axis=(1, 3))
bg = col + BGC * (1 - cov[..., None])              # over the page colour

grid = np.zeros((9, 9, 2))
for r in range(9):
    for fl in range(9):
        grid[r, fl] = project([fl, 0, r])

def with_labels(img, black):
    im = Image.fromarray(np.clip(img + 0.5, 0, 255).astype(np.uint8))
    d = ImageDraw.Draw(im)
    for i in range(8):
        for kind in "fr":
            if kind == "f":
                p = [i + 0.5, 0, -0.30]; txt = chr(ord("a") + (7 - i if black else i))
            else:
                p = [-0.30, 0, i + 0.5]; txt = str(8 - i if black else i + 1)
            x, y = project(p)
            # size from the local scale of a square at that depth
            s = abs(project([p[0] + 0.5, 0, p[2]])[0] - x) * 2
            font = ImageFont.truetype("/usr/share/fonts/dejavu-sans-fonts/DejaVuSans-Bold.ttf",
                                      max(8, int(s * 0.30)))
            bb = d.textbbox((0, 0), txt, font=font)
            d.text((x - (bb[0] + bb[2]) / 2, y - (bb[1] + bb[3]) / 2), txt, font=font,
                   fill=(176, 190, 214))
    return np.asarray(im).astype(np.float64)

bgw = with_labels(bg, False)
bgb = with_labels(bg, True)

# ---------------------------------------------------------------- sprites
SH_RGB = np.array([16.0, 20.0, 34.0])
NLEV = 10
out_sprites = []
for (colour, ptype, facing, sq, x0, y0, w, h, off) in sprites:
    if facing != 0:
        continue
    a, b = frames(off, w, h)
    pcol, pcov, fx0, fy0 = downsample(a, b, x0, y0)
    # contact shadow: a soft disc on the floor, nudged away from the key light
    cx, cz = (sq & 7) + 0.5 + 0.06, (sq >> 3) + 0.5 + 0.05
    corners = [project([cx + dx, 0, cz + dz]) for dx in (-0.5, 0.5) for dz in (-0.5, 0.5)]
    sx0 = int(np.floor(min(c[0] for c in corners))) - 1
    sx1 = int(np.ceil(max(c[0] for c in corners))) + 1
    sy0 = int(np.floor(min(c[1] for c in corners))) - 1
    sy1 = int(np.ceil(max(c[1] for c in corners))) + 1
    ux0, uy0 = min(fx0, sx0), min(fy0, sy0)
    ux1, uy1 = max(fx0 + pcov.shape[1], sx1), max(fy0 + pcov.shape[0], sy1)
    hh, ww = uy1 - uy0, ux1 - ux0
    ys, xs = np.mgrid[uy0:uy1, ux0:ux1] + 0.5
    d = cf + ((xs - ccx) / focal)[..., None] * cr - ((ys - ccy) / focal)[..., None] * cu
    t = -eye[1] / d[..., 1]
    hit = eye + d * t[..., None]
    dist = np.hypot(hit[..., 0] - cx, hit[..., 2] - cz)
    rad = 0.40 if ptype != 1 else 0.33
    sa = 0.42 * np.clip((rad - dist) / 0.16, 0, 1) ** 1.5
    A = np.zeros((hh, ww)); C = np.zeros((hh, ww, 3))
    oy, ox = fy0 - uy0, fx0 - ux0
    A[oy:oy + pcov.shape[0], ox:ox + pcov.shape[1]] = pcov
    C[oy:oy + pcov.shape[0], ox:ox + pcov.shape[1]] = pcol
    # piece over shadow, premultiplied
    outA = A + sa * (1 - A)
    outC = C + SH_RGB * sa[..., None] * (1 - A)[..., None]
    rgb = np.where(outA[..., None] > 0, outC / np.maximum(outA, 1e-9)[..., None], 0)
    alpha = np.clip(np.round(outA * 255), 0, 255).astype(np.uint8)
    rgb = np.clip(np.round(rgb), 0, 255).astype(np.uint8)
    # click outline: left and right edge of the piece itself at NLEV heights
    solid = A >= 0.5
    rows = np.nonzero(solid.any(1))[0]
    click = []
    for i in range(NLEV):
        y = rows[0] + (rows[-1] - rows[0]) * i // (NLEV - 1)
        lo, hi = max(rows[0], y - 2), min(rows[-1], y + 2) + 1
        cols = np.nonzero(solid[lo:hi].any(0))[0]
        click.append((uy0 + y, ux0 + cols[0], ux0 + cols[-1] + 1))
    out_sprites.append(dict(colour=colour, type=ptype, sq=sq, x0=ux0, y0=uy0, rgb=rgb,
                            alpha=alpha, click=click))

# ---------------------------------------------------------------- reference compositor
BAYER = np.array([[0, 8, 2, 10], [12, 4, 14, 6], [3, 11, 1, 9], [15, 7, 13, 5]])
spr_index = {(s["colour"], s["type"], s["sq"]): s for s in out_sprites}

def compose(board, black, tints=()):
    """board: dict square -> (colour, type). Mirrors play.c exactly (integer math)."""
    img = (bgb if black else bgw).round().astype(np.int64).copy()
    for (vsq, tint, ta) in tints:              # vsq is the square as drawn (image position)
        f, r = vsq & 7, vsq >> 3
        P = [grid[r, f], grid[r, f + 1], grid[r + 1, f + 1], grid[r + 1, f]]
        P = [(int(round(p[0] * 8)), int(round(p[1] * 8))) for p in P]
        x0 = max(0, min(p[0] for p in P) // 8); x1 = min(W - 1, max(p[0] for p in P) // 8 + 1)
        y0 = max(0, min(p[1] for p in P) // 8); y1 = min(H - 1, max(p[1] for p in P) // 8 + 1)
        ys, xs = np.mgrid[y0:y1 + 1, x0:x1 + 1]
        px, py = xs * 8 + 4, ys * 8 + 4
        inside = np.ones(px.shape, bool)
        for i in range(4):
            ax, ay = P[i]; bx, by = P[(i + 1) % 4]
            inside &= ((bx - ax) * (py - ay) - (by - ay) * (px - ax)) <= 0
        sub = img[y0:y1 + 1, x0:x1 + 1]
        sub[inside] = (sub[inside] * (256 - ta) + np.array(tint) * ta) >> 8
    def pos(sq):                               # where the sprite sits in white's view
        return 63 - sq if black else sq
    order = sorted(board.items(), key=lambda kv: (-(pos(kv[0]) // 8), -abs((pos(kv[0]) & 7) - 3.5)))
    for sq, (colour, ptype) in order:
        s = spr_index[(colour, ptype, 63 - sq if black else sq)]
        h, w = s["alpha"].shape
        a = s["alpha"].astype(np.int64)[..., None]
        reg = img[s["y0"]:s["y0"] + h, s["x0"]:s["x0"] + w]
        reg[:] = (s["rgb"].astype(np.int64) * a + reg * (255 - a) + 127) // 255
    return img

def quantise(img, lut):
    h, w, _ = img.shape
    d = BAYER[np.arange(h)[:, None] & 3, np.arange(w)[None, :] & 3][..., None] // 2
    c = np.minimum(img + d, 255) >> 3
    return lut[(c[..., 0] << 10) | (c[..., 1] << 5) | c[..., 2]]

START = {}
for f, t in enumerate([4, 2, 3, 5, 6, 3, 2, 4]):
    START[f] = (0, t); START[8 + f] = (0, 1); START[48 + f] = (1, 1); START[56 + f] = (1, t)

# ---------------------------------------------------------------- palette
rng = np.random.default_rng(5)
samples = [compose(START, False), compose(START, True)]
TINTS = [((240, 216, 96), 120), ((110, 200, 100), 130), ((245, 150, 60), 125)]   # = board3d.c
for k in range(6):
    sqs = rng.choice(64, size=rng.integers(8, 28), replace=False)
    board = {int(s): (int(rng.integers(0, 2)), int(rng.integers(1, 7))) for s in sqs}
    tints = [(int(s), *TINTS[i % 3]) for i, s in enumerate(rng.choice(64, size=10, replace=False))]
    samples.append(compose(board, bool(k & 1), tints))
sheet = Image.fromarray(np.concatenate(samples, axis=0).astype(np.uint8))
q = sheet.quantize(colors=256, method=Image.Quantize.MEDIANCUT, dither=Image.Dither.NONE)
pal = np.array(q.getpalette()[:768], dtype=np.int64).reshape(256, 3)
ibg = int(((pal - BGC) ** 2).sum(1).argmin())
pal[ibg] = BGC.astype(np.int64)          # the page colour exactly
grid5 = np.stack(np.meshgrid(np.arange(32), np.arange(32), np.arange(32), indexing="ij"), -1).reshape(-1, 3)
centres = grid5 * 255 / 31 + 255 / 62                 # middle of each 5-bit cell
lut = np.zeros(32768, dtype=np.uint8)
for i in range(0, 32768, 4096):
    dd = ((centres[i:i + 4096, None, :] - pal[None, :, :]) ** 2).sum(-1)
    lut[i:i + 4096] = dd.argmin(1)
cell = (BGC.astype(int) >> 3)
lut[(cell[0] << 10) | (cell[1] << 5) | cell[2]] = ibg

# ---------------------------------------------------------------- write
def spans(s):
    out = bytearray()
    for y in range(s["alpha"].shape[0]):
        nz = np.nonzero(s["alpha"][y])[0]
        if len(nz) == 0:
            out += struct.pack(">HH", 0, 0)
            continue
        a, b = nz[0], nz[-1] + 1
        out += struct.pack(">HH", a, b - a)
        px = np.concatenate([s["rgb"][y, a:b], s["alpha"][y, a:b, None]], axis=1)
        out += px.astype(np.uint8).tobytes()
    return bytes(out)

head = b"LBA1" + struct.pack(">HHHH", W, H, len(out_sprites), 0)
head += pal.astype(np.uint8).tobytes() + lut.tobytes()
head += np.round(grid.reshape(-1, 2) * 8).astype(">i2").tobytes()
head += np.clip(bgw.round(), 0, 255).astype(np.uint8).tobytes()
head += np.clip(bgb.round(), 0, 255).astype(np.uint8).tobytes()
table_at = len(head)
data = bytearray(); table = bytearray()
base = table_at + (16 + 6 * NLEV) * len(out_sprites)
for s in out_sprites:
    h, w = s["alpha"].shape
    table += struct.pack(">BBBBhhHHI", s["colour"], s["type"], s["sq"], 0, s["x0"], s["y0"],
                         w, h, base + len(data))
    table += struct.pack(">" + "hhh" * NLEV, *[v for lev in s["click"] for v in lev])
    data += spans(s)
with open(OUT, "wb") as fo:
    fo.write(head + table + data)
print("sprites", len(out_sprites), "atlas bytes", len(head) + len(table) + len(data),
      "sprite data", len(data))

if PREV:
    os.makedirs(PREV, exist_ok=True)
    for name, img in [("start-w", compose(START, False)), ("start-b", compose(START, True))]:
        idx = quantise(img, lut)
        im = Image.fromarray(idx.astype(np.uint8), "P"); im.putpalette(pal.astype(np.uint8).tobytes())
        im.save(os.path.join(PREV, name + ".gif"))
        Image.fromarray(img.astype(np.uint8)).save(os.path.join(PREV, name + "-rgb.png"))
    np.save(os.path.join(PREV, "lut.npy"), lut)
    np.save(os.path.join(PREV, "pal.npy"), pal)
    # references for test/parity.py
    #  0: White to move after 1.e4 Nc6, Ng1 selected (White's view)
    #  1: Black to move after 1.e4, Nb8 selected (Black's view)
    b = dict(START); del b[12]; b[28] = (0, 1); del b[57]; b[42] = (1, 2)
    tl = [(57, *TINTS[0]), (42, *TINTS[0]), (6, *TINTS[1]), (12, *TINTS[2]), (21, *TINTS[2]), (23, *TINTS[2])]
    np.save(os.path.join(PREV, "ref-0.npy"), quantise(compose(b, False, tl), lut))
    b = dict(START); del b[12]; b[28] = (0, 1)
    tl = [(12, *TINTS[0]), (28, *TINTS[0]), (57, *TINTS[1]), (40, *TINTS[2]), (42, *TINTS[2])]
    np.save(os.path.join(PREV, "ref-1.npy"), quantise(compose(b, True, [(63 - q, c, a) for (q, c, a) in tl]), lut))