"""Correct foot contacts in authored clips, baking IK results without synthesizing new clips.
Pilot exports go to candidates/. Original final GLBs remain untouched until checked.
"""
import bpy,bmesh,sys,json,math,numpy as np
from pathlib import Path
from mathutils import Vector,Matrix
P=Path(sys.argv[sys.argv.index('--')+1]);kind=sys.argv[sys.argv.index('--')+2];ident='enemy-'+kind
bpy.ops.wm.read_factory_settings(use_empty=True);bpy.context.scene.render.fps=30
bpy.ops.import_scene.gltf(filepath=str(P/'glb'/(ident+'.glb')))
rig=next(o for o in bpy.context.scene.objects if o.type=='ARMATURE');meshes=[o for o in bpy.context.scene.objects if o.type=='MESH' and any(m.type=='ARMATURE' for m in o.modifiers)]
for o in meshes:
 bm=bmesh.new();bm.from_mesh(o.data);bad=[f for f in bm.faces if f.calc_area()<1e-14]
 if bad:bmesh.ops.delete(bm,geom=bad,context='FACES_ONLY');bm.to_mesh(o.data)
 bm.free()
for t in list(rig.animation_data.nla_tracks):rig.animation_data.nla_tracks.remove(t)
entry=json.loads((P/'evidence'/(ident+'-entry.json')).read_text());runtime=json.loads((P/'evidence'/(ident+'-runtime.json')).read_text())
quad='IKFrontLeg.L' in rig.pose.bones
feet=[('IKFrontLeg.'+s,'FrontLowerLeg.'+s,2) for s in ['L','R']]+[('IKBackLeg.'+s,'BackLowerLeg.'+s,3) for s in ['L','R']] if quad else [('Foot.'+s,'LowerLeg.'+s,2) for s in ['L','R']]
oldactions={a.name:a for a in bpy.data.actions};logs={}
def activate(a,frame):
 rig.animation_data.action=a;rig.animation_data.action_slot=a.slots[0];bpy.context.scene.frame_set(frame);bpy.context.view_layer.update()
def rigworld(pb):return rig.matrix_world@pb.matrix
# Track actual sole vertices, not IK controls that sit above the feet.
activate(oldactions[entry['clips']['idle']],0)
height=runtime['height_m'];sole_ids={f:[] for f,_,_ in feet}
for o in meshes:
 ev=o.evaluated_get(bpy.context.evaluated_depsgraph_get());me=ev.to_mesh();pts=[ev.matrix_world@v.co for v in me.vertices]
 for f,_,_ in feet:
  head=rig.matrix_world@rig.pose.bones[f].head
  ids=[i for i,v in enumerate(pts) if v.z<height*.13 and (Vector((v.x,v.y))-Vector((head.x,head.y))).length<height*.2 and v.x*head.x>0 and (not quad or v.y*head.y>0)]
  if ids:
   cut=np.percentile([pts[i].z for i in ids],40);ids=[i for i in ids if pts[i].z<=cut];sole_ids[f].append((o,ids))
 ev.to_mesh_clear()
def sole_heights():
 cache={};out={}
 for o in meshes:
  ev=o.evaluated_get(bpy.context.evaluated_depsgraph_get());me=ev.to_mesh();cache[o]=[ev.matrix_world@v.co for v in me.vertices];ev.to_mesh_clear()
 for f,_,_ in feet:out[f]=min(cache[o][i].z for o,ids in sole_ids[f] for i in ids)
 return out
