"""Check mirrored kinematics, inertia, contact distances and launcher conversion."""
import os
os.environ.setdefault('MUJOCO_GL', 'egl')
import json
import numpy as np
import mujoco
from PIL import Image, ImageDraw
from build import ROOT, MODEL, RIGHT, LEFT, OUTPUT


def check():
    right = mujoco.MjModel.from_xml_path(str(RIGHT))
    left = mujoco.MjModel.from_xml_path(str(LEFT))
    live = mujoco.MjModel.from_xml_path(str(MODEL/'linker_hand_o30_left/linker_hand_o30_left.xml'))
    rd, ld, vd = (mujoco.MjData(m) for m in [right, left, live])
    spec = json.loads((OUTPUT/'commands.json').read_text())
    order = spec['left_ros_order']
    reflection = np.diag([1., -1., 1.])
    commands = [v['left_ros_radians'] for v in spec['pinches'].values()]
    lo = np.array([right.jnt_range[right.joint(n).id,0] for n in order])
    hi = np.array([right.jnt_range[right.joint(n).id,1] for n in order])
    commands += [spec['zero']['left_ros_radians'], lo, hi, *np.random.default_rng(309).uniform(lo, hi, (30,20))]
    for name in order:
        assert np.array_equal(right.jnt_range[right.joint(name).id], left.jnt_range[left.joint(name).id])
    max_position = max_rotation = max_runtime = 0.
    contact_rows = []
    for k, command in enumerate(commands):
        for m,d in [(right,rd),(left,ld),(live,vd)]:
            for n,a in zip(order,command):d.qpos[m.jnt_qposadr[m.joint(n).id]]=a
            mujoco.mj_forward(m,d)
        for b in range(1,right.nbody):
            name=right.body(b).name;lb=left.body(name).id;vb=live.body(name).id
            max_position=max(max_position,float(np.max(abs(ld.xpos[lb]-reflection@rd.xpos[b]))))
            max_rotation=max(max_rotation,float(np.max(abs(ld.xmat[lb].reshape(3,3)-reflection@rd.xmat[b].reshape(3,3)@reflection))))
            max_runtime=max(max_runtime,float(np.max(abs(ld.xpos[lb]-vd.xpos[vb]))),float(np.max(abs(ld.xmat[lb]-vd.xmat[vb]))))
            assert abs(right.body_mass[b]-left.body_mass[lb])<1e-12
            for m,bid in [(right,b),(left,lb)]:
                rot=np.empty(9);mujoco.mju_quat2Mat(rot,m.body_iquat[bid]);rot=rot.reshape(3,3)
                tensor=rot@np.diag(m.body_inertia[bid])@rot.T
                if m is right:rt=tensor
                else:assert np.max(abs(tensor-reflection@rt@reflection))<1e-10
        if k<4:
            finger=list(spec['pinches'])[k]
            gaps=[]
            for m,d in [(right,rd),(left,ld),(live,vd)]:
                def geom(n):
                    ids=[g for g in range(m.ngeom) if m.geom_bodyid[g]==m.body(n+'_distal').id]
                    return next((g for g in ids if m.geom_group[g]==3), ids[0])
                gaps.append(float(mujoco.mj_geomDistance(m,d,geom('thumb'),geom(finger),.2,np.zeros(6))*1000))
            assert max(abs(x) for x in gaps)<.2,(finger,gaps)
            contact_rows.append(dict(finger=finger,right_gap_mm=gaps[0],left_gap_mm=gaps[1],runtime_left_gap_mm=gaps[2]))
    assert max_position<2e-8 and max_rotation<2e-8 and max_runtime<2e-8
    result=dict(commands=len(commands),max_mirror_position_error_m=max_position,max_mirror_rotation_error=max_rotation,
                max_runtime_pose_error=max_runtime,contacts=contact_rows,mirrored_inertia=True)
    (OUTPUT/'validation.json').write_text(json.dumps(result,indent=2)+'\n')
    print(json.dumps(result,indent=2))
    # Paired render: equal commands, opposite camera azimuths expose matching palms.
    sheet=Image.new('RGB',(1000,700),'#eef2f6');draw=ImageDraw.Draw(sheet)
    for col,(m,d,azimuth) in enumerate([(right,rd,175),(left,ld,-175)]):
        command=spec['pinches']['middle']['left_ros_radians']
        for n,a in zip(order,command):d.qpos[m.jnt_qposadr[m.joint(n).id]]=a
        mujoco.mj_forward(m,d);m.vis.global_.offwidth=500;m.vis.global_.offheight=650
        camera=mujoco.MjvCamera();camera.lookat[:]=[.045,0,.10];camera.distance=.32;camera.azimuth=azimuth;camera.elevation=-5
        with mujoco.Renderer(m,650,500) as renderer:
            renderer.update_scene(d,camera);sheet.paste(Image.fromarray(renderer.render()),(500*col,40))
        draw.text((500*col+12,12),'Right V10' if col==0 else 'Left mirrored V10',fill='black')
    sheet.save(OUTPUT/'comparison.png')


if __name__=='__main__':check()