import json, math
import numpy as np
from run_experiment import make_graph, collision_count, pairwise_c

SEED = 421
A, _ = make_graph()
n = len(A)
incoming = A.T.astype(np.uint8)
rng = np.random.default_rng(SEED + 20)
c = pairwise_c(A)

def fast_greedy(max_s):
    remaining = np.arange(n)
    chosen = []
    Z = np.zeros((n, 0), dtype=np.uint8)
    for _ in range(min(max_s, n)):
        if collision_count(Z) == 0:
            break
        _, inv = np.unique(Z, axis=0, return_inverse=True)
        groups = [np.flatnonzero(inv == g) for g in range(int(inv.max()) + 1)]
        cand = incoming[:, remaining]
        scores = np.zeros(len(remaining), dtype=np.int64)
        for inds in groups:
            if len(inds) > 1:
                a = cand[inds].sum(axis=0, dtype=np.int64)
                scores += a * (len(inds) - a)
        j = int(np.argmax(scores))
        if scores[j] <= 0:
            break
        q = int(remaining[j])
        chosen.append(q)
        remaining = np.delete(remaining, j)
        Z = np.column_stack((Z, incoming[:, q]))
    return chosen

rows = []
for s in [8, 16, 24, 32, 40, 49, 64, 96, 128, 160]:
    trials = 100 if s < 100 else 50
    vals = []
    for _ in range(trials):
        S = rng.choice(n, s, replace=False)
        vals.append(collision_count(incoming[:, S]))
    Sg = fast_greedy(s)
    bound = math.comb(n, 2) * (math.comb(n-c, s) / math.comb(n, s) if n-c >= s else 0.0)
    rows.append({'s': s, 'uniform_mean_collisions': float(np.mean(vals)),
                 'uniform_zero_fraction': float(np.mean(np.asarray(vals) == 0)),
                 'greedy_landmarks_used': len(Sg),
                 'greedy_collisions': collision_count(incoming[:, Sg]),
                 'bound_using_c': bound})
out = {'n': n, 'c_exact': c,
       'theoretical_s_scale': math.ceil(2*n*math.log(n)/c), 'sweep': rows}
with open('sweep_results.json', 'w') as f:
    json.dump(out, f, indent=2)
print(json.dumps(out, indent=2))