"""
纯投影计算模块
- 仅负责坐标变换、投影和基础视野过滤
- 不包含渲染、颜色或过滤逻辑
"""

from typing import Optional, Tuple

import numpy as np

from lidar_manager.logger import get_module_logger

logger = get_module_logger(__name__)


class LidarToCameraProjector:
    """LiDAR → 相机的纯投影计算器"""

    def __init__(
        self,
        camera_intrinsic: np.ndarray,
        extrinsic_lidar2cam: np.ndarray,
        remove_backward_points: bool = True,
    ):
        self.K = camera_intrinsic.copy()
        self.T = extrinsic_lidar2cam.copy()
        self.remove_backward_points = remove_backward_points

    # ========== 基础参数更新 ==========
    def update_intrinsic(self, new_K: np.ndarray) -> None:
        """更新相机内参"""
        self.K = new_K.copy()

    def update_extrinsic(self, new_T: np.ndarray) -> None:
        """更新外参（LiDAR → Camera）"""
        self.T = new_T.copy()

    def set_remove_backward_points(self, remove: bool) -> None:
        """设置是否移除相机后方的点"""
        self.remove_backward_points = remove

    # ========== 纯计算能力 ==========
    def transform_to_camera(self, points_lidar: np.ndarray) -> np.ndarray:
        """
        将 LiDAR 点云转换到相机坐标系

        参数:
            points_lidar: (N,3) 点云
        返回:
            points_cam: (N,4) 齐次坐标（x,y,z,1）
        """
        points_homo = np.hstack([points_lidar, np.ones((len(points_lidar), 1))])
        return points_homo @ self.T.T

    def filter_visible_points(self, points_cam: np.ndarray, image_size: Tuple[int, int]) -> np.ndarray:
        """
        基础可见性过滤：前方 & 图像范围内

        使用相机坐标系Z轴深度判断点是否在相机前方。

        参数:
            points_cam: (N,4) 相机坐标齐次
            image_size: (width, height)
        返回:
            mask: (N,) 有效点掩码
        """
        if self.remove_backward_points:
            front_mask = points_cam[:, 2] > 1e-6
        else:
            front_mask = np.ones(len(points_cam), dtype=bool)

        if not np.any(front_mask):
            return np.zeros(len(points_cam), dtype=bool)

        points_3d = points_cam[front_mask, :3]
        points_2d = points_3d @ self.K.T
        points_2d = points_2d[:, :2] / points_2d[:, 2:3]
        points_2d_int = np.round(points_2d).astype(int)

        width, height = image_size
        in_image_mask = (
            (points_2d_int[:, 0] >= 0)
            & (points_2d_int[:, 0] < width)
            & (points_2d_int[:, 1] >= 0)
            & (points_2d_int[:, 1] < height)
        )

        full_mask = np.zeros(len(points_cam), dtype=bool)
        full_mask[front_mask] = in_image_mask
        return full_mask

    def project_to_image(
        self, points_cam: np.ndarray, image_size: Tuple[int, int]
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        将相机坐标系点云投影到图像

        返回:
            points_2d: (M,2) 投影像素坐标（仅有效点）
            valid_mask: (N,) 有效点掩码
        """
        valid_mask = self.filter_visible_points(points_cam, image_size)
        if not np.any(valid_mask):
            return np.array([], dtype=int).reshape(0, 2), valid_mask

        points_3d = points_cam[valid_mask, :3]
        points_2d = points_3d @ self.K.T
        points_2d = points_2d[:, :2] / points_2d[:, 2:3]
        points_2d_int = np.round(points_2d).astype(int)
        return points_2d_int, valid_mask


# ========== 畸变模型函数 ==========

def apply_pinhole_distortion(
    points_normalized: np.ndarray,
    D: np.ndarray,
    K: np.ndarray
) -> np.ndarray:
    """应用针孔相机畸变模型（Brown-Conrady模型）。返回 (N, 2) 畸变后的像素坐标。"""
    if len(points_normalized) == 0:
        return np.array([], dtype=float).reshape(0, 2)

    x, y = points_normalized[:, 0], points_normalized[:, 1]
    r2 = x**2 + y**2
    r4 = r2**2
    r6 = r2**3

    k1, k2 = D[0], D[1]
    k3 = D[4] if len(D) >= 5 else 0.0
    radial = 1 + k1*r2 + k2*r4 + k3*r6

    RADIAL_MIN, RADIAL_MAX = 0.1, 10.0
    radial_valid_mask = (radial >= RADIAL_MIN) & (radial <= RADIAL_MAX)
    if not np.any(radial_valid_mask):
        return np.array([], dtype=float).reshape(0, 2)

    x, y = x[radial_valid_mask], y[radial_valid_mask]
    r2, radial = r2[radial_valid_mask], radial[radial_valid_mask]
    p1, p2 = D[2], D[3]
    x_distorted = x * radial + 2*p1*x*y + p2*(r2 + 2*x**2)
    y_distorted = y * radial + p1*(r2 + 2*y**2) + 2*p2*x*y

    fx, fy = K[0, 0], K[1, 1]
    cx, cy = K[0, 2], K[1, 2]
    u = fx * x_distorted + cx
    v = fy * y_distorted + cy
    return np.column_stack([u, v])


def apply_fisheye_distortion(
    points_normalized: np.ndarray,
    D: np.ndarray,
    K: np.ndarray
) -> np.ndarray:
    """应用鱼眼相机畸变模型 (Kannala-Brandt)。返回 (N, 2) 畸变后的像素坐标。"""
    if len(points_normalized) == 0:
        return np.array([], dtype=float).reshape(0, 2)

    x, y = points_normalized[:, 0], points_normalized[:, 1]
    r = np.sqrt(x**2 + y**2)
    theta = np.arctan(r)
    r_safe = np.where(r < 1e-8, 1e-8, r)

    if len(D) < 4:
        D_padded = np.pad(D[:min(4, len(D))], (0, max(0, 4 - len(D))), 'constant')
    else:
        D_padded = D[:4]
    k1, k2, k3, k4 = D_padded[0], D_padded[1], D_padded[2], D_padded[3]
    theta2 = theta**2
    theta_d = theta * (1 + k1*theta2 + k2*theta2**2 + k3*theta2**3 + k4*theta2**4)
    scale = theta_d / r_safe
    x_distorted = scale * x
    y_distorted = scale * y

    fx, fy = K[0, 0], K[1, 1]
    cx, cy = K[0, 2], K[1, 2]
    u = fx * x_distorted + cx
    v = fy * y_distorted + cy
    return np.column_stack([u, v])
