"""Agent-authored first-frame assembly. Generated images only; V5 tail intact."""
from pathlib import Path
import copy
import math
import sys
import numpy as np
from PIL import Image
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
from livideo2d.rig import bone_matrices
from livideo2d.deform import smooth_curve
from experiments.pose_v5.retarget import solve_two_link,angle_of,rotation_of

BASE=ROOT/'experiments/flow_v6';ART=ROOT/'artifacts/flow_v6'
PARTS=ROOT/'assets/generated_v6/parts'
PLACEMENTS={
 'torso_core':('torso',[400,262,139,339]),
 'lapel_left':('torso',[354,282,85,270]),
 'lapel_right':('torso',[493,282,71,262]),
 'coat_left':('torso',[169,477,274,303]),
 'coat_right':('torso',[511,465,169,196]),
 'coat_back':('pelvis',[153,570,271,85]),
 'front_cloth':('torso',[442,508,158,247]),
 'leg_near':('pelvis',[482,553,223,211]),
 'leg_far':('pelvis',[382,609,236,174]),
 'sleeve_left':('upper_left',[337,302,100,164]),
 'sleeve_right':('upper_right',[511,321,55,210]),
 'forearm_right':('fore_right',[521,490,86,115]),
 'hair_back_left':('head',[314,210,77,162]),
 'hair_back_right':('head',[498,177,119,127]),
 'hair_sweep_right':('head',[507,189,96,67]),
 'lock_right':('head',[496,156,41,156]),
}

def part(name):
    bone,box=PLACEMENTS[name];path=PARTS/f'{name}.png'
    cache=ART/'mesh_cache'/f'{name}-{digest(path)[:12]}.json'
    if cache.exists():mesh=read_json(cache)
    else:
        mesh=edge_mesh([Image.open(path)],spacing=18,max_size=400)
        write_json(cache,mesh)
    mesh=copy.deepcopy(mesh);uv=np.asarray(mesh.pop('uv'));stats=mesh.pop('stats')
    mesh['vertices']=np.c_[np.array(box[:2])+uv*np.array(box[2:]),uv].tolist()
    mesh['active_triangles']=list(range(len(mesh['triangles'])))
    return {'id':name,'image':f'../../../assets/generated_v6/parts/{name}.png','type':'mesh',
      'mesh':mesh,'skin':{'bones':[bone],'weights':[[1.] for _ in uv]},
      'topology':{'method':'alpha + Canny + constrained triangulation',**stats}}

def move_similarity(layer,src0,src1,dst0,dst1):
    a,b,c,d=map(lambda v:np.asarray(v,float),[src0,src1,dst0,dst1])
    angle=angle_of(d-c)-angle_of(b-a);scale=np.linalg.norm(d-c)/np.linalg.norm(b-a)
    rotation=np.array([[math.cos(angle),-math.sin(angle)],[math.sin(angle),math.cos(angle)]])*scale
    v=np.array(layer['mesh']['vertices']);v[:,:2]=(v[:,:2]-a)@rotation.T+c;layer['mesh']['vertices']=v.tolist()

