"""Export exact comparison mesh/bind/pose data for Pipeline-authored Unity assets.

Blender Z-up -> Unity Y-up is (x,z,y), so anatomical -Y faces Unity -Z.
All asset creation is subsequently performed inside Unity through Pipeline.
"""
import json
import sys
from pathlib import Path

import bpy
import numpy as np
from math import pi
from mathutils import Matrix

root=Path(sys.argv[sys.argv.index('--')+1])
folder=root/'Tools/output/campguide_studio/rig_mixamo41'
base=root/'AIClient/Assets/Bundles/Characters/Human/test/Boy_mixamo'
out=folder/'unity_review'
out.mkdir(exist_ok=True)
C=Matrix(((1,0,0,0),(0,0,1,0),(0,1,0,0),(0,0,0,1)))

def trs(m):
    p,q,s=m.decompose()
    return list(p)+[q.x,q.y,q.z,q.w]+list(s)

result={'coordinateSystem':'Unity Y-up, -Z front; positions in metres','variants':[],'bones':[],'clips':[]}
names=None
for label,path in [('Mixamo',base/'CampGuideStudio.fbx'),('TokenRigRaw',folder/'tokenrig_mixamo41.blend'),('TokenRigLocal',folder/'tokenrig_mixamo41_local.blend')]:
    bpy.ops.wm.read_factory_settings(use_empty=True)
    if path.suffix=='.blend':bpy.ops.wm.open_mainfile(filepath=str(path))
    else:bpy.ops.import_scene.fbx(filepath=str(path))
    arm=next(o for o in bpy.data.objects if o.type=='ARMATURE')
    mesh=next(o for o in bpy.data.objects if o.type=='MESH')
    pts=np.array([mesh.matrix_world@v.co for v in mesh.data.vertices])
    normalize=Matrix.Rotation(pi/2,4,'X') if np.argmax(np.ptp(pts,axis=0))==1 else Matrix.Identity(4)
    world={b.name:C@normalize@arm.matrix_world@b.matrix_local@C for b in arm.data.bones}
    if names is None:
        names=[b.name for b in arm.data.bones]
        for b in arm.data.bones:
            parent=b.parent.name if b.parent else None
            local=world[parent].inverted()@world[b.name] if parent else world[b.name]
            result['bones'].append({'name':b.name,'parent':names.index(parent) if parent else -1,
                'rest':trs(local),'bindpose':[list(row) for row in world[b.name].inverted()]})
    assert set(names)==set(arm.data.bones.keys())
    weights=[]
    for v in mesh.data.vertices:
        w={mesh.vertex_groups[g.group].name:g.weight for g in v.groups if g.weight>1e-8}
        total=sum(w.values())
        weights.append([[names.index(n),value/total] for n,value in sorted(w.items(),key=lambda x:-x[1])])
    transform=C@normalize@mesh.matrix_world
    normal_transform=transform.to_3x3().inverted().transposed()
    mesh.data.calc_loop_triangles()
    uv_layer=mesh.data.uv_layers.active
    vertices=[];normals=[];uv=[];influences=[];triangles=[];lookup={};first={}
    for triangle in mesh.data.loop_triangles:
        ids=[]
        for loop_index in triangle.loops:
            loop=mesh.data.loops[loop_index]
            vi=loop.vertex_index
            n=normal_transform@mesh.data.corner_normals[loop_index].vector
            n.normalize()
            tex=list(uv_layer.data[loop_index].uv) if uv_layer else [0,0]
            key=(vi,*[round(x,6) for x in tex],*[round(x,6) for x in n])
            if key not in lookup:
                index=len(vertices);lookup[key]=index;first.setdefault(vi,index)
                vertices.append(list(transform@mesh.data.vertices[vi].co));normals.append(list(n));uv.append(tex);influences.append(weights[vi])
            ids.append(lookup[key])
        triangles.extend([ids[0],ids[2],ids[1]])
    # A reflection reverses the winding; the reversed triangle order above is
    # checked visually in Unity and does not alter the comparison edge lengths.
    raw=np.array([normalize@mesh.matrix_world@v.co for v in mesh.data.vertices])
    height=float(np.ptp(raw[:,2]));z=(raw[:,2]-raw[:,2].min())/height
    hand=(np.abs(raw[:,0])>height*.255)&(z>.55)&(z<.67)
    neck=(z>.67)&(z<.78)&(np.abs(raw[:,0])<height*.12)
    head=z>.78
    edges=[]
    for edge in mesh.data.edges:
        a,b=edge.vertices
        length=float(np.linalg.norm(raw[a]-raw[b]))
        if length<=height*.002:continue
        mask=1|(2 if hand[a] and hand[b] else 0)|(4 if neck[a] and neck[b] else 0)|(8 if head[a] and head[b] else 0)
        edges.append([first[a],first[b],length,mask])
    result['variants'].append({'name':label,'vertices':vertices,'normals':normals,'uv':uv,'triangles':triangles,
        'weights':influences,'auditEdges':edges,'originalVertices':len(mesh.data.vertices),'trianglesCount':len(triangles)//3})
    print('MESH',label,len(vertices),len(triangles)//3)

for label in ['Walking','Breathing Idle','Happy Idle']:
    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_scene.fbx(filepath=str(base/('CampGuideStudio@'+label+'.fbx')))
    arm=next(o for o in bpy.data.objects if o.type=='ARMATURE')
    first,last=map(int,arm.animation_data.action.frame_range)
    samples=[];previous={}
    for frame in range(first,last+1):
        bpy.context.scene.frame_set(frame);bpy.context.view_layer.update()
        world={n:C@arm.matrix_world@arm.pose.bones[n].matrix@C for n in names}
        sample=[]
        for i,n in enumerate(names):
            parent=result['bones'][i]['parent']
            local=world[names[parent]].inverted()@world[n] if parent>=0 else world[n]
            values=trs(local)
            if n in previous and sum(a*b for a,b in zip(previous[n],values[3:7]))<0:
                values[3:7]=[-x for x in values[3:7]]
            previous[n]=values[3:7]
            sample.append(values)
        samples.append(sample)
    result['clips'].append({'name':label.replace(' ','_'),'fps':30,'samples':samples})
    print('CLIP',label,len(samples))
(out/'review_data.json').write_text(json.dumps(result,separators=(',',':')),encoding='utf-8')
print('EXPORTED',out/'review_data.json')