#!/usr/bin/env python3
"""V10 constrained zero refinement, with exact active-pinch pose anchors.

The photographs suggest straight, nearly parallel fingers; zero-degree targets
are a geometric prior, not calibrated image measurements. Sweep bounded changes
and reject any contact, endpoint or measured-pose regression beyond stated gates.
"""
import argparse,json,os
os.environ.setdefault('MUJOCO_GL','egl')
from pathlib import Path
import xml.etree.ElementTree as ET
import numpy as np
from scipy.spatial.transform import Rotation
import fit_middle as v9
base=v9.base
OUTPUT=base.MODEL_DIR/'linkerhand_o30_right_zero_v10_calibrated.urdf'
PARAMETERS=base.MODEL_DIR/'o30_right_zero_v10_calibration.json'
REPORT=base.ROOT/'outputs/reports/O30/o30_right_zero_v10_calibration'
ZERO=np.array([0,0,255,197,110,28,0,0,0,0,0,0,0,0,0,0,0,0,0,0],dtype=float)
SELECTED={2:0,3:1,4:2,5:3,7:0,11:0}
GROUPS=[[0,1,6,15],[2,7,11,16],[3,8,12,17],[4,9,13,18],[5,10,14,19]]
PROTECTED_ENDPOINTS=[8,9,10,12,14,15,17,18,19]


def setup():
    fit,previous=v9.setup();_,old=v9.load(previous)
    return fit,old


def anchored_candidate(c,old,fraction):
    q=old.copy();oldzero=c.positions(old,ZERO)
    for i,k in SELECTED.items():
        u=c.commands[k][i]/255;z=ZERO[i]/255
        anchor=c.positions(old,c.commands[k])[i];target=oldzero[i]*(1-fraction)
        span=(anchor-target)/(u-z)
        q[i]=target-z*span-c.lo[i];q[20+i]=span/c.span[i]-1
    return q


def score(fit,old,q):
    c=fit.cal;contacts=fit.contacts(q)
    for row in contacts:
        assert abs(row['mesh_gap_mm'])<.2 and row['penetration_mm']<.2
    for k,cmd in enumerate(c.commands):
        indices=GROUPS[0]+GROUPS[k+1]
        assert np.max(abs(c.positions(q,cmd)[indices]-c.positions(old,cmd)[indices]))<1e-12
    ep=np.full(20,255.)
    assert np.array_equal(c.positions(q,ep)[PROTECTED_ENDPOINTS],c.positions(old,ep)[PROTECTED_ENDPOINTS])
    c.set_pose(q,ZERO);penetration=max([0.]+[-ct.dist*1000 for ct in c.data.contact]);assert penetration<.2
    measured=[]
    for cmd,finger in zip(v9.v7.review.COMMANDS,['middle','thumb']):
        points=[]
        for parameters in [old,q]:
            c.set_pose(parameters,cmd)
            ids=[b for b in range(c.model.nbody) if c.model.body(b).name.startswith(finger+'_')]
            points.append(c.data.xpos[ids].copy())
        change=float(np.max(np.linalg.norm(points[1]-points[0],axis=1))*1000)
        assert change<.5
        measured.append(dict(finger=finger,max_body_origin_change_mm=change))
    angles=np.rad2deg(c.positions(q,ZERO))
    return dict(contacts=contacts,zero_penetration_mm=penetration,
        zero_selected_rms_deg=float(np.sqrt(np.mean(angles[list(SELECTED)]**2))),
        zero_angles_deg=dict(zip(base.NAMES,angles.tolist())),measured_single_joint_pose_changes=measured,
        protected_endpoints_deg={base.NAMES[i]:float(np.rad2deg(c.positions(q,ep)[i])) for i in PROTECTED_ENDPOINTS})


