#!/usr/bin/env python3
"""
preview_undistorted.py — Preview undistorted IMX219 image using calibration YAML.

Usage (ROS2 mode):
  python3 preview_undistorted.py \\
    --camera front \\
    --calibration /path/to/calibration/front/camera_info.yaml

Usage (image file mode):
  python3 preview_undistorted.py \\
    --image /path/to/image.png \\
    --calibration /path/to/calibration/front/camera_info.yaml
"""

import argparse
import sys
from pathlib import Path

import cv2
import numpy as np
import yaml


CAMERA_TOPICS = {
    'front': '/camera/front/image_raw',
    'right': '/camera/right/image_raw',
    'rear':  '/camera/rear/image_raw',
    'left':  '/camera/left/image_raw',
}


def load_camera_info(yaml_path: str) -> dict:
    """Load and parse camera_info YAML.

    Returns dict with keys: K, D, image_size, camera_name, distortion_model
    """
    with open(yaml_path, 'r') as f:
        data = yaml.safe_load(f)

    w = data['image_width']
    h = data['image_height']
    K_data = data['camera_matrix']['data']
    K = np.array(K_data, dtype=np.float64).reshape(3, 3)

    D_data = data['distortion_coefficients']['data']
    D = np.array(D_data, dtype=np.float64).reshape(-1, 1)

    model = data.get('distortion_model', 'equidistant')
    name = data.get('camera_name', 'unknown')

    return {
        'K': K,
        'D': D,
        'image_size': (w, h),
        'camera_name': name,
        'distortion_model': model,
    }


def build_undistort_maps(K: np.ndarray, D: np.ndarray,
                         image_size: tuple, balance: float = 0.0) -> tuple:
    """Compute fisheye undistortion maps.

    Args:
        K: 3x3 camera matrix
        D: (4,1) fisheye distortion coefficients
        image_size: (width, height)
        balance: 0.0 = no black borders (crop), 1.0 = keep all pixels (black borders)

    Returns:
        (map1, map2) suitable for cv2.remap
    """
    new_K = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(
        K, D, image_size, np.eye(3), balance=balance
    )
    map1, map2 = cv2.fisheye.initUndistortRectifyMap(
        K, D, np.eye(3), new_K, image_size, cv2.CV_16SC2
    )
    return map1, map2, new_K


def undistort_image(img: np.ndarray, map1: np.ndarray, map2: np.ndarray) -> np.ndarray:
    """Apply precomputed undistortion maps to image."""
    return cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR,
                     borderMode=cv2.BORDER_CONSTANT)


