#!/usr/bin/env python3
"""
calibrate_camera.py — Run OpenCV fisheye calibration from captured images.

Usage:
  python3 calibrate_camera.py \\
    --images-dir /path/to/cal/cam_front \\
    --output /path/to/calibration/front/camera_info.yaml \\
    --camera-name camera_front \\
    --board-size 8x6 \\
    --square-size 0.025 \\
    --image-size 640x480 \\
    [--show-reprojection]
"""

import argparse
import os
import sys
from pathlib import Path

import cv2
import numpy as np


def parse_size(s: str) -> tuple:
    """Parse 'WxH' string to (width, height) tuple of ints."""
    parts = s.lower().split('x')
    if len(parts) != 2:
        raise argparse.ArgumentTypeError(f'Size must be WxH, got: {s}')
    return (int(parts[0]), int(parts[1]))


def main():
    parser = argparse.ArgumentParser(
        description='Compute IMX219 fisheye camera intrinsics from checkerboard images.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        '--images-dir', required=True,
        help='Directory containing captured calibration images (*.png, *.jpg).',
    )
    parser.add_argument(
        '--output', required=True,
        help='Output path for camera_info.yaml.',
    )
    parser.add_argument(
        '--camera-name', required=True,
        help='Camera name for YAML header (e.g. camera_front).',
    )
    parser.add_argument(
        '--board-size', type=str, default='8x6',
        help='Checkerboard inner corners as WxH (default: 8x6).',
    )
    parser.add_argument(
        '--square-size', type=float, default=0.025,
        help='Physical square size in metres (default: 0.025 = 25mm).',
    )
    parser.add_argument(
        '--image-size', type=str, default='640x480',
        help='Image dimensions as WxH (default: 640x480).',
    )
    parser.add_argument(
        '--show-reprojection', action='store_true',
        help='Display original vs reprojected corners for each image.',
    )
    args = parser.parse_args()

    # Parse board and image sizes
    bw, bh = parse_size(args.board_size)
    board_size = (bw, bh)
    image_size = parse_size(args.image_size)

    # Collect image paths
    images_dir = Path(args.images_dir)
    if not images_dir.exists():
        print(f'[ERROR] Images directory not found: {images_dir}')
        sys.exit(1)

    image_paths = sorted(
        [str(p) for p in images_dir.glob('*.png')] +
        [str(p) for p in images_dir.glob('*.jpg')] +
        [str(p) for p in images_dir.glob('*.jpeg')]
    )
    if not image_paths:
        print(f'[ERROR] No images found in: {images_dir}')
        sys.exit(1)

    print(f'Board size:   {board_size[0]}x{board_size[1]} inner corners')
    print(f'Square size:  {args.square_size * 1000:.0f}mm')
    print(f'Image size:   {image_size[0]}x{image_size[1]}')
    print(f'Images found: {len(image_paths)}')
    print(f'Output:       {args.output}')
    print()

    # Run calibration using camera_calibrator module
    # Add package to path for standalone use
    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.camera_calibrator import (
        calibrate_fisheye,
        write_camera_info_yaml,
        compute_reprojection_errors,
    )

    print('Processing images...')
    result = calibrate_fisheye(image_paths, board_size, args.square_size, image_size)

    K = result['K']
    D = result['D']
    rms = result['rms']
    rvecs = result['rvecs']
    tvecs = result['tvecs']
    obj_pts = result['obj_pts']
    img_pts = result['img_pts']
    n_used = result['n_images_used']
    n_skipped = result['n_images_skipped']

    # Per-image reprojection errors
    errors = compute_reprojection_errors(obj_pts, img_pts, rvecs, tvecs, K, D)

    # Summary table
    print()
    print(f'{"Image":<40}  {"RMS error (px)":>14}')
    print('-' * 57)
    for path, err in zip(
        [p for p in image_paths if Path(p).name not in
         {Path(p).name for p in image_paths if not Path(p).exists()}],
        errors
    ):
        flag = '  <<' if err > 2.0 else ''
        print(f'  {Path(path).name:<38}  {err:>14.4f}{flag}')

    print('-' * 57)
    print(f'  {"Overall RMS":<38}  {rms:>14.4f}')
    print(f'  {"Images used":<38}  {n_used:>14}')
    print(f'  {"Images skipped":<38}  {n_skipped:>14}')
    print()

    # Warn on high reprojection error
    if rms > 1.5:
        print(f'[WARN] RMS error {rms:.4f} px > 1.5 px threshold.')
        print('  Suggestions:')
        print('  - Remove images with corners at extreme edges (strong fisheye distortion)')
        print('  - Ensure board is flat and rigidly printed (not bent)')
        print('  - Increase n_images and vary board orientation more')
        print('  - Check that --image-size matches actual capture resolution')
    else:
        print(f'[OK] RMS error {rms:.4f} px — good calibration.')

    # Intrinsics summary
    print()
    print('Intrinsics:')
    print(f'  fx={K[0,0]:.2f}  fy={K[1,1]:.2f}  cx={K[0,2]:.2f}  cy={K[1,2]:.2f}')
    print(f'  D=[{D[0,0]:.6f}, {D[1,0]:.6f}, {D[2,0]:.6f}, {D[3,0]:.6f}]')
    print()

    # Write YAML
    write_camera_info_yaml(
        output_path=args.output,
        camera_name=args.camera_name,
        image_size=image_size,
        K=K,
        D=D,
        rms=rms,
    )

    # Show reprojection visualisation if requested
    if args.show_reprojection:
        print('Showing reprojection errors. Press any key to advance, q to quit.')
        for i, (path, err) in enumerate(zip(
            [p for p in image_paths],
            errors + [0.0] * (len(image_paths) - len(errors))
        )):
            img = cv2.imread(path)
            if img is None:
                continue

            # Project points back
            projected, _ = cv2.fisheye.projectPoints(
                obj_pts[i] if i < len(obj_pts) else None,
                rvecs[i] if i < len(rvecs) else None,
                tvecs[i] if i < len(tvecs) else None,
                K, D
            ) if i < len(obj_pts) else (None, None)

            vis = img.copy()
            if projected is not None and i < len(img_pts):
                # Draw original corners (green)
                for pt in img_pts[i]:
                    cv2.circle(vis, tuple(pt[0].astype(int)), 5, (0, 255, 0), 2)
                # Draw reprojected corners (red)
                for pt in projected:
                    cv2.circle(vis, tuple(pt[0].astype(int)), 3, (0, 0, 255), -1)

            cv2.putText(vis, f'{Path(path).name}  RMS={err:.3f}px', (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
            cv2.putText(vis, 'Green=detected  Red=reprojected', (10, 60),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (200, 200, 200), 1)

            cv2.imshow('Reprojection', vis)
            key = cv2.waitKey(0) & 0xFF
            if key == ord('q'):
                break

        cv2.destroyAllWindows()

    print('Calibration complete.')


if __name__ == '__main__':
    main()
