"""calibration_loader.py — Load camera_info_manager YAML for IMX219 cameras.

Supports equidistant (fisheye) and plumb_bob (pinhole) distortion models.
Falls back to placeholder IMX219 values if YAML not found.
"""

import yaml
import numpy as np
from pathlib import Path

# Placeholder IMX219 intrinsics at 640×480 (used when calibration YAML not found)
_PLACEHOLDER_K = np.array([
    [229.0,   0.0, 320.0],
    [  0.0, 229.0, 240.0],
    [  0.0,   0.0,   1.0],
], dtype=np.float64)
_PLACEHOLDER_D = np.zeros((4, 1), dtype=np.float64)  # no distortion as fallback


def load_camera_info(yaml_path: str) -> dict:
    """Load camera_info_manager YAML.

    Returns dict with keys:
        K: np.ndarray [3, 3] camera matrix
        D: np.ndarray [4, 1] distortion coefficients (equidistant) or [5,1] (plumb_bob)
        distortion_model: str ('equidistant' or 'plumb_bob')
        image_size: (width, height)
        camera_name: str
        valid: bool (False if using placeholder)
    """
    p = Path(yaml_path)
    if not p.exists():
        return {
            'K': _PLACEHOLDER_K.copy(),
            'D': _PLACEHOLDER_D.copy(),
            'distortion_model': 'equidistant',
            'image_size': (640, 480),
            'camera_name': p.parent.name,
            'valid': False,
        }

    with open(p) as f:
        data = yaml.safe_load(f)

    cm = data['camera_matrix']
    K = np.array(cm['data'], dtype=np.float64).reshape(cm['rows'], cm['cols'])

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

    model = data.get('distortion_model', 'equidistant')
    w = data['image_width']
    h = data['image_height']

    return {
        'K': K, 'D': D, 'distortion_model': model,
        'image_size': (w, h), 'camera_name': data.get('camera_name', ''),
        'valid': True,
    }


def load_all_camera_infos(calibration_base: str) -> dict:
    """Load calibration for all 4 cameras.

    calibration_base: e.g. '/mnt/nvme/saltybot/calibration'
    Returns dict: {cam_name: camera_info_dict}
    """
    cams = ['front', 'right', 'rear', 'left']
    result = {}
    for name in cams:
        yaml_path = str(Path(calibration_base) / name / 'camera_info.yaml')
        info = load_camera_info(yaml_path)
        result[name] = info
        status = 'OK' if info['valid'] else 'PLACEHOLDER'
        print(f'  [{status}] {name}: {yaml_path}')
    return result
