# SPDX-License-Identifier: Apache-2.0
"""Tests for ``examples.datasets.endonerf`` (branch: vnath_gsharp).

Builds a tiny synthetic EndoNeRF-shaped directory on disk and exercises
:class:`EndoNeRFParser` + :class:`EndoNeRFDataset` end-to-end. Seed is set
to 42 by the autouse fixture in ``conftest.py``.
"""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import pytest
import torch

# Pillow is an `[examples]` / `[dev]` extra dep — skip the whole module
# cleanly if a fresh `pip install .` doesn't pull it.
Image = pytest.importorskip("PIL.Image")

# The ``examples`` directory is a sibling of ``gsplat/`` and isn't installed.
# Add the gsplat repo root to ``sys.path`` so ``examples.datasets.endonerf``
# resolves the same way the trainer would.
_REPO_ROOT = Path(__file__).resolve().parent.parent
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))

from examples.datasets.endonerf import EndoNeRFDataset, EndoNeRFParser  # noqa: E402


# ---------------------------------------------------------------------------
# Fixture — tiny synthetic EndoNeRF directory
# ---------------------------------------------------------------------------


def _write_endonerf_fixture(
    root: Path, n_frames: int = 4, height: int = 8, width: int = 8
) -> Path:
    """Write a minimal EndoNeRF-shaped directory under *root* and return *root*."""
    (root / "images").mkdir(parents=True, exist_ok=True)
    (root / "depth").mkdir(parents=True, exist_ok=True)
    (root / "masks").mkdir(parents=True, exist_ok=True)

    # poses_bounds.npy: (N, 17) — 15 = pose [3 x 5] flattened, 2 = near/far bounds.
    poses = np.zeros((n_frames, 3, 5), dtype=np.float32)
    for i in range(n_frames):
        poses[i, :, :3] = np.eye(3)  # identity rotation
        poses[i, :, 3] = [0.0, 0.0, float(i)]  # translation along z so frames differ
        poses[i, :, 4] = [height, width, 100.0]  # [H, W, focal]
    bounds = np.tile([[0.1, 100.0]], (n_frames, 1)).astype(np.float32)
    poses_arr = np.concatenate([poses.reshape(n_frames, 15), bounds], axis=1)
    np.save(root / "poses_bounds.npy", poses_arr)

    for i in range(n_frames):
        # Per-frame greyscale value spread across [0, 255] so the fixture
        # works for any n_frames (the old `i * 60` overflowed uint8 at i>=5
        # under numpy >=2.0, which CI runs).
        gray = int(round(i * 255.0 / max(1, n_frames - 1)))
        rgb = np.full((height, width, 3), gray, dtype=np.uint8)
        Image.fromarray(rgb).save(root / "images" / f"{i:06d}.png")

        depth = np.full((height, width), 10 + i, dtype=np.uint8)
        Image.fromarray(depth).save(root / "depth" / f"{i:06d}.png")

        # Binary mask on-disk (will be inverted by the dataset to tool=1).
        mask = np.full((height, width), 0 if i % 2 == 0 else 255, dtype=np.uint8)
        Image.fromarray(mask).save(root / "masks" / f"{i:06d}.png")

    return root


@pytest.fixture
def endonerf_dir(tmp_path: Path) -> Path:
    return _write_endonerf_fixture(tmp_path)


# ---------------------------------------------------------------------------
# Parser — happy path
# ---------------------------------------------------------------------------


