"""Inspect/export the user's Mixamo LmGirl without changing its skin weights."""
import json
import sys
from pathlib import Path
import math
import bpy
import numpy as np
from mathutils import Matrix, Vector, Quaternion

ROOT=Path(__file__).resolve().parents[2]
OUT=ROOT/'Tools/output/lmgirl_tpose_studio/rig_mixamo41/website'
SOURCE=ROOT/'AIClient/Assets/Bundles/Characters/Human/test/girl/LmGirlTpose.fbx'

def inspect():
    bpy.ops.wm.read_factory_settings(use_empty=True);bpy.ops.import_scene.fbx(filepath=str(SOURCE))
    arm=next(o for o in bpy.context.scene.objects if o.type=='ARMATURE');mesh=next(o for o in bpy.context.scene.objects if o.type=='MESH')
    p=np.array([mesh.matrix_world@v.co for v in mesh.data.vertices])
    normalize=Matrix.Rotation(math.pi/2,4,'X') if np.argmax(np.ptp(p,axis=0))==1 else Matrix.Identity(4)
    p=np.array([normalize@mesh.matrix_world@v.co for v in mesh.data.vertices])
    info={'source':str(SOURCE),'vertices':len(p),'triangles':sum(len(x.vertices)-2 for x in mesh.data.polygons),
          'bounds':[p.min(0).tolist(),p.max(0).tolist()],
          'bones':{b.name:{'parent':b.parent.name if b.parent else None,'matrix':[list(r) for r in normalize@arm.matrix_world@b.matrix_local]} for b in arm.data.bones},
          'maxInfluences':max(len(v.groups) for v in mesh.data.vertices),'hasAnimation':bool(arm.animation_data and arm.animation_data.action)}
    OUT.mkdir(parents=True,exist_ok=True);(OUT/'inspection.json').write_text(json.dumps(info,indent=2),encoding='utf-8')
    print(json.dumps({k:v for k,v in info.items() if k!='bones'}));print('BONES',len(info['bones']),list(info['bones']))

