"""Cut visible surfaces from the NEW generated master, never source frames.

The atlas supplies concealed underlap. Shared generated pixels make the rest
pose coherent; semantic ownership and mapping below are explicitly authored.
"""
from pathlib import Path
import copy
import sys
import cv2
import numpy as np
from PIL import Image,ImageDraw
from scipy.ndimage import distance_transform_edt,binary_fill_holes
sys.path.insert(0,str(Path(__file__).resolve().parents[2]))
from livideo2d.core import ROOT,read_json,write_json,digest,validate_scene
from livideo2d.topology import edge_mesh

BASE=ROOT/'experiments/flow_v6';ART=ROOT/'artifacts/flow_v6';ASSET=ROOT/'assets/generated_v6'
# Coordinates on the generated master resized to height 720, not source-video UVs.
POLYGONS={
 'coat_left':[(333,476),(378,464),(401,491),(388,548),(375,594),(336,720),(90,720),(120,650),(236,606)],
 'coat_right':[(466,469),(509,469),(605,547),(640,587),(567,572),(473,550)],
 'leg_far':[(354,612),(408,611),(456,637),(567,720),(327,720)],
 'leg_near':[(466,534),(519,544),(688,593),(704,650),(640,720),(560,720),(586,658),(474,652)],
 'front_cloth':[(388,510),(465,512),(467,548),(488,574),(532,636),(568,720),(539,720),(463,653),(418,612),(386,585)],
 'lapel_left':[(343,286),(398,292),(382,340),(398,398),(375,407),(349,365)],
 'lapel_right':[(422,285),(460,280),(466,326),(516,334),(500,381),(499,423),(485,455),(464,455),(473,394),(463,361),(443,322)],
 'forearm_right':[(470,508),(488,507),(507,539),(531,561),(503,562),(480,545)],
 'sleeve_right':[(481,396),(504,442),(527,498),(481,526),(463,503),(468,470)],
 'sleeve_left':[(319,297),(370,302),(373,351),(391,390),(376,420),(342,453),(310,459),(285,429),(294,398),(301,350)],
 'forearm_left':[(340,429),(366,423),(399,369),(408,344),(433,322),(456,317),(469,328),(455,345),(455,358),(441,372),(423,393),(406,452),(392,491),(375,496),(360,478)],
 'hand_left':[(397,377),(423,383),(449,368),(470,340),(462,314),(437,315),(407,341),(398,366)],
 'badge':[(333,324),(373,337),(395,346),(395,376),(347,374),(330,360)],
}
HAIR={
 'hair_back_left':[(294,245),(333,260),(351,281),(324,293),(318,305),(319,317),(312,325),(312,350),(293,370),(251,370),(238,328),(275,315),(286,289)],
 'hair_back_right':[(437,250),(457,248),(459,220),(476,219),(484,232),(512,245),(544,250),(566,231),(573,262),(588,280),(588,310),(530,311),(497,307),(483,307),(479,314),(466,314),(465,282),(443,278),(422,271),(421,259)],
 'hair_sweep_right':[(472,229),(507,244),(545,249),(566,236),(570,256),(552,269),(520,266),(486,248)],
 'hair_temple':[(420,245),(464,220),(481,220),(483,265),(494,279),(492,308),(480,320),(420,310)],
}

def poly_mask(size,points):
    im=Image.new('L',size);ImageDraw.Draw(im).polygon(points,fill=255);return np.asarray(im)>0

