"""Tests for the two-panel refraction comparison figure and the sample field."""

from __future__ import annotations

import sys
from pathlib import Path

import pytest

np = pytest.importorskip("numpy")

from orisys_sensor_design.optics import (
    Interface,
    OpticalStack,
    OpticsLayer,
    camera_rays,
    quantify,
    sample_stack,
    summarize,
)
from orisys_sensor_design.optics_plot import (
    DEFAULT_SANITY_MM,
    figure_summary,
    project_pinhole,
    render_comparison,
)

REPO = Path(__file__).resolve().parents[1]
XYNOVA_EXPORT = REPO / "build" / "export" / "finger-xynova" / "r01" / "compute"


def _quad(z: float, half: float = 10.0, name: str = "q") -> Interface:
    vertices = np.array(
        [[-half, -half, z], [half, -half, z], [half, half, z], [-half, half, z]],
        dtype=np.float64,
    )
    faces = np.array([[0, 1, 2], [0, 2, 3]], dtype=np.int64)
    return Interface(name=name, vertices=vertices, faces=faces)


def _air_stack():
    layer = OpticsLayer(name="air", outer="top", inner="bottom", ior=1.0, ior_source="t")
    layers = [(layer, _quad(0.0, name="top"), _quad(-2.0, name="bottom"))]
    marker = _quad(-10.0, half=60.0, name="marker")
    return OpticalStack(marker_surface="marker", layers=(layer,)), layers, marker


def _rays(samples: int = 4) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    K = np.array([[100.0, 0.0, 50.0], [0.0, 100.0, 50.0], [0.0, 0.0, 1.0]])
    # Place the camera 5 units in front of the stack: at the sensor origin it
    # would sit exactly on the first interface and every ray would miss.
    world_to_cam = np.eye(4)
    world_to_cam[2, 3] = -5.0
    return camera_rays(K, world_to_cam, 100, 100, samples)


# ----------------------------------------------------------- sample field


def test_sample_stack_returns_arrays_with_nan_for_unresolved() -> None:
    stack, layers, marker = _air_stack()
    origin, directions, pixel_u, pixel_v = _rays(4)
    field = sample_stack(
        stack=stack,
        layers=layers,
        marker=marker,
        origin=origin,
        directions=directions,
        pitch_mm=0.25,
        measured_thickness=[2.0],
        pixel_u=pixel_u,
        pixel_v=pixel_v,
        image_size=(100, 100),
    )
    assert field.count == 16
    assert field.error_mm.shape == (16,)
    assert field.straight.shape == (16, 3)
    # A plane-facing camera resolves every ray here.
    assert field.resolved == 16
    assert np.all(np.isfinite(field.error_mm))
    assert field.grid_shape == (4, 4)
    assert field.image_size == (100, 100)
    assert np.allclose(field.error_mm, 0.0, atol=1e-9)


def test_sample_stack_marks_misses_as_nan() -> None:
    stack, layers, marker = _air_stack()
    # Rays pointing away from every surface must not become zeros.
    directions = np.tile(np.array([0.0, 0.0, 1.0]), (5, 1))
    field = sample_stack(
        stack=stack,
        layers=layers,
        marker=marker,
        origin=np.array([0.0, 0.0, 5.0]),
        directions=directions,
        pitch_mm=0.25,
        measured_thickness=[2.0],
    )
    assert field.resolved == 0
    assert np.all(np.isnan(field.error_mm))
    assert field.status_counts().get("miss") == 5


def test_summarize_agrees_with_quantify() -> None:
    """The refactor must not change the JSON report's numbers."""
    stack, layers, marker = _air_stack()
    origin, directions, pixel_u, pixel_v = _rays(4)
    kwargs = dict(
        stack=stack,
        layers=layers,
        marker=marker,
        origin=origin,
        directions=directions,
        pitch_mm=0.25,
        measured_thickness=[2.0],
    )
    summary = summarize(sample_stack(**kwargs))
    legacy = quantify(**kwargs)
    assert summary == legacy
    assert summary["statuses"]["ok"] == 16


def test_camera_rays_pixel_grid_matches_the_direction_count() -> None:
    origin, directions, pixel_u, pixel_v = _rays(5)
    assert directions.shape == (25, 3)
    assert pixel_u.shape == (25,)
    assert pixel_v.shape == (25,)
    assert pixel_u.min() > 0.0 and pixel_u.max() < 100.0
    assert origin.shape == (3,)


# ----------------------------------------------------------------- plotting


def test_project_pinhole_places_a_centred_point_at_the_principal_point() -> None:
    K = np.array([[100.0, 0.0, 50.0], [0.0, 100.0, 50.0], [0.0, 0.0, 1.0]])
    u, v, depth = project_pinhole(np.array([[0.0, 0.0, -5.0]]), K, np.eye(4))
    assert u[0] == pytest.approx(50.0)
    assert v[0] == pytest.approx(50.0)
    assert depth[0] == pytest.approx(5.0)


matplotlib = pytest.importorskip("matplotlib")


def _render(tmp_path: Path, field, marker, *, sanity_mm: float = DEFAULT_SANITY_MM):
    K = np.array([[100.0, 0.0, 50.0], [0.0, 100.0, 50.0], [0.0, 0.0, 1.0]])
    return render_comparison(
        samples=field,
        marker=marker,
        K=K,
        world_to_cam=np.eye(4),
        out_path=tmp_path / "figure.png",
        sanity_mm=sanity_mm,
    )


