"""Measure bone trajectories independently of skinning for the shared walk."""
import json
import sys
import math
from pathlib import Path
import bpy
import numpy as np
from mathutils import Matrix,Vector,Quaternion

ROOT=Path(__file__).resolve().parents[2]
FOLDER=ROOT/'Tools/output/lmgirl_tpose_studio/rig_mixamo41'
SOURCE=json.loads((ROOT/'Tools/output/campguide_studio/rig_mixamo41/unity_review/review_data.json').read_text())

def quaternion(q):return Quaternion((q[3],*q[:3]))
def sample(data,frame):
    source_indices={b['name']:i for i,b in enumerate(SOURCE['bones'])}
    values=[SOURCE['clips'][0]['samples'][frame][source_indices[b['name']]] for b in data['bones']]
    source_rot=[];mapped=[];world=[]
    for i,(b,v) in enumerate(zip(data['bones'],values)):
        parent=b['parent'];q=quaternion(v[3:7]);source_rot.append(source_rot[parent]@q if parent>=0 else q)
        mapped.append(source_rot[i]@quaternion(data['bindRotationCorrections'][i]) if data.get('bindRotationCorrections') else source_rot[i])
        q=mapped[parent].inverted()@mapped[i] if parent>=0 else mapped[i]
        p=Vector(v[:3])
        if data.get('translationOffsets'):
            p*=data['translationScale']
            if data.get('translationDeltaRotations'):p=quaternion(data['translationDeltaRotations'][i])@p
            p+=Vector(data['translationOffsets'][i])
        local=Matrix.LocRotScale(p,q,Vector(v[7:10]));world.append(world[parent]@local if parent>=0 else local)
    return {b['name'].split(':')[-1]:world[i] for i,b in enumerate(data['bones'])}

def audit():
    results={}
    for label,data in [('boy',SOURCE),('girl_tokenrig',json.loads((FOLDER/'unity_review_data.json').read_text())),
                       ('girl_mixamo',json.loads((FOLDER/'website/unity_review_data.json').read_text()))]:
        samples=[]
        rest={b['name'].split(':')[-1]:Matrix(b['bindpose']).inverted() for b in data['bones']}
        for f in range(32):
            pose=sample(data,f);hips=pose['Hips'].translation
            lateral=(pose['LeftUpLeg'].translation-pose['RightUpLeg'].translation).normalized()
            feet={side:list(pose[side+'Foot'].translation-hips) for side in ['Left','Right']}
            gap=(pose['LeftFoot'].translation-pose['RightFoot'].translation).dot(lateral)
            knees=(pose['LeftLeg'].translation-pose['RightLeg'].translation).dot(lateral)
            samples.append({'frame':f+1,'footSignedGap':gap,'kneeSignedGap':knees,'feetRelativeHips':feet})
        results[label]={'restHipWidth':(rest['LeftUpLeg'].translation-rest['RightUpLeg'].translation).length,
            'restAnkleWidth':(rest['LeftFoot'].translation-rest['RightFoot'].translation).length,
            'leftRestLegVector':list(rest['LeftFoot'].translation-rest['LeftUpLeg'].translation),
            'minimumFootSignedGap':min(s['footSignedGap'] for s in samples),'minimumKneeSignedGap':min(s['kneeSignedGap'] for s in samples),
            'crossedFrames':[s['frame'] for s in samples if s['footSignedGap']<0],
            'samples':samples}
    (FOLDER/'gait_diagnosis_before.json').write_text(json.dumps(results,indent=2),encoding='utf-8')
    print(json.dumps({k:{n:v for n,v in d.items() if n!='samples'} for k,d in results.items()},indent=2))

