#!/usr/bin/env python3
"""Provisional single-joint endpoint fitting, preserving V4 pinch/zero anchors exactly."""
import argparse
import json
import os
from pathlib import Path
import sys
import xml.etree.ElementTree as ET
if '--render' in sys.argv:
    os.environ.setdefault('MUJOCO_GL','egl')
import numpy as np
from scipy.optimize import least_squares
import fit_pinky as v4
base=v4.base
CONFIG=Path(__file__).with_name('endpoint_constraints.json')
OUTPUT=base.MODEL_DIR/'linkerhand_o30_right_endpoint_calibrated.urdf'
PARAMETERS=base.MODEL_DIR/'o30_right_endpoint_calibration.json'
REPORT=base.ROOT/'outputs/reports/O30/o30_right_endpoint_calibration'


class EndpointFit:
    def __init__(self):
        source=v4.PinkyFit()
        self.previous,self.q=v4.load(source)
        self.cal=source.cal
        self.config=json.loads(CONFIG.read_text())
        self.constraints=self.config['constraints']
        self.names=[row['joint'] for row in self.constraints]
        self.indices=np.array([base.NAMES.index(n) for n in self.names])
        self.offsets=self.q[:20].copy()
        tree=ET.parse(v4.OUTPUT)
        self.limits=np.array([[float(tree.getroot().find(f"joint[@name='{name}']/limit").get(key))
                               for key in ['lower','upper']] for name in base.NAMES])
        self.span=self.limits[:,1]-self.limits[:,0]
        self.targets=np.deg2rad([row['target_deg'] for row in self.constraints])
        self.sigma=np.deg2rad([row['uncertainty_deg'] for row in self.constraints])
        self.old_endpoints=self.limits[self.indices,1]+self.offsets[self.indices]
        # All existing measured pinch commands and the measured zero are protected.
        self.anchors=[np.unique([0.,*[cmd[i] for cmd in self.cal.commands],self.cal.zero_command[i]])
                      for i in self.indices]
        if any(a[-1]>=255 for a in self.anchors):
            raise ValueError('An existing pinch already constrains 255; endpoint needs a joint refit')
        self.scenarios=[]
        for background in ['measured_zero','all_zero']:
            for row in self.constraints:
                command=self.cal.zero_command.copy() if background=='measured_zero' else np.zeros(20)
                for name in row['active']: command[base.NAMES.index(name)]=255
                self.scenarios.append(dict(background=background,joint=row['joint'],command=command))
        self.baseline_penetration=np.array([self.penetration(self.positions(s['command'])) for s in self.scenarios])
        self.evaluations=0

    def positions(self,command,endpoints=None):
        cmd=np.clip(np.asarray(command,dtype=float),0,255)
        q=self.limits[:,0]+self.span*cmd/255+self.offsets
        if endpoints is not None:
            for j,i in enumerate(self.indices):
                x=np.r_[self.anchors[j],255.]
                y=np.r_[self.limits[i,0]+self.span[i]*self.anchors[j]/255+self.offsets[i],endpoints[j]]
                q[i]=np.interp(cmd[i],x,y)
        return q

    def set_pose(self,positions):
        self.cal.data.qpos[self.cal.adr]=positions
        base.mujoco.mj_forward(self.cal.model,self.cal.data)

    def penetration(self,positions):
        self.set_pose(positions)
        return max([0.]+[-c.dist for c in self.cal.data.contact])

    def residual(self,endpoints):
        self.evaluations+=1
        depths=np.array([self.penetration(self.positions(s['command'],endpoints)) for s in self.scenarios])
        return np.r_[(endpoints-self.targets)/self.sigma,
                     np.maximum(depths-self.baseline_penetration-.0002,0)/.0002]

    def fit(self):
        # Several endpoint hypotheses, then local refinement using MuJoCo collisions.
        history=[];best=None
        for start in [self.targets,self.targets-np.deg2rad(5),self.targets+np.deg2rad(5)]:
            result=least_squares(self.residual,start,bounds=(self.targets-self.sigma,self.targets+self.sigma),
                                 max_nfev=120,diff_step=1e-5,ftol=1e-10,xtol=1e-10,gtol=1e-10)
            history.append(dict(cost=float(result.cost),nfev=result.nfev,endpoints_deg=np.rad2deg(result.x).tolist()))
            if best is None or result.cost<best.cost:best=result
        self.check_anchors(best.x)
        return best.x,history

    def check_anchors(self,endpoints):
        for cmd in [*self.cal.commands,self.cal.zero_command]:
            if np.max(abs(self.positions(cmd,endpoints)-self.positions(cmd)))>1e-12:
                raise AssertionError('An existing pinch/zero joint angle changed')
        for j,i in enumerate(self.indices):
            last=self.limits[i,0]+self.span[i]*self.anchors[j][-1]/255+self.offsets[i]
            if endpoints[j]<=last:raise AssertionError('Non-monotone endpoint')

    def export(self,endpoints,history):
        self.check_anchors(endpoints)
        tree=ET.parse(v4.OUTPUT);root=tree.getroot()
        extension=ET.SubElement(root,'o30_sdk_mapping',version='1',interpolation='linear')
        changed={}
        for j,(name,i) in enumerate(zip(self.names,self.indices)):
            upper=endpoints[j]-self.offsets[i]
            root.find(f"joint[@name='{name}']/limit").set('upper',format(upper,'.16g'))
            xs=np.r_[self.anchors[j],255.]
            ys=np.r_[self.limits[i,0]+self.span[i]*self.anchors[j]/255,upper]
            ET.SubElement(extension,'map',joint=name,sdk=' '.join(format(x,'.16g') for x in xs),
                          radians=' '.join(format(y,'.16g') for y in ys))
            changed[name]=dict(previous_endpoint_deg=float(np.rad2deg(self.old_endpoints[j])),
                               endpoint_deg=float(np.rad2deg(endpoints[j])),sdk_knots=xs.tolist(),
                               urdf_radian_knots=ys.tolist())
        ET.indent(tree,space='  ');tree.write(OUTPUT,encoding='utf-8',xml_declaration=True)
        info=dict(status='Provisional photo-estimated endpoints, not measured angles or confirmed full commands',
                  constraints=self.config,parameterization='Piecewise linear SDK-to-radian mapping with exact V4 pinch and zero anchors',
                  endpoints_rad=endpoints.tolist(),changed_joints=changed,history=history,evaluations=self.evaluations,
                  hashes={str(p.relative_to(base.ROOT)):base.sha256(p) for p in [CONFIG,v4.OUTPUT,v4.PARAMETERS,base.STOCK_XML,base.SAMPLES]},
                  output_sha256=base.sha256(OUTPUT),
                  hypotheses=[dict(background=s['background'],joint=s['joint'],command=s['command'].tolist()) for s in self.scenarios])
        PARAMETERS.write_text(json.dumps(info,indent=2,ensure_ascii=False)+'\n')
        return info

    def load(self):
        info=json.loads(PARAMETERS.read_text())
        for path,digest in info['hashes'].items():
            if base.sha256(base.ROOT/path)!=digest:raise ValueError(f'Calibration input changed: {path}')
        if base.sha256(OUTPUT)!=info['output_sha256']:raise ValueError('Output URDF changed')
        endpoints=np.array(info['endpoints_rad']);self.check_anchors(endpoints)
        return info,endpoints

    def validate(self,endpoints):
        from linker_hand_mujoco_ros2.utils import o30_right_joints as mapping
        mapping.configure_o30_right_limits(base.LIVE_XML)
        if set(mapping.O30_RIGHT_SDK_KNOTS)!=set(self.names):raise AssertionError('Model did not load six nonlinear mappings')
        model=base.mujoco.MjModel.from_xml_path(str(base.LIVE_XML));data=base.mujoco.MjData(model)
        commands=[*self.cal.commands,self.cal.zero_command,*[s['command'] for s in self.scenarios],
                  np.zeros(20),np.full(20,255),*np.random.default_rng(305).uniform(0,255,(30,20))]
        error=0.;contact=[];scenarios=[]
        for k,cmd in enumerate(commands):
            self.set_pose(self.positions(cmd,endpoints))
            mapped=mapping.sdk_range_to_full_o30_right_urdf_positions(list(cmd))
            for name,value in mapped.items():
                jid=model.joint(name).id;data.qpos[model.jnt_qposadr[jid]]=value
                if not model.jnt_range[jid,0]-1e-10<=value<=model.jnt_range[jid,1]+1e-10:
                    raise AssertionError('Mapped value outside model limits')
            base.mujoco.mj_forward(model,data)
            for bid in range(1,self.cal.model.nbody):
                other=model.body(self.cal.model.body(bid).name).id
                diff=float(np.linalg.norm(data.xpos[other]-self.cal.data.xpos[bid]));error=max(error,diff)
                if diff>2e-7 or np.max(abs(data.xmat[other]-self.cal.data.xmat[bid]))>2e-6:
                    raise AssertionError('Actual exported kinematics mismatch')
            depth=max([0.]+[-c.dist for c in data.contact])*1000
            if k<4:
                f=self.cal.samples[k]['finger']
                def geom(finger):return next(g for g in range(model.ngeom) if model.geom_group[g]==3 and model.geom_bodyid[g]==model.body(finger+'_distal').id)
                gap=base.mujoco.mj_geomDistance(model,data,geom('thumb'),geom(f),.2,np.zeros(6))*1000
                if abs(gap)>.2 or depth>.2:raise AssertionError('Pinch contact regression')
                contact.append(dict(finger=f,mesh_gap_mm=float(gap),max_penetration_mm=float(depth)))
            elif k==4:
                if depth>.2:raise AssertionError('Zero collision regression')
            elif k<17:
                old=self.baseline_penetration[k-5]*1000
                scenarios.append(dict(joint=self.scenarios[k-5]['joint'],background=self.scenarios[k-5]['background'],
                                      previous_penetration_mm=float(old),penetration_mm=float(depth)))
                if depth>old+.3:raise AssertionError('Endpoint collision regression')
        # Dense commands verify monotonic interpolation, bounds and exact anchors.
        mapped=np.array([list(mapping.sdk_range_to_full_o30_right_urdf_positions([u]*20).values())
                         for u in np.linspace(0,255,1021)])
        if np.any(np.diff(mapped,axis=0)<-1e-12):raise AssertionError('Non-monotone SDK mapping')
        self.check_anchors(endpoints)
        return dict(tested_commands=len(commands),dense_mapping_commands=1021,max_export_position_error_m=error,
                    pinch_and_zero_angles_preserved=True,pinches=contact,endpoint_hypotheses=scenarios)

    def render(self,endpoints):
        from PIL import Image,ImageDraw,ImageFont
        REPORT.mkdir(parents=True,exist_ok=True)
        model=self.cal.model;model.vis.global_.offwidth=480;model.vis.global_.offheight=600
        for g in range(model.ngeom):
            if model.geom_group[g]==2:model.geom_rgba[g]=[.48,.50,.52,1] if model.body(model.geom_bodyid[g]).name.endswith('_distal') else [.68,.70,.72,1]
        font=ImageFont.truetype('/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',16)
        with base.mujoco.Renderer(model,600,480) as renderer:
            for bg in ['measured_zero','all_zero']:
                sheet=Image.new('RGB',(480*6,640*2),'#f1f5f9');draw=ImageDraw.Draw(sheet)
                for row,ep in enumerate([None,endpoints]):
                    for col,scenario in enumerate([s for s in self.scenarios if s['background']==bg]):
                        self.set_pose(self.positions(scenario['command'],ep))
                        camera=base.mujoco.MjvCamera();camera.lookat[:]=[.04,0,.10]
                        camera.distance=.32;camera.azimuth=90;camera.elevation=-5
                        renderer.update_scene(self.cal.data,camera=camera)
                        sheet.paste(Image.fromarray(renderer.render()),(480*col,640*row+40))
                        draw.text((480*col+8,640*row+10),('V4' if row==0 else 'V5 provisional')+' / '+scenario['joint'],font=font,fill='#172033')
                sheet.save(REPORT/f'{bg}_comparison.png');print(REPORT/f'{bg}_comparison.png')


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    for flag in ['fit','check','render']:parser.add_argument('--'+flag,action='store_true')
    parser.add_argument('--preview',choices=['ring_mcp_pitch','ring_dip','pinky_mcp_pitch','pinky_pip','pinky_dip','middle_dip'])
    parser.add_argument('--background',choices=['measured_zero','all_zero'],default='measured_zero')
    args=parser.parse_args()
    if not any([args.fit,args.check,args.render,args.preview]):parser.error('Choose an action')
    fit=EndpointFit()
    if args.preview:
        import time
        import mujoco.viewer
        from linker_hand_mujoco_ros2.utils import o30_right_joints as mapping
        mapping.configure_o30_right_limits(base.LIVE_XML)
        model=base.mujoco.MjModel.from_xml_path(str(base.LIVE_XML));data=base.mujoco.MjData(model)
        command=next(s['command'] for s in fit.scenarios if s['joint']==args.preview and s['background']==args.background)
        for name,value in mapping.sdk_range_to_full_o30_right_urdf_positions(list(command)).items():
            data.qpos[model.jnt_qposadr[model.joint(name).id]]=value
        base.mujoco.mj_forward(model,data)
        print('Background hypothesis:',args.background,'command:',command.tolist())
        with mujoco.viewer.launch_passive(model,data) as viewer:
            viewer.cam.lookat[:]=[.04,0,.10];viewer.cam.distance=.32
            viewer.cam.azimuth=90;viewer.cam.elevation=-5
            while viewer.is_running():viewer.sync();time.sleep(.02)
        return
    if args.fit:
        endpoints,history=fit.fit();info=fit.export(endpoints,history)
        print(json.dumps(info['changed_joints'],indent=2));print('Regenerate launcher XML before --check.')
    else:info,endpoints=fit.load()
    if args.check:
        info['validation']=fit.validate(endpoints);REPORT.mkdir(parents=True,exist_ok=True)
        (REPORT/'validation.json').write_text(json.dumps(info,indent=2,ensure_ascii=False)+'\n')
        print(json.dumps(info['validation'],indent=2))
    if args.render:fit.render(endpoints)


if __name__=='__main__':main()