"""Contact-aware L20 pose refinement with a required, explicit object specification.

Optimizes calibrated finger coordinates and all six wrist correction coordinates.
The --fixture path is an analytic sphere test, NOT the object in the source video.
"""
import os
os.environ.setdefault('MUJOCO_GL','osmesa')
import argparse,json
from pathlib import Path
import xml.etree.ElementTree as ET
import copy
import numpy as np
import mujoco
from scipy.optimize import least_squares
from scipy.spatial.transform import Rotation
from l20_calibrated import ROOT,Calibration,FINGERS,TIPS

OUT=ROOT/'output/l20_2047635068'

def scene_for_object(spec,out):
    tree=ET.parse(OUT/'l20_moving.xml');root=tree.getroot()
    root.find('option').set('gravity','0 0 -9.81')
    root.find('option').set('timestep','0.001')
    root.find('option').set('cone','elliptic')
    # Resolve all mesh references relative to the existing scene.
    for mesh in root.findall('asset/mesh'):
        mesh.set('file',str((OUT/mesh.get('file')).resolve()))
    hand=root.find("worldbody/body[@name='hand_base_link']")
    for i,geom in enumerate(hand.iter('geom')):
        geom.set('name',f'hand_collision_{i}')
        geom.set('contype','1');geom.set('conaffinity','2')
        geom.set('friction',f"{spec.get('friction',.8)} .005 .0001")
        geom.set('condim','4')
    obj=ET.SubElement(root.find('worldbody'),'body',name='grasp_object',pos='0 0 0')
    ET.SubElement(obj,'freejoint',name='object_free')
    if spec['type']=='mjcf':
        source=Path(spec['mjcf_path']).resolve();bottle=ET.parse(source).getroot()
        for asset in bottle.find('asset'):
            item=copy.deepcopy(asset)
            if item.get('file'):item.set('file',str((source.parent/item.get('file')).resolve()))
            root.find('asset').append(item)
        for item in bottle.find('worldbody/body'):
            if item.tag=='freejoint':continue
            child=copy.deepcopy(item)
            if child.tag=='geom':
                child.set('name','bottle_part_'+str(len(obj)))
                if child.get('contype','1')!='0':
                    child.set('contype','2');child.set('conaffinity','1')
                    child.set('friction',f"{spec.get('friction',.6)} .005 .001")
                    child.set('solref','.004 1');child.set('solimp','.95 .99 .001')
            obj.append(child)
        path=out/'contact_scene.xml';tree.write(path)
        return path
    attrs=dict(name='object_geom',type=spec['type'],mass=str(spec.get('mass_kg',.1)),
               contype='2',conaffinity='1',friction=f"{spec.get('friction',.8)} .005 .0001",
               condim='4',rgba='.9 .65 .2 1',solref='.005 1')
    if spec['type']=='mesh':
        mesh_path=Path(spec['mesh_path']).resolve()
        ET.SubElement(root.find('asset'),'mesh',name='object_mesh',file=str(mesh_path),
                      scale=' '.join(map(str,spec.get('mesh_scale',[1,1,1]))))
        attrs['mesh']='object_mesh'
    else:attrs['size']=' '.join(map(str,spec['size']))
    ET.SubElement(obj,'geom',**attrs)
    path=out/'contact_scene.xml';tree.write(path)
    return path

