#!/usr/bin/env python3
"""
calib.py — Odin1 的相机/雷达标定参数与 FishPoly 鱼眼相机模型。

公式来源：odin_ros_driver/src/pointcloud_depth_converter.cpp
          createDistortionMaps()（厂商逐像素实现），这里改写成向量化的连续函数。

坐标约定（已由驱动源码确认）：
    Tcl_0 把【雷达系】的点变到【相机系】：p_cam = R · p_lidar + t
    （证据：驱动里 Kcl_ = K * Tcl，然后 transformPointCloud(cloud, cloud_in_cam, Kcl_)）

相机模型（FishPoly）：
    归一化针孔坐标 (x, y)  ->  r = sqrt(x²+y²)  ->  θ = atan(r)
    θ_d = θ + k2θ² + k3θ³ + k4θ⁴ + k5θ⁵ + k6θ⁶ + k7θ⁷
    畸变后的归一化坐标 = (x/r, y/r) * θ_d
    像素 = (A11·x_d + A12·y_d + u0,  A22·y_d + v0)

注意：这里所有"像素"都指**原始（未去畸变）图像**的像素，也就是 /odin1/image。
"""
import numpy as np

# ---------------- calib.yaml（设备 #O1-N090100083） ----------------
A11 = 730.51652555414909
A12 = -0.68917984297115186
A22 = 730.30384140813805
U0 = 785.10214573595919
V0 = 670.33088344368650

K2 = -8.2510665959742261e-03
K3 = 1.6324575756549702e-02
K4 = -5.6837005957061821e-02
K5 = 4.9991959834060670e-02
K6 = -2.4463726125325536e-02
K7 = 3.8528646040351232e-03

IMAGE_WIDTH = 1600
IMAGE_HEIGHT = 1296
MAX_INCIDENT_DEG = 120.0

# Tcl_0：雷达 -> 相机
TCL = np.array([
    [-0.00177, -0.99998,  0.00586,  0.03809],
    [ 0.00744, -0.00588, -0.99996, -0.00189],
    [ 0.99997, -0.00173,  0.00745, -0.01966],
    [ 0.0,      0.0,      0.0,      1.0    ],
], dtype=np.float64)
R_CAM_LIDAR = TCL[:3, :3]
T_CAM_LIDAR = TCL[:3, 3]

# 安全上限：θ 接近 90° 时 tan(θ) 发散。
# 这颗镜头水平半视角约 63°，图像四角可到 ~87°，所以上限放到 89°。
_THETA_MAX = np.deg2rad(89.0)


def theta_d(theta):
    """把入射角 θ 映射成畸变后的归一化半径。"""
    th = np.asarray(theta, dtype=np.float64)
    return (th + K2 * th**2 + K3 * th**3 + K4 * th**4
            + K5 * th**5 + K6 * th**6 + K7 * th**7)


def dtheta_d(theta):
    """θ_d 对 θ 的导数（牛顿迭代用）。"""
    th = np.asarray(theta, dtype=np.float64)
    return (1.0 + 2 * K2 * th + 3 * K3 * th**2 + 4 * K4 * th**3
            + 5 * K5 * th**4 + 6 * K6 * th**5 + 7 * K7 * th**6)


def solve_theta(rho_d, iters=30):
    """反解：给定畸变归一化半径 rho_d，求入射角 θ。"""
    rho_d = np.asarray(rho_d, dtype=np.float64)
    th = np.minimum(rho_d, _THETA_MAX).astype(np.float64)
    for _ in range(iters):
        f = theta_d(th) - rho_d
        fp = dtheta_d(th)
        fp = np.where(np.abs(fp) < 1e-12, 1e-12, fp)
        step = f / fp
        th = np.clip(th - step, 0.0, _THETA_MAX)
        if np.all(np.abs(step) < 1e-12):
            break
    return th


def pixel_to_ray(u, v):
    """原始(畸变)图像像素 -> 相机系下的单位方向向量。

    标量输入返回 (3,)，数组输入返回 (N,3)。
    """
    u = np.asarray(u, dtype=np.float64)
    v = np.asarray(v, dtype=np.float64)
    yd = (v - V0) / A22
    xd = (u - U0 - A12 * yd) / A11

    rho = np.hypot(xd, yd)
    small = rho < 1e-9
    rho_safe = np.where(small, 1.0, rho)

    theta = solve_theta(rho)
    r = np.tan(theta)

    x = xd * (r / rho_safe)
    y = yd * (r / rho_safe)
    d = np.stack([x, y, np.ones_like(x)], axis=-1)
    d = d / np.linalg.norm(d, axis=-1, keepdims=True)
    return d


