batch

class BatchedDistanceFunction[T](ABC):
    """Protocol for distance functions that support batch computation."""

    @abstractmethod
    def __call__(self, a: T, b: T) -> float:
        """Compute distance between two items (single pair)."""
        ...

    @abstractmethod
    def extract(self, sessions: cg.Array, /) -> np.ndarray:
        """Extract numeric features from sessions as a 2D array."""
        ...

    @abstractmethod
    def pairwise(self, features_a: np.ndarray, features_b: np.ndarray) -> np.ndarray:
        """Compute pairwise distances from pre-extracted features."""
        ...

    def matrix(
        self,
        vals_a: cg.Array,
        vals_b: cg.Array | None = None,
    ) -> np.ndarray:
        """Compute pairwise distance matrix."""
        if vals_b is None:
            vals_b = vals_a

        # Handle empty lists - return empty matrix
        n = len(vals_a)
        m = len(vals_b)
        if n == 0 or m == 0:
            return np.zeros((n, m), dtype=np.float64)

        # Check for missing values
        mask_a = np.array([not cg.is_missing(v) for v in vals_a])
        mask_b = np.array([not cg.is_missing(v) for v in vals_b])

        if not np.any(mask_a) or not np.any(mask_b):
            raise ValueError("All values are missing")

        # Extract features only for non-missing sessions
        valid_a = cg.Array([v for v, m in zip(vals_a, mask_a, strict=True) if m])
        valid_b = cg.Array([v for v, m in zip(vals_b, mask_b, strict=True) if m])

        features_a = self.extract(valid_a)
        features_b = self.extract(valid_b)

        # Compute pairwise distances for valid pairs
        valid_distances = self.pairwise(features_a, features_b)

        # Map back to full matrix (missing pairs get 0 contribution)
        full_distances = np.zeros((len(vals_a), len(vals_b)), dtype=np.float64)
        valid_i = np.where(mask_a)[0]
        valid_j = np.where(mask_b)[0]
        full_distances[np.ix_(valid_i, valid_j)] = valid_distances
        return full_distances

    def prepared(self, candidates: cg.Array) -> "PreparedDistance":
        """Extract ``candidates`` features once for repeated column queries."""
        return _FeatureColumns(self, candidates)

class PreparedDistance(ABC):
    """A distance function bound to a fixed candidate set, extracted once.

    Void & Cluster reads candidate-to-candidate distances one column at a time,
    always against the *same* candidate set. Extracting that set's features
    once turns each column into a single vectorized ``pairwise`` call instead
    of re-walking every session object — the per-column cost that dominates the
    sampler on large grids. Obtain one with [prepare][].
    """

    @abstractmethod
    def column(self, index: int) -> np.ndarray:
        """Distance from every candidate to ``candidates[index]`` (length N)."""
        ...

    @abstractmethod
    def against(self, others: cg.Array) -> np.ndarray:
        """Distance from every candidate to each of ``others`` (N × len(others))."""
        ...

@runtime_checkable
class SupportsMatrix(Protocol):
    """A distance function that can compute a full pairwise matrix in one call.

    Structural, not nominal: both [BatchedDistanceFunction][] and the
    composite ``CombinedBatchDistance`` provide ``matrix`` without sharing a base.
    """

    def matrix(self, vals_a: cg.Array, vals_b: cg.Array | None = None) -> np.ndarray:
        """Compute the full pairwise distance matrix between two value sets."""
        ...