import numpy as np, os, sys
from PIL import Image
T=os.environ["TMP"]
d=np.load(T+"/bl/tuna_tris.npz"); UV=d["uv"]; CO=d["co"]
img=np.asarray(Image.open(T+"/tunatex/img0.jpg").convert("RGB")).astype(np.float32)/255; H,W,_=img.shape
def sample(uv): 
    x=np.clip((uv[...,0]*W).astype(int),0,W-1); y=np.clip(((1-uv[...,1])*H).astype(int),0,H-1); return img[y,x]
def hsv(c):
    mx=c.max(-1); mn=c.min(-1); d=mx-mn+1e-6; r,g,b=c[...,0],c[...,1],c[...,2]
    h=np.where(mx==r,((g-b)/d)%6,np.where(mx==g,(b-r)/d+2,(r-g)/d+4))*60; s=d/(mx+1e-6); return h,s,mx
# classify vertices by texture colour (sample slightly inside the triangle)
cen=UV.mean(1,keepdims=True); c=sample(UV*0.8+cen*0.2); h,s,v=hsv(c)
yellow=(h>30)&(h<62)&(s>0.45)&(v>0.45); black=(v<0.22); purple=(h>265)&(h<335)&(s>0.25)
fur=(yellow|black)
body_tri=fur.all(1)
P=CO[body_tri].reshape(-1,3); ctr=P.mean(0)
eye=CO[purple.any(1)].reshape(-1,3); head=np.median(eye,0)
tail=P[np.argmax(np.linalg.norm(P-head,axis=1))]
ax=(tail-head)/np.linalg.norm(tail-head)
pr=(P-head)@ax; rear=P[pr>np.percentile(pr,88)]; tail=rear.mean(0)
mid=P[(pr>np.percentile(pr,45))&(pr<np.percentile(pr,55))].mean(0)
ax=(tail-head)/np.linalg.norm(tail-head); ctr=head
print("rear centroid",tail.round(2),"mid",mid.round(2))
print("head",head.round(2),"tail",tail.round(2),"purple tris",len(eye)//3)
tp=(P-ctr)@ax; t0,t1=np.percentile(tp,1),np.percentile(tp,99.5)
te=(eye-ctr)@ax; near=np.linalg.norm(eye-head,axis=1)<0.45; neck=np.percentile(te[near],95); print('eye t pct',np.percentile(te,[10,50,90]).round(2))
print("axis",ax.round(3),"t range",round(t0,3),round(t1,3),"eyes up to",round(neck,3))
# palette from the original
Y=sample(UV[yellow.all(1)]).reshape(-1,3); Kc=sample(UV[black.all(1)]).reshape(-1,3)
ycol=np.median(Y,0); kcol=np.median(Kc,0)*0.9; print("yellow",ycol.round(2),"black",kcol.round(2))
NB=float(sys.argv[1]) if len(sys.argv)>1 else 3.0     # number of black rings behind the head
start=neck+0.04
out=img.copy(); mask=np.zeros((H,W),bool)
lum=img.mean(-1)
RMAX=float(sys.argv[2]) if len(sys.argv)>2 else 0.5
allc=CO[body_tri].mean(1); rr=np.linalg.norm((allc-ctr)-np.outer((allc-ctr)@ax,ax),axis=1)
print('radial pct',np.percentile(rr,[50,90,97,99]).round(2))
CORE=float(sys.argv[3]) if len(sys.argv)>3 else 0.47
cc_=sample(UV.mean(1)); hh,ss,vv=hsv(cc_)
redtuft=((hh<22)|(hh>335))&(ss>0.5)&(vv>0.35)
for k in range(len(CO)):
    co=CO[k]; tt=(co-ctr)@ax
    cm=co.mean(0)-ctr; r=np.linalg.norm(cm-(cm@ax)*ax)
    if r>RMAX: continue
    rad=(cm-(cm@ax)*ax)/max(r,1e-6); down=-rad[2]
    leglike=(hh[k]<35)&(ss[k]>0.35)&(vv[k]<0.62)&(not yellow[k].all())
    winglike=(ss[k]<0.2)&(vv[k]>0.6)
    if winglike or (leglike and down>0.35 and r>CORE): continue
    if redtuft[k] and (cm@ax)>0.75*t1: continue
    if tt.max()<start: continue
    uv=UV[k]*[W,H]; uv[:,1]=H-uv[:,1]
    x0,y0=np.floor(uv.min(0)).astype(int); x1,y1=np.ceil(uv.max(0)).astype(int)+1
    x0=max(x0-1,0); y0=max(y0-1,0); x1=min(x1+1,W); y1=min(y1+1,H)
    xs,ys=np.meshgrid(np.arange(x0,x1)+0.5,np.arange(y0,y1)+0.5)
    a,b,cc=uv; den=(b[1]-cc[1])*(a[0]-cc[0])+(cc[0]-b[0])*(a[1]-cc[1])
    if abs(den)<1e-9: continue
    l1=((b[1]-cc[1])*(xs-cc[0])+(cc[0]-b[0])*(ys-cc[1]))/den
    l2=((cc[1]-a[1])*(xs-cc[0])+(a[0]-cc[0])*(ys-cc[1]))/den; l3=1-l1-l2
    ins=(l1>=-0.02)&(l2>=-0.02)&(l3>=-0.02)
    if not ins.any(): continue
    t=l1*tt[0]+l2*tt[1]+l3*tt[2]
    u=(t-start)/(t1-start)                     # 0 at neck, 1 at tail
    ph=u*NB*2                                  # 2 bands per ring
    frac=ph%2
    # soft but crisp edge (1.5% of a band)
    e=0.06; w=np.clip((np.minimum(frac,2-frac)-0.0)/e,0,1)   # unused smoothing helper
    isblack=(frac<1.0)
    edge=np.minimum(np.abs(frac-1.0),np.minimum(frac,2-frac))
    soft=np.clip(edge/0.05,0,1)
    base=np.where(isblack[...,None],kcol,ycol)
    other=np.where(isblack[...,None],ycol,kcol)
    col=base*(0.5+0.5*soft[...,None])+other*(0.5-0.5*soft[...,None])
    col=np.where((u<0)[...,None],img[y0:y1,x0:x1],col)
    det=np.clip(lum[y0:y1,x0:x1]/max(lum.mean(),1e-3),0.6,1.35)**0.2
    col=np.clip(col*det[...,None],0,1)
    sl=out[y0:y1,x0:x1]; sl[ins]=col[ins]; mask[y0:y1,x0:x1]|=ins
Image.fromarray((out*255).astype(np.uint8)).save(T+"/tunatex/striped.jpg",quality=93)
Image.fromarray((mask*255).astype(np.uint8)).resize((512,512)).save(T+"/tunatex/mask.png")
print("painted texels",int(mask.sum()))