def compile_arms(scene):
    bones={b['id']:b for b in scene['rig']['bones']};old={}
    for b in bones.values():old[b['id']]=np.asarray(b['bind'])+old.get(b.get('parent'),np.zeros(2))
    positions={k:v.copy() for k,v in old.items()}
    positions.update(upper_left=np.array([382.,325.]),fore_left=np.array([417.,480.]),hand_left=np.array([459.,379.]),
      upper_right=np.array([520.,326.]),fore_right=np.array([530.,472.]),hand_right=np.array([548.,540.]))
    for side in ['left','right']:
        for b in bones.values():
            if b.get('parent')=='hand_'+side:positions[b['id']]+=positions['hand_'+side]-old['hand_'+side]
    for b in bones.values():b['bind']=(positions[b['id']]-positions.get(b.get('parent'),np.zeros(2))).tolist()
    events=read_json(ROOT/'experiments/pose_v5/authored_motion.json')['events']
    for f in np.arange(0,99.001,.25):
        _,poses=bone_matrices(scene,float(f));parent=poses['torso']
        for side in ['left','right']:
            upper,fore,hand=[bones[n+'_'+side] for n in ['upper','fore','hand']]
            root=(parent@np.r_[upper['bind'],1.])[:2]
            target=smooth_curve([[e['frame'],*(positions[hand['id']]+np.array(e['points'][hand['id']])-events[0]['points'][hand['id']])] for e in events],f)
            direction=positions[hand['id']]-positions[upper['id']];bend=positions[fore['id']]-positions[upper['id']]
            sign=1 if direction[0]*bend[1]-direction[1]*bend[0]>0 else -1
            elbow,wrist,error=solve_two_link(root,target,np.linalg.norm(fore['bind']),np.linalg.norm(hand['bind']),sign)
            aa=angle_of(elbow-root)-angle_of(fore['bind']);bb=angle_of(wrist-elbow)-angle_of(hand['bind'])
            for b,angle in [(upper,aa-rotation_of(parent)),(fore,bb-aa),(hand,rotation_of(parent)-bb)]:
                if f==0:b['keys']=[]
                b['keys'].append([float(f),0,0,angle,1,1])
    for name in ['upper_left','fore_left','hand_left','upper_right','fore_right','hand_right']:
        k=np.array(bones[name]['keys']);k[:,3]=np.unwrap(k[:,3]);bones[name]['keys']=k.tolist()
    return positions

def main():
    scene=read_json(ROOT/'experiments/pose_v5/scenes/04_final.json')
    old={l['id']:l for l in scene['layers']}
    for l in old.values():
        if l['id']!='tail':l.pop('brushes',None)
    compile_arms(scene)
    # Approved face, hood, ear, bangs, hands and generated background remain available.
    for name in ['forearm_left']:
        move_similarity(old[name],[401,445],[453,379],[417,480],[459,379])
    for name in ['palm_left','thumb_left','index_left','middle_left','ring_left','little_left','watch_band','watch_face']:
        move_similarity(old[name],[453,379],[453,329],[459,379],[459,329])
    for name,delta in [('papers',[11,6]),('badge',[8,1]),('bow',[5,22])]:
        v=np.array(old[name]['mesh']['vertices']);v[:,:2]+=delta;old[name]['mesh']['vertices']=v.tolist()
    # Left waist ribbon lies behind the body along the seat; right bow stays at cuff.
    move_similarity(old['ribbon_left'],[408,494],[375,626],[384,555],[263,580])
    move_similarity(old['ribbon_right'],[550,481],[630,618],[551,496],[592,551])
    fresh={n:part(n) for n in PLACEMENTS}
    order=['background','tail','coat_back','ribbon_left','leg_far','coat_right','forearm_right','leg_near',
      'hair_back_left','hair_back_right','hair_sweep_right','sleeve_right','front_cloth','torso_core',
      'coat_left','lapel_left','lapel_right','ribbon_right','sleeve_left','forearm_left','watch_band','palm_left',
      'thumb_left','index_left','middle_left','ring_left','little_left','papers','badge','watch_face','bow',
      'face','ear','hood','fringe_left','bangs','fringe_right','lock_right','crown','tags']
    scene['layers']=[fresh[n] if n in fresh else old[n] for n in order]
    scene['name']='V6 · seated costume and flow-reviewed hair / cloth'
    scene['notes']={'texture_policy':'Only generated V3/V6 imagery; no source video cutouts',
      'first_frame':'New source-faithful seated costume, visible bent legs and four broad hair layers; no added long skirt or stockings',
      'tail':'Exact V5 tail layer and pelvis animation retained',
      'motion_policy':'V5 main events, regenerated arm bind, flow evidence reviewed separately for secondary keys',
      'known_limits':['Planar face turn and imperfect source/illustration correspondence remain.']}
    write_json(BASE/'placements.json',PLACEMENTS)
    write_json(BASE/'scenes/00_assembly.json',validate_scene(scene))
    print('Assembled',len(scene['layers']),'layers;',sum(len(l['mesh']['vertices']) for l in scene['layers']),'vertices')

if __name__=='__main__':main()