#!/usr/bin/env python
"""
rig_pivots_render.py — 풀캔버스 리그 파츠 검증 + 관절 피벗 자동 산출 + 배경춤 프리뷰 렌더.

사용법:
  python rig_pivots_render.py <Profile_dir> <prefix>
예:
  python rig_pivots_render.py "D:\\...\\Noeul_Profile" noeul_part_

하는 일:
  1) <Profile>/03_Assets/Parts/Images/<prefix><bone>.png (마스터+16파츠) 검증
     - 전부 520x900 / 32-bit alpha / 16파츠 스택 == 마스터(alpha union) missed/extra 출력
  2) 관절 피벗 자동 산출: pivot(bone) = centroid(opaque(bone) & opaque(parent))
  3) <Profile>/05_Animation/dance_idle.json 로 배경춤 프레임 렌더(풀캔버스 FK)
  출력: <Profile>/04_Rig/_pivots.json (rig.json에 반영), <Profile>/04_Rig/_dance_preview.png

주의: PowerShell은 대소문자 무시 변수 충돌이 잦으니 리그 계산은 이 Python 도구를 쓸 것.
"""
import sys, os, json, math
import numpy as np
from PIL import Image

CW, CH, TH = 520, 900, 32
BONES = [  # (name, parent, z)
 ("pelvis",None,6),("chest","pelvis",8),("neck","chest",9),("head","neck",10),
 ("upperarm_r","chest",5),("forearm_r","upperarm_r",5),("hand_r","forearm_r",5),
 ("upperarm_l","chest",12),("forearm_l","upperarm_l",12),("hand_l","forearm_l",13),
 ("thigh_r","pelvis",4),("shin_r","thigh_r",3),("foot_r","shin_r",2),
 ("thigh_l","pelvis",4),("shin_l","thigh_l",3),("foot_l","shin_l",2)]

def main(profile, prefix):
    pdir = os.path.join(profile, "03_Assets", "Parts", "Images")
    anim_path = os.path.join(profile, "05_Animation", "dance_idle.json")
    rigdir = os.path.join(profile, "04_Rig")
    names = ["master_apose"] + [n for n,_,_ in BONES]
    for nm in names:
        p = os.path.join(pdir, f"{prefix}{nm}.png")
        assert os.path.exists(p), f"MISSING {p}"
        assert Image.open(p).size == (CW, CH), f"size!=520x900: {nm}"
    print("all 17 present, all 520x900: OK")

    img = {n: Image.open(os.path.join(pdir, f"{prefix}{n}.png")).convert("RGBA") for n,_,_ in BONES}
    alpha = {n: (np.asarray(img[n])[:,:,3] > TH) for n,_,_ in BONES}
    master = np.asarray(Image.open(os.path.join(pdir, f"{prefix}master_apose.png")).convert("RGBA"))[:,:,3] > TH

    pivot = {}
    for n,p,_ in BONES:
        par = p if p else "chest"
        ys,xs = np.nonzero(alpha[n] & alpha[par])
        if len(xs): pivot[n] = (round(float(xs.mean()),1), round(float(ys.mean()),1))
        else:
            ys2,xs2 = np.nonzero(alpha[n]); pivot[n] = (round(float(xs2.mean()),1), round(float(ys2.min()),1))
    print("== pivots (paste into rig.json) ==")
    for n,_,_ in BONES: print(f"  {n:12s} {pivot[n]}")

    union = np.zeros_like(master)
    for n,_,_ in BONES: union |= alpha[n]
    mo = int(master.sum())
    print(f"stack union vs master: opaque={mo} missed={int((master & ~union).sum())} extra={int((union & ~master).sum())}")

    os.makedirs(rigdir, exist_ok=True)
    json.dump({n:list(pivot[n]) for n,_,_ in BONES}, open(os.path.join(rigdir,"_pivots.json"),"w"), ensure_ascii=False, indent=2)

    anim = json.load(open(anim_path, encoding="utf-8"))
    ease = lambda u: -(math.cos(math.pi*u)-1)/2
    def samp(keys,t):
        if not keys: return 0.0
        if t<=keys[0]["t"]: return float(keys[0]["v"])
        if t>=keys[-1]["t"]: return float(keys[-1]["v"])
        for i in range(len(keys)-1):
            a,b=keys[i],keys[i+1]
            if a["t"]<=t<=b["t"]:
                sp=b["t"]-a["t"]
                return float(a["v"]) if sp<=0 else float(a["v"])+(float(b["v"])-float(a["v"]))*ease((t-a["t"])/sp)
        return 0.0
    T=lambda x,y: np.array([[1,0,x],[0,1,y],[0,0,1]],float)
    def R(d):
        r=math.radians(d);c,s=math.cos(r),math.sin(r);return np.array([[c,-s,0],[s,c,0],[0,0,1]],float)
    def render(t):
        W={}
        for n,p,_ in BONES:
            tr=anim["tracks"].get(n,{}); rot=samp(tr.get("rot"),t);tx=samp(tr.get("tx"),t);ty=samp(tr.get("ty"),t)
            jx,jy=pivot[n]; ml=T(tx,ty)@T(jx,jy)@R(rot)@T(-jx,-jy); W[n]=(W[p]@ml) if p else ml
        fr=Image.new("RGBA",(CW,CH),(0,0,0,0))
        for n,_,_ in sorted(BONES,key=lambda b:b[2]):
            Mi=np.linalg.inv(W[n]); a,c,e=Mi[0]; b,d,f=Mi[1]
            fr=Image.alpha_composite(fr,img[n].transform((CW,CH),Image.AFFINE,(a,c,e,b,d,f),resample=Image.BICUBIC))
        return fr
    times=[0.0,0.5,1.0,1.5]; sc=0.42; fw,fh=int(CW*sc),int(CH*sc)
    mont=Image.new("RGBA",(fw*len(times),fh),(28,32,40,255))
    for i,t in enumerate(times): mont.alpha_composite(render(t).resize((fw,fh),Image.BICUBIC),(fw*i,0))
    out=os.path.join(rigdir,"_dance_preview.png"); mont.convert("RGB").save(out)
    print("saved preview:", out)

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print(__doc__); sys.exit(1)
    main(sys.argv[1], sys.argv[2])