#!/usr/bin/env python3
"""Jointly fit zero/travel and uncalibrated cameras to five annotated real photos.

The photo coordinates are manual estimates with uncertainty, not measured 3-D
poses. Four contact constraints still use actual MuJoCo forward kinematics.
"""
from __future__ import annotations

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 calibrate as base

ANNOTATIONS = Path(__file__).with_name('photo_landmarks.json')
V1_JSON = base.MODEL_DIR / 'o30_right_pinch_calibration_v1.json'
PHOTO_JSON = base.MODEL_DIR / 'o30_right_photo_calibration.json'
PHOTO_URDF = base.MODEL_DIR / 'linkerhand_o30_right_photo_calibrated.urdf'
REPORT = base.ROOT / 'outputs/reports/O30/o30_right_photo_calibration'
BASE_COUNT = 56  # 40 shared joint parameters + 4*4 pad contact coordinates


class PhotoCalibration(base.Calibration):
    def __init__(self):
        super().__init__()
        self.annotations = json.loads(ANNOTATIONS.read_text())
        self.zero_command = np.array(self.annotations['zero_command'], dtype=float)
        self.photos = self.annotations['photos']
        if self.n != 4 or [p['name'] for p in self.photos] != ['index', 'middle', 'ring', 'pinky', 'zero']:
            raise ValueError('Expected four ordered pinches and one zero photo')
        if self.zero_command.shape != (20,) or not np.all(np.isfinite(self.zero_command)):
            raise ValueError('Invalid zero command')
        if not np.all((self.zero_command >= 0) & (self.zero_command <= 255)):
            raise ValueError('Zero command outside SDK range')
        self.photo_commands = [*self.commands, self.zero_command]
        self.width, self.height = self.annotations['image_size']
        self.center = np.array([self.width/2, self.height/2])
        landmarks = {name: (int(self.model.jnt_bodyid[self.model.joint(name).id]), np.zeros(3))
                     for name in base.NAMES}
        for finger in base.FINGERS:
            raw = (base.MODEL_DIR/'meshes'/f'{finger}_distal.STL').read_bytes()
            dtype = np.dtype([('normal', '<f4', 3), ('v', '<f4', (3, 3)), ('attr', '<u2')])
            triangles = np.frombuffer(raw[84:], dtype=dtype, count=int.from_bytes(raw[80:84], 'little'))
            vertices = triangles['v'].reshape(-1, 3)
            axis = 1 if finger == 'thumb' else 2
            tip = vertices[vertices[:, axis] > vertices[:, axis].max()-.0005].mean(0)
            landmarks[finger+'_tip'] = (self.model.body(finger+'_distal').id, tip)
        self.photo_data = []
        for photo in self.photos:
            points = photo['points']
            body = np.array([landmarks[p[0]][0] for p in points])
            local = np.array([landmarks[p[0]][1] for p in points])
            uv = np.array([p[1:3] for p in points], dtype=float)
            sigma = np.array([p[3] for p in points], dtype=float)
            if np.any(sigma <= 0) or not np.all(np.isfinite(uv)):
                raise ValueError('Invalid photo annotations')
            self.photo_data.append((body, local, uv, sigma))
        self.v1 = np.array(json.loads(V1_JSON.read_text())['parameter_vector'])

    def world_points(self, photo_index, q):
        self.data.qpos[self.adr] = self.positions(q, self.photo_commands[photo_index])
        base.mujoco.mj_kinematics(self.model, self.data)
        body, local, _, _ = self.photo_data[photo_index]
        return self.data.xpos[body] + np.einsum('ijk,ik->ij', self.data.xmat[body].reshape(-1, 3, 3), local)

    def project(self, world, camera, focal):
        point = world @ Rotation.from_rotvec(camera[:3]).as_matrix().T + camera[3:]
        return point[:, :2]/point[:, 2, None]*focal+self.center

    @staticmethod
    def robust_photo_residual(error):
        # Reduce the influence of occlusion and uncertain manually picked centres.
        return error/(1+(error/3)**2)**.25

    def camera_initial(self, q):
        rotation = Rotation.from_matrix(np.array([[0, 1, 0], [0, 0, -1], [-1, 0, 0]])).as_rotvec()
        cameras = []
        for i in range(5):
            world = self.world_points(i, q)
            _, _, observed, sigma = self.photo_data[i]
            result = least_squares(lambda v: ((self.project(world, v, 1000)-observed)/sigma[:, None]).ravel(),
                                   np.r_[rotation, 0, .08, .5], max_nfev=100)
            cameras.append(result.x)
        return np.r_[np.array(cameras).ravel(), np.log(1000)]

    def camera_bounds(self):
        return (np.r_[np.tile([-6, -6, -6, -.3, -.2, .2], 5), np.log(500)],
                np.r_[np.tile([6, 6, 6, .3, .3, 1.5], 5), np.log(2000)])

    def photo_residual(self, q, camera_parameters):
        cameras = camera_parameters[:30].reshape(5, 6)
        focal = np.exp(camera_parameters[-1])
        out = []
        for i in range(5):
            uv = self.project(self.world_points(i, q), cameras[i], focal)
            _, _, observed, sigma = self.photo_data[i]
            out.extend(self.robust_photo_residual((uv-observed)/sigma[:, None]).ravel())
        out.append((camera_parameters[-1]-np.log(1000))/.35)
        return out

    def joint_residual(self, v):
        q = v[:BASE_COUNT]
        out = list(super().residual(q))
        # Slightly strengthen opposing-pad orientation after adding noisy photos.
        for i in range(4):
            self.set_pose(q, self.commands[i])
            _, thumb = self.world_pad(q, i, 'thumb')
            _, finger = self.world_pad(q, i, self.samples[i]['finger'])
            out.extend((thumb+finger)/.065)
        out.extend(self.photo_residual(q, v[BASE_COUNT:]))
        # Weak anatomical prior from the straight zero photo, not measured angles.
        sigma_degrees = np.full(20, 3.)
        sigma_degrees[[0, 1, 6, 15]] = 5.
        out.extend(self.positions(q, self.zero_command)/np.deg2rad(sigma_degrees))
        return np.array(out)

    def fit_cameras(self, q, initial=None):
        initial = self.camera_initial(q) if initial is None else initial
        return least_squares(lambda v: self.photo_residual(q, v), initial,
                             bounds=self.camera_bounds(), max_nfev=300, x_scale='jac').x

    def photo_metrics(self, q, camera_parameters):
        cameras = camera_parameters[:30].reshape(5, 6)
        focal = np.exp(camera_parameters[-1])
        result, all_errors = [], []
        for i in range(5):
            uv = self.project(self.world_points(i, q), cameras[i], focal)
            _, _, observed, sigma = self.photo_data[i]
            errors = np.linalg.norm(uv-observed, axis=1)
            all_errors.extend(errors)
            result.append(dict(photo=self.photos[i]['name'], points=len(errors),
                               rms_px=float(np.sqrt(np.mean(errors**2))),
                               points_detail=[dict(name=p[0], observed_px=obs.tolist(), predicted_px=pred.tolist(),
                                                   sigma_px=float(s), error_px=float(e))
                                              for p, obs, pred, s, e in zip(self.photos[i]['points'], observed, uv, sigma, errors)]))
        return dict(aggregate_rms_px=float(np.sqrt(np.mean(np.array(all_errors)**2))),
                    focal_px=float(focal), photos=result)

    def fit_photos(self, initial, starts, max_nfev):
        rng = np.random.default_rng(20260921)
        low, high = self.bounds()
        clow, chigh = self.camera_bounds()
        bounds = np.r_[low, clow], np.r_[high, chigh]
        best, history = None, []
        for i in range(starts):
            start = initial.copy() if best is None else best.x.copy()
            if i:
                start[:40] += rng.normal(0, .025, 40)
            result = least_squares(self.joint_residual, np.clip(start, bounds[0]+1e-10, bounds[1]-1e-10),
                                   bounds=bounds, max_nfev=max_nfev, x_scale='jac', diff_step=1e-4,
                                   ftol=1e-7, xtol=1e-9)
            rows = self.metrics(result.x[:BASE_COUNT])
            try:
                base.check_metrics(rows)
                accepted = True
            except RuntimeError:
                accepted = False
            history.append(dict(start=i+1, cost=float(result.cost), nfev=result.nfev,
                                accepted=accepted, success=bool(result.success), message=result.message))
            print(json.dumps(history[-1]), flush=True)
            if accepted and (best is None or result.cost < best.cost):
                best = result
        if best is None:
            raise RuntimeError('No candidate passed contact acceptance; existing model was not replaced')
        return best, history