def test_render_comparison_writes_a_png(tmp_path: Path) -> None:
    stack, layers, marker = _air_stack()
    origin, directions, pixel_u, pixel_v = _rays(4)
    field = sample_stack(
        stack=stack,
        layers=layers,
        marker=marker,
        origin=origin,
        directions=directions,
        pitch_mm=0.25,
        measured_thickness=[2.0],
        pixel_u=pixel_u,
        pixel_v=pixel_v,
        image_size=(100, 100),
    )
    path = _render(tmp_path, field, marker)
    assert path.is_file() and path.stat().st_size > 0
    assert path.read_bytes()[:8] == b"\x89PNG\r\n\x1a\n"


def test_render_comparison_survives_an_all_unresolved_field(tmp_path: Path) -> None:
    stack, layers, marker = _air_stack()
    directions = np.tile(np.array([0.0, 0.0, 1.0]), (4, 1))
    field = sample_stack(
        stack=stack,
        layers=layers,
        marker=marker,
        origin=np.array([0.0, 0.0, 5.0]),
        directions=directions,
        pitch_mm=0.25,
        measured_thickness=[2.0],
        pixel_u=np.array([1.0, 2.0, 3.0, 4.0]),
        pixel_v=np.array([1.0, 2.0, 3.0, 4.0]),
        image_size=(100, 100),
    )
    path = _render(tmp_path, field, marker)
    assert path.is_file() and path.stat().st_size > 0


def test_render_comparison_handles_an_all_zero_field(tmp_path: Path) -> None:
    """An air-only stack has zero displacement; the colour range must not collapse."""
    stack, layers, marker = _air_stack()
    origin, directions, pixel_u, pixel_v = _rays(4)
    field = sample_stack(
        stack=stack,
        layers=layers,
        marker=marker,
        origin=origin,
        directions=directions,
        pitch_mm=0.25,
        measured_thickness=[2.0],
        pixel_u=pixel_u,
        pixel_v=pixel_v,
        image_size=(100, 100),
    )
    assert np.allclose(field.error_mm, 0.0, atol=1e-9)
    assert _render(tmp_path, field, marker).is_file()


def test_figure_summary_counts_implausible_samples() -> None:
    stack, layers, marker = _air_stack()
    origin, directions, pixel_u, pixel_v = _rays(4)
    field = sample_stack(
        stack=stack,
        layers=layers,
        marker=marker,
        origin=origin,
        directions=directions,
        pitch_mm=0.25,
        measured_thickness=[2.0],
        pixel_u=pixel_u,
        pixel_v=pixel_v,
        image_size=(100, 100),
    )
    counts = figure_summary(field, sanity_mm=5.0)
    assert counts["resolved"] == 16
    assert counts["implausible"] == 0  # air only: no displacement at all
    assert counts["sanity_mm"] == 5.0


# ---------------------------------------------------------------------- CLI


@pytest.mark.skipif(
    not (XYNOVA_EXPORT / "mesh.npz").is_file(), reason="export not built"
)
def test_cli_writes_both_the_report_and_the_figure(tmp_path: Path) -> None:
    from orisys_sensor_design import cli

    revision = REPO / "sensors" / "finger-xynova" / "r01"
    args = type(
        "Args",
        (),
        {
            "profile": "compute",
            "output_dir": XYNOVA_EXPORT,
            "samples": 8,
            "out": tmp_path / "report.json",
            "plot": tmp_path / "figure.png",
            "no_plot": False,
            "sanity_mm": 5.0,
        },
    )()
    assert cli.cmd_optics(revision, args) == 0
    assert (tmp_path / "report.json").is_file()
    assert (tmp_path / "figure.png").is_file()


@pytest.mark.skipif(
    not (XYNOVA_EXPORT / "mesh.npz").is_file(), reason="export not built"
)
def test_cli_no_plot_skips_the_figure(tmp_path: Path) -> None:
    from orisys_sensor_design import cli

    revision = REPO / "sensors" / "finger-xynova" / "r01"
    args = type(
        "Args",
        (),
        {
            "profile": "compute",
            "output_dir": XYNOVA_EXPORT,
            "samples": 8,
            "out": tmp_path / "report.json",
            "plot": tmp_path / "figure.png",
            "no_plot": True,
            "sanity_mm": 5.0,
        },
    )()
    assert cli.cmd_optics(revision, args) == 0
    assert (tmp_path / "report.json").is_file()
    assert not (tmp_path / "figure.png").exists()


@pytest.mark.skipif(
    not (XYNOVA_EXPORT / "mesh.npz").is_file(), reason="export not built"
)
def test_cli_without_matplotlib_exits_two_without_tracing(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """A missing plotting dependency must fail fast, before the ray trace."""
    from orisys_sensor_design import cli

    monkeypatch.setitem(sys.modules, "matplotlib", None)
    revision = REPO / "sensors" / "finger-xynova" / "r01"
    args = type(
        "Args",
        (),
        {
            "profile": "compute",
            "output_dir": XYNOVA_EXPORT,
            "samples": 8,
            "out": tmp_path / "report.json",
            "plot": tmp_path / "figure.png",
            "no_plot": False,
            "sanity_mm": 5.0,
        },
    )()
    assert cli.cmd_optics(revision, args) == 2
    assert not (tmp_path / "report.json").exists()