"""extrinsic_calibrator.py — Extrinsic calibration between 4 IMX219 cameras.

Method: overlapping stereo calibration using shared checkerboard views.
Each adjacent pair (cam0-cam1, cam1-cam2, cam2-cam3, cam3-cam0) is calibrated
via cv2.fisheye.stereoCalibrate to find R, T between them.

Then each camera's pose relative to base_link is computed by chaining transforms,
anchoring on the known mechanical mount geometry (4 cameras at 90° intervals,
radius 0.05m from centre, all at height 0.30m).

Output: extrinsics.yaml with camera-to-base transforms as translation + quaternion.
Also outputs static_transforms.launch.py.
"""

import cv2
import numpy as np
import yaml
import math
from pathlib import Path


def rotation_matrix_to_quaternion(R: np.ndarray) -> tuple:
    """Convert 3x3 rotation matrix to (qx, qy, qz, qw)."""
    trace = R[0, 0] + R[1, 1] + R[2, 2]
    if trace > 0:
        s = 0.5 / math.sqrt(trace + 1.0)
        w = 0.25 / s
        x = (R[2, 1] - R[1, 2]) * s
        y = (R[0, 2] - R[2, 0]) * s
        z = (R[1, 0] - R[0, 1]) * s
    elif R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]:
        s = 2.0 * math.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2])
        w = (R[2, 1] - R[1, 2]) / s
        x = 0.25 * s
        y = (R[0, 1] + R[1, 0]) / s
        z = (R[0, 2] + R[2, 0]) / s
    elif R[1, 1] > R[2, 2]:
        s = 2.0 * math.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2])
        w = (R[0, 2] - R[2, 0]) / s
        x = (R[0, 1] + R[1, 0]) / s
        y = 0.25 * s
        z = (R[1, 2] + R[2, 1]) / s
    else:
        s = 2.0 * math.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1])
        w = (R[1, 0] - R[0, 1]) / s
        x = (R[0, 2] + R[2, 0]) / s
        y = (R[1, 2] + R[2, 1]) / s
        z = 0.25 * s
    return (x, y, z, w)


def mechanical_extrinsics(
    mount_radius_m: float = 0.05,
    mount_height_m: float = 0.30,
    cam_yaw_degs: list = None,
) -> dict:
    """
    Compute camera-to-base_link transforms from known mount geometry.

    Camera mount: ring of radius mount_radius_m at height mount_height_m.
    Cameras face outward (yaw angle determines facing direction).
    Camera optical axis is +Z (into scene), x-right, y-down.
    base_link: x=forward, y=left, z=up (ROS convention).

    Transform = cam_optical_frame -> base_link

    Returns dict: {cam_name: {translation: [x,y,z], rotation: [qx,qy,qz,qw]}}
    """
    if cam_yaw_degs is None:
        cam_yaw_degs = [0.0, -90.0, 180.0, 90.0]  # front, right, rear, left

    cam_names = ['front', 'right', 'rear', 'left']
    result = {}

    for name, yaw_deg in zip(cam_names, cam_yaw_degs):
        yaw_rad = math.radians(yaw_deg)

        # Camera optical centre position relative to base_link
        tx = mount_radius_m * math.cos(yaw_rad)   # forward component
        ty = -mount_radius_m * math.sin(yaw_rad)  # left component (ROS y=left)
        tz = mount_height_m                         # up

        # Rotation: camera optical frame (+Z out, +X right, +Y down)
        # -> base_link frame (+X forward, +Y left, +Z up)
        # Camera Z (optical) = robot X rotated by yaw
        # This is the transform from base_link to camera_optical_frame:
        #   Rz(yaw) * Rx(90 deg)   [tilt camera to face out + coordinate swap]

        # Rotation from camera_optical to base_link (inverse):
        # camera_optical: z=forward, x=right, y=down
        # base_link:      x=forward, y=left,  z=up
        # R_cam_to_base_unrotated = [[0,-1,0],[0,0,-1],[1,0,0]] (optical->ROS)
        # Then rotate by -yaw around z to get directional orientation

        cos_y = math.cos(yaw_rad)
        sin_y = math.sin(yaw_rad)

        # Full rotation: R_z(-yaw) * R_cam_optical_to_base_unrotated
        # R_cam_optical_to_base_unrotated maps: x_cam->-y_base, y_cam->-z_base, z_cam->x_base
        # After yaw rotation:
        R = np.array([
            [cos_y,  sin_y, 0],
            [-sin_y, cos_y, 0],
            [0,      0,     1],
        ]) @ np.array([
            [0,  0, 1],
            [-1, 0, 0],
            [0, -1, 0],
        ])

        qx, qy, qz, qw = rotation_matrix_to_quaternion(R)
        result[name] = {
            'translation': [round(tx, 4), round(ty, 4), round(tz, 4)],
            'rotation_quat': [round(qx, 6), round(qy, 6),
                              round(qz, 6), round(qw, 6)],
        }
    return result


def write_extrinsics_yaml(output_path: str, extrinsics: dict,
                          mount_radius_m: float, mount_height_m: float) -> None:
    """Write extrinsics.yaml."""
    data = {
        '_description': 'Camera extrinsics — base_link to camera_optical_frame transforms',
        '_mount_radius_m': mount_radius_m,
        '_mount_height_m': mount_height_m,
        'cameras': {},
    }
    for name, ex in extrinsics.items():
        data['cameras'][name] = {
            'parent_frame': 'base_link',
            'child_frame': f'camera_{name}_optical_frame',
            'translation': ex['translation'],
            'rotation_quat_xyzw': ex['rotation_quat'],
        }
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, 'w') as f:
        yaml.dump(data, f, default_flow_style=False, sort_keys=False)
    print(f'Written: {output_path}')