def save(cal, vector, history):
    q = vector[:BASE_COUNT]
    base.CAL_URDF, base.CAL_JSON = PHOTO_URDF, PHOTO_JSON
    info = base.save_solution(cal, q, dict(method='Joint bounded photogeometric fitting', history=history,
                                          objective_evaluations=cal.evaluations, seed=20260921))
    # Silver housing and grey pads match the photos; this does not change physics.
    tree = ET.parse(PHOTO_URDF)
    for link in tree.getroot().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')
    tree.getroot().insert(0, ET.Comment(
        ' Photo/zero fitted candidate, 2026-09-21. Zero offsets are baked into '
        'joint origins; SDK 0..255 maps to limit lower..upper. '
        'See o30_right_photo_calibration.json; not a measured full-stroke calibration. '))
    ET.indent(tree, space='  ')
    tree.write(PHOTO_URDF, encoding='utf-8', xml_declaration=True)
    baseline_cameras = cal.fit_cameras(cal.v1)
    info.update(status='five-photo fitted candidate; hardware validation pending',
                method='MuJoCo contact constraints plus manual 2-D landmark reprojection, jointly optimized with cameras',
                limitations='Single-view photographs with manually estimated landmarks and unknown lens/camera pose. '
                            'Small image residual does not establish unique 3-D angles or real-world accuracy. '
                            'Full stroke, unseen poses, and PD dynamics remain unmeasured.',
                annotations_sha256=base.sha256(ANNOTATIONS), v1_sha256=base.sha256(V1_JSON),
                calibrated_urdf_sha256=base.sha256(PHOTO_URDF),
                zero_command=cal.zero_command.tolist(), camera_parameters=vector[BASE_COUNT:].tolist(),
                baseline_camera_parameters=baseline_cameras.tolist(),
                previous_v1=cal.photo_metrics(cal.v1, baseline_cameras),
                photo_fit=cal.photo_metrics(q, vector[BASE_COUNT:]),
                zero_model_angles_deg=dict(zip(base.NAMES, np.rad2deg(cal.positions(q, cal.zero_command)).tolist())),
                travel_bounds_hit=[base.NAMES[i] for i in range(20) if abs(abs(q[20+i])-.5) < 1e-4])
    PHOTO_JSON.write_text(json.dumps(info, indent=2, ensure_ascii=False)+'\n')
    return info


