from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np
import numpy.typing as npt

from k1link.device_plugins.xgrids_k1.calibration_schema import (
    ADMITTED_MAIN_STREAM_RESOLUTION,
    MAIN_CAMERA_SLOT_BY_SOURCE,
    K1FactoryCalibration,
)

FloatArray = npt.NDArray[np.float64]
IntArray = npt.NDArray[np.int64]


class CalibratedProjectionError(ValueError):
    """Raised when a calibrated projection input violates the reviewed contract."""


@dataclass(frozen=True, slots=True)
class Kb4ProjectionProfile:
    source_id: str
    calibration_slot: str
    width: int
    height: int
    intrinsic_fx_fy_cx_cy: tuple[float, float, float, float]
    distortion_kb4: tuple[float, float, float, float]
    t_camera_from_lidar: FloatArray

    @classmethod
    def from_factory_calibration(
        cls,
        calibration: K1FactoryCalibration,
        source_id: str,
    ) -> Kb4ProjectionProfile:
        try:
            slot = MAIN_CAMERA_SLOT_BY_SOURCE[source_id]
        except KeyError as exc:
            raise CalibratedProjectionError("source is not an admitted K1 main camera") from exc
        camera = calibration.camera(slot)
        width, height = ADMITTED_MAIN_STREAM_RESOLUTION
        scale_x = width / camera.image_width
        scale_y = height / camera.image_height
        intrinsic = (
            camera.intrinsic[0] * scale_x,
            camera.intrinsic[1] * scale_y,
            camera.intrinsic[2] * scale_x,
            camera.intrinsic[3] * scale_y,
        )
        transform = np.asarray(calibration.t_camera_from_lidar(slot), dtype=np.float64)
        transform.setflags(write=False)
        return cls(
            source_id=source_id,
            calibration_slot=slot,
            width=width,
            height=height,
            intrinsic_fx_fy_cx_cy=intrinsic,
            distortion_kb4=camera.distortion,
            t_camera_from_lidar=transform,
        )


@dataclass(frozen=True, slots=True)
class ProjectedPointCloud:
    pixels_xy: FloatArray
    depths_m: FloatArray
    source_indices: IntArray
    source_point_count: int
    camera_front_point_count: int

    @property
    def projected_point_count(self) -> int:
        return int(self.pixels_xy.shape[0])


def quaternion_xyzw_to_rotation_matrix(
    orientation_xyzw: tuple[float, float, float, float],
) -> FloatArray:
    """Return R_map_from_lidar for a finite, non-zero xyzw quaternion."""

    quaternion = np.asarray(orientation_xyzw, dtype=np.float64)
    if quaternion.shape != (4,) or not np.isfinite(quaternion).all():
        raise CalibratedProjectionError("pose quaternion must contain four finite values")
    norm = float(np.linalg.norm(quaternion))
    if not math.isfinite(norm) or norm < 1e-9:
        raise CalibratedProjectionError("pose quaternion has no usable norm")
    x, y, z, w = quaternion / norm
    rotation = np.asarray(
        [
            [
                1.0 - 2.0 * (y * y + z * z),
                2.0 * (x * y - z * w),
                2.0 * (x * z + y * w),
            ],
            [
                2.0 * (x * y + z * w),
                1.0 - 2.0 * (x * x + z * z),
                2.0 * (y * z - x * w),
            ],
            [
                2.0 * (x * z - y * w),
                2.0 * (y * z + x * w),
                1.0 - 2.0 * (x * x + y * y),
            ],
        ],
        dtype=np.float64,
    )
    rotation.setflags(write=False)
    return rotation


def map_points_to_lidar(
    points_map_xyz: npt.ArrayLike,
    *,
    position_map_xyz: tuple[float, float, float],
    orientation_map_from_lidar_xyzw: tuple[float, float, float, float],
) -> FloatArray:
    """Invert the K1 T_map_from_lidar pose for row-vector map points."""

    points = _finite_points(points_map_xyz)
    position = np.asarray(position_map_xyz, dtype=np.float64)
    if position.shape != (3,) or not np.isfinite(position).all():
        raise CalibratedProjectionError("pose position must contain three finite values")
    rotation_map_from_lidar = quaternion_xyzw_to_rotation_matrix(
        orientation_map_from_lidar_xyzw
    )
    # Column-vector form is R.T @ (p_map - t). With row vectors this is
    # (p_map - t) @ R.
    return (points - position) @ rotation_map_from_lidar