def semantic_hair_masks(image):
    """Authored material domains come first; hue only cleans their boundary.

    Gold lapels, badges, skin and the earring share hair-like colors. A broad
    color window cannot distinguish them and caused floating animated shards.
    """
    rgba=np.array(image);hsv=cv2.cvtColor(rgba[:,:,:3],cv2.COLOR_RGB2HSV)
    orange=(hsv[:,:,0]<30)&(hsv[:,:,1]>35)&(rgba[:,:,0].astype(int)-rgba[:,:,1]>12)
    domains={k:poly_mask(image.size,v) for k,v in HAIR.items()}
    domains['hair_temple'] &= domains['hair_back_right']
    domains['hair_sweep_right'] &= domains['hair_back_right']
    masks={}
    for name,domain in domains.items():
        support=cv2.morphologyEx(np.uint8(domain&orange&(rgba[:,:,3]>3)),cv2.MORPH_CLOSE,np.ones((3,3),np.uint8))
        masks[name]=binary_fill_holes(support)&domain
    # Material domain covers every native-alpha strand tip. Partition that union
    # once, so the torso cannot inherit unassigned static hair slivers.
    masks['hair_temple'] &= masks['hair_back_right']
    masks['hair_sweep_right'] &= masks['hair_back_right']&~masks['hair_temple']
    masks['hair_back_right'] &= ~(masks['hair_sweep_right']|masks['hair_temple'])
    # This lower curl attaches through the temple lock, not the wide rear blade.
    # Keep its entire generated-alpha component with that attachment.
    _,components=cv2.connectedComponents(np.uint8(masks['hair_back_right']&(rgba[:,:,3]>3)))
    areas=np.bincount(components.ravel())
    if len(areas)>1:
        rear_label=1+int(areas[1:].argmax())
        attached_scraps=(components>0)&(components!=rear_label)
        masks['hair_temple'] |= attached_scraps
        masks['hair_back_right'] &= ~attached_scraps
    return masks

def new_layer(name,image,mask,bone):
    rgba=np.array(image);rgba[:,:,3]*=mask.astype(np.uint8)
    y,x=np.where(rgba[:,:,3]>0);box=[max(0,int(x.min())-2),max(0,int(y.min())-2),min(image.width,int(x.max())+3),min(image.height,int(y.max())+3)]
    tile=Image.fromarray(rgba[box[1]:box[3],box[0]:box[2]])
    folder=ASSET/'surface';folder.mkdir(exist_ok=True);path=folder/f'{name}.png';tile.save(path)
    cache=ART/'mesh_cache'/f'surface-{name}-{digest(path)[:12]}.json'
    if cache.exists():mesh=read_json(cache)
    else:
        mesh=edge_mesh([tile],spacing=15,max_size=400);write_json(cache,mesh)
    mesh=copy.deepcopy(mesh);uv=np.array(mesh.pop('uv'));stats=mesh.pop('stats')
    xy=uv*np.array([tile.width,tile.height])+box[:2]
    # Two authored pose anchors: neck and nearer knee. No pixel-error search.
    xy[:,0]=xy[:,0]*.936+76.25;xy[:,1]=xy[:,1]-4+np.maximum(0,xy[:,1]-620)*.16
    mesh['vertices']=np.c_[xy,uv].tolist();mesh['active_triangles']=list(range(len(mesh['triangles'])))
    layer={'id':name,'image':f'../../../assets/generated_v6/surface/{name}.png','type':'mesh','mesh':mesh,
      'skin':{'bones':[bone],'weights':[[1.] for _ in uv]},'topology':{'method':'generated alpha + Canny + constrained triangulation',**stats}}
    return layer,{'id':name,'image':str(path.relative_to(ROOT)),'generated_master_crop':box,'sha256':digest(path),'bone':bone}