def load(cal):
    info = json.loads(PHOTO_JSON.read_text())
    for key, path in [('stock_urdf_sha256', base.STOCK_URDF), ('stock_xml_sha256', base.STOCK_XML),
                      ('samples_sha256', base.SAMPLES), ('annotations_sha256', ANNOTATIONS), ('v1_sha256', V1_JSON)]:
        if info[key] != base.sha256(path):
            raise ValueError(f'Calibration input changed: {path}')
    if 'calibrated_urdf_sha256' in info and info['calibrated_urdf_sha256'] != base.sha256(PHOTO_URDF):
        raise ValueError('Calibrated URDF was edited after fitting')
    if info['zero_command'] != cal.zero_command.tolist():
        raise ValueError('Zero command changed')
    return info, np.r_[info['parameter_vector'], info['camera_parameters']]


def validate(cal, vector):
    q = vector[:BASE_COUNT]
    result = base.validate_live(cal, q)
    cal.set_pose(q, cal.zero_command)
    zero_penetration = max([0.] + [-c.dist for c in cal.data.contact])*1000
    if zero_penetration > .2:
        raise AssertionError(f'Zero pose penetration {zero_penetration} mm')
    result['zero_penetration_mm'] = float(zero_penetration)
    result['zero_command'] = cal.zero_command.tolist()
    return result