def project_map_points_kb4(
    points_map_xyz: npt.ArrayLike,
    *,
    position_map_xyz: tuple[float, float, float],
    orientation_map_from_lidar_xyzw: tuple[float, float, float, float],
    profile: Kb4ProjectionProfile,
) -> ProjectedPointCloud:
    """Project one world-frame K1 lio_pcl frame into an admitted camera image."""

    points_map = _finite_points(points_map_xyz)
    points_lidar = map_points_to_lidar(
        points_map,
        position_map_xyz=position_map_xyz,
        orientation_map_from_lidar_xyzw=orientation_map_from_lidar_xyzw,
    )
    transform = profile.t_camera_from_lidar
    if transform.shape != (4, 4) or not np.isfinite(transform).all():
        raise CalibratedProjectionError("camera transform must be a finite 4x4 matrix")
    points_camera = points_lidar @ transform[:3, :3].T + transform[:3, 3]
    x, y, z = points_camera.T
    front = z > 1e-6
    front_indices = np.flatnonzero(front)
    front_points = points_camera[front]
    if front_points.size == 0:
        return ProjectedPointCloud(
            pixels_xy=np.empty((0, 2), dtype=np.float64),
            depths_m=np.empty((0,), dtype=np.float64),
            source_indices=np.empty((0,), dtype=np.int64),
            source_point_count=int(points_map.shape[0]),
            camera_front_point_count=0,
        )

    x, y, z = front_points.T
    radial = np.hypot(x, y)
    theta = np.arctan2(radial, z)
    theta_squared = theta * theta
    k1, k2, k3, k4 = profile.distortion_kb4
    theta_distorted = theta * (
        1.0
        + k1 * theta_squared
        + k2 * theta_squared**2
        + k3 * theta_squared**3
        + k4 * theta_squared**4
    )
    radial_scale = np.divide(
        theta_distorted,
        radial,
        out=np.zeros_like(theta_distorted),
        where=radial > 1e-12,
    )
    fx, fy, cx, cy = profile.intrinsic_fx_fy_cx_cy
    u = fx * x * radial_scale + cx
    v = fy * y * radial_scale + cy
    in_frame = (
        np.isfinite(u)
        & np.isfinite(v)
        & (u >= 0.0)
        & (u < profile.width)
        & (v >= 0.0)
        & (v < profile.height)
    )
    pixels = np.column_stack((u[in_frame], v[in_frame])).astype(np.float64, copy=False)
    depths = z[in_frame].astype(np.float64, copy=False)
    indices = front_indices[in_frame].astype(np.int64, copy=False)
    return ProjectedPointCloud(
        pixels_xy=pixels,
        depths_m=depths,
        source_indices=indices,
        source_point_count=int(points_map.shape[0]),
        camera_front_point_count=int(front_points.shape[0]),
    )


def unproject_pixels_kb4(
    pixels_xy: npt.ArrayLike,
    *,
    profile: Kb4ProjectionProfile,
) -> FloatArray:
    """Return unit camera-frame rays for finite pixels under the KB4 model."""

    pixels = np.asarray(pixels_xy, dtype=np.float64)
    if (
        pixels.ndim != 2
        or pixels.shape[1:] != (2,)
        or not np.isfinite(pixels).all()
    ):
        raise CalibratedProjectionError("camera pixels must have shape (N, 2)")
    if pixels.size == 0:
        return np.empty((0, 3), dtype=np.float64)
    fx, fy, cx, cy = profile.intrinsic_fx_fy_cx_cy
    if not all(math.isfinite(value) and value > 0.0 for value in (fx, fy)):
        raise CalibratedProjectionError("camera focal lengths must be positive")
    distorted_x = (pixels[:, 0] - cx) / fx
    distorted_y = (pixels[:, 1] - cy) / fy
    theta_distorted = np.hypot(distorted_x, distorted_y)
    theta = theta_distorted.copy()
    k1, k2, k3, k4 = profile.distortion_kb4
    for _ in range(12):
        squared = theta * theta
        polynomial = (
            1.0
            + k1 * squared
            + k2 * squared**2
            + k3 * squared**3
            + k4 * squared**4
        )
        derivative = (
            1.0
            + 3.0 * k1 * squared
            + 5.0 * k2 * squared**2
            + 7.0 * k3 * squared**3
            + 9.0 * k4 * squared**4
        )
        if np.any(np.abs(derivative) < 1e-9):
            raise CalibratedProjectionError("KB4 inverse derivative became singular")
        theta -= (theta * polynomial - theta_distorted) / derivative
    if not np.isfinite(theta).all() or np.any(theta < 0.0) or np.any(theta >= math.pi):
        raise CalibratedProjectionError("KB4 inverse produced invalid angles")
    scale = np.divide(
        np.sin(theta),
        theta_distorted,
        out=np.ones_like(theta),
        where=theta_distorted > 1e-12,
    )
    directions = np.column_stack(
        (
            distorted_x * scale,
            distorted_y * scale,
            np.cos(theta),
        )
    )
    norms = np.linalg.norm(directions, axis=1)
    if not np.isfinite(norms).all() or np.any(norms < 1e-9):
        raise CalibratedProjectionError("KB4 inverse produced a zero camera ray")
    return np.asarray(directions / norms[:, None], dtype=np.float64)


def depth_colors(depths_m: npt.ArrayLike) -> npt.NDArray[np.uint8]:
    """Return deterministic blue→cyan→green→yellow→red diagnostic colors."""

    depths = np.asarray(depths_m, dtype=np.float64)
    if depths.ndim != 1 or not np.isfinite(depths).all() or np.any(depths <= 0.0):
        raise CalibratedProjectionError("projected depths must be a positive finite vector")
    if depths.size == 0:
        return np.empty((0, 3), dtype=np.uint8)
    lower, upper = np.percentile(depths, [5.0, 95.0])
    span = max(float(upper - lower), 1e-9)
    normalized = np.clip((depths - lower) / span, 0.0, 1.0)
    colors = np.column_stack(
        (
            255.0 * (1.0 - normalized),
            255.0 * (1.0 - np.abs(2.0 * normalized - 1.0)),
            255.0 * normalized,
        )
    )
    return np.rint(colors).astype(np.uint8)


def _finite_points(points_xyz: npt.ArrayLike) -> FloatArray:
    points = np.asarray(points_xyz, dtype=np.float64)
    if points.ndim != 2 or points.shape[1:] != (3,):
        raise CalibratedProjectionError("point cloud must have shape (N, 3)")
    if not np.isfinite(points).all():
        raise CalibratedProjectionError("point cloud contains a non-finite coordinate")
    return points
