from __future__ import annotations

import math
from pathlib import Path

import numpy as np
import pytest

from k1link.device_plugins.xgrids_k1.analyze.calibrated_projection import (
    CalibratedProjectionError,
    Kb4ProjectionProfile,
    depth_colors,
    map_points_to_lidar,
    project_map_points_kb4,
    quaternion_xyzw_to_rotation_matrix,
)
from k1link.device_plugins.xgrids_k1.calibration_schema import (
    parse_k1_factory_calibration,
)

FIXTURES = Path(__file__).parent / "fixtures" / "k1" / "calibration"


def _identity_profile() -> Kb4ProjectionProfile:
    transform = np.eye(4, dtype=np.float64)
    transform.setflags(write=False)
    return Kb4ProjectionProfile(
        source_id="sensor.camera.right",
        calibration_slot="camera_1",
        width=800,
        height=600,
        intrinsic_fx_fy_cx_cy=(100.0, 100.0, 400.0, 300.0),
        distortion_kb4=(0.0, 0.0, 0.0, 0.0),
        t_camera_from_lidar=transform,
    )


def test_pose_inverse_maps_world_points_back_into_lidar_frame() -> None:
    half = math.sqrt(0.5)
    points_lidar = map_points_to_lidar(
        [[10.0, 21.0, 30.0]],
        position_map_xyz=(10.0, 20.0, 30.0),
        orientation_map_from_lidar_xyzw=(0.0, 0.0, half, half),
    )
    np.testing.assert_allclose(points_lidar, [[1.0, 0.0, 0.0]], atol=1e-12)


def test_quaternion_rotation_is_orthonormal_after_normalization() -> None:
    rotation = quaternion_xyzw_to_rotation_matrix((0.0, 0.0, 2.0, 2.0))
    assert rotation @ rotation.T == pytest.approx(np.eye(3), abs=1e-12)
    assert np.linalg.det(rotation) == pytest.approx(1.0, abs=1e-12)


def test_kb4_projection_uses_theta_polynomial_and_rejects_behind_camera() -> None:
    projected = project_map_points_kb4(
        [[0.0, 0.0, 1.0], [1.0, 0.0, 1.0], [0.0, 0.0, -1.0]],
        position_map_xyz=(0.0, 0.0, 0.0),
        orientation_map_from_lidar_xyzw=(0.0, 0.0, 0.0, 1.0),
        profile=_identity_profile(),
    )
    assert projected.source_point_count == 3
    assert projected.camera_front_point_count == 2
    assert projected.projected_point_count == 2
    assert projected.pixels_xy[0] == pytest.approx([400.0, 300.0])
    assert projected.pixels_xy[1] == pytest.approx(
        [400.0 + 100.0 * math.pi / 4.0, 300.0]
    )
    assert projected.source_indices.tolist() == [0, 1]


def test_factory_profile_binds_right_main_camera_and_scales_intrinsics() -> None:
    calibration = parse_k1_factory_calibration(
        (FIXTURES / "camera.yaml").read_bytes(),
        (FIXTURES / "extrinsic_camera_lidar.yaml").read_bytes(),
    )
    profile = Kb4ProjectionProfile.from_factory_calibration(
        calibration,
        "sensor.camera.right",
    )
    camera = calibration.camera("camera_1")
    assert profile.calibration_slot == "camera_1"
    assert (profile.width, profile.height) == (800, 600)
    assert profile.intrinsic_fx_fy_cx_cy == pytest.approx(
        tuple(value * 0.2 for value in camera.intrinsic)
    )
    assert profile.t_camera_from_lidar == pytest.approx(
        np.asarray(calibration.t_camera_from_lidar("camera_1"))
    )


def test_projection_rejects_unknown_camera_and_invalid_quaternion() -> None:
    calibration = parse_k1_factory_calibration(
        (FIXTURES / "camera.yaml").read_bytes(),
        (FIXTURES / "extrinsic_camera_lidar.yaml").read_bytes(),
    )
    with pytest.raises(CalibratedProjectionError, match="admitted K1 main camera"):
        Kb4ProjectionProfile.from_factory_calibration(calibration, "sensor.camera.unknown")
    with pytest.raises(CalibratedProjectionError, match="no usable norm"):
        quaternion_xyzw_to_rotation_matrix((0.0, 0.0, 0.0, 0.0))


def test_depth_colors_are_bounded_and_repeatable() -> None:
    first = depth_colors([1.0, 2.0, 3.0])
    second = depth_colors([1.0, 2.0, 3.0])
    assert first.dtype == np.uint8
    assert first.shape == (3, 3)
    assert np.array_equal(first, second)
    assert first[0, 0] > first[0, 2]
    assert first[-1, 2] > first[-1, 0]