import json
import numpy as np
from scipy.linalg import eigh

SEED = 2891
rng = np.random.default_rng(SEED)


def make_graph(n=8):
    # Connected, nonuniform weighted graph (ring plus reproducible chords).
    A = np.zeros((n, n), dtype=float)
    for i in range(n):
        w = 0.25 + 0.10 * ((i * 3) % 4)
        j = (i + 1) % n
        A[i, j] = A[j, i] = w
    for i, j, w in [(0, 3, .90), (1, 5, .18), (2, 6, .55), (3, 7, .28), (0, 6, .12)]:
        A[i, j] = A[j, i] = w
    L = np.diag(A.sum(1)) - A
    return A, L


def gap(L, pins, p):
    P = np.zeros(L.shape[0])
    P[list(pins)] = p
    return float(eigh(L + np.diag(P), eigvals_only=True, subset_by_index=[0, 0])[0])


def greedy_pins(L, budget, p):
    selected = []
    current = 0.0
    gains = []
    for _ in range(budget):
        candidates = [i for i in range(L.shape[0]) if i not in selected]
        vals = [(gap(L, selected + [i], p), i) for i in candidates]
        bestval, best = max(vals, key=lambda z: (z[0], -z[1]))
        best = int(best)
        gains.append(bestval - current)
        selected.append(best)
        current = bestval
    return selected, gains


def rate_from_simulation(L, pins, p, c, Gamma, x0):
    """Exact eigendecomposition simulation, then fit the late log-energy slope."""
    n, d = L.shape[0], Gamma.shape[0]
    P = np.zeros(n); P[list(pins)] = p
    M = np.kron(L + np.diag(P), Gamma)
    vals, vecs = eigh(M)
    z0 = vecs.T @ x0.reshape(-1)
    # Enough time for the slowest excited mode to dominate.
    lam_pos = vals[vals > 1e-12][0]
    times = np.linspace(0.0, max(30.0, 12.0 / (c * lam_pos)), 500)
    coeff = z0[:, None] * np.exp(-c * vals[:, None] * times[None, :])
    xt = vecs @ coeff
    V = 0.5 * np.sum(xt * xt, axis=0)
    # fit last half, avoiding numerical underflow
    take = slice(250, 470)
    slope = float(np.polyfit(times[take], np.log(np.maximum(V[take], 1e-300)), 1)[0])
    predicted = -2.0 * c * float(vals[0])
    return slope, predicted, float(V[0]), float(V[-1])


def main():
    A, L = make_graph()
    n, d = L.shape[0], 3
    Gamma = np.array([[1.0, .15, 0.0], [.15, 1.7, .08], [0.0, .08, 2.4]])
    gamma_min = float(eigh(Gamma, eigvals_only=True, subset_by_index=[0, 0])[0])
    # Kronecker ordering is module-major, each module has d coordinates.
    x0 = rng.normal(size=(n, d))
    p = 1.0
    budget = 2
    greedy, greedy_gains = greedy_pins(L, budget, p)
    degrees = A.sum(1)
    high_degree = [int(i) for i in np.argsort(-degrees)[:budget]]
    # Deterministic random baseline plus all one/two node possibilities for context.
    random_sets = [[int(i) for i in x] for x in [(0, 1), (2, 7), (1, 6), (3, 5), (0, 4)]]
    random_gaps = [gap(L, s, p) for s in random_sets]
    gd = gap(L, greedy, p)
    hd = gap(L, high_degree, p)

    # Prediction 1: rate = -2*c*lambda_min(Lg)*lambda_min(Gamma), across c.
    scaling = []
    for c in [0.25, 0.5, 1.0, 2.0]:
        obs, pred, _, _ = rate_from_simulation(L, greedy, p, c, Gamma, x0)
        scaling.append({'c': c, 'observed_slope': obs, 'predicted_slope': pred,
                        'ratio_abs_observed_to_predicted': abs(obs / (pred * gamma_min))})

    # Prediction 2: stronger pinning raises the grounded gap with diminishing marginal gain.
    pin_sweep = []
    for strength in [0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 4.0]:
        g = gap(L, greedy, strength)
        obs, _, _, _ = rate_from_simulation(L, greedy, strength, 1.0, Gamma, x0)
        predicted_rate = -2.0 * g * gamma_min
        pin_sweep.append({'p': strength, 'lambda_min': g, 'observed_slope': obs,
                          'predicted_slope': predicted_rate,
                          'ratio': abs(obs / predicted_rate)})

    # Prediction 3: spectral greedy produces the largest gap among candidate choices.
    comparisons = {'greedy': {'pins': greedy, 'gap': gd},
                   'high_degree': {'pins': high_degree, 'gap': hd},
                   'random_sets': [{'pins': s, 'gap': g} for s, g in zip(random_sets, random_gaps)]}
    for name, pins in [('greedy', greedy), ('high_degree', high_degree)]:
        obs, pred0, _, _ = rate_from_simulation(L, pins, p, 1.0, Gamma, x0)
        comparisons[name]['observed_slope'] = obs
        comparisons[name]['predicted_slope'] = pred0
    random_slopes = []
    for s in random_sets:
        obs, pred0, _, _ = rate_from_simulation(L, s, p, 1.0, Gamma, x0)
        random_slopes.append(obs)
    comparisons['random_mean_gap'] = float(np.mean(random_gaps))
    comparisons['random_mean_observed_slope'] = float(np.mean(random_slopes))

    all_ratios = [x['ratio_abs_observed_to_predicted'] for x in scaling] + [x['ratio'] for x in pin_sweep]
    output = {
        'seed': SEED, 'graph_nodes': n, 'module_dimension': d,
        'lambda_min_Gamma': gamma_min, 'degrees': degrees.tolist(),
        'predictions': {
            'decay_rate': 'late log(V) slope = -2*c*lambda_min(L+P)*lambda_min(Gamma)',
            'pinning': 'lambda_min and decay magnitude increase with p, with diminishing gains',
            'selection': 'greedy marginal gap maximization beats random choices on this graph'
        },
        'greedy_pins': greedy, 'greedy_marginal_gains': greedy_gains,
        'scaling_sweep': scaling, 'pin_strength_sweep': pin_sweep,
        'selection_comparison': comparisons,
        'max_rate_prediction_ratio_error': float(max(abs(r - 1.0) for r in all_ratios)),
        'all_rate_ratios_within_25pct': bool(all(0.75 <= r <= 1.25 for r in all_ratios)),
        'greedy_beats_random_mean_gap': bool(gd > np.mean(random_gaps)),
        'greedy_beats_high_degree_gap': bool(gd > hd),
    }
    with open('results.json', 'w') as f:
        json.dump(output, f, indent=2, default=lambda x: x.item() if hasattr(x, 'item') else x)
    print(json.dumps(output, indent=2))


if __name__ == '__main__':
    main()