#!/usr/bin/env python3
"""Pinky-only refinement of V3: three flexion travels, fixed zeros and other joints."""
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_multiview as multi
base = multi.base
HERE = Path(__file__).resolve().parent
ANNOTATIONS = HERE/'pinky_landmarks.json'
OUTPUT = base.MODEL_DIR/'linkerhand_o30_right_pinky_calibrated.urdf'
PARAMETERS = base.MODEL_DIR/'o30_right_pinky_calibration.json'
REPORT = base.ROOT/'outputs/reports/O30/o30_right_pinky_refinement'
JOINTS = ['pinky_mcp_pitch', 'pinky_pip', 'pinky_dip']


class PinkyFit:
    def __init__(self):
        self.cal = multi.MultiviewCalibration()
        self.previous, _ = multi.load(self.cal)
        self.q0 = np.array(self.previous['parameter_vector'])
        self.annotation = json.loads(ANNOTATIONS.read_text())
        self.command = np.array(self.annotation['command'])
        if not np.array_equal(self.command, self.cal.commands[3]):
            raise ValueError('Pinky command changed')
        self.width, self.height = self.annotation['image_size']
        self.center = np.array([self.width/2, self.height/2])
        self.names = [p[0] for p in self.annotation['points']]
        self.obs = np.array([p[1:3] for p in self.annotation['points']], dtype=float)
        self.sigma = np.array([p[3] for p in self.annotation['points']], dtype=float)
        self.bids = np.array([self.cal.model.jnt_bodyid[self.cal.model.joint(n).id] for n in self.names])
        self.indices = np.array([20+base.NAMES.index(n) for n in JOINTS])
        self.count = 0

    def unpack(self, v):
        q = self.q0.copy()
        q[self.indices] = v[:3]
        return q

    def project(self, q, camera):
        self.cal.set_pose(q, self.command)
        points = self.cal.data.xpos[self.bids]@Rotation.from_rotvec(camera[:3]).as_matrix().T+camera[3:6]
        return points[:, :2]/points[:, 2, None]*np.exp(camera[6])+self.center

    def contact(self, q):
        c = self.cal
        c.set_pose(q, self.command)
        closest = np.zeros(6)
        gap = base.mujoco.mj_geomDistance(c.model, c.data, c.geoms['thumb'], c.geoms['pinky'], .2, closest)
        local = {}
        for i, finger in enumerate(['thumb', 'pinky']):
            bid = c.bodies[finger]
            local[finger] = c.data.xmat[bid].reshape(3, 3).T@(closest[3*i:3*i+3]-c.data.xpos[bid])
        penetration = max([0.] + [-contact.dist for contact in c.data.contact])
        return gap, penetration, local

    def camera_fit(self):
        # Fit unchanged anatomy and fixed MCP origin first, excluding the two target hinges.
        initial = np.r_[self.previous['camera_parameters'][48:54], np.log(2000)]
        self.camera_bounds = ([-6,-6,-6,-.3,-.3,.15,np.log(1200)],
                              [6,6,6,.3,.3,2,np.log(8000)])
        result = least_squares(lambda v: np.r_[
            ((self.project(self.q0, v)[:6]-self.obs[:6])/self.sigma[:6,None]).ravel(),
            (v[6]-np.log(2500))/.5], initial, bounds=self.camera_bounds, max_nfev=400)
        return result.x

    def residual(self, v):
        self.count += 1
        q = self.unpack(v)
        gap, penetration, local = self.contact(q)
        p, t = local['pinky'], local['thumb']
        # Rounded fingertip pad region. Prevent an apparent fit through contact on the back or housing.
        region = np.array([max(0., -.002-p[0]), max(0., .014-p[2]), max(0., p[2]-.024),
                           max(0., abs(p[1])-.006), max(0., .028-t[1]), max(0., t[1]-.046),
                           max(0., -t[2]), max(0., abs(t[0])-.008)])
        uv = self.project(q, v[3:])
        return np.r_[gap/.00005, max(0., penetration-.00003)/.00005,
                     region/.0001, ((uv-self.obs)/self.sigma[:,None]).ravel()*3,
                     (v[-1]-self.camera0[-1])/.15, (v[:3]-self.q0[self.indices])/.6*.3]

    def assess(self, q, camera):
        gap, penetration, local = self.contact(q)
        uv = self.project(q, camera)
        error = np.linalg.norm(uv-self.obs, axis=1)
        # Fair pose comparison: same camera for V3 and candidate; also report V3's anchor-fit camera.
        old_same = np.linalg.norm(self.project(self.q0, camera)-self.obs, axis=1)
        old_anchor = np.linalg.norm(self.project(self.q0, self.camera0)-self.obs, axis=1)
        return dict(mesh_gap_mm=float(gap*1000), max_penetration_mm=float(penetration*1000),
                    closest_points_local_m={k:v.tolist() for k,v in local.items()},
                    hinge_rms_px=float(np.sqrt(np.mean(error[6:]**2))),
                    previous_same_camera_hinge_rms_px=float(np.sqrt(np.mean(old_same[6:]**2))),
                    previous_anchor_camera_hinge_rms_px=float(np.sqrt(np.mean(old_anchor[6:]**2))),
                    anchors_rms_px=float(np.sqrt(np.mean(error[:6]**2))),
                    previous_anchor_rms_px=float(np.sqrt(np.mean(old_anchor[:6]**2))),
                    points=[dict(name=n, observed_px=o.tolist(), predicted_px=p.tolist(), error_px=float(e))
                            for n,o,p,e in zip(self.names,self.obs,uv,error)])

    def accept(self, q, camera):
        result = self.assess(q, camera)
        p, t = result['closest_points_local_m']['pinky'], result['closest_points_local_m']['thumb']
        if abs(result['mesh_gap_mm']) > .2 or result['max_penetration_mm'] > .2:
            raise ValueError('Contact/collision tolerance exceeded')
        if not (p[0] >= -.003 and .014 <= p[2] <= .024 and abs(p[1]) <= .006 and
                .028 <= t[1] <= .046 and t[2] >= 0 and abs(t[0]) <= .008):
            raise ValueError('Contact outside fingertip pad region')
        if result['hinge_rms_px'] >= min(result['previous_same_camera_hinge_rms_px'],result['previous_anchor_camera_hinge_rms_px']):
            raise ValueError('Target hinges did not improve')
        if result['anchors_rms_px'] > result['previous_anchor_rms_px']+10:
            raise ValueError('Camera fit compromised unchanged anatomy')
        untouched = np.ones(56,dtype=bool);untouched[self.indices]=False
        if not np.array_equal(q[untouched],self.q0[untouched]):
            raise ValueError('Unexpected parameter change')
        for command in [*self.cal.commands[:3],self.cal.zero_command]:
            if not np.array_equal(self.cal.positions(q,command),self.cal.positions(self.q0,command)):
                raise ValueError('Other pinch or zero pose changed')
        return result

    def fit(self, starts, max_nfev):
        self.camera0 = self.camera_fit()
        low = np.r_[np.full(3,-.5),self.camera_bounds[0]]
        high = np.r_[np.full(3,.5),self.camera_bounds[1]]
        initial = np.r_[self.q0[self.indices],self.camera0]
        rng = np.random.default_rng(3004)
        best, history = None, []
        for i in range(starts):
            start = initial.copy() if best is None else best.x.copy()
            if i: start[:3] += rng.normal(0,.05,3)
            result = least_squares(self.residual,np.clip(start,low+1e-9,high-1e-9),bounds=(low,high),
                                   max_nfev=max_nfev,diff_step=1e-4,ftol=1e-9,xtol=1e-10)
            try:
                score=self.accept(self.unpack(result.x),result.x[3:]); accepted=True;reason=''
            except ValueError as exc:
                accepted=False;score=None;reason=str(exc)
            row=dict(start=i+1,cost=float(result.cost),nfev=result.nfev,accepted=accepted,reason=reason,score=score)
            history.append(row)
            print(json.dumps({k:v for k,v in row.items() if k!='score'}),flush=True)
            if accepted and (best is None or result.cost<best.cost): best=result
        if best is None: raise RuntimeError('No candidate passed; active V3 untouched')
        return best.x,history


