#!/usr/bin/env python3
"""
calibrate_extrinsics.py — Compute camera-to-base_link extrinsic transforms.

Two modes:
  1. --mode mechanical: use known mount geometry (radius, height, yaw angles)
     -> produces analytically exact transforms from CAD measurements
  2. --mode stereo: use cv2.fisheye.stereoCalibrate on paired checkerboard images
     -> requires shared-view images from adjacent camera pairs

Usage (mechanical mode, recommended first pass):
  python3 calibrate_extrinsics.py \\
    --mode mechanical \\
    --mount-radius 0.05 \\
    --mount-height 0.30 \\
    --output /path/to/calibration/extrinsics.yaml \\
    --generate-launch /path/to/jetson/launch/static_transforms.launch.py
"""

import argparse
import math
import sys
from pathlib import Path


def print_extrinsics_table(extrinsics: dict) -> None:
    """Print formatted extrinsics table to stdout."""
    print()
    print(f'{"Camera":<8}  {"tx":>8}  {"ty":>8}  {"tz":>8}  '
          f'{"qx":>10}  {"qy":>10}  {"qz":>10}  {"qw":>10}')
    print('-' * 86)
    for name, ex in extrinsics.items():
        tx, ty, tz = ex['translation']
        qx, qy, qz, qw = ex['rotation_quat']
        print(f'  {name:<6}  {tx:>8.4f}  {ty:>8.4f}  {tz:>8.4f}  '
              f'{qx:>10.6f}  {qy:>10.6f}  {qz:>10.6f}  {qw:>10.6f}')
    print()


def generate_static_transforms_launch(extrinsics: dict, output_path: str) -> None:
    """Generate static_transforms.launch.py from extrinsics dict."""
    # Add package to path
    script_dir = Path(__file__).resolve().parent
    pkg_dir = script_dir.parent
    if str(pkg_dir) not in sys.path:
        sys.path.insert(0, str(pkg_dir))

    from saltybot_calibration.extrinsic_calibrator import write_extrinsics_yaml

    lines = [
        '"""',
        'static_transforms.launch.py — Auto-generated by calibrate_extrinsics.py',
        'DO NOT EDIT manually — regenerate via:',
        '  python3 calibrate_extrinsics.py --mode mechanical --generate-launch <this_file>',
        '"""',
        '',
        'from launch import LaunchDescription',
        'from launch_ros.actions import Node',
        '',
        '',
        'def generate_launch_description():',
        '    return LaunchDescription([',
    ]

    for name, ex in extrinsics.items():
        tx, ty, tz = ex['translation']
        qx, qy, qz, qw = ex['rotation_quat']
        child_frame = f'camera_{name}_optical_frame'
        lines += [
            f'        # camera_{name}: base_link -> {child_frame}',
            f'        Node(',
            f'            package=\'tf2_ros\',',
            f'            executable=\'static_transform_publisher\',',
            f'            name=\'tf_base_to_camera_{name}\',',
            f'            arguments=[',
            f'                \'{tx}\', \'{ty}\', \'{tz}\',',
            f'                \'{qx}\', \'{qy}\', \'{qz}\', \'{qw}\',',
            f'                \'base_link\', \'{child_frame}\',',
            f'            ],',
            f'        ),',
        ]

    lines += [
        '    ])',
        '',
    ]

    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, 'w') as f:
        f.write('\n'.join(lines))
    print(f'Generated launch file: {output_path}')


def run_mechanical_mode(args) -> None:
    """Run mechanical extrinsics calculation."""
    script_dir = Path(__file__).resolve().parent
    pkg_dir = script_dir.parent
    if str(pkg_dir) not in sys.path:
        sys.path.insert(0, str(pkg_dir))

    from saltybot_calibration.extrinsic_calibrator import (
        mechanical_extrinsics,
        write_extrinsics_yaml,
    )

    print(f'Mode:          mechanical')
    print(f'Mount radius:  {args.mount_radius * 100:.1f}cm')
    print(f'Mount height:  {args.mount_height * 100:.1f}cm')
    print(f'Camera yaw angles: front=0 deg, right=-90 deg, rear=180 deg, left=90 deg')
    print()

    extrinsics = mechanical_extrinsics(
        mount_radius_m=args.mount_radius,
        mount_height_m=args.mount_height,
    )

    print_extrinsics_table(extrinsics)

    write_extrinsics_yaml(
        output_path=args.output,
        extrinsics=extrinsics,
        mount_radius_m=args.mount_radius,
        mount_height_m=args.mount_height,
    )

    if args.generate_launch:
        generate_static_transforms_launch(extrinsics, args.generate_launch)


