# _crowdpass.py — turn the two crowd ROLLS into the six sheets the game ships, plus the table of
# who in them can stand up.
#
#   uv run --with pillow --with numpy --with scipy python assets/_crowdpass.py            # install
#   uv run --with pillow --with numpy --with scipy python assets/_crowdpass.py --review   # to _refcheck
#
# INPUT   assets/_refcheck/crowd_seat.png + crowd_clap.png  (see `_crowdroll.py`)
# OUTPUT  assets/x_crowd{,2}.png · x_crowd_c{,2}.png · x_crowd_h{,2}.png · src/crowd.js
#
# THE PROBLEM THIS SOLVES. Ask a generator for "the same picture with some arms raised" and about a
# third of the pixels come back different — faces, seat backs, the railing, everything. Swap any of
# that in and the stand re-hashes itself rather than applauding, which is exactly what the shipped
# sheets used to do: Oliver, "the crowd wobbles with a client side animation right now.. that feels
# odd." So the drift is thrown away here and only the arms are kept.
#
# AND "THE ARMS" IS MEANT LITERALLY. The clap frame is the seated frame plus the pixels that were
# STAND in the seated roll and became PERSON in the clap roll — arms and hands raised into the empty
# wall above a row, and nothing else. Not the raiser's face, not his shirt, not his seat. It cannot
# drift, because 98.8% of the sheet is not copied at all: it IS the seated sheet.
#
# TWO INVARIANTS, and render.js leans on both:
#   1. the clap sheet is the seated sheet EXCEPT inside the boxes listed in src/crowd.js;
#   2. every one of those boxes is a WHOLE PERSON, so the game animates people ONE AT A TIME — a
#      raiser's box drawn over the seated stand moves that spectator and nobody else.
# Boxes that overlap are merged into one entry: two neighbours who would always rise together are
# one person as far as the animation is concerned, and pretending otherwise reads as a glitch.
#
# THE THREE COURTS ARE ONE ROLL, GRADED. Each court's look is a per-channel histogram match onto the
# OLD sheet for that court, so terracotta clay, night grass and blue hard court all come back — but
# the crowd underneath is one drawing, so anything added to it shows up on every court. (Before this
# the three were separate rolls; Oliver: "the whole thing somehow does not translate well to the
# other courts, there the difference is barely visible".) THE SAME LUT IS APPLIED TO BOTH FRAMES OF
# A PAIR — grading them apart would move identical pixels to different colours and put the wobble
# straight back.
import os, subprocess, sys
import numpy as np
from PIL import Image
from scipy import ndimage

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.join(HERE, '..')
REFCHECK = os.path.join(HERE, '_refcheck')
PAL = 128              # colours a sheet is quantised to on save
# The court LOOKS are matched to the sheets as they were BEFORE any of this — a pinned commit, not
# HEAD, so re-running the tool after installing its own output cannot drift the palettes.
GRADE_REF = '76ea1d0'

COURTS = [('x_crowd', 'x_crowd2'), ('x_crowd_c', 'x_crowd_c2'), ('x_crowd_h', 'x_crowd_h2')]

MIN_AREA, MAX_AREA = 150, 3000   # an arm blob: one pair of raised arms, or one arm
ARM_H, ARM_W = (10, 70), (8, 90)  # ...and its shape, in px of the 1408x768 sheet
BAND = 8               # rows the background colour is sampled over
PAD = 3                # px of margin around a raiser's box


def old_sheet(name):
    """The pre-roll sheet for a court, read out of git — the grade's target, and nothing else."""
    p = subprocess.run(['git', 'show', f'{GRADE_REF}:assets/{name}.png'], cwd=ROOT, capture_output=True)
    if p.returncode:
        sys.exit(f'cannot read {GRADE_REF}:assets/{name}.png — {p.stderr.decode()[:200]}')
    tmp = os.path.join(REFCHECK, f'_old_{name}.png')
    open(tmp, 'wb').write(p.stdout)
    return np.asarray(Image.open(tmp).convert('RGB'))


def modal(px):
    a = px.reshape(-1, 3).astype(int)
    v = int(np.bincount((a[:, 0] << 16) + (a[:, 1] << 8) + a[:, 2]).argmax())
    return np.array([v >> 16, (v >> 8) & 255, v & 255])


def stand(im):
    """The STAND itself — wall, steps, seat backs — as whatever each band of rows is mostly made of.
    Being able to say "this pixel is empty stand" is the whole trick: a raise is empty stand that
    turned into a person."""
    m = np.zeros(im.shape[:2], bool)
    for y in range(0, im.shape[0], BAND):
        band = im[y:y + BAND]
        m[y:y + BAND] = np.abs(band.astype(int) - modal(band)).sum(2) < 70
    return m