def corrected_motion(data,label):
    names=[b['name'].split(':')[-1] for b in data['bones']];indices={n:i for i,n in enumerate(names)}
    parents=[b['parent'] for b in data['bones']]
    rest={b['name'].split(':')[-1]:Matrix(b['bindpose']).inverted() for b in data['bones']}
    source_rest={b['name'].split(':')[-1]:Matrix(b['bindpose']).inverted() for b in SOURCE['bones']}
    target_width=(rest['LeftUpLeg'].translation-rest['RightUpLeg'].translation).length
    source_width=(source_rest['LeftUpLeg'].translation-source_rest['RightUpLeg'].translation).length
    scale=Vector((target_width/source_width,data['translationScale'],data['translationScale']))
    samples=[];tracks=[];previous={}
    for frame in range(32):
        source=sample(SOURCE,frame);initial=sample(data,frame)
        local=[]
        for i,n in enumerate(names):
            m=initial[names[parents[i]]].inverted()@initial[n] if parents[i]>=0 else initial[n]
            p,q,s=m.decompose();local.append([p.copy(),q.copy(),s.copy()])
        def fk():
            result=[]
            for i,(p,q,s) in enumerate(local):
                m=Matrix.LocRotScale(p,q,s);result.append(result[parents[i]]@m if parents[i]>=0 else m)
            return result
        world=fk();target_hips=world[indices['Hips']].translation.copy();source_hips=source['Hips'].translation.copy()
        goals={};errors={}
        for side in ['Left','Right']:
            upper=indices[side+'UpLeg'];lower=indices[side+'Leg'];foot=indices[side+'Foot']
            source_foot=source[side+'Foot'].translation-source_hips;source_knee=source[side+'Leg'].translation-source_hips
            goal=target_hips+Vector(tuple(source_foot[i]*scale[i] for i in range(3)))
            hint=target_hips+Vector(tuple(source_knee[i]*scale[i] for i in range(3)))
            h=world[upper].translation.copy();k=world[lower].translation.copy();f=world[foot].translation.copy()
            foot_q=world[foot].to_quaternion().copy();a=(k-h).length;b=(f-k).length
            direction=goal-h;wanted_distance=direction.length;direction.normalize()
            distance=max(abs(a-b)+1e-7,min(wanted_distance,a+b-1e-7))
            reachable=h+direction*distance
            bend=hint-h-direction*(hint-h).dot(direction)
            if bend.length<1e-7:bend=k-h-direction*(k-h).dot(direction)
            if bend.length<1e-7:bend=Vector((0,0,-1))-direction*direction.dot(Vector((0,0,-1)))
            bend.normalize();along=(a*a+distance*distance-b*b)/(2*distance)
            desired_knee=h+direction*along+bend*math.sqrt(max(0,a*a-along*along))
            upper_q=(k-h).rotation_difference(desired_knee-h)@world[upper].to_quaternion()
            local[upper][1]=world[parents[upper]].to_quaternion().inverted()@upper_q
            world=fk()
            lower_q=(world[foot].translation-world[lower].translation).rotation_difference(reachable-world[lower].translation)@world[lower].to_quaternion()
            local[lower][1]=world[parents[lower]].to_quaternion().inverted()@lower_q
            world=fk();local[foot][1]=world[parents[foot]].to_quaternion().inverted()@foot_q;world=fk()
            goals[side]=list(goal);errors[side]=(world[foot].translation-goal).length
        values=[]
        for i,(p,q,s) in enumerate(local):
            if i in previous and previous[i].dot(q)<0:q.negate()
            previous[i]=q.copy();values.append(list(p)+[q.x,q.y,q.z,q.w]+list(s))
        samples.append(values)
        lateral=(world[indices['LeftUpLeg']].translation-world[indices['RightUpLeg']].translation).normalized()
        gap=(world[indices['LeftFoot']].translation-world[indices['RightFoot']].translation).dot(lateral)
        tracks.append({'frame':frame+1,'footSignedGap':gap,'goalErrors':errors,'goals':goals,
            'positions':{n:list(world[i].translation) for n,i in indices.items() if n in ['Hips','LeftUpLeg','RightUpLeg','LeftLeg','RightLeg','LeftFoot','RightFoot']}})
    return {'name':'Walking_LmGirl'+('Mixamo33' if label=='mixamo' else '41'),'fps':30,'samples':samples}, {
        'footMotionScale':list(scale),'minimumFootSignedGap':min(t['footSignedGap'] for t in tracks),
        'crossedFrames':[t['frame'] for t in tracks if t['footSignedGap']<0],
        'maxUnreachableGoalError':max(max(t['goalErrors'].values()) for t in tracks),'tracks':tracks}