def hashes():
    annotation=json.loads(ANNOTATIONS.read_text())
    return {str(p.relative_to(base.ROOT)):base.sha256(p) for p in [ANNOTATIONS,
            HERE/annotation['hardware_image'],HERE/annotation['simulation_image'],
            multi.OUTPUT_JSON,multi.OUTPUT_URDF,base.STOCK_XML,base.STOCK_URDF,base.SAMPLES]}


def save(fit, vector, history):
    q=fit.unpack(vector);result=fit.accept(q,vector[3:])
    tree=ET.parse(multi.OUTPUT_URDF)
    for name,index in zip(JOINTS,fit.indices):
        joint=tree.getroot().find(f"joint[@name='{name}']")
        i=index-20
        joint.find('limit').set('upper',format(fit.cal.lo[i]+fit.cal.span[i]*(1+q[index]),'.12g'))
    ET.indent(tree,space='  ');tree.write(OUTPUT,encoding='utf-8',xml_declaration=True)
    info=dict(schema_version=1,method='Pinky-only bounded multistart image/mesh-contact fitting',
              limitations='Manual 2-D estimates and unknown camera; not measured full-stroke or 3-D hardware calibration. '
                           'V3 pad normal/point nuisance parameters are not reused for pinky contact acceptance.',
              input_hashes=hashes(),output_sha256=base.sha256(OUTPUT),parameter_vector=q.tolist(),
              camera=vector[3:].tolist(),anchor_camera=fit.camera0.tolist(),
              objective_evaluations=fit.count,history=history,score=result,
              zero_command=fit.cal.zero_command.tolist(),command=fit.command.tolist(),
              changed_joints={name:dict(previous_scale=float(1+fit.q0[idx]),scale=float(1+q[idx]),
                   previous_pose_deg=float(np.rad2deg(fit.cal.positions(fit.q0,fit.command)[idx-20])),
                   pose_deg=float(np.rad2deg(fit.cal.positions(q,fit.command)[idx-20])))
                   for name,idx in zip(JOINTS,fit.indices)})
    PARAMETERS.write_text(json.dumps(info,indent=2,ensure_ascii=False)+'\n')
    return info,q


