"""Tests for the geometry oracle (position -> material/layer_id + escape)."""

from pathlib import Path
from unittest.mock import patch

import numpy as np
import pytest

from giant import geometry as g

# scikit-learn is an optional extra; skip the whole module if it's missing.
pytest.importorskip("sklearn")


def _box_batch(n, rng):
    """A labelled point cloud: inside a 100mm box -> PbWO4/0, else AIR/-1."""
    pos = rng.uniform(-200, 200, (n, 3)).astype(np.float32)
    inside = (np.abs(pos) < 100).all(axis=1)
    mat = np.where(inside, "G4_PbWO4", "G4_AIR").astype(object)
    lay = np.where(inside, 0, -1).astype(np.int64)
    return pos, mat, lay


def _build(subsample=30000, method="knn", escape_factor=5.0):
    rng = np.random.default_rng(0)
    batches = [_box_batch(20000, rng) for _ in range(3)]
    with patch.object(g, "_iter_point_batches", lambda p: iter(batches)):
        return g.build_geometry_oracle(
            [Path("x")], method=method, subsample=subsample,
            escape_factor=escape_factor,
        )


def test_classes_discovered():
    orc = _build()
    assert set(orc.classes) == {("G4_PbWO4", 0), ("G4_AIR", -1)}


def test_query_labels_inside_and_outside():
    orc = _build()
    pos = np.array([[0.0, 0.0, 0.0], [150.0, 150.0, 150.0]])
    material, layer_id, _ = orc.query(pos)
    assert material[0] == "G4_PbWO4" and layer_id[0] == 0
    assert material[1] == "G4_AIR" and layer_id[1] == -1


def test_escape_flag_fires_far_from_data():
    orc = _build()
    pos = np.array([[0.0, 0.0, 0.0], [1e5, 0.0, 0.0]])
    _, _, escaped = orc.query(pos)
    assert not escaped[0]
    assert escaped[1]


def test_query_empty():
    orc = _build()
    material, layer_id, escaped = orc.query(np.empty((0, 3)))
    assert len(material) == len(layer_id) == len(escaped) == 0


def test_query_bad_shape_raises():
    orc = _build()
    with pytest.raises(ValueError):
        orc.query(np.zeros((4, 2)))


def test_save_load_roundtrip(tmp_path):
    orc = _build()
    p = tmp_path / "oracle.pkl"
    orc.save(p)
    loaded = g.GeometryOracle.load(p)

    pos = np.array([[0.0, 0.0, 0.0], [150.0, 150.0, 150.0], [1e5, 0.0, 0.0]])
    m0, l0, e0 = orc.query(pos)
    m1, l1, e1 = loaded.query(pos)
    assert (m0 == m1).all() and (l0 == l1).all() and (e0 == e1).all()
    assert loaded.escape_threshold == orc.escape_threshold
    assert loaded.classes == orc.classes


def test_svm_method_has_escape_tree():
    orc = _build(subsample=4000, method="svm")
    # SVM cannot answer NN-distance, so a reference tree backs the escape test.
    assert orc._ref_tree is not None
    _, _, escaped = orc.query(np.array([[1e5, 0.0, 0.0]]))
    assert escaped[0]