def fit_zero(fit,old):
    history=[];best=None
    for fraction in np.linspace(0,1,21):
        q=anchored_candidate(fit.cal,old,fraction)
        try:
            metrics=score(fit,old,q)
            history.append(dict(fraction=float(fraction),accepted=True,rms_deg=metrics['zero_selected_rms_deg']))
            if best is None or metrics['zero_selected_rms_deg']<best[0]:best=(metrics['zero_selected_rms_deg'],q)
        except AssertionError:
            history.append(dict(fraction=float(fraction),accepted=False))
    if best is None:raise RuntimeError('No candidate satisfies prior constraints')
    return best[1],history


def save(fit,old,q,history):
    tree=ET.parse(v9.OUTPUT)
    for i in SELECTED:
        joint=tree.getroot().find(f"joint[@name='{base.NAMES[i]}']");origin=joint.find('origin')
        axis=np.fromstring(joint.find('axis').get('xyz'),sep=' ')
        rot=Rotation.from_euler('xyz',np.fromstring(origin.get('rpy'),sep=' '))*Rotation.from_rotvec(axis*(q[i]-old[i]))
        origin.set('rpy',' '.join(format(v,'.16g') for v in rot.as_euler('xyz')))
        joint.find('limit').set('upper',format(fit.cal.lo[i]+fit.cal.span[i]*(1+q[i+20]),'.16g'))
    ET.indent(tree,space='  ');tree.write(OUTPUT,encoding='utf-8',xml_declaration=True)
    info=dict(parameter_vector=q.tolist(),zero_command=ZERO.tolist(),changed_joints=[base.NAMES[i] for i in SELECTED],
        baseline=score(fit,old,old),result=score(fit,old,q),history=history,
        inputs={str(p.relative_to(base.ROOT)):base.sha256(p) for p in [v9.OUTPUT,v9.PARAMETERS,v9.SAMPLES]},
        output_sha256=base.sha256(OUTPUT),
        limitations='Straight/parallel zero is a visual geometric prior, not calibrated pixel/angle measurement. '
        'Only active pinching fingers and thumb are exactly pose-preserved; resting fingers can change. '
        'Index DIP zero is frozen because its pinch and zero commands are both 0. Left stays V9.')
    PARAMETERS.write_text(json.dumps(info,indent=2,ensure_ascii=False)+'\n');return info


def load(old):
    info=json.loads(PARAMETERS.read_text())
    for p,h in info['inputs'].items():assert base.sha256(base.ROOT/p)==h,p
    assert base.sha256(OUTPUT)==info['output_sha256']
    q=np.array(info['parameter_vector']);mask=np.ones(56,dtype=bool)
    mask[list(SELECTED)+[20+i for i in SELECTED]]=False
    assert np.array_equal(q[mask],old[mask])
    return info,q


