"""Actuated L20 + freely simulated metric bottle, with a prescribed moving wrist."""
import os
os.environ.setdefault('MUJOCO_GL','osmesa')
import argparse,json
import xml.etree.ElementTree as ET
from pathlib import Path
import numpy as np
import mujoco
from scipy.spatial.transform import Rotation
from scipy.optimize import lsq_linear
from l20_calibrated import ROOT,Calibration

BASE=ROOT/'output/l20_2047635068'
OUT=BASE/'bottle_grasp'

def feedback_torque(model,data,cal,addr,objects,object_body,target_pos,target_velocity):
    selected=[c for c in data.contact if any(g in objects for g in c.geom)]
    if not selected:return None,None
    center=data.xipos[object_body].copy();velocity=np.zeros(6)
    mujoco.mj_objectVelocity(model,data,mujoco.mjtObj.mjOBJ_BODY,object_body,velocity,0)
    desiredforce=np.array([0,0,model.body_mass[object_body]*9.81])+100*(target_pos-center)+2*(target_velocity-velocity[3:])
    rotation=Rotation.from_matrix(data.xmat[object_body].reshape(3,3))
    desiredtorque=-.08*rotation.as_rotvec()-.008*velocity[:3]
    target=np.r_[desiredforce,desiredtorque/.03]
    columns=[];jacobians=[];rays_all=[];preferred=[]
    for c in selected:
        object_first=c.geom[0] in objects
        normal=-c.frame[:3]*(1 if object_first else -1)
        tangent=np.cross(normal,[0,0,1])
        if np.linalg.norm(tangent)<.01:tangent=np.cross(normal,[1,0,0])
        tangent/=np.linalg.norm(tangent);second=np.cross(normal,tangent)
        rays=np.array([normal+.6*(np.cos(a)*tangent+np.sin(a)*second) for a in np.arange(4)*np.pi/2]).T
        rays_all.append(rays);columns.append(np.vstack([rays,np.cross((c.pos-center)[None,:],rays.T).T/.03]))
        body=int(model.geom_bodyid[c.geom[1] if object_first else c.geom[0]])
        preferred.append(6 if model.body(body).name.startswith('thumb') else 1.5)
        jac=np.zeros((3,model.nv));mujoco.mj_jac(model,data,jac,None,c.pos,body)
        ja=np.stack([jac[:,model.jnt_dofadr[model.joint(n).id]] for n in cal.active],axis=1)
        for n,(parent,_,_) in cal.curves.items():
            slope=cal.passive_value(n,data.qpos[addr[parent]])[1]
            ja[:,cal.active.index(parent)]+=slope*jac[:,model.jnt_dofadr[model.joint(n).id]]
        jacobians.append(ja)
    A=np.concatenate(columns,axis=1);reg=np.zeros((len(selected),len(selected)*4))
    for i in range(len(selected)):reg[i,i*4:i*4+4]=.025
    sol=lsq_linear(np.vstack([A,reg]),np.r_[target,.025*np.array(preferred)],bounds=(0,10),tol=1e-6,max_iter=50)
    tau=sum(j.T@(r@sol.x[i*4:i*4+4]) for i,(j,r) in enumerate(zip(jacobians,rays_all)))
    for i,n in enumerate(cal.active):tau[i]+=data.qfrc_bias[model.jnt_dofadr[model.joint(n).id]]
    return tau,float(np.linalg.norm(A@sol.x-target))

def build_physics(kp=30):
    tree=ET.parse(BASE/'contact_object/contact_scene.xml');root=tree.getroot()
    hand=root.find("worldbody/body[@name='hand_base_link']")
    hand.remove(hand.find('freejoint'));hand.set('mocap','true')
    option=root.find('option');option.set('timestep','.001');option.set('integrator','implicitfast');option.set('iterations','80')
    root.find('default/joint').set('damping','.08')
    eq=root.find('equality')
    if eq is None:eq=ET.SubElement(root,'equality')
    actuator=ET.SubElement(root,'actuator')
    cal=Calibration()
    for n in cal.active:
        ET.SubElement(actuator,'position',name='motor_'+n,joint=n,kp=str(kp),kv='.3',
              forcelimited='true',forcerange='-1 1')
    for n,(parent,x,y) in cal.curves.items():
        ET.SubElement(eq,'joint',name='lut_'+n,joint1=n,joint2=parent,
            polycoef='0 1 0 0 0',solref='.003 1',solimp='.95 .99 .001')
    path=OUT/'physics_scene.xml';tree.write(path)
    return path

