from __future__ import annotations

import numpy as np
import pytest

import k1link.sessions.canonical_lab_spatial as spatial_module
from k1link.sessions.canonical_lab_spatial import (
    _bounded_local_slam,
    _CanonicalSpatialIndex,
    _estimate_local_sensor_height,
    _estimate_sensor_height,
    _gravity_stable_basis_map_from_body,
    _ground_origin_map,
    _TimedPoints,
    _TimedPoses,
    canonical_lab_spatial_playback_points,
)


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]), 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)


def test_gravity_stable_body_frame_converts_rfu_to_forward_left_up() -> None:
    times = (0, 1_000_000_000, 2_000_000_000)
    poses = _TimedPoses(
        times_ns=times,
        translations=(
            np.asarray([0.0, 0.0, 0.4]),
            np.asarray([0.0, 1.0, 0.5]),
            np.asarray([0.0, 2.0, 0.3]),
        ),
        quaternions_xyzw=tuple(
            np.asarray([0.25, 0.0, 0.0, np.sqrt(1.0 - 0.25**2)]) for _ in times
        ),
    )

    basis, source = _gravity_stable_basis_map_from_body(poses, 1_000_000_000)

    assert source == "smoothed-pose-trajectory-tangent"
    assert basis[:, 0].tolist() == pytest.approx([0.0, 1.0, 0.0], abs=1e-7)
    assert basis[:, 1].tolist() == pytest.approx([-1.0, 0.0, 0.0], abs=1e-7)
    assert basis[:, 2].tolist() == pytest.approx([0.0, 0.0, 1.0], abs=1e-7)
    assert np.linalg.det(basis) == pytest.approx(1.0, abs=1e-7)


def test_ground_origin_is_projected_only_along_map_gravity() -> None:
    origin = _ground_origin_map(np.asarray([4.0, -2.0, 1.25]), 0.32)
    assert origin.tolist() == pytest.approx([4.0, -2.0, 0.93], abs=1e-9)


def test_sensor_height_tracks_current_source_window_instead_of_fixed_mount() -> None:
    times = tuple(index * 500_000_000 for index in range(8))
    points = _TimedPoints(
        times_ns=times,
        values=tuple(
            _calibration_cloud(0.18 if index < 4 else 1.05, index)
            for index in range(8)
        ),
    )
    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),
    )

    low, low_samples, _, low_source = _estimate_local_sensor_height(
        points, poses, 500_000_000, 0.5,
    )
    high, high_samples, _, high_source = _estimate_local_sensor_height(
        points, poses, 3_000_000_000, 0.5,
    )

    assert low == pytest.approx(0.18, abs=0.03)
    assert high == pytest.approx(1.05, abs=0.03)
    assert low_samples >= 3 and high_samples >= 3
    assert low_source == high_source == "local-source-cloud-ground-quantile-median"


def test_playback_track_binds_sparse_map_increments_to_dense_camera_timeline(
    tmp_path,
    monkeypatch,
) -> None:
    recording = tmp_path / "recording.rrd"
    recording.write_bytes(b"sealed")
    points = _TimedPoints(
        times_ns=(10, 20),
        values=(
            np.asarray([[1.0, 2.0, 3.0]], dtype=np.float32),
            np.asarray([[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], dtype=np.float32),
        ),
    )
    empty_poses = _TimedPoses(times_ns=(), translations=(), quaternions_xyzw=())
    index = _CanonicalSpatialIndex(
        points=points,
        poses=empty_poses,
        trajectories=_TimedPoints(times_ns=(), values=()),
        sensor_height_m=0.4,
        sensor_height_sample_count=0,
        sensor_height_mad_m=0.0,
    )
    monkeypatch.setattr(spatial_module, "_load_index", lambda *_args: index)

    track, offsets = canonical_lab_spatial_playback_points(
        recording,
        "a" * 64,
        (10, 15, 20, 25),
    )

    assert offsets == (0, 1, 1, 3, 3)
    assert track.tolist() == [
        [1.0, 2.0, 3.0],
        [4.0, 5.0, 6.0],
        [7.0, 8.0, 9.0],
    ]
    assert track.dtype == np.dtype("<f4")
    assert not track.flags.writeable