from __future__ import annotations

import importlib.util
import math
import sys
from pathlib import Path
from types import ModuleType

import numpy as np
import pytest

WORKER_ROOT = (
    Path(__file__).parents[1] / "experiments" / "perception" / "worker"
)
PROFILE_PATH = WORKER_ROOT / "e3_k1_camera1_profile.json"


def _worker_module() -> ModuleType:
    path = WORKER_ROOT / "run_e3_rectified_segmentation.py"
    sys.path.insert(0, str(WORKER_ROOT))
    spec = importlib.util.spec_from_file_location("e3_rectified_segmentation", path)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def test_e3_profile_is_pinned_to_camera1_kb4_and_exact_model() -> None:
    worker = _worker_module()

    profile, digest = worker._profile(PROFILE_PATH)

    assert len(digest) == 64
    assert profile["source"]["calibration_slot"] == "camera_1"
    assert profile["source"]["distortion_kb4"] == pytest.approx(
        [
            -0.023164451386679667,
            -0.0014974198594105452,
            -0.001039213149441563,
            -0.000035237331915978814,
        ]
    )
    assert profile["model"]["revision"] == "8d6b6d1a3f7b50d441afd7d247c2ed10db186e8f"
    assert profile["model"]["files"]["model.safetensors"]["sha256"] == (
        "c265da9a74f58f5c3f4826d23ca4ca78beac0b106cca5842beca61580de5b782"
    )


def test_kb4_inverse_round_trips_valid_field_angles() -> None:
    worker = _worker_module()
    profile, _digest = worker._profile(PROFILE_PATH)
    coefficients = np.asarray(profile["source"]["distortion_kb4"])
    theta = np.linspace(0.0, math.radians(96.4), 512)
    squared = theta * theta
    distorted = theta * (
        1.0
        + coefficients[0] * squared
        + coefficients[1] * squared**2
        + coefficients[2] * squared**3
        + coefficients[3] * squared**4
    )

    recovered = worker._invert_kb4(distorted, coefficients)

    assert recovered == pytest.approx(theta, abs=1e-10)


def test_five_view_rectification_covers_the_k1_valid_circle() -> None:
    worker = _worker_module()
    profile, _digest = worker._profile(PROFILE_PATH)
    width, height = profile["source"]["resolution"]
    _fx, _fy, cx, cy = profile["source"]["intrinsic_fx_fy_cx_cy"]
    y, x = np.mgrid[:height, :width]
    valid_mask = np.hypot(x - cx, y - cy) <= 293.0

    maps = worker._rectification_maps(profile, valid_mask)

    assert maps["coverage"]["coverage_fraction"] == 1.0
    assert maps["coverage"]["covered_pixel_count"] == int(valid_mask.sum())
    assert maps["coverage"]["overlap_count"]["min"] >= 1
    assert maps["coverage"]["overlap_count"]["max"] <= 3


def test_tile_fusion_uses_the_selected_view_and_masks_outside_fov() -> None:
    worker = _worker_module()
    valid_mask = np.asarray([[True, True], [False, True]])
    maps = {
        "selected_tile": np.asarray([[0, 1], [-1, 1]], dtype=np.int8),
        "tiles": [
            {
                "tile_map_x": np.zeros((2, 2), dtype=np.float32),
                "tile_map_y": np.zeros((2, 2), dtype=np.float32),
            },
            {
                "tile_map_x": np.ones((2, 2), dtype=np.float32),
                "tile_map_y": np.ones((2, 2), dtype=np.float32),
            },
        ],
    }
    semantics = [
        np.asarray([[4, 4], [4, 4]], dtype=np.uint8),
        np.asarray([[7, 7], [7, 7]], dtype=np.uint8),
    ]

    result = worker._fuse_tiles(semantics, maps, valid_mask)

    assert result.tolist() == [[4, 7], [0, 7]]