def export():
    from mathutils.kdtree import KDTree
    inspect()
    arm=next(o for o in bpy.context.scene.objects if o.type=='ARMATURE');mesh=next(o for o in bpy.context.scene.objects if o.type=='MESH')
    arm.animation_data_clear()
    points=np.array([mesh.matrix_world@v.co for v in mesh.data.vertices])
    normalize=Matrix.Rotation(math.pi/2,4,'X') if np.argmax(np.ptp(points,axis=0))==1 else Matrix.Identity(4)
    C=Matrix(((1,0,0,0),(0,0,1,0),(0,1,0,0),(0,0,0,1)))
    transform=C@normalize@mesh.matrix_world
    points=np.array([transform@v.co for v in mesh.data.vertices])
    names=[b.name for b in arm.data.bones]
    world={b.name:C@normalize@arm.matrix_world@b.matrix_local@C for b in arm.data.bones}
    source=json.loads((ROOT/'Tools/output/campguide_studio/rig_mixamo41/unity_review/review_data.json').read_text())
    source_index={b['name']:i for i,b in enumerate(source['bones'])}
    source_world={b['name']:Matrix(b['bindpose']).inverted() for b in source['bones']}
    assert set(names).issubset(source_index)
    bones=[];offsets=[];delta_rotations=[];corrections=[]
    def trs(m):
        p,q,s=m.decompose();return list(p)+[q.x,q.y,q.z,q.w]+list(s)
    def xyzw(q):return [q.x,q.y,q.z,q.w]
    target_span=world['mixamorig:Hips'].translation.y-sum(world['mixamorig:'+s+'Foot'].translation.y for s in ['Left','Right'])/2
    source_span=source_world['mixamorig:Hips'].translation.y-sum(source_world['mixamorig:'+s+'Foot'].translation.y for s in ['Left','Right'])/2
    ratio=target_span/source_span
    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]
        values=trs(local);src=source['bones'][source_index[b.name]]['rest']
        delta=(world[parent].to_quaternion().inverted()@source_world[parent].to_quaternion()) if parent else Quaternion()
        correction=source_world[b.name].to_quaternion().inverted()@world[b.name].to_quaternion()
        bones.append({'name':b.name,'parent':names.index(parent) if parent else -1,'rest':values,'bindpose':[list(r) for r in world[b.name].inverted()]})
        offsets.append(list(Vector(values[:3])-(delta@Vector(src[:3]))*ratio));delta_rotations.append(xyzw(delta));corrections.append(xyzw(correction))
    weights=[]
    for v in mesh.data.vertices:
        row={mesh.vertex_groups[g.group].name:g.weight for g in v.groups if g.weight>1e-8}
        assert all(n in names for n in row)
        total=sum(row.values());weights.append([[names.index(n),float(w/total)] for n,w in sorted(row.items(),key=lambda x:-x[1])])
    # Confirm identical geometry to the exact static/TokenRig input.
    target=json.loads((OUT.parent/'unity_review_data.json').read_text())
    reference=target['variants'][0]['vertices'];tree=KDTree(len(reference))
    for i,p in enumerate(reference):tree.insert(Vector(p),i)
    tree.balance();maximum=max(tree.find(Vector(p))[2] for p in points);assert maximum<1e-5,maximum
    normal_matrix=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 li in triangle.loops:
            vi=mesh.data.loops[li].vertex_index;n=normal_matrix@mesh.data.corner_normals[li].vector;n.normalize();tex=list(uv_layer.data[li].uv)
            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(points[vi].tolist());normals.append(list(n));uv.append(tex);influences.append(weights[vi])
            ids.append(lookup[key])
        triangles.extend([ids[0],ids[2],ids[1]])
    height=float(np.ptp(points[:,1]));edges=np.array([list(e.vertices) for e in mesh.data.edges])
    lengths=np.linalg.norm(points[edges[:,0]]-points[edges[:,1]],axis=1);valid=lengths>height*.002
    hands=np.abs(points[:,0])>.32;neck=(points[:,1]>.91)&(points[:,1]<.99)&(np.abs(points[:,0])<.1)
    masks={'all':valid,'hands':valid&hands[edges[:,0]]&hands[edges[:,1]],'neck':valid&neck[edges[:,0]]&neck[edges[:,1]]}
    audit_edges=[[first[a],first[b],float(lengths[i]),1|(2 if masks['hands'][i] else 0)|(4 if masks['neck'][i] else 0)] for i,(a,b) in enumerate(edges) if valid[i]]
    variant={'name':'mixamo','vertices':vertices,'normals':normals,'uv':uv,'triangles':triangles,'weights':influences,
             'auditEdges':audit_edges,'originalVertices':len(points),'trianglesCount':len(triangles)//3}
    data={'bones':bones,'variants':[variant],'height':height,'translationScale':ratio,'translationOffsets':offsets,
          'translationDeltaRotations':delta_rotations,'bindRotationCorrections':corrections,'rotationParents':[b['parent'] for b in bones]}
    (OUT/'unity_review_data.json').write_text(json.dumps(data,separators=(',',':')),encoding='utf-8')
    # Independent matrix-based LBS reference for native Unity verification.
    dense=np.zeros((len(points),len(names)))
    for i,row in enumerate(weights):
        for j,w in row:dense[i,j]=w
    p4=np.column_stack((points,np.ones(len(points))));maxs={k:0 for k in masks};over={k:0 for k in masks};samples=[]
    for frame,sample in enumerate(source['clips'][0]['samples']):
        source_pose=[]
        for i,values in enumerate(sample):
            local=Matrix.LocRotScale(Vector(values[:3]),Quaternion((values[6],*values[3:6])),Vector(values[7:10]))
            parent=source['bones'][i]['parent'];source_pose.append(source_pose[parent]@local if parent>=0 else local)
        target_pose=[];mapped_q=[]
        for i,b in enumerate(bones):
            n=b['name'];si=source_index[n];values=sample[si];parent=b['parent']
            rot=source_pose[si].to_quaternion()@Quaternion((corrections[i][3],*corrections[i][:3]));mapped_q.append(rot)
            local_q=mapped_q[parent].inverted()@rot if parent>=0 else rot
            delta=Quaternion((delta_rotations[i][3],*delta_rotations[i][:3]))
            local_p=delta@(Vector(values[:3])*ratio)+Vector(offsets[i])
            local=Matrix.LocRotScale(local_p,local_q,Vector(values[7:10]));target_pose.append(target_pose[parent]@local if parent>=0 else local)
        deformed=np.zeros_like(points)
        for i,n in enumerate(names):deformed+=dense[:,i,None]*(p4@np.array(target_pose[i]@world[n].inverted()).T)[:,:3]
        stretch=np.linalg.norm(deformed[edges[:,0]]-deformed[edges[:,1]],axis=1)/np.maximum(lengths,1e-10)
        for k,mask in masks.items():maxs[k]=max(maxs[k],float(stretch[mask].max()));over[k]=max(over[k],int((stretch[mask]>2).sum()))
        samples.append({'frame':frame+1,'maxStretch':float(stretch[valid].max())})
    report={'source':str(SOURCE),'bones':len(names),'triangles':4999,'maxInfluences':max(len(w) for w in weights),
            'geometryNearestError':maximum,'translationScale':ratio,'maxStretch':maxs,'maxEdgesOver2x':over,'samples':samples,
            'weights':'Original website weights, normalized only; no top4 truncation or cleanup',
            'motion':'Same shared Walking; source bind-relative world rotations mapped to website rest axes; target bind translations preserved',
            'missingSourceBones':sorted(set(source_index)-set(names))}
    (OUT/'walking_comparison.json').write_text(json.dumps(report,indent=2),encoding='utf-8')
    print('WEBSITE_EXPORTED',json.dumps({k:v for k,v in report.items() if k!='samples'}))

if __name__=='__main__':
    export() if '--export' in sys.argv else inspect()