def run(seconds=3.,lift=False,kp=30.,preload=0.,force_preload=False,force_control=False,feedback=False):
    path=build_physics(kp);model=mujoco.MjModel.from_xml_path(str(path));data=mujoco.MjData(model)
    cal=Calibration();old=mujoco.MjModel.from_xml_path(str(BASE/'contact_object/contact_scene.xml'))
    fitted=np.load(BASE/'contact_object/contact_result.npz')['qpos']
    names=list(cal.tables)
    addr={n:model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)] for n in names}
    for n in names:
        oldaddr=old.jnt_qposadr[mujoco.mj_name2id(old,mujoco.mjtObj.mjOBJ_JOINT,n)]
        data.qpos[addr[n]]=fitted[oldaddr]
    ow=old.jnt_qposadr[mujoco.mj_name2id(old,mujoco.mjtObj.mjOBJ_JOINT,'wrist_free')]
    oo=old.jnt_qposadr[mujoco.mj_name2id(old,mujoco.mjtObj.mjOBJ_JOINT,'object_free')]
    obj=model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,'object_free')]
    # Rigidly reorient the entire registration so bottle +Z is world up.
    rotation=Rotation.from_quat(fitted[oo+3:oo+7][[1,2,3,0]]).as_matrix().T
    handpos=rotation@(fitted[ow:ow+3]-fitted[oo:oo+3])+[0,0,.3]
    handrot=rotation@Rotation.from_quat(fitted[ow+3:ow+7][[1,2,3,0]]).as_matrix()
    data.mocap_pos[0]=handpos;data.mocap_quat[0]=Rotation.from_matrix(handrot).as_quat()[[3,0,1,2]]
    data.qpos[obj:obj+7]=[0,0,.3,1,0,0,0]
    controls=np.array([data.qpos[addr[n]] for n in cal.active])
    reference_controls=controls.copy()
    for i,n in enumerate(cal.active):
        if n.endswith(('pitch','pip','mcp')):controls[i]+=preload
        controls[i]=np.clip(controls[i],*cal.bounds[n])
    data.ctrl[:]=controls
    objects=np.flatnonzero(model.geom_contype==2)
    object_body=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_BODY,'grasp_object')
    preload_report={}
    if force_preload:
        for n,(parent,x,y) in cal.curves.items():
            a=data.qpos[addr[parent]];v,s=cal.passive_value(n,a)
            ei=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_EQUALITY,'lut_'+n)
            model.eq_data[ei,:5]=[v-s*a,s,0,0,0]
        mujoco.mj_forward(model,data)
        selected=[c for c in data.contact if any(g in objects for g in c.geom)]
        columns=[];jacobians=[];cone=[]
        center=data.xipos[object_body].copy()
        for c in selected:
            object_first=c.geom[0] in objects
            outward=c.frame[:3]*(1 if object_first else -1)
            normal=-outward
            tangent=np.cross(normal,[0,0,1])
            if np.linalg.norm(tangent)<.01:tangent=np.cross(normal,[1,0,0])
            tangent/=np.linalg.norm(tangent);second=np.cross(normal,tangent)
            rays=np.array([normal+.6*(np.cos(a)*tangent+np.sin(a)*second) for a in np.arange(4)*np.pi/2]).T
            cone.append(rays)
            columns.append(np.vstack([rays,np.cross((c.pos-center)[None,:],rays.T).T/.03]))
            body=model.geom_bodyid[c.geom[1] if object_first else c.geom[0]]
            jac=np.zeros((3,model.nv));mujoco.mj_jac(model,data,jac,None,c.pos,int(body))
            activejac=np.stack([jac[:,model.jnt_dofadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)]] for n in cal.active],axis=1)
            for n,(parent,x,y) in cal.curves.items():
                slope=cal.passive_value(n,data.qpos[addr[parent]])[1]
                di=model.jnt_dofadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)]
                activejac[:,cal.active.index(parent)]+=slope*jac[:,di]
            jacobians.append(activejac)
        if not selected:raise RuntimeError('No contacts for force preload')
        A=np.concatenate(columns,axis=1)
        regular=np.zeros((len(selected),len(selected)*4))
        for i in range(len(selected)):regular[i,i*4:i*4+4]=.025
        target=np.r_[0,0,model.body_mass[object_body]*9.81,0,0,0]
        preferred=np.array([6.]+[1.5]*(len(selected)-1))
        solution=lsq_linear(np.vstack([A,regular]),np.r_[target,.025*preferred],bounds=(0,10),tol=1e-10)
        tau=np.zeros(16)
        for i,jac in enumerate(jacobians):tau+=jac.T@(cone[i]@solution.x[i*4:i*4+4])
        for i,n in enumerate(cal.active):
            dof=model.jnt_dofadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)]
            tau[i]+=data.qfrc_bias[dof]
        controls+=np.clip(tau,-1,1)/kp
        for i,n in enumerate(cal.active):controls[i]=np.clip(controls[i],*cal.bounds[n])
        data.ctrl[:]=controls
        preload_report=dict(wrench_residual=float(np.linalg.norm(A@solution.x-target)),desired_motor_torque_Nm=tau.tolist(),contact_normal_loads_N=(regular@solution.x/.025).tolist())
    rows=[];forces=[];contacts=[];relative=[];torques=[];wrist_path=[]
    initial_rel=np.array([0,0,.3])-handpos
    feedback_errors=[]
    for k in range(int(seconds/model.opt.timestep)):
        t=k*model.opt.timestep
        for n,(parent,x,y) in cal.curves.items():
            active_angle=data.qpos[addr[parent]];passive,slope=cal.passive_value(n,active_angle)
            eid=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_EQUALITY,'lut_'+n)
            model.eq_data[eid,:5]=[passive-slope*active_angle,slope,0,0,0]
        displacement=np.zeros(3)
        if lift:
            phase=np.clip((t-.5)/1.,0,1);smooth=phase*phase*(3-2*phase)
            displacement=np.array([.02*smooth,0,.10*smooth])
        data.mocap_pos[0]=handpos+displacement
        if feedback and k%10==0:
            phase=np.clip((t-.5),0,1)
            target_velocity=np.array([.02,0,.10])*6*phase*(1-phase) if lift else np.zeros(3)
            target_center=np.array([0,0,.3663])+displacement
            newtau,residual=feedback_torque(model,data,cal,addr,objects,object_body,target_center,target_velocity)
            if newtau is not None:tau=newtau;feedback_errors.append(residual)
        if force_control:
            if not force_preload:raise ValueError('force-control requires force-preload')
            angles=np.array([data.qpos[addr[n]] for n in cal.active])
            data.ctrl[:]=angles+(tau+1.0*(reference_controls-angles))/kp
        mujoco.mj_step(model,data)
        if not np.isfinite(data.qpos).all():raise RuntimeError('Nonfinite simulation')
        if k%33==0:
            rows.append(data.qpos.copy());wrist_path.append(np.r_[data.mocap_pos[0],data.mocap_quat[0]])
            selected=[i for i,c in enumerate(data.contact) if any(g in objects for g in c.geom)]
            normalforce=0
            for ci in selected:
                f=np.zeros(6);mujoco.mj_contactForce(model,data,ci,f);normalforce+=max(0,f[0])
            forces.append(normalforce);contacts.append(len(selected));torques.append(data.actuator_force.copy())
            relative.append(np.linalg.norm(data.qpos[obj:obj+3]-data.mocap_pos[0]-initial_rel))
    rows=np.array(rows);relative=np.array(relative)
    report=dict(seconds=seconds,lift=lift,kp=kp,preload_rad=preload,force_preload=preload_report,force_control=force_control,feedback=feedback,
        feedback_wrench_residual_mean=float(np.mean(feedback_errors)) if feedback_errors else None,
        max_relative_drift_m=float(relative.max()),final_relative_drift_m=float(relative[-1]),
        object_vertical_displacement_m=float(data.qpos[obj+2]-.3),
        contact_frames=int(np.sum(np.array(contacts)>0)),total_frames=len(rows),
        normal_force_mean_N=float(np.mean(forces)),peak_actuator_torque_Nm=float(np.max(np.abs(torques))),
        retained=bool(relative.max()<.02),
        mass_kg=float(model.body_mass[object_body]),friction=.6,
        boundaries=['Prescribed moving wrist, no arm dynamics','16 force-limited active joints; 5 nonlinear LUT equality constraints','No object weld or position reset','Object registration is estimated from hand pose, not tri-camera object tracking'])
    np.savez_compressed(OUT/'physics_motion.npz',qpos=rows,wrist=np.array(wrist_path),fps=1/(33*model.opt.timestep),
             object_qpos_address=obj,contact_count=contacts,normal_force=forces,relative_drift=relative,actuator_force=torques,
             source_to_upright_rotation=rotation,controls=controls)
    (OUT/'physics_validation.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2),flush=True)
    return report

if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('--seconds',type=float,default=3);p.add_argument('--lift',action='store_true');p.add_argument('--kp',type=float,default=30);p.add_argument('--preload',type=float,default=0);p.add_argument('--force-preload',action='store_true');p.add_argument('--force-control',action='store_true');p.add_argument('--feedback',action='store_true');a=p.parse_args()
    run(a.seconds,a.lift,a.kp,a.preload,a.force_preload,a.force_control,a.feedback)