def correct():
    result={'method':'Pelvis-relative source ankle paths scaled by target hip width and hip-to-ankle span; analytic two-bone IK with source knee hints, fixed segment lengths and preserved mapped foot world rotations.',
            'sourceClipUnmodified':True,'sourceClips':'CampGuide native Walking','motions':{},'reports':{}}
    for label,path in [('tokenrig',FOLDER/'unity_review_data.json'),('mixamo',FOLDER/'website/unity_review_data.json')]:
        result['motions'][label],result['reports'][label]=corrected_motion(json.loads(path.read_text()),label)
    result['deformation']={}
    for label,path in [('tokenrig',FOLDER/'unity_review_data.json'),('mixamo',FOLDER/'website/unity_review_data.json')]:
        data=json.loads(path.read_text());motion=result['motions'][label];parents=[b['parent'] for b in data['bones']]
        for variant in data['variants']:
            points=np.array(variant['vertices']);p4=np.column_stack((points,np.ones(len(points))))
            weights=np.zeros((len(points),len(parents)))
            for i,row in enumerate(variant['weights']):
                for j,w in row:weights[i,j]=w
            edges=np.array(variant['auditEdges']);ids=edges[:,:2].astype(int);lengths=edges[:,2];masks=edges[:,3].astype(int)
            maxima={name:0 for name in ['all','hands','neck']};over=maxima.copy();floors=[]
            for values in motion['samples']:
                world=[]
                for i,v in enumerate(values):
                    local=Matrix.LocRotScale(Vector(v[:3]),quaternion(v[3:7]),Vector(v[7:10]));world.append(world[parents[i]]@local if parents[i]>=0 else local)
                deformed=np.zeros_like(points)
                for i,b in enumerate(data['bones']):deformed+=weights[:,i,None]*(p4@np.array(world[i]@Matrix(b['bindpose'])).T)[:,:3]
                ratios=np.linalg.norm(deformed[ids[:,0]]-deformed[ids[:,1]],axis=1)/lengths;floors.append(float(deformed[:,1].min()))
                for name,mask in [('all',1),('hands',2),('neck',4)]:
                    region=ratios[(masks&mask)!=0];maxima[name]=max(maxima[name],float(region.max()));over[name]=max(over[name],int((region>2).sum()))
            result['deformation'][variant['name']]={'maxStretch':maxima,'maxEdgesOver2x':over,'floorMinRange':[min(floors),max(floors)]}
    (FOLDER/'gait_correction.json').write_text(json.dumps(result,indent=2),encoding='utf-8')
    print(json.dumps({k:{n:v for n,v in d.items() if n!='tracks'} for k,d in result['reports'].items()},indent=2))