def raisers(A, B):
    """The arms. Blobs of stand-turned-person, shaped like a pair of raised arms, sitting directly
    on top of somebody who is in the seated sheet — which is what stops a re-rolled hat or a shifted
    railing highlight being mistaken for applause."""
    new = stand(A) & ~stand(B)
    lab, n = ndimage.label(ndimage.binary_closing(new, np.ones((3, 3))), np.ones((3, 3)))
    sizes = ndimage.sum(np.ones_like(lab), lab, range(1, n + 1))
    bgA = stand(A)
    out = []
    for i, sl in enumerate(ndimage.find_objects(lab)):
        h, w = sl[0].stop - sl[0].start, sl[1].stop - sl[1].start
        if not (MIN_AREA <= sizes[i] <= MAX_AREA):
            continue
        if not (ARM_H[0] <= h <= ARM_H[1] and ARM_W[0] <= w <= ARM_W[1]):
            continue
        # IT HAS TO BE ON SOMEBODY. The strip just below the blob must be a spectator in the SEATED
        # sheet; arms that rise out of empty stand belong to nobody and read as floating hands.
        if bgA[sl[0].stop:min(A.shape[0], sl[0].stop + 6), sl[1]].mean() > 0.55:
            continue
        out.append([sl[0].start, sl[1].start, sl[0].stop, sl[1].stop, i + 1])
    return out, lab


def merge(boxes):
    """Two raisers whose boxes overlap are ONE entry: the game draws a box at a time, so a box that
    contains half of a neighbour's arms would raise that neighbour too, on the wrong clock."""
    boxes = [b[:4] for b in boxes]
    changed = True
    while changed:
        changed = False
        for i in range(len(boxes)):
            for j in range(len(boxes) - 1, i, -1):
                a, b = boxes[i], boxes[j]
                if a[0] - PAD < b[2] and b[0] - PAD < a[2] and a[1] - PAD < b[3] and b[1] - PAD < a[3]:
                    boxes[i] = [min(a[0], b[0]), min(a[1], b[1]), max(a[2], b[2]), max(a[3], b[3])]
                    boxes.pop(j)
                    changed = True
    return boxes


def grade(src, ref):
    """Per-channel histogram match: the roll's colours redistributed onto a court's own look.
    Returns a function, so the pair can be graded with ONE transform."""
    luts = []
    for ch in range(3):
        s, r = src[:, :, ch].ravel(), ref[:, :, ch].ravel()
        sv, sc = np.unique(s, return_counts=True)
        rv, rc = np.unique(r, return_counts=True)
        mapped = np.interp(np.cumsum(sc) / s.size, np.cumsum(rc) / r.size, rv)
        # colours the seated frame never uses (the clap frame has a few) borrow their neighbours'
        luts.append(np.interp(np.arange(256), sv, mapped).astype(np.uint8))
    return lambda im: np.stack([luts[c][im[:, :, c]] for c in range(3)], 2)


def save(seated, pa, clap, pb):
    """ONE shared palette for the pair — see the invariant at the top."""
    sa = Image.fromarray(seated).quantize(colors=PAL, method=Image.MEDIANCUT)
    Image.fromarray(clap).quantize(palette=sa, dither=Image.NONE).save(pb, optimize=True)
    sa.save(pa, optimize=True)


def write_table(boxes, h, w):
    """src/crowd.js — the raisers as fractions of the sheet, sorted left to right so a wave can walk
    them in order."""
    rows = sorted(((x0 - PAD) / w, (y0 - PAD) / h, (x1 - x0 + 2 * PAD) / w, (y1 - y0 + 2 * PAD) / h)
                  for y0, x0, y1, x1 in boxes)
    body = ',\n  '.join('[' + ', '.join(f'{v:.4f}' for v in r) + ']' for r in rows)
    open(os.path.join(ROOT, 'src', 'crowd.js'), 'w').write(
        '// GENERATED by assets/_crowdpass.py — do not edit by hand, re-run the tool.\n'
        '// The spectators who can stand up: [x, y, w, h] as fractions of the crowd sheet, sorted\n'
        '// left to right. Each box holds ONE person whose arms are raised in <crowd>2 and nobody\n'
        '// else; everything outside these boxes is identical in the two sheets, which is what lets\n'
        '// render.js raise them one at a time with no seam. See "THE CLAP FRAME" in PROJECT.md.\n'
        f'export const RAISERS = [\n  {body},\n];\n')
    return len(rows)


def main(review):
    A = np.asarray(Image.open(os.path.join(REFCHECK, 'crowd_seat.png')).convert('RGB'))
    B = np.asarray(Image.open(os.path.join(REFCHECK, 'crowd_clap.png')).convert('RGB'))
    arms, lab = raisers(A, B)
    h, w = A.shape[:2]

    clap = A.copy()
    for y0, x0, y1, x1, i in arms:                     # the BLOB, not its box: only the arms move
        m = lab[y0:y1, x0:x1] == i
        clap[y0:y1, x0:x1][m] = B[y0:y1, x0:x1][m]
    boxes = merge(arms)
    was = (np.abs(A.astype(int) - B.astype(int)).sum(2) > 30).mean() * 100
    now = (np.abs(A.astype(int) - clap.astype(int)).sum(2) > 30).mean() * 100
    n = write_table(boxes, h, w)
    print(f'{n} raisers kept; the clap frame moves {now:.1f}% of the sheet (the roll moved {was:.1f}%)')

    for a_name, b_name in COURTS:
        lut = grade(A, old_sheet(a_name))
        pa, pb = (os.path.join(HERE, x + '.png') for x in (a_name, b_name))
        if review:
            pa, pb = (os.path.join(REFCHECK, x + '_new.png') for x in (a_name, b_name))
        save(lut(A), pa, lut(clap), pb)
        print(f'  {os.path.basename(pa)} + {os.path.basename(pb)}')


if __name__ == '__main__':
    main('--review' in sys.argv)