# _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/_sheet_crowd3.png (seated) + _sheet_crowd3b.png (arms high) + _sheet_crowd3c.png
#         (hands together, mid-clap) — the three rolls (see `_crowdroll.py`), tracked, because
#         _refcheck is scratch and these are the only copy of the art
# OUTPUT  assets/x_crowd.png · x_crowd_c.png · x_crowd_h.png — ONE STACKED SPRITE SHEET PER COURT,
#         three frames tall — plus 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.
#
# THE SHEET IS AN ANIMATION SHEET: frame 0 seated, frame 1 arms high, frame 2 hands together in
# front of the face. A spectator with material in BOTH raised frames can be cycled between them —
# that is a clap, not an appearance — and 26 of them can. The rest alternate with frame 0, which is
# arms going up and down: also motion, which is the point. Oliver: "when clapping I barely see any
# movement at all now", and then "can't we make an animation sprite sheet of the new crowd?"
#
# TWO INVARIANTS, and render.js leans on both:
#   1. every raised frame is the SEATED frame plus that pose's arms, and nothing else;
#   2. every box in src/crowd.js 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 (and their frame masks OR'd together): 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'

# sheet -> how much its grade TARGET is lifted before the histogram match (a gamma; 1.0 is the old
# sheet as it was). Centre court's old crowd sat at a mean luminance of 39 against clay's 96 — dark
# enough that the detail in the new roll was being thrown away, and then dimmed on top of that:
# Oliver, "it is too dim to see anything on the centre court and the dark court". The night court
# borrows the hard-court sheet and does its own dimming in `COURTS.dim`, so it is fixed there.
COURTS = [('x_crowd', 0.72), ('x_crowd_c', 1.0), ('x_crowd_h', 1.0)]

# WHAT THE HUD EATS. A raiser drawn under the scoreboard is a draw call nobody sees, and worse: the
# "how many join in" fraction spends part of every cheer on them. Oliver: "in the upper parts of the
# crowd there seems to be a lot more action going on, but I can barely see it because the hud is
# covering it." Measured on the sheet before this rule, 31 of 134 raisers sat behind the board or the
# blood panel. They are dropped here, which also makes the visible applause denser for free.
# In LOGICAL px of the PORTRAIT layout, which is the primary target: the board sits at (6,6) 150x42
# with its game/sets panel under it, the blood panel is up in the right corner, and the stand band is
# about 195 px tall (render.js: standsBottom - AD_H). Landscape has a shorter band, so the same
# fractions cover more of the sheet there — this rule is conservative, never over-eager.
HUD_W, HUD_BAND = 270, 195
HUD = [(6, 6, 156, 62),          # scoreboard + the GAME/SETS panel under it
       (170, 6, 264, 30)]        # the BLOOD panel, clear of the quit chip

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
    carry a frame mask, and merging ORs it — which is also how a spectator ends up owning both
    raised poses and therefore a real clap."""
    boxes = [b[:4] + [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]),
                                a[4] | b[4]]
                    boxes.pop(j)
                    changed = True
    return boxes


def visible(boxes, h, w):
    """Drop the raisers the HUD sits on top of — see HUD above."""
    out = []
    for y0, x0, y1, x1, m in boxes:
        cx, cy = (x0 + x1) / 2 / w * HUD_W, (y0 + y1) / 2 / h * HUD_BAND
        if any(a <= cx <= c and b <= cy <= d for a, b, c, d in HUD):
            continue
        out.append([y0, x0, y1, x1, m])
    return out


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, frames):
    """src/crowd.js — the raisers as fractions of ONE FRAME of the sheet, sorted left to right so a
    wave can walk them in order, each with the mask of the frames its arms appear in."""
    rows = sorted(((x0 - PAD) / w, (y0 - PAD) / h, (x1 - x0 + 2 * PAD) / w, (y1 - y0 + 2 * PAD) / h, m)
                  for y0, x0, y1, x1, m in boxes)
    body = ',\n  '.join('[' + ', '.join(f'{v:.4f}' for v in r[:4]) + f', {r[4]}]' 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 crowd sheet is FRAMES frames stacked vertically: 0 seated (the backdrop), 1 arms\n'
        '// high, 2 hands together in front of the face. RAISERS lists the spectators who can move:\n'
        '// [x, y, w, h] as fractions of ONE FRAME, then a bitmask of the raised frames that person\n'
        '// actually appears in (1 = arms high, 2 = mid-clap, 3 = both, and 3 is a real clap cycle).\n'
        '// Each box holds ONE person and nobody else, and outside these boxes every frame is\n'
        '// pixel-identical to frame 0 — which is what lets render.js animate them one at a time\n'
        '// with no seam. See "THE CROWD SHEET" in PROJECT.md.\n'
        f'export const FRAMES = {frames};\n'
        f'export const RAISERS = [\n  {body},\n];\n')
    return rows


def main(review):
    A = np.asarray(Image.open(os.path.join(HERE, '_sheet_crowd3.png')).convert('RGB'))
    poses = [np.asarray(Image.open(os.path.join(HERE, n)).convert('RGB'))
             for n in ('_sheet_crowd3b.png', '_sheet_crowd3c.png')]
    h, w = A.shape[:2]

    frames, boxes = [A], []
    for bit, P in enumerate(poses):
        arms, lab = raisers(A, P)
        f = 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
            f[y0:y1, x0:x1][m] = P[y0:y1, x0:x1][m]
        frames.append(f)
        boxes += [b[:4] + [1 << bit] for b in arms]   # b[4] is the label id; the box carries a FRAME MASK
        drift = (np.abs(A.astype(int) - P.astype(int)).sum(2) > 30).mean() * 100
        moved = (np.abs(A.astype(int) - f.astype(int)).sum(2) > 30).mean() * 100
        print(f'  pose {bit + 1}: {len(arms)} arm blobs, {moved:.1f}% of the sheet '
              f'(the roll itself drifted {drift:.1f}%)')
    boxes = merge(boxes)
    seen = visible(boxes, h, w)
    print(f'  {len(boxes) - len(seen)} of {len(boxes)} raisers dropped: the HUD covers them')
    rows = write_table(seen, h, w, len(frames))
    both = sum(1 for r in rows if r[4] == 3)
    print(f'{len(rows)} spectators can move, {both} of them with both poses (a real clap)')

    stack = np.concatenate(frames, 0)
    for a_name, lift in COURTS:
        ref = old_sheet(a_name)
        if lift != 1.0:
            ref = (255 * (ref / 255.0) ** lift).astype(np.uint8)
        lut = grade(A, ref)                            # the LOOK is judged on the seated frame...
        pa = os.path.join(REFCHECK if review else HERE, a_name + ('_new.png' if review else '.png'))
        # ...and applied to the whole stack at once, so the frames cannot drift apart in colour
        Image.fromarray(lut(stack)).quantize(colors=PAL, method=Image.MEDIANCUT).save(pa, optimize=True)
        print(f'  {os.path.basename(pa)} ({stack.shape[1]}x{stack.shape[0]}, {len(frames)} frames)')


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