#!/usr/bin/env python3
"""Extract individual sprites from a Spriters-Resource-style sheet.

v2 - projection based:
  1. Remove background (border-connected pixels near the border palette).
  2. Remove the flat ground-shadow color(s) (default SF2 green (0,96,0)).
  3. Segment into ROW BANDS (runs of rows with foreground, split by empty rows).
  4. Within each band, segment into COLUMN RUNS (split by empty-column gutters).
     Each column run -> one sprite. Shadows removed => sprites separate cleanly.
  5. Drop thin text labels (short) and tiny specks.
  6. Save each trimmed transparent sprite + atlas.json + a numbered contact sheet.

Atlas frames carry x,y,w,h on the sheet plus the trimmed size, so the engine can
align by bottom-center (foot baseline).
"""
import sys, os, json
import numpy as np
from PIL import Image, ImageDraw
from scipy import ndimage

SHADOW_COLORS = [(0, 96, 0)]      # SF2 ground-shadow green

def build_fg(src, bg_tol=26, shadow_tol=22):
    im = Image.open(src).convert("RGBA")
    arr = np.asarray(im).astype(np.int16)
    H, W = arr.shape[:2]
    rgb = arr[:, :, :3]

    # bg palette = border-ring colors UNION globally most-frequent flat colors
    ring = np.concatenate([rgb[0], rgb[-1], rgb[:, 0], rgb[:, -1]], axis=0).reshape(-1, 3)
    cand = [tuple(int(x) for x in c) for c in ring]
    flat = rgb.reshape(-1, 3)[::17]
    colors, counts = np.unique(flat, axis=0, return_counts=True)
    order = np.argsort(-counts)
    total = flat.shape[0]
    for j in order[:8]:
        if counts[j] > 0.02 * total:                 # a big flat area == background
            cand.append(tuple(int(x) for x in colors[j]))
    palette = []
    for c in cand:
        if not any(np.abs(np.array(p) - c).sum() < 12 for p in palette):
            palette.append(c)

    bg_cand = np.zeros((H, W), bool)
    for p in palette:
        bg_cand |= (np.abs(rgb - np.array(p)).sum(axis=2) < bg_tol)
    lbl, _ = ndimage.label(bg_cand)
    border = set(np.unique(np.concatenate([lbl[0], lbl[-1], lbl[:, 0], lbl[:, -1]])))
    border.discard(0)
    bg = np.isin(lbl, list(border))

    # also treat original-alpha==0 as background (Ryu sheet has real transparency)
    bg |= (arr[:, :, 3] < 8)

    shadow = np.zeros((H, W), bool)
    for s in SHADOW_COLORS:
        shadow |= (np.abs(rgb - np.array(s)).sum(axis=2) < shadow_tol)

    fg = (~bg) & (~shadow)
    out = arr.copy()
    out[~fg, 3] = 0                      # transparent bg + shadow
    return out.astype(np.uint8), fg, palette, W, H

def runs(mask_1d, min_gap):
    """Return (start,end) runs of True separated by >= min_gap False."""
    idx = np.where(mask_1d)[0]
    if len(idx) == 0:
        return []
    splits = np.where(np.diff(idx) > min_gap)[0]
    segs, start = [], idx[0]
    prev = idx[0]
    for k in idx[1:]:
        if k - prev > min_gap:
            segs.append((start, prev + 1)); start = k
        prev = k
    segs.append((start, prev + 1))
    return segs

def extract(src, outdir, row_gap=4, col_gap=3, min_h=22, min_w=10, min_area=200):
    os.makedirs(outdir, exist_ok=True)
    out, fg, palette, W, H = build_fg(src)
    frames = []
    for (ry0, ry1) in runs(fg.any(axis=1), row_gap):
        band = fg[ry0:ry1]
        for (cx0, cx1) in runs(band.any(axis=0), col_gap):
            sub = band[:, cx0:cx1]
            rows = np.where(sub.any(axis=1))[0]
            cols = np.where(sub.any(axis=0))[0]
            y0, y1 = ry0 + rows[0], ry0 + rows[-1] + 1
            x0, x1 = cx0 + cols[0], cx0 + cols[-1] + 1
            w, h, area = x1 - x0, y1 - y0, int(sub.sum())
            if h < min_h or w < min_w or area < min_area:
                continue
            if h < 30 and w > 4 * h:      # text label
                continue
            frames.append((x0, y0, w, h, area))

    atlas = []
    for idx, (x, y, w, h, area) in enumerate(frames):
        Image.fromarray(out[y:y+h, x:x+w], "RGBA").save(os.path.join(outdir, f"f{idx:03d}.png"))
        atlas.append({"i": idx, "file": f"f{idx:03d}.png", "x": int(x), "y": int(y),
                      "w": int(w), "h": int(h), "area": int(area)})
    json.dump({"source": os.path.basename(src), "sheet_w": W, "sheet_h": H,
               "palette": palette, "frames": atlas},
              open(os.path.join(outdir, "atlas.json"), "w"), indent=1)
    print(f"{os.path.basename(src)}: palette={palette} -> {len(atlas)} frames")
    return atlas, outdir

def contact(atlas, outdir, cols=14, cell=104, out="contact.png"):
    n = len(atlas); rows = (n + cols - 1) // cols
    sheet = Image.new("RGBA", (cols * cell, rows * cell), (38, 38, 46, 255))
    d = ImageDraw.Draw(sheet)
    for k, fr in enumerate(atlas):
        r, c = divmod(k, cols); cx, cy = c * cell, r * cell
        im = Image.open(os.path.join(outdir, fr["file"])).convert("RGBA")
        s = min((cell - 20) / im.width, (cell - 20) / im.height, 1.0)
        im = im.resize((max(1, int(im.width * s)), max(1, int(im.height * s))))
        sheet.paste(im, (cx + (cell - im.width) // 2, cy + (cell - im.height) // 2 + 6), im)
        d.rectangle([cx, cy, cx + cell - 1, cy + cell - 1], outline=(88, 88, 108))
        d.text((cx + 3, cy + 2), f"{fr['i']}:{fr['w']}x{fr['h']}", fill=(255, 230, 120))
    sheet.save(out)
    print("contact ->", out, sheet.size)

if __name__ == "__main__":
    src, outdir = sys.argv[1], sys.argv[2]
    cout = sys.argv[3] if len(sys.argv) > 3 else os.path.join(outdir, "contact.png")
    atlas, od = extract(src, outdir)
    contact(atlas, od, out=cout)