"""camera_calibrator.py — OpenCV fisheye intrinsic calibration for IMX219 cameras.

Uses cv2.fisheye model (Scaramuzza/Kannala-Brandt 4-param distortion):
  k1, k2, k3, k4 — equidistant distortion coefficients

Outputs camera_info_manager compatible YAML:
  image_width, image_height, camera_name
  camera_matrix:  3x3 K
  distortion_model: equidistant
  distortion_coefficients: [k1, k2, k3, k4]
  rectification_matrix: 3x3 I
  projection_matrix: 3x4 P (same as K padded)
"""

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


def find_checkerboard(img_bgr: np.ndarray, board_size: tuple,
                      flags: int = None) -> tuple:
    """Detect checkerboard corners in image.

    Args:
        img_bgr: BGR image
        board_size: (cols, rows) inner corners, e.g. (8, 6)
        flags: cv2.findChessboardCorners flags

    Returns:
        (found, corners) — corners is [N, 1, 2] float32 or None
    """
    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
    if flags is None:
        flags = cv2.CALIB_CB_ADAPTIVE_THRESH | cv2.CALIB_CB_NORMALIZE_IMAGE
    found, corners = cv2.findChessboardCorners(gray, board_size, flags)
    if found:
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 50, 0.001)
        corners = cv2.cornerSubPix(gray, corners, (7, 7), (-1, -1), criteria)
    return found, corners


def calibrate_fisheye(
    image_paths: list,
    board_size: tuple,
    square_size_m: float,
    image_size: tuple,  # (width, height)
) -> dict:
    """Run cv2.fisheye calibration.

    Args:
        image_paths: list of calibration image file paths
        board_size: (cols, rows) inner corners
        square_size_m: physical square size in metres (e.g. 0.025)
        image_size: (width, height) of images

    Returns:
        dict with keys: K, D, rvecs, tvecs, rms, n_images_used
    """
    obj_p = np.zeros((board_size[0] * board_size[1], 1, 3), dtype=np.float64)
    obj_p[:, 0, :2] = np.mgrid[0:board_size[0], 0:board_size[1]].T.reshape(-1, 2)
    obj_p *= square_size_m

    obj_pts = []
    img_pts = []
    used = 0
    skipped = 0

    for path in image_paths:
        img = cv2.imread(path)
        if img is None:
            print(f'  [WARN] Cannot read: {path}')
            skipped += 1
            continue
        found, corners = find_checkerboard(img, board_size)
        if not found:
            print(f'  [SKIP] No corners: {os.path.basename(path)}')
            skipped += 1
            continue
        obj_pts.append(obj_p)
        img_pts.append(corners)
        used += 1
        print(f'  [OK]   {os.path.basename(path)}')

    if used < 5:
        raise RuntimeError(
            f'Only {used} valid images (need >=5). Check lighting and board visibility.'
        )

    K = np.zeros((3, 3))
    D = np.zeros((4, 1))
    flags = (
        cv2.fisheye.CALIB_RECOMPUTE_EXTRINSIC |
        cv2.fisheye.CALIB_CHECK_COND |
        cv2.fisheye.CALIB_FIX_SKEW
    )
    rms, K, D, rvecs, tvecs = cv2.fisheye.calibrate(
        obj_pts, img_pts, image_size, K, D, flags=flags
    )
    print(f'  Calibration RMS reprojection error: {rms:.4f} px ({used} images)')
    return {
        'K': K, 'D': D, 'rvecs': rvecs, 'tvecs': tvecs,
        'rms': rms, 'n_images_used': used, 'n_images_skipped': skipped,
        'obj_pts': obj_pts, 'img_pts': img_pts,
    }


def write_camera_info_yaml(
    output_path: str,
    camera_name: str,
    image_size: tuple,  # (width, height)
    K: np.ndarray,
    D: np.ndarray,
    rms: float,
) -> None:
    """Write camera_info_manager compatible YAML.

    Format compatible with camera_info_manager CameraInfoManager.loadCameraInfo().
    distortion_model = 'equidistant' (ROS2 fisheye standard).
    D has 4 coefficients [k1, k2, k3, k4] for fisheye model.
    """
    w, h = image_size
    fx = float(K[0, 0])
    fy = float(K[1, 1])
    cx = float(K[0, 2])
    cy = float(K[1, 2])

    data = {
        'image_width':  w,
        'image_height': h,
        'camera_name':  camera_name,
        'camera_matrix': {
            'rows': 3, 'cols': 3,
            'data': [fx, 0.0, cx, 0.0, fy, cy, 0.0, 0.0, 1.0],
        },
        'distortion_model': 'equidistant',
        'distortion_coefficients': {
            'rows': 1, 'cols': 4,
            'data': [float(D[0, 0]), float(D[1, 0]),
                     float(D[2, 0]), float(D[3, 0])],
        },
        'rectification_matrix': {
            'rows': 3, 'cols': 3,
            'data': [1.0, 0.0, 0.0,
                     0.0, 1.0, 0.0,
                     0.0, 0.0, 1.0],
        },
        'projection_matrix': {
            'rows': 3, 'cols': 4,
            'data': [fx, 0.0, cx, 0.0,
                     0.0, fy, cy, 0.0,
                     0.0, 0.0, 1.0, 0.0],
        },
        # Metadata (not part of camera_info_manager spec, but informational)
        '_calibration_rms_px': round(rms, 4),
    }

    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}')


def compute_reprojection_errors(
    obj_pts: list, img_pts: list, rvecs: list, tvecs: list,
    K: np.ndarray, D: np.ndarray,
) -> list:
    """Compute per-image reprojection errors for fisheye model."""
    errors = []
    for i, (op, ip) in enumerate(zip(obj_pts, img_pts)):
        projected, _ = cv2.fisheye.projectPoints(op, rvecs[i], tvecs[i], K, D)
        err = cv2.norm(ip, projected, cv2.NORM_L2) / np.sqrt(len(ip))
        errors.append(err)
    return errors