def test_endonerf_parser_loads_fixture(endonerf_dir: Path):
    parser = EndoNeRFParser(data_dir=endonerf_dir)

    assert parser.height == 8
    assert parser.width == 8
    assert parser.focal == pytest.approx(100.0)
    assert parser.K.shape == (3, 3)
    assert parser.K[0, 0] == pytest.approx(100.0)
    assert parser.K[1, 1] == pytest.approx(100.0)

    assert parser.camtoworlds.shape == (4, 4, 4)
    # The fixture writes an identity rotation in LLFF raw format ([down,
    # right, backwards]). The parser applies the LLFF → OpenGL/Nerfstudio
    # axis permutation: `c2w[:, :, [1, 0, 2, 3]] * [1, -1, 1, 1]`.
    # Identity LLFF maps to the column-permuted matrix below; translation
    # is unchanged.
    llff_identity_after_permutation = np.array(
        [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]],
        dtype=np.float32,
    )
    np.testing.assert_allclose(
        parser.camtoworlds[0, :3, :3], llff_identity_after_permutation, atol=1e-6
    )
    np.testing.assert_allclose(parser.camtoworlds[2, :3, 3], [0.0, 0.0, 2.0])
    np.testing.assert_allclose(parser.camtoworlds[:, 3, :], [[0.0, 0.0, 0.0, 1.0]] * 4)

    assert len(parser.image_paths) == 4
    assert len(parser.depth_paths) == 4
    assert len(parser.mask_paths) == 4
    np.testing.assert_allclose(parser.times, [0.0, 0.25, 0.5, 0.75])


def test_endonerf_parser_applies_llff_axis_permutation(endonerf_dir: Path):
    """Parser must apply ``c2w[:, :, [1, 0, 2, 3]] * [1, -1, 1, 1]`` to
    match the OpenGL / Nerfstudio camera convention (G-SHARP ``endo_loader.py``).

    Verified by writing a non-trivial LLFF pose and asserting the parser's
    output matches the hand-computed permutation."""
    # Custom fixture: scaled-identity LLFF rotation per frame.
    import shutil

    custom_root = endonerf_dir.parent / "endonerf_axis_test"
    if custom_root.exists():
        shutil.rmtree(custom_root)
    custom_root = _write_endonerf_fixture(custom_root)

    poses = np.load(custom_root / "poses_bounds.npy")[:, :15].reshape(-1, 3, 5)
    llff_c2w_3x4 = poses[..., :4].copy()

    # Expected post-permutation: apply the same transform the parser does.
    expected = llff_c2w_3x4[:, :, [1, 0, 2, 3]] * np.array(
        [1.0, -1.0, 1.0, 1.0], dtype=np.float32
    )

    parser = EndoNeRFParser(data_dir=custom_root)
    np.testing.assert_allclose(parser.camtoworlds[:, :3, :4], expected, atol=1e-6)


def test_endonerf_parser_split_indices_obey_test_every():
    """test_every=2: indices with (i-1)%2 == 0 → test; others → train."""
    # n_frames=8, test_every=2 → test = {1, 3, 5, 7}, train = {0, 2, 4, 6}.
    with pytest.MonkeyPatch.context() as _:
        pass
    # Use a fresh temp dir to control n_frames.
    import tempfile

    with tempfile.TemporaryDirectory() as tmp:
        root = Path(tmp)
        _write_endonerf_fixture(root, n_frames=8)
        parser = EndoNeRFParser(data_dir=root, test_every=2)
        assert parser.train_idxs == [0, 2, 4, 6]
        assert parser.test_idxs == [1, 3, 5, 7]
        assert parser.video_idxs == [0, 1, 2, 3, 4, 5, 6, 7]


# ---------------------------------------------------------------------------
# Parser — error paths
# ---------------------------------------------------------------------------


def test_endonerf_parser_scared_routing_raises():
    """SCARED dataset_type currently raises a clear NotImplementedError."""
    import tempfile

    with tempfile.TemporaryDirectory() as tmp:
        with pytest.raises(NotImplementedError, match="SCARED"):
            EndoNeRFParser(data_dir=Path(tmp), dataset_type="scared")


def test_endonerf_parser_missing_data_dir_raises(tmp_path: Path):
    with pytest.raises(FileNotFoundError, match="data_dir"):
        EndoNeRFParser(data_dir=tmp_path / "does-not-exist")


