"""Read TokenRig GLB weights and match quantized joint heads to supplied bones.

Writes JSON only; does not rewrite binary assets. A global one-to-one assignment
avoids greedy nearest-bone matching errors around closely spaced fingers.
"""
import json
import struct
import sys
from pathlib import Path

import numpy as np
from scipy.optimize import linear_sum_assignment
from scipy.spatial import cKDTree


class GLB:
    def __init__(self,path):
        data=Path(path).read_bytes()
        assert data[:4]==b'glTF'
        self.bin=b''
        offset=12
        while offset<len(data):
            size,kind=struct.unpack_from('<II',data,offset)
            chunk=data[offset+8:offset+8+size]
            if kind==0x4E4F534A:
                self.doc=json.loads(chunk)
            elif kind==0x004E4942:
                self.bin=chunk
            offset+=8+size
        self.world={}
        def visit(index,parent):
            node=self.doc['nodes'][index]
            if 'matrix' in node:
                local=np.array(node['matrix']).reshape(4,4,order='F')
            else:
                x,y,z,w=node.get('rotation',[0,0,0,1])
                rotation=np.array([[1-2*(y*y+z*z),2*(x*y-z*w),2*(x*z+y*w)],
                                   [2*(x*y+z*w),1-2*(x*x+z*z),2*(y*z-x*w)],
                                   [2*(x*z-y*w),2*(y*z+x*w),1-2*(x*x+y*y)]])
                local=np.eye(4)
                local[:3,:3]=rotation @ np.diag(node.get('scale',[1,1,1]))
                local[:3,3]=node.get('translation',[0,0,0])
            self.world[index]=parent@local
            for child in node.get('children',[]):
                visit(child,self.world[index])
        for node in self.doc['scenes'][self.doc.get('scene',0)]['nodes']:
            visit(node,np.eye(4))

    def accessor(self,index):
        a=self.doc['accessors'][index]
        v=self.doc['bufferViews'][a['bufferView']]
        dtype=np.dtype({5120:'i1',5121:'u1',5122:'<i2',5123:'<u2',5125:'<u4',5126:'<f4'}[a['componentType']])
        count={'SCALAR':1,'VEC2':2,'VEC3':3,'VEC4':4,'MAT4':16}[a['type']]
        stride=v.get('byteStride',dtype.itemsize*count)
        result=np.ndarray((a['count'],count),dtype=dtype,buffer=self.bin,
                          offset=v.get('byteOffset',0)+a.get('byteOffset',0),strides=(stride,dtype.itemsize)).copy()
        if a.get('normalized'):
            result=result.astype(float)/np.iinfo(dtype).max
        return result

    def mesh(self):
        index,node=next((i,n) for i,n in enumerate(self.doc['nodes']) if 'mesh' in n)
        primitive=self.doc['meshes'][node['mesh']]['primitives'][0]
        p=self.accessor(primitive['attributes']['POSITION'])
        p4=np.column_stack((p,np.ones(len(p))))
        if 'skin' in node:
            # For skinned glTF nodes the mesh-node world transform cancels out
            # of the skin transform. Evaluate the supplied inverse bind matrices
            # rather than applying the armature object's scale a second time.
            skin=self.doc['skins'][node['skin']]
            inverse=self.accessor(skin['inverseBindMatrices'])
            transforms=np.array([self.world[j] @ inverse[k].reshape(4,4,order='F') for k,j in enumerate(skin['joints'])])
            ids=self.accessor(primitive['attributes']['JOINTS_0'])
            weights=self.accessor(primitive['attributes']['WEIGHTS_0'])
            p=np.zeros_like(p4,dtype=float)
            for slot in range(ids.shape[1]):
                p+=weights[:,slot,None]*np.einsum('nij,nj->ni',transforms[ids[:,slot]],p4)
            p=p[:,:3]
        else:
            p=(p4 @ self.world[index].T)[:,:3]
        return p,primitive,node


folder=Path(sys.argv[1])
raw_name=sys.argv[2] if len(sys.argv)>2 else 'rigged_raw.glb'
prefix=sys.argv[3] if len(sys.argv)>3 else ''
source=GLB(folder/'source.glb')
raw=GLB(folder/raw_name)
src_skin=source.doc['skins'][0]
raw_skin=raw.doc['skins'][0]
assert len(src_skin['joints'])==len(raw_skin['joints'])==41
src_heads=np.array([source.world[i][:3,3] for i in src_skin['joints']])
raw_heads=np.array([raw.world[i][:3,3] for i in raw_skin['joints']])
cost=np.linalg.norm(raw_heads[:,None,:]-src_heads[None,:,:],axis=2)
rows,cols=linear_sum_assignment(cost)
names=[source.doc['nodes'][i]['name'] for i in src_skin['joints']]
mapping={int(i):names[j] for i,j in zip(rows,cols)}
report={'joint_count':41,'max_joint_quantization_drift_m':float(cost[rows,cols].max()),
    'joint_mapping':[{'output_skin_index':int(i),'output_name':raw.doc['nodes'][raw_skin['joints'][i]]['name'],
                     'input_name':names[j],'distance_m':float(cost[i,j]),
                     'row_second_best_distance_m':float(np.sort(cost[i])[1])} for i,j in zip(rows,cols)]}
assert report['max_joint_quantization_drift_m']<.015, report
positions,primitive,_=raw.mesh()
source_positions,_,_=source.mesh()
distances,_=cKDTree(positions).query(source_positions)
report['max_vertex_nearest_distance_m']=float(distances.max())
assert distances.max()<1e-5
reverse_distances,_=cKDTree(source_positions).query(positions)
report['max_reverse_vertex_nearest_distance_m']=float(reverse_distances.max())
report['source_render_vertices']=len(source_positions)
assert reverse_distances.max()<1e-5
attributes=primitive['attributes']
joints=raw.accessor(attributes['JOINTS_0'])
weights=raw.accessor(attributes['WEIGHTS_0'])
# Transfer may split render vertices at normal/UV corners. Both-direction
# geometric correspondence above is authoritative; render counts may differ.
decoded=[]
for p,ids,w in zip(positions,joints,weights):
    named={mapping[int(i)]:float(value) for i,value in zip(ids,w) if value>0}
    decoded.append({'position_y_up':p.tolist(),'weights':named})
report['vertices']=len(decoded)
report['weight_sum_range']=[float(weights.sum(1).min()),float(weights.sum(1).max())]
(folder/(prefix+'joint_mapping.json')).write_text(json.dumps(report,indent=2),encoding='utf-8')
(folder/(prefix+'decoded_weights.json')).write_text(json.dumps(decoded,separators=(',',':')),encoding='utf-8')
print(json.dumps({k:v for k,v in report.items() if k!='joint_mapping'},indent=2))