#!/usr/bin/env python3
"""Portable pure-URDF linear calibration: compromise endpoints, zero and pinch shape."""
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
from scipy.spatial.transform import Rotation
import fit_endpoints as endpoints
base=endpoints.base
OUTPUT=base.MODEL_DIR/'linkerhand_o30_right_linear_calibrated.urdf'
PARAMETERS=base.MODEL_DIR/'o30_right_linear_calibration.json'
REPORT=base.ROOT/'outputs/reports/O30/o30_right_linear_calibration'


class LinearCalibration:
    def __init__(self):
        self.source=endpoints.EndpointFit()
        self.cal=self.source.cal
        self.previous=self.source.q.copy()
        self.old_zero=self.cal.positions(self.previous,self.cal.zero_command)
        self.old_poses=np.array([self.cal.positions(self.previous,cmd) for cmd in self.cal.commands])
        self.photo_parameters=np.array(json.loads(endpoints.v4.multi.OUTPUT_JSON.read_text())['camera_parameters'])
        self.pinky=endpoints.v4.PinkyFit()
        self.pinky_camera=np.array(self.source.previous['camera'])
        self.evaluations=0
        self.contact_sigma=.002

    def unpack(self,vector):
        q=self.previous.copy();q[:40]=vector
        return q

    def contacts(self,q):
        result=[]
        for sample in self.cal.samples:
            self.cal.set_pose(q,sample['command'])
            closest=np.zeros(6)
            gap=base.mujoco.mj_geomDistance(self.cal.model,self.cal.data,self.cal.geoms['thumb'],
                                           self.cal.geoms[sample['finger']],.2,closest)
            local={}
            for k,finger in enumerate(['thumb',sample['finger']]):
                bid=self.cal.bodies[finger]
                local[finger]=self.cal.data.xmat[bid].reshape(3,3).T@(closest[3*k:3*k+3]-self.cal.data.xpos[bid])
            result.append(dict(finger=sample['finger'],mesh_gap_mm=float(gap*1000),
                               penetration_mm=float(max([0.]+[-ct.dist for ct in self.cal.data.contact])*1000),
                               local_contact_m={name:value.tolist() for name,value in local.items()}))
        return result

    def residual(self,vector):
        self.evaluations+=1;q=self.unpack(vector);res=[]
        for k,sample in enumerate(self.cal.samples):
            self.cal.set_pose(q,sample['command']);closest=np.zeros(6)
            gap=base.mujoco.mj_geomDistance(self.cal.model,self.cal.data,self.cal.geoms['thumb'],
                                           self.cal.geoms[sample['finger']],.2,closest)
            res.append(gap/self.contact_sigma)
            res.extend(np.maximum(self.cal.inter_finger_penetrations()-.00003,0)/self.contact_sigma)
            for j,finger in enumerate(['thumb',sample['finger']]):
                bid=self.cal.bodies[finger]
                p=self.cal.data.xmat[bid].reshape(3,3).T@(closest[3*j:3*j+3]-self.cal.data.xpos[bid])
                region=([max(0,-.003-p[0]),max(0,.009-p[2]),max(0,p[2]-.025),max(0,abs(p[1])-.006)] if j else
                        [max(0,.027-p[1]),max(0,p[1]-.046),max(0,-.001-p[2]),max(0,abs(p[0])-.008)])
                res.extend(np.array(region)/.0003)
            res.extend((self.cal.positions(q,sample['command'])-self.old_poses[k])/np.deg2rad(15))
        res.extend((self.cal.positions(q,self.cal.zero_command)-self.old_zero)/np.deg2rad(2))
        res.extend((self.cal.positions(q,np.full(20,255))[self.source.indices]-self.source.targets)/np.deg2rad(3))
        res.extend(np.array(self.cal.photo_residual(q,self.photo_parameters))*.15)
        res.extend(((self.pinky.project(q,self.pinky_camera)-self.pinky.obs)/self.pinky.sigma[:,None]).ravel()*.15)
        res.extend((vector[20:40]-self.previous[20:40])/.5*.3)
        self.cal.set_pose(q,self.cal.zero_command)
        res.extend(np.maximum(self.cal.inter_finger_penetrations()-.00003,0)/.0001)
        return np.array(res)

    def score(self,q):
        self.cal.set_pose(q,self.cal.zero_command)
        zero_penetration=max([0.]+[-ct.dist for ct in self.cal.data.contact])*1000
        zero_delta=np.rad2deg(self.cal.positions(q,self.cal.zero_command)-self.old_zero)
        pose_delta=np.rad2deg(np.array([self.cal.positions(q,cmd) for cmd in self.cal.commands])-self.old_poses)
        return dict(contacts=self.contacts(q),zero_penetration_mm=float(zero_penetration),
                    max_zero_change_deg=float(np.max(abs(zero_delta))),
                    zero_changes_deg=dict(zip(base.NAMES,zero_delta.tolist())),
                    max_pinch_joint_change_deg=float(np.max(abs(pose_delta))),
                    pinch_joint_changes_deg={s['finger']:dict(zip(base.NAMES,row.tolist())) for s,row in zip(self.cal.samples,pose_delta)},
                    endpoints_deg=dict(zip(self.source.names,np.rad2deg(self.cal.positions(q,np.full(20,255))[self.source.indices]).tolist())))

    def accept(self,q):
        score=self.score(q)
        if score['max_zero_change_deg']>3 or score['zero_penetration_mm']>.2:
            raise ValueError('Zero pose regressed beyond tolerance')
        if score['max_pinch_joint_change_deg']>22:raise ValueError('Pinch shape changed too much')
        if max(abs(value-90) for value in score['endpoints_deg'].values())>8:
            raise ValueError('Endpoint outside photo-estimate tolerance')
        for row in score['contacts']:
            if abs(row['mesh_gap_mm'])>.2 or row['penetration_mm']>.2:
                raise ValueError('Pinch contact/collision tolerance exceeded')
            thumb=row['local_contact_m']['thumb'];finger=row['local_contact_m'][row['finger']]
            if not (.026<=thumb[1]<=.047 and thumb[2]>=-.0015 and abs(thumb[0])<=.009 and
                    finger[0]>=-.004 and .008<=finger[2]<=.026 and abs(finger[1])<=.007):
                raise ValueError('Contact outside near-tip region')
        return score

    def fit(self,starts=3):
        low=np.r_[np.maximum(self.previous[:20]-.12,-.4),np.maximum(self.previous[20:40]-.6,-.7)]
        high=np.r_[np.minimum(self.previous[:20]+.12,.4),np.minimum(self.previous[20:40]+.6,.7)]
        best=None;history=[];rng=np.random.default_rng(306)
        for attempt in range(starts):
            start=self.previous[:40].copy()
            start[20+self.source.indices]=(self.source.targets-self.cal.lo[self.source.indices]-start[self.source.indices])/self.cal.span[self.source.indices]-1
            if attempt:start+=rng.normal(0,.012,40)
            for self.contact_sigma in [.002,.0005,.00015,.00005]:
                result=least_squares(self.residual,np.clip(start,low+1e-9,high-1e-9),bounds=(low,high),
                                     max_nfev=200,diff_step=1e-4,ftol=1e-8,xtol=1e-10,x_scale='jac')
                start=result.x
                row=dict(start=attempt+1,contact_sigma_m=self.contact_sigma,cost=float(result.cost),nfev=result.nfev)
                history.append(row);print(json.dumps(row),flush=True)
            q=self.unpack(result.x)
            try: self.accept(q);accepted=True
            except ValueError as exc:accepted=False;history[-1]['rejection']=str(exc)
            history[-1]['accepted']=accepted
            if accepted and (best is None or result.cost<best.cost):best=result
        if best is None:raise RuntimeError('No acceptable linear compromise; V5 remains available')
        return self.unpack(best.x),history

    def inputs(self):
        paths=[endpoints.CONFIG,endpoints.v4.OUTPUT,endpoints.v4.PARAMETERS,base.STOCK_URDF,base.STOCK_XML,
               base.SAMPLES,endpoints.v4.multi.OUTPUT_JSON,endpoints.v4.ANNOTATIONS,
               endpoints.v4.multi.SIDE,endpoints.v4.multi.photo.ANNOTATIONS]
        return {str(p.relative_to(base.ROOT)):base.sha256(p) for p in paths}

    def save(self,q,history):
        score=self.accept(q)
        base.CAL_URDF=OUTPUT;base.export_urdf(self.cal,q)
        tree=ET.parse(OUTPUT);root=tree.getroot()
        for link in root.findall('link'):
            color=link.find('visual/material/color')
            if color is not None:color.set('rgba','0.48 0.50 0.52 1' if link.get('name','').endswith('_distal') else '0.68 0.70 0.72 1')
        root.insert(0,ET.Comment(' V6 pure URDF. Zero offsets baked into origin; ordinary linear SDK 0..255 to joint lower..upper. No custom mapping. '))
        # Relative meshes make the asset directory relocatable.
        for mesh in root.findall('.//mesh'):
            mesh.set('filename','meshes/'+Path(mesh.get('filename')).name)
        ET.indent(tree,space='  ');tree.write(OUTPUT,encoding='utf-8',xml_declaration=True)
        info=dict(status='Pure-linear compromise; endpoint targets remain provisional photo estimates',
                  parameter_vector=q.tolist(),inputs=self.inputs(),output_sha256=base.sha256(OUTPUT),
                  objective_evaluations=self.evaluations,history=history,score=score,
                  limitations='Linear mapping cannot retain every V5 intermediate angle and endpoint. '
                              'Pinch contact is preserved approximately; intermediate joint shapes change. '
                              'Unknown full endpoint commands and approximate 90-degree targets remain uncertain.',
                  mapping='q = lower + (upper-lower)*clip(sdk,0,255)/255; no extra offset or lookup table')
        PARAMETERS.write_text(json.dumps(info,indent=2,ensure_ascii=False)+'\n')
        return info

    def load(self):
        info=json.loads(PARAMETERS.read_text())
        if info['inputs']!=self.inputs() or info['output_sha256']!=base.sha256(OUTPUT):raise ValueError('Inputs or URDF changed')
        q=np.array(info['parameter_vector']);self.accept(q)
        return info,q

    def validate(self,q):
        from linker_hand_mujoco_ros2.utils import o30_right_joints as mapping
        mapping.configure_o30_right_limits(base.LIVE_XML)
        if mapping.O30_RIGHT_SDK_KNOTS:raise AssertionError('Linear model has custom mapping')
        urdf=ET.parse(OUTPUT).getroot()
        if urdf.find('o30_sdk_mapping') is not None:raise AssertionError('Custom URDF mapping present')
        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));native_data=base.mujoco.MjData(native)
        commands=[*self.cal.commands,self.cal.zero_command,*[s['command'] for s in self.source.scenarios],
                  np.zeros(20),np.full(20,255),*np.random.default_rng(306).uniform(0,255,(30,20))]
        max_error=0.;portable_error=0.;native_error=0.;scenarios=[]
        # Independent standard-URDF forward kinematics: no calibration vectors or SDK code.
        def portable_fk(angles):
            joints=urdf.findall('joint')
            children={joint.find('child').get('link') for joint in joints}
            roots={link.get('name') for link in urdf.findall('link')}-children
            if len(roots)!=1:raise AssertionError('URDF does not have one root')
            transforms={roots.pop():np.eye(4)}
            pending=list(joints)
            while pending:
                progressed=False
                for joint in pending[:]:
                    parent=joint.find('parent').get('link');child=joint.find('child').get('link')
                    if parent not in transforms:continue
                    origin=joint.find('origin');relative=np.eye(4)
                    if origin is not None:
                        relative[:3,3]=np.fromstring(origin.get('xyz','0 0 0'),sep=' ')
                        relative[:3,:3]=Rotation.from_euler('xyz',np.fromstring(origin.get('rpy','0 0 0'),sep=' ')).as_matrix()
                    if joint.get('type')=='revolute':
                        axis=np.fromstring(joint.find('axis').get('xyz'),sep=' ');axis/=np.linalg.norm(axis)
                        relative[:3,:3]=relative[:3,:3]@Rotation.from_rotvec(axis*angles[joint.get('name')]).as_matrix()
                    elif joint.get('type')!='fixed':raise AssertionError('Unsupported test joint type')
                    transforms[child]=transforms[parent]@relative
                    pending.remove(joint);progressed=True
                if not progressed:raise AssertionError('URDF tree cannot be resolved')
            return transforms
        for k,cmd in enumerate(commands):
            self.cal.set_pose(q,cmd)
            mapped=mapping.sdk_range_to_full_o30_right_urdf_positions(list(cmd))
            generic_angles={}
            for i,name in enumerate(base.NAMES):
                limit=urdf.find(f"joint[@name='{name}']/limit")
                low,high=float(limit.get('lower')),float(limit.get('upper'))
                # Independent foreign consumer: only standard URDF limits and a linear interpolation.
                angle=low+(high-low)*float(np.clip(cmd[i],0,255))/255
                generic_angles[name]=angle
                if abs(angle-mapped[name])>1e-10:raise AssertionError('Runtime differs from generic linear consumer')
                data.qpos[model.jnt_qposadr[model.joint(name).id]]=angle
                native_data.qpos[native.jnt_qposadr[native.joint(name).id]]=angle
            base.mujoco.mj_forward(model,data)
            base.mujoco.mj_forward(native,native_data)
            for bid in range(1,native.nbody):
                other=model.body(native.body(bid).name).id
                delta=float(np.linalg.norm(native_data.xpos[bid]-data.xpos[other]));native_error=max(native_error,delta)
                if delta>2e-8 or np.max(abs(native_data.xmat[bid]-data.xmat[other]))>2e-8:
                    raise AssertionError('Native URDF load differs from runtime')
            transforms=portable_fk(generic_angles)
            for name,transform in transforms.items():
                bid=model.body(name).id
                delta=float(np.linalg.norm(transform[:3,3]-data.xpos[bid]));portable_error=max(portable_error,delta)
                if delta>2e-8 or np.max(abs(transform[:3,:3]-data.xmat[bid].reshape(3,3)))>2e-8:
                    raise AssertionError('Standard URDF kinematics differs from runtime')
            for bid in range(1,self.cal.model.nbody):
                other=model.body(self.cal.model.body(bid).name).id
                delta=float(np.linalg.norm(data.xpos[other]-self.cal.data.xpos[bid]));max_error=max(max_error,delta)
                if delta>2e-7 or np.max(abs(data.xmat[other]-self.cal.data.xmat[bid]))>2e-6:
                    raise AssertionError('Export kinematics mismatch')
            if 5<=k<17:
                s=self.source.scenarios[k-5]
                scenarios.append(dict(joint=s['joint'],background=s['background'],
                     penetration_mm=float(max([0.]+[-ct.dist for ct in data.contact])*1000)))
        self.accept(q)
        return dict(tested_commands=len(commands),generic_linear_consumer_matches=True,custom_mapping_absent=True,
                    max_export_position_error_m=max_error,independent_urdf_fk_error_m=portable_error,
                    native_urdf_load_error_m=native_error,endpoint_scenarios=scenarios)

    def render(self,q):
        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:
            groups=[('pinches_left',[(s['finger'],s['command']) for s in self.cal.samples]+[('zero',self.cal.zero_command)],90),
                    ('pinches_front',[(s['finger'],s['command']) for s in self.cal.samples]+[('zero',self.cal.zero_command)],175),
                    ('endpoints',[(s['joint'],s['command']) for s in self.source.scenarios if s['background']=='measured_zero'],90)]
            for group,poses,azimuth in groups:
                sheet=Image.new('RGB',(480*len(poses),1280),'#f1f5f9');draw=ImageDraw.Draw(sheet)
                for row in range(2):
                    for col,(name,cmd) in enumerate(poses):
                        if row:self.cal.set_pose(q,cmd)
                        else:
                            # V5 equals V4 at the pinches/zero; endpoints use V5's actual piecewise mapping.
                            position=self.source.positions(cmd,self.source.targets)
                            self.cal.data.qpos[self.cal.adr]=position;base.mujoco.mj_forward(model,self.cal.data)
                        camera=base.mujoco.MjvCamera();camera.lookat[:]=[.04,0,.1]
                        camera.distance=.38 if azimuth==175 else .32;camera.azimuth=azimuth;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),('V5 piecewise' if row==0 else 'V6 linear')+' / '+name,font=font,fill='#172033')
                sheet.save(REPORT/f'{group}.png');print(REPORT/f'{group}.png')


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    for flag in ['fit','check','render']:parser.add_argument('--'+flag,action='store_true')
    parser.add_argument('--starts',type=int,default=3)
    args=parser.parse_args()
    if not any([args.fit,args.check,args.render]):parser.error('Choose an action')
    if args.starts<1:parser.error('Starts must be positive')
    fit=LinearCalibration()
    if args.fit:q,history=fit.fit(args.starts);info=fit.save(q,history)
    else:info,q=fit.load()
    if args.check:
        info['validation']=fit.validate(q);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(q)


if __name__=='__main__':main()