void_and_cluster

def select_sessions(
    potential_sessions: cg.Array,
    previous_sessions: cg.Array,
    distance_fn: Callable[[Any, Any], float],
    energy_fn: Callable[[np.ndarray], np.ndarray] = _DEFAULT_ENERGY_FN,
    energy_mode: str = "sum",
    selections: int = 10,
) -> cg.Array:
    """Select diverse sessions with a Void & Cluster sampler.

    Each iteration fills the largest "void": it computes every unselected
    candidate's energy — its similarity to the previous sessions plus the
    selections made so far — and selects the candidate with the lowest
    energy. It then finds the tightest "cluster" — the selected point with
    the highest energy — and, unless that is the point just added (a stable
    set), releases it so a later iteration can place it in a better void.
    This relaxation lets late picks evict early picks that became crowded.
    The first pick is the candidate least covered by `previous_sessions`;
    with no previous sessions it is uniformly random. Total removals are
    bounded by `len(potential_sessions)` to guarantee termination.

    Energy functions are applied to whole numpy distance matrices, so
    custom `energy_fn` callables must be numpy-vectorized. Ties are broken
    with the global numpy RNG; seed `np.random` for reproducibility.

    Args:
        potential_sessions: Candidate sessions to select from. A candidate the
            metric cannot compare — every dimension missing — is never offered.
        previous_sessions: Already-captured sessions. New selections are spread
            away from these; one the metric cannot compare covers nothing.
        distance_fn: Statistical distance between two sessions. A
            `cgsh.distance.*` metric is evaluated vectorized; a plain
            `(session_a, session_b) -> float` callable is evaluated pair by pair.
        energy_fn: Converts distance to coverage energy. Default is Gaussian
            with sigma=1.
        energy_mode: How to combine energies from multiple sources.
            - "sum": Total energy is the sum of all contributions.
            - "max": Total energy is the largest contribution.
        selections: Number of sessions to select. Clamped to the number of
            comparable candidates.

    Returns:
        Array of selected sessions, in `potential_sessions` index order
        (not selection order).

    Raises:
        ValueError: If energy_mode is not "sum" or "max".
    """
    # Each mode: the ufunc that folds contributions, with its identity (the
    # energy of "no coverage").
    combiners: dict[str, tuple[np.ufunc, float]] = {
        "sum": (np.add, 0.0),
        "max": (np.maximum, -np.inf),
    }
    if energy_mode not in combiners:
        raise ValueError(f"Unknown energy mode: {energy_mode}")
    combine, nil_value = combiners[energy_mode]

    # A candidate the metric cannot compare is not a slot worth prompting for.
    metric = as_distance(distance_fn)
    comparable = metric.defined(potential_sessions)
    candidates = cg.Array(
        [session for session, ok in zip(potential_sessions, comparable, strict=True) if ok]
    )
    selections = min(selections, len(candidates))

    def rand_argmin(arr: np.ndarray) -> int:
        """Return the index of the minimum value, breaking ties uniformly at random."""
        return np.random.choice(np.flatnonzero(arr == np.min(arr)))

    def rand_argmax(arr: np.ndarray) -> int:
        """Return the index of the maximum value, breaking ties uniformly at random."""
        return np.random.choice(np.flatnonzero(arr == np.max(arr)))

    def energy(distances: np.ndarray) -> np.ndarray:
        """Coverage energy; a pair with nothing in common covers nothing."""
        return np.where(np.isnan(distances), nil_value, energy_fn(distances))

    # The capture history is small, so its whole energy matrix is built at once.
    if len(previous_sessions) == 0:
        base_energy = np.full(len(candidates), nil_value)
    else:
        base_energy = combine.reduce(
            energy(metric.matrix(candidates, previous_sessions)),
            axis=1,
        )

    # Per-column, on-demand cache — the full N × N energy matrix is never built.
    columns: dict[int, np.ndarray] = {}

    def column(index: int) -> np.ndarray:
        """Energy from every candidate to ``candidates[index]``, cached.

        A point never covers itself, so its own entry holds the nil value.
        """
        cached = columns.get(index)
        if cached is None:
            cached = energy(metric.matrix(candidates, cg.Array([candidates[index]])).reshape(-1))
            cached[index] = nil_value
            columns[index] = cached
        return cached

    def total_energy(mask: np.ndarray) -> np.ndarray:
        """Every candidate's coverage by ``base_energy`` plus the selected points.

        ``combine`` folds each selected point's cached column into a running
        total — additively for "sum", as a running maximum for "max".
        """
        contribution = np.full(len(candidates), nil_value)
        for index in np.flatnonzero(mask):
            contribution = combine(contribution, column(int(index)))
        return combine(base_energy, contribution)

    selected = np.zeros(len(candidates), dtype=bool)
    # Bounded removals guarantee termination (the cluster step can oscillate).
    swap_budget = len(candidates)

    while np.sum(selected) < selections:
        # First selection: pick the point furthest from existing sessions.
        if np.sum(selected) == 0:
            void = rand_argmin(base_energy)
            selected[void] = True
            continue

        # Void step: add the least-covered candidate (the largest void).
        energies = total_energy(selected)
        energies[selected] = np.inf
        void = rand_argmin(energies)
        selected[void] = True

        if swap_budget <= 0:
            continue

        # Cluster step: release the tightest selected point; releasing the
        # point just added means the set is stable.
        new_energy = total_energy(selected)
        new_energy[~selected] = -np.inf
        cluster = rand_argmax(new_energy)

        if cluster != void:
            selected[cluster] = False
            swap_budget -= 1

    return cg.Array([candidates[int(i)] for i in np.where(selected)[0]])