"""Geometry oracle: learn a position -> (material, layer_id) map from step data.

The surrogate conditions on `material` and `layer_id`, but does not predict
them — during a shower rollout they must be looked up from the new position.
There is no in-repo detector geometry (it lives in external miniCaloSim), so we
learn it from positions sampled from a real steps dataset.

miniCaloSim's detector is a stack of planar layer slabs along one axis (see
`physics/detector-design/minicalosim-geometry.md`), so `material`/`layer_id`
are a pure function of depth. The default ("slab") method exploits this: fit a
1D lookup table of depth-axis segment boundaries and do an exact O(log
#segments) binary search per query, with escape decided by depth/transverse
bounds — far cheaper per call than a nearest-neighbour search over hundreds of
thousands of reference points, which matters because this oracle is queried on
every autoregressive step of a shower rollout. "knn"/"svm" remain as generic
fallbacks (a classifier over 3D positions, escape decided by distance to the
nearest reference point) for geometries that aren't simple slab stacks.

scikit-learn / joblib are an optional dependency (the `geometry` extra) needed
by "knn"/"svm" and by `save`/`load`; they are imported lazily so the core
install stays lean.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterable

import numpy as np
import pandas as pd
import pyarrow.parquet as pq

_INSTALL_HINT = (
    "the geometry oracle needs scikit-learn — install it with "
    "`uv sync --extra cpu --extra geometry`"
)


def _require_sklearn():
    try:
        import joblib  # noqa: F401
        from sklearn.neighbors import KNeighborsClassifier  # noqa: F401
        from sklearn.svm import SVC  # noqa: F401
    except ImportError as exc:  # pragma: no cover - exercised only without extra
        raise ImportError(_INSTALL_HINT) from exc


@dataclass
class _SlabLookup:
    """Fast path for a detector that is a stack of planar layer slabs along one
    axis (miniCaloSim's actual geometry — see `giant/geometry.py` module docstring
    and `physics/detector-design/minicalosim-geometry.md`). `material`/`layer_id`
    are then a pure function of depth, found by binary search over `z_edges`
    instead of a nearest-neighbour search over the whole reference point cloud —
    O(log(#segments)) instead of O(log(#reference points)), with a far smaller
    constant factor, and exact rather than approximate.
    """

    axis: int  # which of the 3 position components is the depth axis
    z_edges: np.ndarray  # (n_segments + 1,) sorted boundaries between segments
    materials: np.ndarray  # (n_segments,) object, material of each segment
    layer_ids: np.ndarray  # (n_segments,) int64, layer_id of each segment
    radius_max: float  # largest transverse radius seen in training data

    def query(
        self, pos: np.ndarray, margin: float
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        other = [i for i in range(3) if i != self.axis]
        z = pos[:, self.axis]
        radius = np.sqrt(pos[:, other[0]] ** 2 + pos[:, other[1]] ** 2)

        idx = np.searchsorted(self.z_edges, z, side="right") - 1
        idx = np.clip(idx, 0, len(self.materials) - 1)
        material = self.materials[idx]
        layer_id = self.layer_ids[idx]

        escaped = (
            (z < self.z_edges[0] - margin)
            | (z > self.z_edges[-1] + margin)
            | (radius > self.radius_max + margin)
        )
        return material, layer_id, escaped


@dataclass
class GeometryOracle:
    """Maps world-frame position -> (material, layer_id, escaped).

    Two lookup strategies are supported (`metadata["method"]`):
    - `"slab"`: exact O(log #segments) binary search exploiting the known
      layered-slab detector geometry (see `_SlabLookup`). Fast and preferred.
    - `"knn"` / `"svm"`: a generic sklearn classifier over 3D positions,
      predicting a class index into `classes` (a list of (material, layer_id)
      pairs). Kept as a fallback for geometries that aren't simple slab stacks.

    `escape_threshold` is a distance in position units (mm). For knn/svm it's
    compared against the nearest training reference point. For slab it's the
    slack allowed beyond the observed depth range / transverse radius before a
    point is flagged `escaped`.
    """

    estimator: Any
    classes: list[tuple[str, int]]
    escape_threshold: float
    metadata: dict
    # Only populated for non-neighbour estimators (SVM) to answer the escape
    # distance query; KNeighborsClassifier answers it directly.
    _ref_tree: Any = field(default=None)
    # Only populated when metadata["method"] == "slab".
    _slab: _SlabLookup | None = field(default=None)

    def query(self, pos: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Return (material (N,) str, layer_id (N,) int, escaped (N,) bool)."""
        pos = np.ascontiguousarray(np.asarray(pos, dtype=np.float64))
        if pos.ndim != 2 or pos.shape[1] != 3:
            raise ValueError(f"pos must be (N, 3), got {pos.shape}")
        if len(pos) == 0:
            return (
                np.empty(0, dtype=object),
                np.empty(0, dtype=np.int64),
                np.empty(0, dtype=bool),
            )

        if self._slab is not None:
            return self._slab.query(pos, self.escape_threshold)

        # kneighbors gives the distance to the nearest reference point, which is
        # what the escape test needs; it exists on both KNeighborsClassifier and
        # (via a stored reference tree) our SVM wrapper below.
        dist = self._nearest_distance(pos)
        escaped = dist > self.escape_threshold

        cls_idx = self.estimator.predict(pos).astype(np.int64)
        material = np.array([self.classes[i][0] for i in cls_idx], dtype=object)
        layer_id = np.array([self.classes[i][1] for i in cls_idx], dtype=np.int64)
        return material, layer_id, escaped

    def _nearest_distance(self, pos: np.ndarray) -> np.ndarray:
        from sklearn.neighbors import KNeighborsClassifier

        if isinstance(self.estimator, KNeighborsClassifier):
            dist, _ = self.estimator.kneighbors(pos, n_neighbors=1)
            return dist[:, 0]
        # SVM (or any non-neighbour estimator): use the separately stored
        # NearestNeighbors index purely for the escape distance.
        dist, _ = self._ref_tree.kneighbors(pos, n_neighbors=1)
        return dist[:, 0]

    def save(self, path: str | Path) -> None:
        _require_sklearn()
        import joblib

        joblib.dump(
            {
                "estimator": self.estimator,
                "classes": self.classes,
                "escape_threshold": self.escape_threshold,
                "metadata": self.metadata,
                "ref_tree": getattr(self, "_ref_tree", None),
                "slab": getattr(self, "_slab", None),
            },
            path,
        )

    @classmethod
    def load(cls, path: str | Path) -> "GeometryOracle":
        _require_sklearn()
        import joblib

        d = joblib.load(path)
        obj = cls(
            estimator=d["estimator"],
            classes=[tuple(c) for c in d["classes"]],
            escape_threshold=float(d["escape_threshold"]),
            metadata=d.get("metadata", {}),
        )
        obj._ref_tree = d.get("ref_tree")
        obj._slab = d.get("slab")
        return obj


_PRE_COLS = ["pre_x", "pre_y", "pre_z"]
_POST_COLS = ["post_x", "post_y", "post_z"]
_LABEL_COLS = ["material", "layer_id"]


def _iter_point_batches(path: Path, batch_size: int = 1_000_000):
    """Yield (pos (M,3), material (M,), layer_id (M,)) from any parquet with a
    position + material + layer_id schema (raw steps *or* predict output).

    Only the needed columns are read. post_pos points are included when present
    (they share their step's label) so boundary regions are densely sampled.
    """
    pf = pq.ParquetFile(path)
    have = set(pf.schema_arrow.names)
    missing = [c for c in (*_PRE_COLS, *_LABEL_COLS) if c not in have]
    if missing:
        raise ValueError(
            f"{path} is missing columns {missing} needed to build a geometry "
            "oracle (expected pre_x/y/z, material, layer_id)"
        )
    has_post = all(c in have for c in _POST_COLS)
    cols = [*_PRE_COLS, *_LABEL_COLS] + (_POST_COLS if has_post else [])

    for batch in pf.iter_batches(batch_size=batch_size, columns=cols):
        d = batch.to_pydict()
        pre = np.array([d[c] for c in _PRE_COLS], dtype=np.float32).T
        mat = np.array([str(m) for m in d["material"]], dtype=object)
        lay = np.asarray(d["layer_id"], dtype=np.int64)
        if has_post:
            post = np.array([d[c] for c in _POST_COLS], dtype=np.float32).T
            yield (
                np.concatenate([pre, post], axis=0),
                np.concatenate([mat, mat], axis=0),
                np.concatenate([lay, lay], axis=0),
            )
        else:
            yield pre, mat, lay


def _collect_points(
    files: Iterable[Path],
    subsample: int,
    seed: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Stream files, reservoir-sample (pos, material, layer_id) points.

    Reservoir sampling keeps memory bounded regardless of total file size.
    """
    rng = np.random.default_rng(seed)
    res_pos = np.empty((subsample, 3), dtype=np.float32)
    res_mat = np.empty(subsample, dtype=object)
    res_lay = np.empty(subsample, dtype=np.int64)
    seen = 0

    for path in files:
        for pos, mat, lay in _iter_point_batches(path):
            m = len(pos)

            if seen < subsample:
                take = min(subsample - seen, m)
                res_pos[seen : seen + take] = pos[:take]
                res_mat[seen : seen + take] = mat[:take]
                res_lay[seen : seen + take] = lay[:take]
                seen += take
                if take < m:
                    # Reservoir is now full; run the standard replacement rule
                    # on the remainder of this chunk.
                    _reservoir_replace(
                        res_pos,
                        res_mat,
                        res_lay,
                        pos[take:],
                        mat[take:],
                        lay[take:],
                        seen,
                        rng,
                    )
                    seen += m - take
            else:
                _reservoir_replace(res_pos, res_mat, res_lay, pos, mat, lay, seen, rng)
                seen += m

    n = min(seen, subsample)
    return res_pos[:n], res_mat[:n], res_lay[:n]


def _reservoir_replace(res_pos, res_mat, res_lay, pos, mat, lay, seen, rng) -> None:
    """Vectorized reservoir replacement for a batch of incoming points."""
    m = len(pos)
    k = res_pos.shape[0]
    # For incoming global index j (seen..seen+m-1), keep with prob k/(j+1),
    # replacing a uniformly-chosen reservoir slot.
    idx = seen + np.arange(m)
    j = rng.integers(0, idx + 1)  # j in [0, global_index]
    keep = j < k
    slots = j[keep]
    res_pos[slots] = pos[keep]
    res_mat[slots] = mat[keep]
    res_lay[slots] = lay[keep]


def _fit_slab_lookup(
    pos: np.ndarray,
    mat: np.ndarray,
    lay: np.ndarray,
    axis: int,
    n_bins: int,
) -> tuple[_SlabLookup, dict]:
    """Fit a `_SlabLookup` assuming material/layer_id are a function of depth
    (`pos[:, axis]`) alone — true for a stack of planar layer slabs.

    Bins the depth axis into `n_bins` equal-width bins, takes the majority
    (material, layer_id) label per bin (robust to the handful of points near a
    boundary whose true label is ambiguous at bin resolution), fills any empty
    bins from the nearest populated bin, then run-length-encodes consecutive
    bins sharing a label into segments. Binary search over the segment
    boundaries then answers a query in O(log #segments).
    """
    other = [i for i in range(3) if i != axis]
    z = pos[:, axis].astype(np.float64)
    radius = np.sqrt(
        pos[:, other[0]].astype(np.float64) ** 2
        + pos[:, other[1]].astype(np.float64) ** 2
    )

    z_min, z_max = float(z.min()), float(z.max())
    if z_min == z_max:
        raise ValueError(
            "all points share the same depth-axis coordinate — pick a "
            "different `depth_axis` or use method='knn'/'svm'"
        )
    edges = np.linspace(z_min, z_max, n_bins + 1)
    bin_idx = np.clip(np.searchsorted(edges, z, side="right") - 1, 0, n_bins - 1)

    counts = (
        pd.DataFrame({"bin": bin_idx, "material": mat, "layer_id": lay})
        .groupby(["bin", "material", "layer_id"])
        .size()
        .to_frame("n")
        .reset_index()
        .sort_values("n", ascending=False)
        .drop_duplicates("bin")
    )

    bin_material = np.full(n_bins, "", dtype=object)
    bin_layer = np.full(n_bins, -1, dtype=np.int64)
    has_data = np.zeros(n_bins, dtype=bool)
    idx = counts["bin"].to_numpy()
    bin_material[idx] = counts["material"].to_numpy()
    bin_layer[idx] = counts["layer_id"].to_numpy()
    has_data[idx] = True

    # Forward/backward-fill bins with no samples from the nearest populated one.
    fill_from = np.where(has_data, np.arange(n_bins), -1)
    for b in range(1, n_bins):
        if fill_from[b] == -1:
            fill_from[b] = fill_from[b - 1]
    for b in range(n_bins - 2, -1, -1):
        if fill_from[b] == -1:
            fill_from[b] = fill_from[b + 1]
    bin_material = bin_material[fill_from]
    bin_layer = bin_layer[fill_from]

    # Run-length-encode consecutive bins sharing a label into segments.
    changed = (
        np.flatnonzero(
            (bin_material[1:] != bin_material[:-1]) | (bin_layer[1:] != bin_layer[:-1])
        )
        + 1
    )
    seg_starts = np.concatenate([[0], changed])
    z_edges = np.concatenate([edges[seg_starts], edges[-1:]])
    materials = bin_material[seg_starts]
    layer_ids = bin_layer[seg_starts]

    unique_z = np.unique(z)
    median_spacing = float(np.median(np.diff(unique_z))) if len(unique_z) > 1 else 1.0
    radius_max = float(radius.max())

    slab = _SlabLookup(
        axis=axis,
        z_edges=z_edges,
        materials=materials,
        layer_ids=layer_ids,
        radius_max=radius_max,
    )
    info = {
        "n_segments": int(len(materials)),
        "z_range": (z_min, z_max),
        "median_z_spacing": median_spacing,
        "radius_max": radius_max,
    }
    return slab, info


def build_geometry_oracle(
    files: list[Path],
    method: str = "slab",
    k: int = 1,
    subsample: int = 500_000,
    escape_factor: float = 5.0,
    seed: int = 0,
    depth_axis: int = 2,
    n_bins: int = 2000,
) -> GeometryOracle:
    """Fit a position -> (material, layer_id) classifier from steps files.

    method: "slab" (default-recommended fast path exploiting the known
        layered-slab detector geometry — see `_SlabLookup`), "knn"
        (KNeighborsClassifier), or "svm" (SVC). "slab" is O(log #segments) per
        query and exact; "knn"/"svm" are generic fallbacks for geometries that
        aren't simple slab stacks, at the cost of a much slower query (a
        nearest-neighbour or kernel evaluation against up to `subsample`
        reference points) and, for "svm", occasional misclassification.
    k: neighbours for the knn classifier (ignored otherwise).
    subsample: max reference points held in memory / used for the fit.
    escape_factor: escape_threshold = escape_factor * median spacing of the
        reference points along the relevant axis/axes, so it scales with the
        sampling density of the data.
    depth_axis: index (0/1/2 -> x/y/z) of the position component that layers
        stack along. Only used by method="slab"; default 2 (z) matches
        miniCaloSim's beam-axis-aligned layer stack.
    n_bins: depth-axis resolution for method="slab" — should be finer than the
        thinnest layer.
    """
    pos, mat, lay = _collect_points(files, subsample, seed)
    if len(pos) == 0:
        raise ValueError("no points collected — are these steps parquet files?")

    # Combined (material, layer_id) class label, used for `classes` regardless
    # of method (informational for slab; the actual classifier index for
    # knn/svm).
    pairs = list(zip((str(m) for m in mat), (int(v) for v in lay)))
    classes = sorted(set(pairs))

    if method == "slab":
        slab, info = _fit_slab_lookup(pos, mat, lay, axis=depth_axis, n_bins=n_bins)
        escape_threshold = escape_factor * info["median_z_spacing"]
        oracle = GeometryOracle(
            estimator=None,
            classes=classes,
            escape_threshold=escape_threshold,
            metadata={
                "method": "slab",
                "depth_axis": depth_axis,
                "n_bins": n_bins,
                "n_reference_points": int(len(pos)),
                "escape_factor": escape_factor,
                "n_files": len(files),
                **info,
            },
        )
        oracle._slab = slab
        return oracle

    _require_sklearn()
    from sklearn.neighbors import KNeighborsClassifier, NearestNeighbors
    from sklearn.svm import SVC

    class_to_idx = {c: i for i, c in enumerate(classes)}
    y = np.array([class_to_idx[p] for p in pairs], dtype=np.int64)

    X = pos.astype(np.float64)

    if method == "knn":
        estimator = KNeighborsClassifier(n_neighbors=k)
        estimator.fit(X, y)
        ref_tree = None
    elif method == "svm":
        estimator = SVC(kernel="rbf")
        estimator.fit(X, y)
        # SVM cannot answer nearest-neighbour distance queries, so keep a light
        # reference tree alongside it purely for the escape test.
        ref_tree = NearestNeighbors(n_neighbors=1).fit(X)
    else:
        raise ValueError(f"unknown method {method!r}; use 'slab', 'knn', or 'svm'")

    # Escape threshold from the reference point spacing. Sample a subset for the
    # median 2-NN distance (the 1st neighbour of a training point is itself).
    nn = NearestNeighbors(n_neighbors=2).fit(X)
    probe = (
        X
        if len(X) <= 20_000
        else X[np.random.default_rng(seed).choice(len(X), 20_000, replace=False)]
    )
    d2, _ = nn.kneighbors(probe, n_neighbors=2)
    median_nn = float(np.median(d2[:, 1]))
    escape_threshold = escape_factor * median_nn

    oracle = GeometryOracle(
        estimator=estimator,
        classes=classes,
        escape_threshold=escape_threshold,
        metadata={
            "method": method,
            "k": k,
            "n_reference_points": int(len(X)),
            "median_nn_dist": median_nn,
            "escape_factor": escape_factor,
            "n_files": len(files),
        },
    )
    oracle._ref_tree = ref_tree
    return oracle