def project_cam_to_pixel(p_cam, margin=0.0):
    """相机系点 (N,3) -> 原始图像像素 (u, v, valid)。"""
    p = np.atleast_2d(np.asarray(p_cam, dtype=np.float64))
    Z = p[:, 2]
    valid = np.isfinite(p).all(axis=1) & (Z > 1e-6)
    Zs = np.where(valid, Z, 1.0)

    x = p[:, 0] / Zs
    y = p[:, 1] / Zs
    r = np.hypot(x, y)
    theta = np.arctan(r)
    thd = theta_d(theta)

    s = np.where(r > 1e-9, thd / np.where(r > 1e-9, r, 1.0), 1.0)
    xd = x * s
    yd = y * s

    u = A11 * xd + A12 * yd + U0
    v = A22 * yd + V0
    valid &= np.isfinite(u) & np.isfinite(v)
    valid &= (u >= -margin) & (u <= IMAGE_WIDTH - 1 + margin)
    valid &= (v >= -margin) & (v <= IMAGE_HEIGHT - 1 + margin)
    return u, v, valid


def lidar_to_cam(pts):
    """雷达系点 (N,3) -> 相机系点 (N,3)。"""
    pts = np.atleast_2d(np.asarray(pts, dtype=np.float64))
    return pts @ R_CAM_LIDAR.T + T_CAM_LIDAR


def cam_to_lidar(pts):
    """相机系点 (N,3) -> 雷达系点 (N,3)。"""
    pts = np.atleast_2d(np.asarray(pts, dtype=np.float64))
    return (pts - T_CAM_LIDAR) @ R_CAM_LIDAR


def cam_to_base_link(pts, cam_xyz=(0.0, 0.0, 0.0), rpy_deg=(0.0, 0.0, 0.0)):
    """相机系 -> 车体系。

    坐标约定：
        相机系：x 右, y 下, z 前
        车体系：x 前, y 左, z 上（ROS REP-103, base_link）

    参数：
        cam_xyz : 相机光心在车体系下的位置 (m)，例如 (0.1, 0.0, 0.3)
        rpy_deg : 相机相对车体的姿态 (roll, pitch, yaw)，单位度
                  pitch > 0 表示镜头向下俯（常见安装）
    """
    pts = np.atleast_2d(np.asarray(pts, dtype=np.float64))
    R = _base_from_cam_rotation(rpy_deg)
    return pts @ R.T + np.asarray(cam_xyz, dtype=np.float64)


def base_link_to_cam(pts, cam_xyz=(0.0, 0.0, 0.0), rpy_deg=(0.0, 0.0, 0.0)):
    """车体系 -> 相机系（cam_to_base_link 的逆）。"""
    pts = np.atleast_2d(np.asarray(pts, dtype=np.float64))
    R = _base_from_cam_rotation(rpy_deg)
    return (pts - np.asarray(cam_xyz, dtype=np.float64)) @ R


def _base_from_cam_rotation(rpy_deg=(0.0, 0.0, 0.0)):
    """构造 R_base_cam。零姿态时：相机 z(前)->车体 x(前)，x(右)->车体 -y，y(下)->车体 -z。"""
    roll, pitch, yaw = np.deg2rad(rpy_deg)
    Rx = np.array([[1, 0, 0],
                   [0, np.cos(roll), -np.sin(roll)],
                   [0, np.sin(roll), np.cos(roll)]])
    Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)],
                   [0, 1, 0],
                   [-np.sin(pitch), 0, np.cos(pitch)]])
    Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0],
                   [np.sin(yaw), np.cos(yaw), 0],
                   [0, 0, 1]])
    R0 = np.array([[0.0, 0.0, 1.0],
                   [-1.0, 0.0, 0.0],
                   [0.0, -1.0, 0.0]])
    return Rz @ Ry @ Rx @ R0


if __name__ == '__main__':
    # 往返一致性自测：像素 -> 射线 -> 投回像素
    print('Tcl determinant (should be +1):', np.linalg.det(R_CAM_LIDAR))
    print('camera position in lidar frame (m):', cam_to_lidar(np.zeros((1, 3)))[0])
    print()
    for u, v in [(800, 648), (100, 100), (1500, 1200), (400, 900), (1599, 0)]:
        d = pixel_to_ray(u, v)
        p = d * 3.0
        u2, v2, ok = project_cam_to_pixel(p)
        print('pixel (%4d,%4d) -> ray %s -> back (%8.2f,%8.2f) valid=%s  err=%.2e'
              % (u, v, np.round(d, 4), u2[0], v2[0], ok[0],
                 np.hypot(u2[0] - u, v2[0] - v)))