from __future__ import annotations

import numpy as np
import pytest

from k1link.sessions.canonical_lab_spatial import (
    _TimedPoints,
    _TimedPoses,
    _bounded_local_slam,
    _estimate_sensor_height,
    _ground_origin_map,
)


def _calibration_cloud(height_m: float, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    xy = rng.uniform(-5.5, 5.5, size=(500, 2)).astype(np.float32)
    radius = np.linalg.norm(xy, axis=1)
    xy = xy[(radius >= 1.0) & (radius <= 5.5)][:360]
    ground = np.column_stack((
        xy,
        rng.normal(-height_m, 0.006, size=xy.shape[0]),
    )).astype(np.float32)
    vegetation = np.column_stack((
        rng.uniform(-5, 5, size=(300, 2)),
        rng.uniform(0.0, 1.2, size=300),
    )).astype(np.float32)
    return np.concatenate((ground, vegetation), axis=0)


def test_session_sensor_height_is_derived_from_initial_source_cloud() -> None:
    times = tuple(index * 500_000_000 for index in range(12))
    points = _TimedPoints(
        times_ns=times,
        values=tuple(_calibration_cloud(0.32, index) for index in range(12)),
    )
    poses = _TimedPoses(
        times_ns=times,
        translations=tuple(np.zeros(3) for _ in times),
        quaternions_xyzw=tuple(np.asarray([0.0, 0.0, 0.0, 1.0]) for _ in times),
    )

    height, sample_count, mad = _estimate_sensor_height(points, poses)

    assert height == pytest.approx(0.32, abs=0.02)
    assert sample_count == 12
    assert mad < 0.02


def test_local_slam_accumulates_source_increments_in_ground_body_frame() -> None:
    points = _TimedPoints(
        times_ns=(0, 1_000_000_000, 2_000_000_000),
        values=(
            np.asarray([[1.0, 0.0, -0.32]], dtype=np.float32),
            np.asarray([[2.0, 0.0, -0.32]], dtype=np.float32),
            np.asarray([[3.0, 0.0, -0.32]], dtype=np.float32),
        ),
    )
    basis = np.eye(3)
    ground_origin = _ground_origin_map(np.asarray([0.0, 0.0, 0.0]), basis, 0.32)

    local, frame_count, source_count = _bounded_local_slam(
        points,
        2_000_000_000,
        ground_origin,
        basis,
    )

    assert frame_count == 3
    assert source_count == 3
    assert local[:, 2].tolist() == pytest.approx([0.0, 0.0, 0.0], abs=1e-6)