class ContactRefiner:
    def __init__(self,spec,out):
        self.spec=spec;self.cal=Calibration()
        self.model=mujoco.MjModel.from_xml_path(str(scene_for_object(spec,out)))
        self.data=mujoco.MjData(self.model)
        self.names=[n for n in self.cal.tables]
        self.addr={n:self.model.jnt_qposadr[mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_JOINT,n)] for n in self.names}
        self.wrist=self.model.jnt_qposadr[mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_JOINT,'wrist_free')]
        self.obj=self.model.jnt_qposadr[mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_JOINT,'object_free')]
        self.object_geoms=np.flatnonzero(self.model.geom_contype==2)
        self.object_geom=int(self.object_geoms[0])
        self.hand_geoms=np.flatnonzero(self.model.geom_contype==1)
        self.tip_geoms=[]
        contact_links=spec.get('contact_links',[f+'_distal' for f in spec.get('contact_fingers',['thumb','index','middle'])])
        for link in contact_links:
            bid=mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_BODY,link)
            self.tip_geoms.append(next(g for g in self.hand_geoms if self.model.geom_bodyid[g]==bid))
        self.lower=np.array([self.cal.bounds[n][0] for n in self.cal.active]+[-.06]*3+[-.5]*3)
        self.upper=np.array([self.cal.bounds[n][1] for n in self.cal.active]+[.06]*3+[.5]*3)

    def expand(self,active):
        q=dict(zip(self.cal.active,active))
        for n in self.cal.passive:q[n]=self.cal.passive_value(n,q[self.cal.curves[n][0]])[0]
        return q

    def set_pose(self,x,wrist_pos,wrist_quat,obj_pos,obj_quat):
        q=self.expand(x[:16])
        for n,v in q.items():self.data.qpos[self.addr[n]]=v
        self.data.qpos[self.wrist:self.wrist+3]=wrist_pos+x[16:19]
        rotation=Rotation.from_quat(wrist_quat[[1,2,3,0]])*Rotation.from_rotvec(x[19:22])
        self.data.qpos[self.wrist+3:self.wrist+7]=rotation.as_quat()[[3,0,1,2]]
        self.data.qpos[self.obj:self.obj+3]=obj_pos
        self.data.qpos[self.obj+3:self.obj+7]=obj_quat
        mujoco.mj_kinematics(self.model,self.data)

    def distances(self):
        return np.array([min(mujoco.mj_geomDistance(self.model,self.data,int(g),int(o),.5,None) for o in self.object_geoms) for g in self.hand_geoms])

    def solve(self,reference,wrist_pos,wrist_quat,obj_pos,obj_quat,initial=None):
        initial=np.r_[reference,np.zeros(6)] if initial is None else np.asarray(initial).copy()
        tip_index=[list(self.hand_geoms).index(g) for g in self.tip_geoms]
        def residual(x):
            self.set_pose(x,wrist_pos,wrist_quat,obj_pos,obj_quat)
            dist=self.distances()
            compression=self.spec.get('compression_m',.0002)
            return np.r_[(dist[tip_index]+compression)/.004,
                          np.maximum(-dist-max(.0003,compression),0)/.0015,
                          .15*(x[:16]-reference),.2*x[16:19]/.03,.2*x[19:22]/.25]
        before=residual(initial);before_distance=self.distances()
        result=least_squares(residual,np.clip(initial,self.lower+1e-9,self.upper-1e-9),
             bounds=(self.lower,self.upper),diff_step=1e-4,max_nfev=100,ftol=1e-5,xtol=1e-5,gtol=1e-5)
        after=residual(result.x);after_distance=self.distances()
        report=dict(converged=bool(result.success),nfev=result.nfev,
              cost_before=float(before@before),cost_after=float(after@after),
              mean_selected_surface_gap_before_mm=float(np.abs(before_distance[tip_index]).mean()*1000),
              mean_selected_surface_gap_after_mm=float(np.abs(after_distance[tip_index]).mean()*1000),
              max_penetration_before_mm=float(max(0,-before_distance.min())*1000),
              max_penetration_after_mm=float(max(0,-after_distance.min())*1000),
              wrist_correction_m=result.x[16:19].tolist(),wrist_correction_rotvec=result.x[19:22].tolist())
        return result.x,self.data.qpos.copy(),report

    def dynamic_hold(self,qpos,seconds=2.):
        # A diagnostic with prescribed hand coordinates and a freely simulated object.
        # It does not validate actuator effort, compliant tracking, or force closure.
        mujoco.mj_resetData(self.model,self.data);self.data.qpos[:]=qpos
        object_initial=qpos[self.obj:self.obj+3].copy()
        hand_addr=[self.addr[n] for n in self.names]+list(range(self.wrist,self.wrist+7))
        object_dof=self.model.jnt_dofadr[mujoco.mj_name2id(self.model,mujoco.mjtObj.mjOBJ_JOINT,'object_free')]
        controlled_dof=[i for i in range(self.model.nv) if not object_dof<=i<object_dof+6]
        frames=[];contacts=[]
        steps=int(seconds/self.model.opt.timestep)
        for k in range(steps):
            self.data.qpos[hand_addr]=qpos[hand_addr];self.data.qvel[controlled_dof]=0
            mujoco.mj_step(self.model,self.data)
            if not np.isfinite(self.data.qpos).all():raise RuntimeError('Non-finite physics state')
            if k%33==0:
                frames.append(self.data.qpos.copy())
                contacts.append(sum(any(g in self.object_geoms for g in c.geom) for c in self.data.contact))
        displacement=float(np.linalg.norm(self.data.qpos[self.obj:self.obj+3]-object_initial))
        return np.array(frames),dict(seconds=seconds,object_displacement_m=displacement,
             contact_frames=int(np.sum(np.array(contacts)>0)),total_frames=len(frames),
             retained_within_2cm=displacement<.02,
             boundary='Prescribed hand, free object under gravity. No object weld or pose resets. Hand actuator/arm dynamics not validated.')

def main():
    parser=argparse.ArgumentParser();parser.add_argument('--spec',type=Path);parser.add_argument('--fixture',action='store_true')
    parser.add_argument('--frame',type=int,default=666);args=parser.parse_args()
    motion=np.load(OUT/'motion.npz');t=args.frame
    if not args.fixture and args.spec is None:parser.error('Provide an object specification, or explicitly select the synthetic --fixture test.')
    out=OUT/('contact_fixture' if args.fixture else 'contact_object');out.mkdir(exist_ok=True)
    if args.fixture:
        pts=motion['actual'][t]
        center=(pts[4]+pts[12])*.5
        radius=np.linalg.norm(pts[4]-pts[12])*.42
        r=Rotation.from_quat(motion['wrist_quat_wxyz'][t,[1,2,3,0]]).as_matrix()
        spec=dict(type='sphere',size=[float(radius)],mass_kg=.05,friction=.8,
                  position=(r@center+motion['wrist_pos'][t]).tolist(),quat_wxyz=[1,0,0,0],
                  contact_fingers=['thumb','index','middle'],purpose='Synthetic interface/physics fixture; not source-video reconstruction')
    else:spec=json.loads(args.spec.read_text())
    (out/'object_spec.json').write_text(json.dumps(spec,indent=2))
    solver=ContactRefiner(spec,out)
    active=np.array([motion['qpos'][t,list(motion['joint_names']).index(n)] for n in solver.cal.active])
    x,q,report=solver.solve(active,motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],np.array(spec['position']),np.array(spec['quat_wxyz']))
    path,dynamic=solver.dynamic_hold(q)
    report['dynamic_hold']=dynamic;report['frame']=t;report['fixture']=args.fixture
    np.savez_compressed(out/'contact_result.npz',optimization_variables=x,qpos=q,dynamic_qpos=path)
    (out/'validation.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2),flush=True)

if __name__=='__main__':main()