def load(fit):
    info=json.loads(PARAMETERS.read_text())
    if info['input_hashes'] != hashes() or info['output_sha256'] != base.sha256(OUTPUT):
        raise ValueError('Calibration input or output changed')
    fit.camera0=np.array(info['anchor_camera']);q=np.array(info['parameter_vector'])
    fit.accept(q,np.array(info['camera']))
    return info,q


def validate(fit,q):
    from linker_hand_mujoco_ros2.utils.o30_right_joints import configure_o30_right_limits, sdk_range_to_full_o30_right_urdf_positions
    model=base.mujoco.MjModel.from_xml_path(str(base.LIVE_XML));data=base.mujoco.MjData(model)
    configure_o30_right_limits(base.LIVE_XML)
    c=fit.cal
    commands=[*c.commands,c.zero_command,np.zeros(20),np.full(20,255),*np.random.default_rng(30).uniform(0,255,(30,20))]
    max_error=0.;rows=[]
    for k,command in enumerate(commands):
        c.set_pose(q,command)
        for name,value in sdk_range_to_full_o30_right_urdf_positions(list(command)).items():
            jid=model.joint(name).id
            if not model.jnt_range[jid,0]-1e-10 <= value <= model.jnt_range[jid,1]+1e-10:
                raise AssertionError('Joint limit exceeded')
            data.qpos[model.jnt_qposadr[jid]]=value
        base.mujoco.mj_forward(model,data)
        for bid in range(1,c.model.nbody):
            other=model.body(c.model.body(bid).name).id
            error=float(np.linalg.norm(data.xpos[other]-c.data.xpos[bid]));max_error=max(max_error,error)
            if error>2e-7 or np.max(abs(data.xmat[other]-c.data.xmat[bid]))>2e-6:
                raise AssertionError('Exported model differs from optimized kinematics')
        if k<5:
            penetration=max([0.]+[-ct.dist for ct in data.contact])*1000
            if penetration>.2: raise AssertionError('Known pose penetration')
            if k<4:
                finger=c.samples[k]['finger']
                def geom(f):
                    return next(g for g in range(model.ngeom) if model.geom_group[g]==3 and model.geom_bodyid[g]==model.body(f+'_distal').id)
                gap=base.mujoco.mj_geomDistance(model,data,geom('thumb'),geom(finger),.2,np.zeros(6))*1000
                if abs(gap)>.2: raise AssertionError('Known pose contact gap')
                rows.append(dict(finger=finger,mesh_gap_mm=float(gap),max_penetration_mm=float(penetration)))
            else: zero_penetration=penetration
    # Structural guarantee: only three limit upper values differ from accepted V3.
    previous=ET.parse(multi.OUTPUT_URDF).getroot();current=ET.parse(OUTPUT).getroot()
    for name in JOINTS:
        old=previous.find(f"joint[@name='{name}']/limit")
        current.find(f"joint[@name='{name}']/limit").set('upper',old.get('upper'))
    def structure(element):
        # Ignore pretty-print whitespace, but compare every element and attribute.
        return (element.tag, sorted(element.attrib.items()), (element.text or '').strip(),
                [structure(child) for child in element])
    if structure(previous)!=structure(current): raise AssertionError('Unexpected URDF structural change')
    return dict(tested_commands=len(commands),max_export_position_error_m=max_error,actual_mjcf=rows,
                zero_penetration_mm=float(zero_penetration),other_three_pinches_and_zero_exactly_preserved=True,
                only_three_urdf_limit_upper_values_changed=True)


