"""Fit a generated realistic mesh to the licensed, animated Quaternius wolf rig.
Never generates replacement walk/run clips. Transfer weights in rest space and retain source clips.
"""
import bpy,sys,json,math,numpy as np
from pathlib import Path
from mathutils import Vector,Matrix
from mathutils.kdtree import KDTree
P=Path(sys.argv[sys.argv.index('--')+1]);mode=sys.argv[sys.argv.index('--')+2] if len(sys.argv)>sys.argv.index('--')+2 else 'bind'
bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.context.scene.render.fps=30
bpy.ops.import_scene.gltf(filepath=str(P/'sources/generated-wolf/raw.glb'))
real=next(o for o in bpy.context.scene.objects if o.type=='MESH')
coords=np.array([list(real.matrix_world@v.co) for v in real.data.vertices]);cen=coords[:,:2].mean(axis=0);eigval,eigvec=np.linalg.eigh(np.cov(coords[:,:2].T));forward=eigvec[:,-1]
headxy=coords[coords[:,2]>np.percentile(coords[:,2],95),:2].mean(axis=0)
if np.dot(headxy-cen,forward)<0:forward=-forward
ang=-math.pi/2-math.atan2(forward[1],forward[0]);rot=Matrix.Rotation(ang,4,'Z')
real.parent=None;real.matrix_world=Matrix.Identity(4)
for v,co in zip(real.data.vertices,coords):v.co=rot@Vector(co)
lo=Vector([min(v.co[i] for v in real.data.vertices) for i in range(3)]);hi=Vector([max(v.co[i] for v in real.data.vertices) for i in range(3)])
scale=.95/(hi.z-lo.z);offset=Vector((-(lo.x+hi.x)/2,-(lo.y+hi.y)/2,-lo.z))
for v in real.data.vertices:v.co=(v.co+offset)*scale
real.name='RealisticWolf';real.data.name='RealisticWolfMesh'
for p in real.data.polygons:p.use_smooth=True
for im in bpy.data.images:
 if im.type=='IMAGE' and im.size[0]>1024:im.scale(1024,1024);im.pack()
# Keep a static final-frame model for comparing the rig transfer to the supplied anatomy.
bpy.ops.object.select_all(action='DESELECT');real.select_set(True);bpy.context.view_layer.objects.active=real
bpy.ops.export_scene.gltf(filepath=str(P/'sources/generated-wolf/aligned.glb'),export_format='GLB',use_selection=True,export_yup=True,export_image_format='JPEG',export_jpeg_quality=88)
if mode=='align':
 (P/'evidence/generated-wolf-alignment.json').write_text(json.dumps({'rotation_z':ang,'scale':scale,'bounds_before':[list(lo),list(hi)]},indent=2))