for lane in ['walk','run']:
 if lane not in entry['clips']:continue
 name=entry['clips'][lane];source=oldactions[name];start,end=map(round,source.frame_range);N=end-start;dt=1/30
 original=[];heads={f:[] for f,_,_ in feet};tails={f:[] for f,_,_ in feet};soles={f:[] for f,_,_ in feet}
 for frame in range(start,end+1):
  activate(source,frame);original.append({b.name:b.matrix.copy() for b in rig.pose.bones});sh=sole_heights()
  for f,_,_ in feet:soles[f].append(sh[f])
  for f,lower,_ in feet:heads[f].append(np.array(rig.matrix_world@rig.pose.bones[f].head));tails[f].append(np.array(rig.matrix_world@rig.pose.bones[lower].tail))
 speed=runtime['speeds'][lane]['speed'];desired={};contact={}
 for f,lower,chain in feet:
  pos=np.array(heads[f]);zs=pos[:N,2];lo,hi=zs.min(),zs.max();vel=np.gradient(pos[:,1],dt)
  mask=(zs<=lo+.28*(hi-lo))&(vel[:N]>.02)
  # Circular stance components; make a crossing-zero stance continuous.
  starts=[i for i in range(N) if mask[i] and not mask[(i-1)%N]]
  segments=[]
  for i in starts:
   seq=[];j=i
   while mask[j%N] and len(seq)<N:seq.append(j);j+=1
   if len(seq)>=3:segments.append(seq)
  tar=pos.copy();contacts=[]
  for seq in segments:
   idx=np.array(seq)%N;mid=(seq[0]+seq[-1])/2;xm=pos[idx,0].mean();ym=pos[idx,1].mean();zbase=lo
   for j in seq:
    k=j%N;w=min(1.,(j-seq[0]+1)/2,(seq[-1]-j+1)/2)
    planted=np.array([xm,ym+speed*(j-mid)*dt,pos[k,2]-soles[f][k]]);tar[k]=pos[k]*(1-w)+planted*w
    if w>.99:contacts.append(k)
  for k in range(N):tar[k,2]=max(tar[k,2],pos[k,2]-min(0,soles[f][k]))
  tar[N]=tar[0];desired[f]=tar;contact[f]=contacts
 targets={};constraints=[]
 for f,lower,chain in feet:
  target=bpy.data.objects.new('ContactTarget_'+f,None);bpy.context.collection.objects.link(target);targets[f]=target
  c=rig.pose.bones[lower].constraints.new('IK');c.target=target;c.chain_count=chain;c.iterations=64;c.use_stretch=False;constraints.append((rig.pose.bones[lower],c))
 baked=[]
 for idx,frame in enumerate(range(start,end+1)):
  # Disable IK while restoring the authored pose, then solve only the contact adjustment.
  for _,c in constraints:c.influence=0
  activate(source,frame)
  for f,lower,_ in feet:
   delta=desired[f][idx]-heads[f][idx];targets[f].location=Vector(tails[f][idx]+delta)
   mat=original[idx][f].copy();mat.translation+=rig.matrix_world.to_3x3().inverted()@Vector(delta);rig.pose.bones[f].matrix=mat
  for _,c in constraints:c.influence=1
  bpy.context.view_layer.update()
  for iteration in range(2):
   sh=sole_heights()
   for f,lower,_ in feet:
    correction=-sh[f] if idx%N in contact[f] else max(0,-sh[f])
    if abs(correction)>.0005:
     targets[f].location.z+=correction;mat=rig.pose.bones[f].matrix.copy();mat.translation+=rig.matrix_world.to_3x3().inverted()@Vector((0,0,correction));rig.pose.bones[f].matrix=mat
   bpy.context.view_layer.update()
  ev=rig.evaluated_get(bpy.context.evaluated_depsgraph_get());mats={b.name:b.matrix.copy() for b in ev.pose.bones};channels={}
  for pb in rig.pose.bones:
   kw={'parent_matrix':mats[pb.parent.name],'parent_matrix_local':pb.parent.bone.matrix_local} if pb.parent else {}
   basis=pb.bone.convert_local_to_pose(mats[pb.name],pb.bone.matrix_local,invert=True,**kw);channels[pb.name]=basis.decompose()
  baked.append(channels)
 for pb,c in constraints:pb.constraints.remove(c)
 for t in targets.values():bpy.data.objects.remove(t,do_unlink=True)
 source.name='Original_'+name;new=bpy.data.actions.new(name);rig.animation_data.action=new
 for idx,frame in enumerate(range(start,end+1)):
  for pb in rig.pose.bones:
   loc,rot,sca=baked[idx][pb.name];pb.rotation_mode='QUATERNION';pb.location=loc;pb.rotation_quaternion=rot;pb.scale=sca
   pb.keyframe_insert('location',frame=frame,group=pb.name);pb.keyframe_insert('rotation_quaternion',frame=frame,group=pb.name);pb.keyframe_insert('scale',frame=frame,group=pb.name)
 bpy.data.actions.remove(source)
 logs[lane]={'speed_from_source_m_s':speed,'contact_frames':contact,'frames':N+1,'method':'source gait, linear ground-relative stance with 2-frame feather; leg IK solved and baked'}
 print('PLANTED',lane,logs[lane],flush=True)
# Export same geometry, portable texture paths, and all the original clip roles.
activate(bpy.data.actions[entry['clips']['idle']],0)
bpy.ops.object.select_all(action='DESELECT')
for o in bpy.context.scene.objects:
 if o.type in ['ARMATURE','EMPTY'] or o in meshes or (o.type=='MESH' and o.parent_type=='BONE'):o.select_set(True)
(P/'candidates').mkdir(exist_ok=True)
bpy.ops.export_scene.gltf(filepath=str(P/'candidates'/(ident+'.glb')),export_format='GLB',use_selection=True,export_yup=True,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,export_image_format='JPEG' if kind=='wolf-realistic' else 'AUTO',export_jpeg_quality=88)
entry.update(id=ident+'-footplant',file='../candidates/'+ident+'.glb')
(P/'evidence'/(ident+'-footplant-entry.json')).write_text(json.dumps(entry,indent=2));(P/'evidence'/(ident+'-footplant-method.json')).write_text(json.dumps(logs,indent=2))