def export_corrected():
    result=json.loads((FOLDER/'gait_correction.json').read_text());C=Matrix(((1,0,0,0),(0,0,1,0),(0,1,0,0),(0,0,0,1)))
    receipts=[]
    for tag in ['raw','post','clean','clean_eps','local']:
        bpy.ops.wm.open_mainfile(filepath=str(FOLDER/('lmgirl_'+tag+'.blend')))
        arm=next(o for o in bpy.context.scene.objects if o.type=='ARMATURE');arm.animation_data_clear();arm.data.pose_position='POSE'
        names=[b.name for b in arm.data.bones];data=json.loads((FOLDER/'unity_review_data.json').read_text());parents=[b['parent'] for b in data['bones']]
        for b in arm.pose.bones:b.rotation_mode='QUATERNION'
        for f,values in enumerate(result['motions']['tokenrig']['samples'],1):
            bpy.context.scene.frame_set(f);world=[]
            for i,v in enumerate(values):
                m=Matrix.LocRotScale(Vector(v[:3]),quaternion(v[3:7]),Vector(v[7:10]));world.append(world[parents[i]]@m if parents[i]>=0 else m)
            for i,n in enumerate(names):
                b=arm.pose.bones[n];b.matrix=arm.matrix_world.inverted()@C@world[i]@C;bpy.context.view_layer.update()
                for attr in ['location','rotation_quaternion','scale']:b.keyframe_insert(attr,frame=f)
        bpy.context.scene.frame_start=1;bpy.context.scene.frame_end=32;bpy.context.scene.render.fps=30
        path=FOLDER/('Walking_'+tag+'_corrected.glb')
        bpy.ops.export_scene.gltf(filepath=str(path),export_format='GLB',export_skins=True,export_animations=True,
            export_animation_mode='ACTIVE_ACTIONS',export_force_sampling=True,export_def_bones=False)
        receipts.append({'file':str(path),'frames':[1,32],'source':result['motions']['tokenrig']['name']})
    (FOLDER/'gait_export_manifest.json').write_text(json.dumps(receipts,indent=2),encoding='utf-8')

def verify_exports():
    from mathutils.kdtree import KDTree
    data=json.loads((FOLDER/'unity_review_data.json').read_text());result=json.loads((FOLDER/'gait_correction.json').read_text());receipts=[]
    for tag in ['raw','post','clean','clean_eps','local']:
        variant=next(v for v in data['variants'] if v['name']==tag)
        bpy.ops.wm.read_factory_settings(use_empty=True);bpy.context.scene.render.fps=30
        bpy.ops.import_scene.gltf(filepath=str(FOLDER/('Walking_'+tag+'_corrected.glb')))
        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')
        # Mesh raw geometry is Blender Z-up; convert to native review Y-up.
        points=np.array([mesh.matrix_world@v.co for v in mesh.data.vertices]);points=points[:,[0,2,1]]
        tree=KDTree(len(points))
        for i,p in enumerate(points):tree.insert(Vector(p),i)
        tree.balance();mapping=[]
        for p in variant['vertices']:
            _,index,distance=tree.find(Vector(p));assert distance<1e-5;mapping.append(index)
        edges=np.array(variant['auditEdges']);ids=edges[:,:2].astype(int);max_stretch=0;min_gap=100
        for frame in range(1,33):
            bpy.context.scene.frame_set(frame);bpy.context.view_layer.update()
            obj=mesh.evaluated_get(bpy.context.evaluated_depsgraph_get());geometry=obj.to_mesh()
            deformed=np.array([obj.matrix_world@v.co for v in geometry.vertices])[mapping];obj.to_mesh_clear()
            max_stretch=max(max_stretch,float((np.linalg.norm(deformed[ids[:,0]]-deformed[ids[:,1]],axis=1)/edges[:,2]).max()))
            p={n:arm.matrix_world@arm.pose.bones['mixamorig:'+n].head for n in ['LeftUpLeg','RightUpLeg','LeftFoot','RightFoot']}
            lateral=(p['LeftUpLeg']-p['RightUpLeg']).normalized();min_gap=min(min_gap,(p['LeftFoot']-p['RightFoot']).dot(lateral))
        error=abs(max_stretch-result['deformation'][tag]['maxStretch']['all']);assert error<.003,(tag,error)
        assert min_gap>.08,(tag,min_gap)
        receipts.append({'variant':tag,'maxStretch':max_stretch,'error':error,'minimumFootSignedGap':min_gap})
    (FOLDER/'gait_export_verification.json').write_text(json.dumps({'passed':True,'results':receipts},indent=2),encoding='utf-8')
    print('GAIT_EXPORTS_VERIFIED',len(receipts))

if __name__=='__main__':
    if '--correct' in sys.argv:correct()
    elif '--export' in sys.argv:export_corrected()
    elif '--verify-exports' in sys.argv:verify_exports()
    else:audit()