def main():
    raw=Image.open(ASSET/'raw/canonical.png').convert('RGBA');image=raw.resize((round(raw.width*720/raw.height),720),Image.Resampling.LANCZOS)
    rgba=np.array(image);h,w=rgba.shape[:2]
    yy,xx=np.mgrid[:h,:w]
    # Native alpha already defines the outside; ownership only separates generated parts.
    hair_masks=semantic_hair_masks(image)
    neck=poly_mask(image.size,[(374,247),(390,254),(409,259),(422,256),(423,269),(447,275),(464,278),(472,304),(458,318),(381,302),(371,286)])
    body=(yy>=285)|neck
    for mask in hair_masks.values():body &= ~mask
    labels=np.zeros((h,w),np.uint8);ids=['torso_core',*POLYGONS]
    for i,(name,polygon) in enumerate(POLYGONS.items(),1):labels[poly_mask(image.size,polygon)]=i
    scene=read_json(BASE/'scenes/00_assembly.json');old={l['id']:l for l in scene['layers']};fresh={};records=[];underlaps=[]
    for i,name in enumerate(ids):
        bone={'leg_near':'pelvis','leg_far':'pelvis','forearm_right':'fore_right','forearm_left':'fore_left','hand_left':'hand_left',
          'sleeve_left':'upper_left','sleeve_right':'upper_right'}.get(name,'torso')
        fresh[name],record=new_layer(name,image,body&(labels==i),bone);records.append(record)
        # Concealed padding extends already-generated material at an internal cut.
        # It is drawn behind ALL visible surfaces, so the first frame remains exact.
        visible=body&(labels==i)&(rgba[:,:,3]>16)
        distance,indices=distance_transform_edt(~visible,return_indices=True)
        radius=38 if name in ['leg_near','leg_far','torso_core'] else 10
        mask=(distance<radius)&(distance>0)&body&(rgba[:,:,3]>16)
        if mask.any():
            extended=rgba[indices[0],indices[1]].copy();extended[:,:,3]*=mask
            padding,padding_record=new_layer(name+'_underlap',Image.fromarray(extended),mask,bone)
            padding_record['operation']='Concealed nearest-color padding of generated surface; drawn behind every visible surface'
            records.append(padding_record)
            underlaps.append(padding)
    hair_padding=[]
    for name,mask in hair_masks.items():
        fresh[name],record=new_layer(name,image,mask,'head');records.append(record)
        distance,indices=distance_transform_edt(~(mask&(rgba[:,:,3]>16)),return_indices=True)
        root=[321,273] if name=='hair_back_left' else [464,264]
        near_root=((xx-root[0])/38)**2+((yy-root[1])/38)**2<1
        overlap=(distance>0)&(distance<20)&(rgba[:,:,3]>16)&near_root
        if overlap.any():
            extended=rgba[indices[0],indices[1]].copy();extended[:,:,3]*=overlap
            padding,_=new_layer(name+'_root_underlap',Image.fromarray(extended),overlap,'head')
            hair_padding.append(padding)
    # Regenerated atlas supplies the large torso area concealed by the raised arm.
    back=copy.deepcopy(old['torso_core']);back['id']='torso_underlap'
    im=Image.open(ASSET/'parts/torso_core.png').convert('RGBA');a=np.array(im);u=np.arange(im.width)/im.width
    a[:,:,3]*=((400+u*139>=424)&(400+u*139<=509))[None,:]
    a[:70,:,3]=0
    (ASSET/'surface').mkdir(exist_ok=True);Image.fromarray(a).save(ASSET/'surface/torso_underlap.png')
    back['image']='../../../assets/generated_v6/surface/torso_underlap.png'
    # Atlas roots remain concealed under visible generated surfaces.
    order=['background','tail','hair_back_left','hair_back_right','hair_sweep_right','hair_temple','forearm_right',
      'leg_far','leg_near','coat_right','torso_underlap','torso_core','front_cloth','coat_left',
      'lapel_left','lapel_right','sleeve_right','sleeve_left','forearm_left','hand_left','badge',
      'face','ear','hood','fringe_left','bangs','fringe_right','crown','tags']
    fresh['torso_underlap']=back
    scene['layers']=[fresh[n] if n in fresh else old[n] for n in order]
    # Root padding was rejected visually: its extrapolated edge became exposed
    # during the large swing. Complete copper surfaces attach behind the head.
    at=next(i for i,l in enumerate(scene['layers']) if l['id']=='forearm_right')
    scene['layers'][at:at]=underlaps
    from experiments.flow_v6.weights import body_weights
    for layer in scene['layers']:
        if layer['id'] in ids or (layer['id'].endswith('_underlap') and not layer['id'].startswith('hair_')):
            layer['skin']=body_weights(np.array(layer['mesh']['vertices'])[:,:2])
    scene['notes']['visible_surface']='Semantic polygons split newly generated canonical artwork; original video pixels are never cut or rendered. Atlas material supplies hidden torso underlap.'
    scene['notes']['rejected_padding']='Hair root nearest-color padding is not in the selected scene; exposed edges looked blocky under large swings.'
    write_json(BASE/'surface_parts.json',records)
    write_json(BASE/'scenes/01_surface.json',validate_scene(scene))
    print('Built coherent generated surfaces:',len(scene['layers']))

if __name__=='__main__':main()