def test_endonerf_parser_missing_poses_bounds_raises(tmp_path: Path):
    # Make all dirs but skip writing poses_bounds.npy.
    (tmp_path / "images").mkdir()
    (tmp_path / "depth").mkdir()
    (tmp_path / "masks").mkdir()
    with pytest.raises(FileNotFoundError, match="poses_bounds"):
        EndoNeRFParser(data_dir=tmp_path)


def test_endonerf_parser_non_binary_mask_rejected(endonerf_dir: Path):
    """Overwrite the first mask with a non-binary value → parser rejects."""
    bad_mask = np.full((8, 8), 127, dtype=np.uint8)  # half-grey, not binary
    Image.fromarray(bad_mask).save(endonerf_dir / "masks" / "000000.png")
    with pytest.raises(ValueError, match="non-binary"):
        EndoNeRFParser(data_dir=endonerf_dir)


def test_endonerf_parser_unknown_dataset_type_raises(endonerf_dir: Path):
    with pytest.raises(ValueError, match="unknown dataset_type"):
        EndoNeRFParser(data_dir=endonerf_dir, dataset_type="foo")  # type: ignore[arg-type]


def test_endonerf_parser_count_mismatch_raises(endonerf_dir: Path):
    """Delete one mask → count mismatch error at parse time."""
    (endonerf_dir / "masks" / "000003.png").unlink()
    with pytest.raises(ValueError, match="masks/"):
        EndoNeRFParser(data_dir=endonerf_dir)


# ---------------------------------------------------------------------------
# Dataset
# ---------------------------------------------------------------------------


def test_endonerf_dataset_default_split_is_train(endonerf_dir: Path):
    parser = EndoNeRFParser(data_dir=endonerf_dir)
    ds = EndoNeRFDataset(parser)
    assert ds.split == "train"
    assert len(ds) == len(parser.train_idxs)


def test_endonerf_dataset_unknown_split_raises(endonerf_dir: Path):
    parser = EndoNeRFParser(data_dir=endonerf_dir)
    with pytest.raises(ValueError, match="unknown split"):
        EndoNeRFDataset(parser, split="bogus")


def test_endonerf_dataset_getitem_returns_expected_keys(endonerf_dir: Path):
    parser = EndoNeRFParser(data_dir=endonerf_dir)
    ds = EndoNeRFDataset(parser, split="video")
    item = ds[0]
    for key in ("image", "depth", "mask", "camtoworld", "K", "time"):
        assert key in item, f"missing key {key!r}"
    assert item["image"].shape == (8, 8, 3)
    assert item["depth"].shape == (8, 8)
    assert item["mask"].shape == (8, 8)
    assert item["camtoworld"].shape == (4, 4)
    assert item["K"].shape == (3, 3)
    assert item["time"].ndim == 0


def test_endonerf_dataset_mask_is_tissue_include_convention(endonerf_dir: Path):
    """EndoNeRF / G-SHARP convention: on-disk 0 = tissue, 255 = tool.

    After ``1 - mask/255`` in the dataset, the returned mask is a
    tissue-include mask: ``1`` where we should keep (tissue), ``0`` where
    the pixel is a tool (drop). Frame 0 on-disk is all 0 → returned all 1
    (keep everything); frame 1 on-disk is all 255 → returned all 0 (drop
    everything).
    """
    parser = EndoNeRFParser(data_dir=endonerf_dir)
    ds = EndoNeRFDataset(parser, split="video")
    assert torch.allclose(ds[0]["mask"], torch.ones(8, 8), atol=1e-6)
    assert torch.allclose(ds[1]["mask"], torch.zeros(8, 8), atol=1e-6)


def test_endonerf_dataset_image_in_unit_range(endonerf_dir: Path):
    parser = EndoNeRFParser(data_dir=endonerf_dir)
    ds = EndoNeRFDataset(parser, split="video")
    for i in range(len(ds)):
        img = ds[i]["image"]
        assert img.min() >= 0.0 - 1e-6
        assert img.max() <= 1.0 + 1e-6