def render(fit,q,camera):
    from PIL import Image,ImageDraw,ImageFont
    c=fit.cal;REPORT.mkdir(parents=True,exist_ok=True)
    width,height=700,724
    c.model.vis.global_.offwidth=width;c.model.vis.global_.offheight=height
    for g in range(c.model.ngeom):
        if c.model.geom_group[g]==2:
            c.model.geom_rgba[g]=[.48,.50,.52,1] if c.model.body(c.model.geom_bodyid[g]).name.endswith('_distal') else [.68,.70,.72,1]
    rotation=Rotation.from_rotvec(camera[:3]).as_matrix();eye=-rotation.T@camera[3:6]
    forward=rotation.T@np.array([0,0,1]);up=rotation.T@np.array([0,-1,0]);focal=np.exp(camera[6])/2
    font=ImageFont.truetype('/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',19)
    # Wider framing than the input photo to keep the complete contact region visible.
    margin=100
    c.model.vis.global_.offwidth=width+2*margin
    sheet=Image.new('RGB',((width+2*margin)*2,height+40),'#f1f5f9')
    with base.mujoco.Renderer(c.model,height,width+2*margin) as renderer:
        for col,parameters in enumerate([fit.q0,q]):
            c.set_pose(parameters,fit.command);renderer.update_scene(c.data)
            for cam in renderer.scene.camera:
                cam.pos[:]=eye;cam.forward[:]=forward;cam.up[:]=up
                cam.frustum_near=.01;cam.frustum_far=5.;cam.frustum_top=.01*height/2/focal
                cam.frustum_bottom=-cam.frustum_top;cam.frustum_center=0;cam.frustum_width=.01*(width+2*margin)/2/focal
            for light in renderer.scene.lights[:renderer.scene.nlight]:
                light.pos[:]=eye;light.dir[:]=forward;light.ambient[:]=[.25]*3;light.diffuse[:]=[.65]*3
            image=Image.fromarray(renderer.render());sheet.paste(image,(col*(width+2*margin),40))
            ImageDraw.Draw(sheet).text((col*(width+2*margin)+12,10),'V3 before' if col==0 else 'V4 pinky-only',font=font,fill='#172033')
            image.save(REPORT/('before.png' if col==0 else 'after.png'))
            if col:
                draw=ImageDraw.Draw(image)
                for obs,pred in zip(fit.obs,fit.project(q,camera)):
                    a=obs/2+[margin,0];b=pred/2+[margin,0]
                    draw.line([tuple(a),tuple(b)],fill='yellow',width=2)
                    draw.ellipse([a[0]-4,a[1]-4,a[0]+4,a[1]+4],outline='cyan',width=2)
                    draw.ellipse([b[0]-3,b[1]-3,b[0]+3,b[1]+3],fill='red')
                draw.text((10,10),'Cyan: manual photo points / red: model',font=font,fill='white')
                image.save(REPORT/'landmarks.png')
    sheet.save(REPORT/'comparison.png')
    print(REPORT/'comparison.png')


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--fit',action='store_true');parser.add_argument('--check',action='store_true')
    parser.add_argument('--render',action='store_true');parser.add_argument('--starts',type=int,default=6)
    parser.add_argument('--max-nfev',type=int,default=350)
    parser.add_argument('--preview',choices=['index','middle','ring','pinky','zero'])
    args=parser.parse_args()
    if not any([args.fit,args.check,args.render,args.preview]): parser.error('Choose an action')
    if args.starts<1 or args.max_nfev<1: parser.error('Budgets must be positive')
    fit=PinkyFit()
    if args.preview:
        multi.photo.preview(fit.cal,args.preview);return
    if args.fit:
        vector,history=fit.fit(args.starts,args.max_nfev);info,q=save(fit,vector,history)
        print(json.dumps(info['changed_joints'],indent=2))
        print('Exported. Regenerate launcher model before --check.')
    else: info,q=load(fit)
    if args.check:
        info['live_validation']=validate(fit,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['live_validation'],indent=2))
    if args.render: render(fit,q,np.array(info['camera']))


if __name__=='__main__': main()