"""SpatialIndex — Fast neighbor queries in N-dimensional space."""

import numpy as np
from typing import Optional


class SpatialIndex:
    """
    N-dimensional spatial index for neighbor queries.

    Uses vectorized NumPy distance calculations. For very large populations,
    could be extended to use scipy.spatial.KDTree.
    """

    def __init__(self, n_dims: int):
        self.n_dims = n_dims

    @staticmethod
    def distance(a: np.ndarray, b: np.ndarray) -> float:
        return float(np.linalg.norm(a - b))

    @staticmethod
    def distances_from_point(point: np.ndarray, positions: np.ndarray) -> np.ndarray:
        """Compute distances from a single point to all positions. (N,) array."""
        return np.linalg.norm(positions - point, axis=1)

    @staticmethod
    def pairwise_distances(positions: np.ndarray) -> np.ndarray:
        """Compute all pairwise distances. Returns (N, N) matrix."""
        diff = positions[:, None, :] - positions[None, :, :]
        return np.linalg.norm(diff, axis=2)
