#!/usr/bin/env python3
"""Prepare (do not activate) two-joint linear candidates and quantify pinch tradeoffs."""
import os
os.environ.setdefault('MUJOCO_GL','egl')
import json
import xml.etree.ElementTree as ET
import numpy as np
from scipy.spatial.transform import Rotation
from PIL import Image,ImageDraw,ImageFont
import fit_linear as linear
base=linear.base
REPORT=base.ROOT/'outputs/reports/O30/o30_right_two_joint_review'
COMMANDS=[
 [172,0,169,169,162,170,0,0,255,0,0,0,0,0,0,0,0,0,0,0],
 [172,0,169,169,162,170,0,0,0,0,0,0,0,0,0,255,0,0,0,0],
]


def export(q,old,names,path):
    tree=ET.parse(linear.OUTPUT)
    for name in names:
        i=base.NAMES.index(name);joint=tree.getroot().find(f"joint[@name='{name}']")
        origin=joint.find('origin');axis=np.fromstring(joint.find('axis').get('xyz'),sep=' ');axis/=np.linalg.norm(axis)
        rotation=Rotation.from_euler('xyz',np.fromstring(origin.get('rpy'),sep=' '))*Rotation.from_rotvec(axis*(q[i]-old[i]))
        if abs(q[i]-old[i])>1e-12:origin.set('rpy',' '.join(format(x,'.16g') for x in rotation.as_euler('xyz')))
        # All selected lower limits are zero, so upper equals stock span * fitted scale.
        stock=ET.parse(base.STOCK_URDF).getroot().find(f"joint[@name='{name}']/limit")
        low,high=float(stock.get('lower')),float(stock.get('upper'))
        joint.find('limit').set('upper',format(low+(high-low)*(1+q[20+i]),'.16g'))
    ET.indent(tree,space='  ');tree.write(path,encoding='utf-8',xml_declaration=True)
    before=ET.parse(linear.OUTPUT).getroot();after=ET.parse(path).getroot()
    def structure(e):return (e.tag,sorted(e.attrib.items()),(e.text or '').strip(),[structure(x) for x in e])
    for n in names:
        oldjoint=before.find(f"joint[@name='{n}']");newjoint=after.find(f"joint[@name='{n}']")
        for tag in ['origin','limit']:newjoint.find(tag).attrib.clear();newjoint.find(tag).attrib.update(oldjoint.find(tag).attrib)
    assert structure(before)==structure(after),'Unexpected change outside selected origins/limits'


def main():
    REPORT.mkdir(parents=True,exist_ok=True)
    fit=linear.LinearCalibration();_,old=fit.load();c=fit.cal;q=old.copy()
    endpoint=np.deg2rad(90.)
    i=8;u=78/255;anchor=c.positions(old,c.commands[1])[i]
    zero=(anchor-u*endpoint)/(1-u)
    q[i]=zero-c.lo[i];q[20+i]=(endpoint-zero)/c.span[i]-1
    middle=q.copy()
    i=15;q[20+i]=(endpoint-c.lo[i]-q[i])/c.span[i]-1
    paths={}
    for label,value,names in [('middle_only_candidate',middle,['middle_mcp_pitch']),
                              ('two_joint_candidate',q,['middle_mcp_pitch','thumb_ip'])]:
        path=base.MODEL_DIR/f'linkerhand_o30_right_{label}.urdf';export(value,old,names,path);paths[label]=str(path)
    info=dict(status='Comparison candidates relative to saved V6; this script does not modify launch configuration',
              exact_test_commands=COMMANDS,target_angles_are_photo_estimates=True,
              reference_urdf_sha256=base.sha256(linear.OUTPUT),candidates=paths,
              middle_zero_deg=float(np.rad2deg(zero)),
              middle_only_contacts=fit.contacts(middle),two_joint_contacts=fit.contacts(q))
    # Verify generic native-URDF consumers reproduce the candidate poses.
    models=[base.mujoco.MjModel.from_xml_path(str(linear.OUTPUT)),base.mujoco.MjModel.from_xml_path(paths['two_joint_candidate'])]
    data=[base.mujoco.MjData(m) for m in models]
    for m in models:
        m.vis.global_.offwidth=600;m.vis.global_.offheight=720
    sheets=Image.new('RGB',(1200,1520),'#eef2f6');draw=ImageDraw.Draw(sheets)
    font=ImageFont.truetype('/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',20)
    max_error=0.
    for row,(model,d,params) in enumerate(zip(models,data,[old,q])):
        with base.mujoco.Renderer(model,720,600) as renderer:
            for col,(name,cmd,azimuth) in enumerate(zip(['middle root 255','thumb tip 255'],COMMANDS,[90,175])):
                for i,joint in enumerate(base.NAMES):
                    jid=model.joint(joint).id;low,high=model.jnt_range[jid]
                    d.qpos[model.jnt_qposadr[jid]]=low+(high-low)*cmd[i]/255
                base.mujoco.mj_forward(model,d);c.set_pose(params,cmd)
                for bid in range(1,model.nbody):
                    source=c.model.body(model.body(bid).name).id
                    max_error=max(max_error,float(np.linalg.norm(d.xpos[bid]-c.data.xpos[source])))
                    assert np.max(abs(d.xmat[bid]-c.data.xmat[source]))<2e-6
                camera=base.mujoco.MjvCamera();camera.lookat[:]=[.045,0,.1]
                camera.distance=.34;camera.azimuth=azimuth;camera.elevation=-5
                renderer.update_scene(d,camera=camera)
                sheets.paste(Image.fromarray(renderer.render()),(600*col,760*row+40))
                draw.text((600*col+10,760*row+10),('V6 current' if row==0 else 'Candidate ~90 deg')+' / '+name,font=font,fill='#172033')
    assert max_error<2e-7
    info['native_urdf_export_error_m']=max_error
    (REPORT/'review.json').write_text(json.dumps(info,indent=2,ensure_ascii=False)+'\n')
    sheets.save(REPORT/'comparison.png')
    print(json.dumps(info,indent=2));print(REPORT/'comparison.png')


if __name__=='__main__':main()