#!/usr/bin/env python3
"""Render a labeled montage of specific extracted frames for pose identification,
and flag frames that likely contain baked-in text labels.

Usage:
  montage.py <char> <out.png> [i0-i1 | i,i,i ...]     # default: all frames
  e.g. montage.py ryu /tmp/m.png 4-45
Text flag heuristic: a short horizontal run of near-white/near-yellow pixels that
sits detached (a gap of transparent rows) from the main sprite mass -> probable label.
"""
import sys, os, json, re
import numpy as np
from PIL import Image, ImageDraw

def parse_sel(args):
    out = []
    for a in args:
        m = re.match(r'^(\d+)-(\d+)$', a)
        if m: out += list(range(int(m.group(1)), int(m.group(2))+1))
        elif ',' in a: out += [int(x) for x in a.split(',') if x!='']
        else: out.append(int(a))
    return out

def text_flag(im):
    """Return True if the frame likely has a detached text-label band."""
    a = np.asarray(im.convert("RGBA"))
    alpha = a[:,:,3] > 20
    if alpha.sum() < 30: return False
    rgb = a[:,:,:3].astype(int)
    # near-white or near-yellow, opaque
    whiteish = (rgb.min(axis=2) > 170) & alpha
    yellowish = (rgb[:,:,0]>190)&(rgb[:,:,1]>190)&(rgb[:,:,2]<140)&alpha
    textish = whiteish | yellowish
    rows_any = alpha.any(axis=1)
    ys = np.where(rows_any)[0]
    if len(ys) < 5: return False
    # find transparent gaps splitting the sprite vertically
    gaps = []
    run = 0
    for y in range(ys[0], ys[-1]+1):
        if not rows_any[y]: run += 1
        else:
            if run >= 3: gaps.append(y)
            run = 0
    # a detached lower band that is mostly text-colored & thin
    for gy in gaps:
        band = slice(gy, ys[-1]+1)
        bmask = alpha[band]
        if bmask.sum() == 0: continue
        h = ys[-1]+1 - gy
        frac_text = textish[band].sum() / max(1, bmask.sum())
        if h < 22 and frac_text > 0.45:
            return True
    return False

def main():
    char, out = sys.argv[1], sys.argv[2]
    d = f"public/sprites/{char}"
    atlas = json.load(open(d+"/atlas.json"))
    meta = {f["i"]: f for f in atlas["frames"]}
    sel = parse_sel(sys.argv[3:]) if len(sys.argv) > 3 else [f["i"] for f in atlas["frames"]]
    cols = 10
    cell = 150
    rows = (len(sel)+cols-1)//cols
    sheet = Image.new("RGBA", (cols*cell, rows*cell), (32,32,40,255))
    dr = ImageDraw.Draw(sheet)
    flags = []
    for k, idx in enumerate(sel):
        if idx not in meta: continue
        r, c = divmod(k, cols); cx, cy = c*cell, r*cell
        im = Image.open(f"{d}/f{idx:03d}.png").convert("RGBA")
        flagged = text_flag(im)
        if flagged: flags.append(idx)
        s = min((cell-30)/im.width, (cell-30)/im.height, 1.6)
        rim = im.resize((max(1,int(im.width*s)), max(1,int(im.height*s))))
        # checkerboard behind for transparency clarity
        for yy in range(0, cell, 12):
            for xx in range(0, cell, 12):
                if ((xx//12)+(yy//12))%2==0:
                    dr.rectangle([cx+xx,cy+yy,cx+xx+11,cy+yy+11], fill=(48,48,58))
        sheet.paste(rim, (cx+(cell-rim.width)//2, cy+(cell-rim.height)//2+8), rim)
        col = (255,90,90) if flagged else (255,230,120)
        dr.rectangle([cx,cy,cx+cell-1,cy+cell-1], outline=(90,90,110))
        dr.text((cx+3,cy+2), f"{idx} {meta[idx]['w']}x{meta[idx]['h']}"+(" TXT" if flagged else ""), fill=col)
    sheet.save(out)
    print("montage ->", out, sheet.size)
    print("text-flagged frames:", flags)

if __name__ == "__main__":
    main()