"""Independent narrow-phase contact, coupling, floor and pose verification."""
import os
os.environ.setdefault('MUJOCO_GL','osmesa')
from pathlib import Path
import json,xml.etree.ElementTree as E
import numpy as np,mujoco
OUT=Path(__file__).resolve().parents[1]/'output/l20_full_replay'
x=np.load(OUT/'motion.npz');m=mujoco.MjModel.from_xml_path(str(OUT/'scene.xml'));d=mujoco.MjData(m)
penetration=[];counts=[];floor_count=0;coupling=0.;limit=0.
for row in x['qpos']:
 d.qpos[:]=row;mujoco.mj_forward(m,d)
 ds=[]
 for c in d.contact:
  types={int(m.geom_contype[c.geom1]),int(m.geom_contype[c.geom2])}
  if types=={1,2}:ds.append(float(c.dist))
  if 4 in types:floor_count+=1
 penetration.append(max(0,-min(ds,default=0)));counts.append(len(ds))
 for eq in range(m.neq):
  j1,j2=m.eq_obj1id[eq],m.eq_obj2id[eq];v=row[m.jnt_qposadr[j2]]
  expected=sum(m.eq_data[eq,k]*v**k for k in range(5))
  coupling=max(coupling,abs(row[m.jnt_qposadr[j1]]-expected))
 for j in range(m.njnt):
  if m.jnt_limited[j]:
   v=row[m.jnt_qposadr[j]];limit=max(limit,m.jnt_range[j,0]-v,v-m.jnt_range[j,1])
r=json.loads((OUT/'validation.json').read_text());r['geom_distance_max_penetration_mm']=float(np.max(x['penetration_m'])*1000)
r.update(max_all_penetration_mm=float(max(penetration)*1000),max_holding_penetration_mm=float(max(penetration[160:])*1000),holding_frames_with_contact=int(np.count_nonzero(counts[160:])),holding_frames=len(counts)-160,floor_contacts=floor_count,linear_mimic_max_error_rad=float(coupling),joint_limit_violation_rad=float(limit),contact_verification='mj_forward contact manifold independently checked every frame; geomDistance has isolated inconsistent spikes')
(OUT/'validation.json').write_text(json.dumps(r,indent=2));np.savez_compressed(OUT/'contact_verification.npz',penetration_m=penetration,contact_count=counts)
assert np.isfinite(x['qpos']).all() and coupling<1e-8 and limit<1e-8 and floor_count==0
print(json.dumps(r,indent=2))