def validate(fit,old,q):
    from linker_hand_mujoco_ros2.utils import o30_right_joints as mapping
    mapping.configure_o30_right_limits(base.LIVE_XML);assert not mapping.O30_RIGHT_SDK_KNOTS
    r=ET.parse(OUTPUT).getroot();original=ET.parse(v9.OUTPUT).getroot();copy=ET.fromstring(ET.tostring(r))
    for i in SELECTED:
        for tag in ['origin','limit']:
            path=f"joint[@name='{base.NAMES[i]}']/{tag}"
            copy.find(path).attrib.clear();copy.find(path).attrib.update(original.find(path).attrib)
    assert ET.canonicalize(ET.tostring(copy,encoding='unicode'),strip_text=True)==ET.canonicalize(ET.tostring(original,encoding='unicode'),strip_text=True)
    model=base.mujoco.MjModel.from_xml_path(str(base.LIVE_XML));data=base.mujoco.MjData(model)
    native=base.mujoco.MjModel.from_xml_path(str(OUTPUT));nd=base.mujoco.MjData(native)
    commands=[*fit.cal.commands,ZERO,fit.cal.zero_command,*v9.v7.review.COMMANDS,np.zeros(20),np.full(20,255),*np.random.default_rng(310).uniform(0,255,(30,20))]
    maxerr=0.;actual=[]
    for k,cmd in enumerate(commands):
        cmd=np.asarray(cmd,dtype=float)
        mapped=mapping.sdk_range_to_full_o30_right_urdf_positions(cmd.tolist())
        for i,n in enumerate(base.NAMES):
            lim=r.find(f"joint[@name='{n}']/limit");lo,hi=float(lim.get('lower')),float(lim.get('upper'))
            angle=lo+(hi-lo)*float(np.clip(cmd[i],0,255))/255
            assert abs(mapped[n]-angle)<1e-10
            for m,d in [(model,data),(native,nd)]:d.qpos[m.jnt_qposadr[m.joint(n).id]]=angle
        base.mujoco.mj_forward(model,data);base.mujoco.mj_forward(native,nd);fit.cal.set_pose(q,cmd)
        for b in range(1,native.nbody):
            n=native.body(b).name;j=model.body(n).id;s=fit.cal.model.body(n).id
            e=max(np.max(abs(nd.xpos[b]-data.xpos[j])),np.max(abs(nd.xmat[b]-data.xmat[j])),np.max(abs(nd.xpos[b]-fit.cal.data.xpos[s])))
            maxerr=max(maxerr,float(e));assert e<2e-8
        if k<4:
            finger=fit.cal.samples[k]['finger']
            def geom(n):return next(g for g in range(model.ngeom) if model.geom_group[g]==3 and model.geom_bodyid[g]==model.body(n+'_distal').id)
            gap=float(base.mujoco.mj_geomDistance(model,data,geom('thumb'),geom(finger),.2,np.zeros(6))*1000)
            penetration=max([0.]+[-ct.dist*1000 for ct in data.contact]);assert abs(gap)<.2 and penetration<.2
            actual.append(dict(finger=finger,gap_mm=gap,penetration_mm=penetration))
        if k in [4,5]:assert max([0.]+[-ct.dist*1000 for ct in data.contact])<.2
    return dict(tested_commands=len(commands),only_selected_origins_and_limits_changed=True,no_piecewise_mapping=True,
        max_native_runtime_reference_error=maxerr,actual_contacts=actual,score=score(fit,old,q))


def render(fit,old,q):
    from PIL import Image,ImageDraw
    c=fit.cal;m=c.model;m.vis.global_.offwidth=600;m.vis.global_.offheight=780
    for g in range(m.ngeom):
        if m.geom_group[g]==2:m.geom_rgba[g]=[.65,.67,.69,1]
    sheet=Image.new('RGB',(1200,1640),'#eef2f6');draw=ImageDraw.Draw(sheet)
    with base.mujoco.Renderer(m,780,600) as renderer:
        for row,param in enumerate([old,q]):
            c.set_pose(param,ZERO)
            for col,azi in enumerate([180,90]):
                cam=base.mujoco.MjvCamera();cam.lookat[:]=[.045,0,.09];cam.distance=.4;cam.azimuth=azi;cam.elevation=-5
                renderer.update_scene(c.data,cam);sheet.paste(Image.fromarray(renderer.render()),(col*600,row*820+40))
                draw.text((col*600+12,row*820+12),('V9 before' if row==0 else 'V10 zero refinement')+' / '+('front' if col==0 else 'side'),fill='black')
    sheet.save(REPORT/'comparison.png')


def main():
    p=argparse.ArgumentParser();p.add_argument('--fit',action='store_true');p.add_argument('--check',action='store_true');p.add_argument('--render',action='store_true');a=p.parse_args()
    fit,old=setup();REPORT.mkdir(parents=True,exist_ok=True)
    if a.fit:q,h=fit_zero(fit,old);info=save(fit,old,q,h);print(json.dumps(info['result'],indent=2))
    else:info,q=load(old)
    if a.check:
        result=validate(fit,old,q);(REPORT/'validation.json').write_text(json.dumps(result,indent=2)+'\n');print(json.dumps(result,indent=2))
    if a.render:render(fit,old,q)

if __name__=='__main__':main()