def run_stereo_mode(args) -> None:
    """Stereo extrinsics calibration — stub for future implementation.

    TODO: Implement cv2.fisheye.stereoCalibrate for adjacent camera pairs.

    This mode requires:
      - Shared-view checkerboard images from each adjacent camera pair:
          {images_dir}/pair_front_right/*.png  (visible to both cam0 and cam1)
          {images_dir}/pair_right_rear/*.png
          {images_dir}/pair_rear_left/*.png
          {images_dir}/pair_left_front/*.png
      - Per-camera intrinsics (camera_info.yaml) for each camera
      - Each pair: cv2.fisheye.stereoCalibrate -> R, T between cameras
      - Chain transforms: cam0->cam1->cam2->cam3, anchor to base_link via
        mechanical position of cam0

    Recommended workflow:
      1. First run --mode mechanical to get approximate transforms
      2. Mount robot in place, place checkerboard in overlap zones
      3. Capture simultaneous images from adjacent pairs (requires hardware sync
         or approximate temporal sync)
      4. Run --mode stereo to refine R, T
      5. Run generate_static_transforms.py to update launch file

    For now, use --mode mechanical which gives correct transforms from CAD
    dimensions (typically accurate to ~5mm / ~1 deg for a rigid mount).
    """
    print('[ERROR] Stereo mode is not yet implemented.')
    print()
    print('Use --mode mechanical for now (accurate from CAD measurements).')
    print()
    print('Future implementation plan:')
    print('  1. Capture simultaneous checkerboard images from adjacent camera pairs')
    print('     in their overlapping FoV regions (approx 70 deg overlap per pair)')
    print('  2. For each pair: cv2.fisheye.stereoCalibrate with per-camera intrinsics')
    print('  3. Chain transforms cam0->cam1->cam2->cam3, anchor to base_link')
    print('  4. Refine against mechanical prior using least-squares optimisation')
    print()
    print('Required directory structure for stereo mode:')
    print('  {images_dir}/pair_front_right/  — images visible to both front+right')
    print('  {images_dir}/pair_right_rear/')
    print('  {images_dir}/pair_rear_left/')
    print('  {images_dir}/pair_left_front/')
    sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
        description='Compute camera-to-base_link extrinsic transforms for 4x IMX219 cameras.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        '--mode', choices=['mechanical', 'stereo'], default='mechanical',
        help='Calibration mode: mechanical (from CAD) or stereo (from images). '
             'Default: mechanical.',
    )
    parser.add_argument(
        '--mount-radius', type=float, default=0.05,
        help='Camera mount ring radius in metres (default: 0.05).',
    )
    parser.add_argument(
        '--mount-height', type=float, default=0.30,
        help='Camera mount height above base_link in metres (default: 0.30).',
    )
    parser.add_argument(
        '--output', required=True,
        help='Output path for extrinsics.yaml.',
    )
    parser.add_argument(
        '--generate-launch', type=str, default=None,
        metavar='LAUNCH_FILE',
        help='Also generate static_transforms.launch.py at this path.',
    )
    # Stereo mode arguments (for future use)
    parser.add_argument(
        '--images-dir', type=str, default=None,
        help='[stereo mode] Directory containing paired calibration images.',
    )
    parser.add_argument(
        '--intrinsics-dir', type=str, default=None,
        help='[stereo mode] Directory containing per-camera camera_info.yaml files.',
    )

    args = parser.parse_args()

    print('=== IMX219 Extrinsic Calibration ===')
    print()

    if args.mode == 'mechanical':
        run_mechanical_mode(args)
    elif args.mode == 'stereo':
        run_stereo_mode(args)

    print()
    print('Extrinsics calibration complete.')
    print(f'Output: {args.output}')
    if args.generate_launch:
        print(f'Launch: {args.generate_launch}')
    print()
    print('Next steps:')
    print('  1. Update saltybot_surround/config/surround_vision_params.yaml')
    print('     with camera_info_url paths pointing to calibrated YAMLs')
    print('  2. Run generate_static_transforms.py (or use --generate-launch above)')
    print('  3. Restart saltybot_cameras and saltybot_surround nodes')


if __name__ == '__main__':
    main()
