"""Laplacian-coherence coreset sampler using adjacency lists only."""
import math, random

class LaplacianColumns:
    def __init__(self, adjacency):
        self.adj = adjacency
        self.degree = [sum(w for _, w in nbrs) for nbrs in adjacency]
        self.norm2 = [self.degree[i] ** 2 + sum(w*w for _, w in adjacency[i])
                      for i in range(len(adjacency))]
    def column(self, i):
        d = {i: self.degree[i]}
        for j, w in self.adj[i]:
            d[j] = d.get(j, 0.0) - w
        return d
    def dot(self, i, j):
        # Iterate through the smaller support; this is the sparse operation.
        a, b = self.column(i), self.column(j)
        if len(a) > len(b): a, b = b, a
        return sum(x * b.get(k, 0.0) for k, x in a.items())
    def coherence(self, i, j, eps=1e-12):
        return abs(self.dot(i, j)) / (math.sqrt(self.norm2[i]*self.norm2[j]) + eps)

def coherent_sample(columns, candidates, batch_size, rng):
    """Greedy min-max coherence selection from one candidate pool."""
    if batch_size > len(candidates):
        raise ValueError("batch_size cannot exceed candidate pool")
    pool = list(candidates)
    first = pool[rng.randrange(len(pool))]
    selected = [first]
    remaining = set(pool); remaining.remove(first)
    while len(selected) < batch_size:
        # Refreshing candidates can be added by calling this function repeatedly;
        # this implementation matches one candidate pool per minibatch.
        best = min(remaining, key=lambda i: max(columns.coherence(i, j)
                                                for j in selected))
        selected.append(best); remaining.remove(best)
    return selected

def uniform_sample(n, batch_size, rng):
    return rng.sample(range(n), batch_size)