def render(cal, vector):
    from PIL import Image, ImageDraw, ImageFont
    REPORT.mkdir(parents=True, exist_ok=True)
    model, data = cal.model, cal.data
    model.vis.global_.offwidth, model.vis.global_.offheight = cal.width, cal.height
    # Neutral materials for both rows keep the pose comparison fair.
    for g in range(model.ngeom):
        if model.geom_group[g] == 2:
            name = model.body(model.geom_bodyid[g]).name
            model.geom_rgba[g] = [.48, .50, .52, 1] if name.endswith('_distal') else [.68, .70, .72, 1]
    renderer = base.mujoco.Renderer(model, cal.height, cal.width)
    cameras, focal = vector[BASE_COUNT:BASE_COUNT+30].reshape(5, 6), np.exp(vector[-1])
    font = ImageFont.truetype('/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf', 19)
    row_height = cal.height+50
    sheet = Image.new('RGB', (cal.width*5, row_height*2), '#f1f5f9')
    header = ImageDraw.Draw(sheet)
    for row, parameters in enumerate([cal.v1, vector[:BASE_COUNT]]):
        for i in range(5):
            cal.set_pose(parameters, cal.photo_commands[i])
            renderer.update_scene(data)
            rotation = Rotation.from_rotvec(cameras[i, :3]).as_matrix()
            eye = -rotation.T@cameras[i, 3:]
            forward, up = rotation.T@np.array([0, 0, 1]), rotation.T@np.array([0, -1, 0])
            for camera in renderer.scene.camera:
                camera.pos[:], camera.forward[:], camera.up[:] = eye, forward, up
                camera.frustum_near, camera.frustum_far = .01, 5.
                camera.frustum_top = .01*cal.height/2/focal
                camera.frustum_bottom = -camera.frustum_top
                camera.frustum_center = 0
                camera.frustum_width = .01*cal.width/2/focal  # MuJoCo GL half-width
            for light in renderer.scene.lights[:renderer.scene.nlight]:
                light.pos[:], light.dir[:] = eye, forward
                light.ambient[:] = [.25, .25, .25]
                light.diffuse[:] = [.65, .65, .65]
            image = Image.fromarray(renderer.render())
            sheet.paste(image, (i*cal.width, row*row_height+50))
            header.text((i*cal.width+14, row*row_height+16),
                        ('Previous contact fit' if row == 0 else 'Photo + zero fit')+' / '+cal.photos[i]['name'],
                        fill='#172033', font=font)
            if row == 1:
                image.save(REPORT/f'{cal.photos[i]["name"]}_pose.png')
                overlay = image.copy()
                draw = ImageDraw.Draw(overlay)
                predicted = cal.project(cal.world_points(i, parameters), cameras[i], focal)
                for observed, point in zip(cal.photo_data[i][2], predicted):
                    draw.line([tuple(observed), tuple(point)], fill='#fbbf24', width=2)
                    draw.ellipse([observed[0]-4, observed[1]-4, observed[0]+4, observed[1]+4], outline='#22d3ee', width=2)
                    draw.ellipse([point[0]-2, point[1]-2, point[0]+2, point[1]+2], fill='#f87171')
                draw.text((10, 8), 'Cyan: manual photo points / red: model', fill='white', font=font)
                overlay.save(REPORT/f'{cal.photos[i]["name"]}_landmarks.png')
    renderer.close()
    sheet.save(REPORT/'photo_fit_comparison.png')
    print(REPORT/'photo_fit_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=3)
    parser.add_argument('--max-nfev', type=int, default=250)
    parser.add_argument('--preview', choices=['index', 'middle', 'ring', 'pinky', 'zero'])
    args = parser.parse_args()
    cal = PhotoCalibration()
    if args.preview:
        preview(cal, args.preview)
        return
    if not (args.fit or args.check or args.render):
        parser.error('Choose --fit, --check, --render or --preview')
    if args.starts < 1 or args.max_nfev < 1:
        parser.error('Optimization budgets must be positive')
    if args.fit:
        if PHOTO_JSON.exists():
            _, initial = load(cal)
        else:
            initial = np.r_[cal.v1, cal.camera_initial(cal.v1)]
        best, history = cal.fit_photos(initial, args.starts, args.max_nfev)
        vector = best.x
        info = save(cal, vector, history)
        print('Exported photo URDF. Regenerate with O30_RIGHT_PREPARE_ONLY=1 ./hands/O30/scripts/run_o30_right_mujoco_plan.sh')
    else:
        info, vector = load(cal)
    if args.check:
        info['live_validation'] = validate(cal, vector)
        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(cal, vector)


def preview(cal, name):
    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 = cal.zero_command if name == 'zero' else next(s['command'] for s in cal.samples if s['finger'] == name)
    for joint, value in mapping.sdk_range_to_full_o30_right_urdf_positions(list(command)).items():
        data.qpos[model.jnt_qposadr[model.joint(joint).id]] = value
    base.mujoco.mj_forward(model, data)
    with mujoco.viewer.launch_passive(model, data) as viewer:
        viewer.cam.lookat[:] = [.045, 0, .10]
        viewer.cam.distance, viewer.cam.azimuth, viewer.cam.elevation = .3, 175, -15
        while viewer.is_running():
            viewer.sync()
            time.sleep(.02)


if __name__ == '__main__':
    main()