else:
 # Import source rig at rest, at the same physical height as the real mesh.
 bpy.ops.import_scene.gltf(filepath=str(P/'glb/enemy-wolf.glb'))
 rig=next(o for o in bpy.context.scene.objects if o.type=='ARMATURE')
 donor=next(o for o in bpy.context.scene.objects if o.type=='MESH' and o!=real and any(m.type=='ARMATURE' for m in o.modifiers))
 rig.data.pose_position='REST';rig.animation_data.action=None
 for t in rig.animation_data.nla_tracks:t.mute=True
 bpy.context.view_layer.update()
 # All donor vertex coordinates and weights in final metres.
 pts=[donor.matrix_world@v.co for v in donor.data.vertices];weights=[]
 for v in donor.data.vertices:weights.append({donor.vertex_groups[g.group].name:g.weight for g in v.groups if g.weight>1e-6})
 # Diagnose fitted paws against the donor rig in ground-space metres.
 diagnostics={}
 for label,vs in [('real',np.array([list(v.co) for v in real.data.vertices])),('donor',np.array([list(p) for p in pts]))]:
  diagnostics[label]={}
  for xsign in [-1,1]:
   for ysign in [-1,1]:
    subset=vs[(vs[:,2]<.10)&(vs[:,0]*xsign>0)&(vs[:,1]*ysign>0)]
    diagnostics[label][f'{xsign},{ysign}']={'mean':subset.mean(axis=0).tolist(),'min':subset.min(axis=0).tolist(),'max':subset.max(axis=0).tolist()}
 (P/'evidence/wolf-paw-fit.json').write_text(json.dumps(diagnostics,indent=2))
 print('PAWS',json.dumps(diagnostics),flush=True)
 # Fit the real torso/limb spacing to the donor: the donor has a long raised tail,
 # so bounding-box centring alone shifts its whole torso relative to a tail-down wolf.
 rf=np.mean([diagnostics['real'][f'{x},-1']['mean'][1] for x in [-1,1]])
 rb=np.mean([diagnostics['real'][f'{x},1']['mean'][1] for x in [-1,1]])
 df=np.mean([diagnostics['donor'][f'{x},-1']['mean'][1] for x in [-1,1]])
 db=np.mean([diagnostics['donor'][f'{x},1']['mean'][1] for x in [-1,1]])
 along=(db-df)/(rb-rf);shift=df-along*rf
 xshift=-np.mean([v['mean'][0] for v in diagnostics['real'].values()])
 for v in real.data.vertices:v.co.y=v.co.y*along+shift;v.co.x+=xshift
 diagnostics['fit']={'length_scale':float(along),'forward_offset':float(shift),'side_offset':float(xshift)}
 (P/'evidence/wolf-paw-fit.json').write_text(json.dumps(diagnostics,indent=2))
 kd=KDTree(len(pts))
 for i,p in enumerate(pts):kd.insert(p,i)
 kd.balance()
 bonepos={b.name:{'head':rig.matrix_world@b.head_local,'tail':rig.matrix_world@b.tail_local} for b in rig.data.bones}
 # Assign donor skin smoothly, constrained by anatomical regions so legs cannot borrow torso/other-leg weights.
 for vg in list(real.vertex_groups):real.vertex_groups.remove(vg)
 for b in rig.data.bones:real.vertex_groups.new(name=b.name)
 diststats=[]
 for v in real.data.vertices:
  nearest=kd.find_n(v.co,6);ww={};den=0
  for p,idx,dist in nearest:
   fac=1/max(dist,.008)**2;den+=fac
   for name,w in weights[idx].items():ww[name]=ww.get(name,0)+w*fac
  ww={k:w/den for k,w in ww.items()}
  # The hanging tail must never borrow hind-leg weights. Tail1 drives a restrained
  # rigid wag at the anatomical root rather than folding the tail with the hock.
  tail_front=.43-.30*(v.co.z-.4)
  if abs(v.co.x)<.14 and v.co.y>tail_front and v.co.z>.12:
   blend=min(1.,max(0.,(v.co.y-tail_front)/.035))
   ww={k:w*(1-blend) for k,w in ww.items()};ww['Tail1']=ww.get('Tail1',0)+blend
  top=sorted(ww.items(),key=lambda x:-x[1])[:4];total=sum(w for _,w in top)
  for name,w in top:real.vertex_groups[name].add([v.index],w/total,'REPLACE')
  diststats.append(nearest[0][2])
 # Keep geometry in world metres: the rig's parent normalisation is already baked into its matrix.
 # Mesh need not parent to the rig; the Armature modifier uses their world matrices.
 mod=real.modifiers.new('QuadrupedSkin','ARMATURE');mod.object=rig
 world=real.matrix_world.copy();real.parent=rig;real.matrix_world=world
 bpy.data.objects.remove(donor,do_unlink=True)
 rig.data.pose_position='POSE'
 actions={a.name:a for a in bpy.data.actions};rig.animation_data.action=actions['Idle'];rig.animation_data.action_slot=actions['Idle'].slots[0]
 bpy.context.scene.frame_set(0);bpy.context.view_layer.update()
 bpy.ops.object.select_all(action='DESELECT')
 for o in bpy.context.scene.objects:
  if o.type in ['ARMATURE','MESH','EMPTY']:o.select_set(True)
 out=P/'glb/enemy-wolf-realistic.glb'
 bpy.ops.export_scene.gltf(filepath=str(out),export_format='GLB',use_selection=True,export_yup=True,export_image_format='JPEG',export_jpeg_quality=88,export_animations=True,export_animation_mode='ACTIONS',export_anim_single_armature=True,export_force_sampling=True,export_optimize_animation_size=True,export_influence_nb=4,export_all_influences=False,export_leaf_bone=False)
 entry=json.loads((P/'evidence/enemy-wolf-entry.json').read_text());entry.update(id='enemy-wolf-realistic',file=out.name,about='A textured grey wolf with natural proportions and eyes in its head; walks, gallops, lunges to bite and falls onto its side.',style='realistic',tint=[],licence='Forge-generated mesh; CC0-1.0 animation rig',credit='Forge, generated reference and Pixal3D mesh; rig and clips: Quaternius Ultimate Animated Animals, https://quaternius.com/packs/ultimateanimatedanimals.html',mood=['natural','wilderness','dark-fantasy'],speeds={})
 (P/'evidence/enemy-wolf-realistic-entry.json').write_text(json.dumps(entry,indent=2))
 (P/'evidence/enemy-wolf-realistic-normalization.json').write_text(json.dumps({'source':'Forge-generated image and Pixal3D mesh, licensed Quaternius rig and clips','rotation_z':ang,'scale':scale,'rig_transfer':'six nearest donor surface vertices, inverse squared distance, four highest weights, normalized','surface_distance_p50':float(np.percentile(diststats,50)),'surface_distance_p95':float(np.percentile(diststats,95)),'bones':{n:{k:list(v) for k,v in val.items()} for n,val in bonepos.items()},'missing_clips':[],'flyer':False},indent=2))
 print('BOUND',out.stat().st_size,'donor distances',np.percentile(diststats,[50,95]),flush=True)