def show_side_by_side(original: np.ndarray, undistorted: np.ndarray,
                      title: str = '', rms_info: str = '') -> np.ndarray:
    """Create side-by-side comparison image."""
    # Resize to same height if needed
    h1, w1 = original.shape[:2]
    h2, w2 = undistorted.shape[:2]
    if h1 != h2:
        undistorted = cv2.resize(undistorted, (w2 * h1 // h2, h1))

    divider = np.zeros((h1, 4, 3), dtype=np.uint8)
    divider[:] = (80, 80, 80)

    combined = np.hstack([original, divider, undistorted])

    # Labels
    cv2.putText(combined, 'ORIGINAL (fisheye)', (10, 28),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 220, 0), 2)
    cv2.putText(combined, 'UNDISTORTED', (w1 + 14, 28),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 220, 255), 2)
    if title:
        cv2.putText(combined, title, (10, combined.shape[0] - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1)
    if rms_info:
        cv2.putText(combined, rms_info, (w1 + 14, combined.shape[0] - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1)
    return combined


def preview_image_file(image_path: str, cal_info: dict, balance: float) -> None:
    """Show undistorted preview of a single image file."""
    img = cv2.imread(image_path)
    if img is None:
        print(f'[ERROR] Cannot read image: {image_path}')
        sys.exit(1)

    K = cal_info['K']
    D = cal_info['D']
    image_size = cal_info['image_size']
    name = cal_info['camera_name']

    print(f'Camera: {name}')
    print(f'Image size: {image_size[0]}x{image_size[1]}')
    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.flatten().tolist()}')

    # Resize input if needed
    ih, iw = img.shape[:2]
    if (iw, ih) != image_size:
        print(f'[WARN] Image size {iw}x{ih} != calibration size {image_size}. Resizing.')
        img = cv2.resize(img, image_size)

    map1, map2, new_K = build_undistort_maps(K, D, image_size, balance=balance)
    undist = undistort_image(img, map1, map2)

    print(f'New K after undistort (balance={balance}):')
    print(f'  fx={new_K[0,0]:.2f}  fy={new_K[1,1]:.2f}  cx={new_K[0,2]:.2f}  cy={new_K[1,2]:.2f}')

    combined = show_side_by_side(img, undist, title=str(Path(image_path).name))
    cv2.imshow('Undistorted Preview (q=quit)', combined)
    print('Press any key to close...')
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def preview_ros_live(cam_name: str, cal_info: dict, balance: float) -> None:
    """Subscribe to ROS2 camera topic and show live undistorted preview."""
    try:
        import rclpy
        from rclpy.node import Node
        from sensor_msgs.msg import Image
        from cv_bridge import CvBridge
    except ImportError as e:
        print(f'[ERROR] ROS2 not available: {e}')
        sys.exit(1)

    K = cal_info['K']
    D = cal_info['D']
    image_size = cal_info['image_size']
    topic = CAMERA_TOPICS[cam_name]

    print(f'Subscribing to: {topic}')
    print(f'Image size: {image_size[0]}x{image_size[1]}')
    print('Press q in preview window to quit.')

    map1, map2, new_K = build_undistort_maps(K, D, image_size, balance=balance)

    rclpy.init()
    node = Node('preview_undistorted')
    bridge = CvBridge()
    latest = [None]

    def cb(msg):
        try:
            latest[0] = bridge.imgmsg_to_cv2(msg, 'bgr8')
        except Exception as e:
            node.get_logger().warn(f'cv_bridge: {e}')

    node.create_subscription(Image, topic, cb, 1)

    import threading
    threading.Thread(target=rclpy.spin, args=(node,), daemon=True).start()

    window = f'Undistorted — {cam_name}'
    cv2.namedWindow(window, cv2.WINDOW_NORMAL)

    try:
        while True:
            frame = latest[0]
            if frame is not None:
                fh, fw = frame.shape[:2]
                if (fw, fh) != image_size:
                    frame = cv2.resize(frame, image_size)
                undist = undistort_image(frame, map1, map2)
                combined = show_side_by_side(frame, undist)
                cv2.imshow(window, combined)

            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                break
    finally:
        cv2.destroyAllWindows()
        node.destroy_node()
        rclpy.shutdown()


def main():
    parser = argparse.ArgumentParser(
        description='Preview undistorted IMX219 fisheye image from calibration YAML.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    mode_group = parser.add_mutually_exclusive_group(required=True)
    mode_group.add_argument(
        '--camera', choices=['front', 'right', 'rear', 'left'],
        help='Camera name for ROS2 live preview mode.',
    )
    mode_group.add_argument(
        '--image', type=str,
        help='Path to a single image file for static preview.',
    )
    parser.add_argument(
        '--calibration', required=True,
        help='Path to camera_info.yaml produced by calibrate_camera.py.',
    )
    parser.add_argument(
        '--balance', type=float, default=0.0,
        help='Undistort balance: 0.0=crop to valid pixels, 1.0=keep all (black borders). '
             'Default: 0.0.',
    )
    args = parser.parse_args()

    cal_path = args.calibration
    if not Path(cal_path).exists():
        print(f'[ERROR] Calibration file not found: {cal_path}')
        sys.exit(1)

    print(f'Loading calibration: {cal_path}')
    cal_info = load_camera_info(cal_path)

    if cal_info['distortion_model'] not in ('equidistant', 'fisheye'):
        print(f'[WARN] Distortion model "{cal_info["distortion_model"]}" — '
              f'expected "equidistant". Results may be incorrect.')

    if args.image:
        preview_image_file(args.image, cal_info, args.balance)
    else:
        preview_ros_live(args.camera, cal_info, args.balance)


if __name__ == '__main__':
    main()
