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: List of candidate sessions to select from.
            Each session must have attributes compatible with the distance_fn.
        previous_sessions: List of already-captured sessions. New selections
            will be spread away from these.
        distance_fn: Function `(session_a, session_b) -> float` returning
            statistical distance between two sessions. A batchable
            `cgsh.distance.*` metric has every candidate's features extracted
            once (each column becomes one vectorized pairwise call); a plain
            callable still works through per-pair evaluation.
        energy_fn: Function `distance -> energy` that converts distance to
            coverage energy. Default is Gaussian with sigma=1.
        energy_mode: How to combine energies from multiple sources.
            - "sum": Total energy is sum of all contributions.
            - "max": Total energy is maximum contribution.
        selections: Number of sessions to select. Clamped to
            `len(potential_sessions)`.

    Returns:
        Array of selected session dicts, 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]

    selections = min(selections, len(potential_sessions))

    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)))

    prepared = prepare(distance_fn, potential_sessions)

    # Energy from each candidate to the already-captured sessions: a thin
    # N × len(previous) matrix (previous is the capture history, not the grid).
    if len(previous_sessions) == 0:
        base_energy = np.full(len(potential_sessions), nil_value)
    else:
        base_energy = combine.reduce(
            energy_fn(prepared.against(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 ``potential_sessions[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_fn(prepared.column(index))
            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(potential_sessions), nil_value)
        for index in np.flatnonzero(mask):
            contribution = combine(contribution, column(int(index)))
        return combine(base_energy, contribution)

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

    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).
        energy = total_energy(selected)
        energy[selected] = np.inf
        void = rand_argmin(energy)
        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([potential_sessions[int(i)] for i in np.where(selected)[0]])