import json, math
from pathlib import Path
import numpy as np

SEED = 1522


def cycle_laplacian(n):
    L = 2*np.eye(n)
    for i in range(n):
        L[i, (i-1) % n] = -1
        L[i, (i+1) % n] = -1
    return L


def run_rollout(L, alpha, gamma, sigma, T, d=3, burn=0, reps=1, seed=0, initial=None):
    """Linear contractive recurrent GNN: H[t+1]=alpha(I-gamma L)H[t]+sigma Xi[t]."""
    local = np.random.default_rng(seed)
    n = L.shape[0]
    P = alpha * (np.eye(n) - gamma * L)
    vals = []
    for _ in range(reps):
        H = np.zeros((n, d)) if initial is None else initial.copy()
        es = []
        for t in range(T):
            H = P @ H + sigma * local.standard_normal((n, d))
            if t >= burn:
                es.append(0.5 * np.trace(H.T @ L @ H))
        vals.append(np.mean(es))
    return float(np.mean(vals)), float(np.std(vals, ddof=1) / math.sqrt(reps)) if reps > 1 else 0.0


def stationary_energy_formula(eigs, alpha, gamma, sigma, d):
    q = alpha * (1 - gamma * eigs)
    return float(0.5 * d * np.sum(eigs[1:] * sigma**2 / (1-q[1:]**2)))


def main():
    alpha, gamma, d = 0.82, 0.20, 3
    n = 12
    L = cycle_laplacian(n)
    eigs = np.linalg.eigvalsh(L)
    rho = np.max(np.abs(alpha * (1 - gamma * eigs)))
    assert rho < 1, rho

    # Prediction 1: stationary energy is proportional to sigma^2, and agrees
    # with the exact sum over graph Fourier modes.
    sigmas = np.array([0.01, 0.02, 0.04, 0.08])
    rows_sigma = []
    for s in sigmas:
        observed, se = run_rollout(L, alpha, gamma, s, T=3500, burn=1000, d=d, reps=80, seed=100+int(s*10000))
        predicted = stationary_energy_formula(eigs, alpha, gamma, s, d)
        rows_sigma.append({'sigma': float(s), 'sigma2_lambda2': float(s*s*eigs[1]), 'observed': observed, 'predicted': predicted, 'relative_error': abs(observed-predicted)/predicted, 'se': se})
    x = sigmas**2
    y = np.array([r['observed'] for r in rows_sigma])
    slope = float(np.dot(x, y) / np.dot(x, x))
    slope_r2 = float(1 - np.sum((y-slope*x)**2) / np.sum((y-y.mean())**2))

    # Prediction 2: a pure lambda_2 mode decays as q_2^(2t), giving the
    # half-life log(1/2)/log(q_2^2).
    v2 = np.linalg.eigh(L)[1][:, 1]
    initial = np.outer(v2, np.ones(d))
    q2 = alpha * (1 - gamma * eigs[1])
    pred_half = math.log(0.5) / math.log(q2*q2)
    P = alpha * (np.eye(n) - gamma * L)
    H = initial.copy(); energies = []
    for t in range(80):
        energies.append(0.5*np.trace(H.T@L@H)); H = P@H
    target = energies[0]/2
    obs_half = next((i for i,e in enumerate(energies) if e <= target), 80)

    # Prediction 3: for one fixed graph topology, uniformly scaling edge
    # weights increases lambda_2 and increases the stationary energy. This
    # avoids confounding gap with graph size/mode count.
    rows_gap = []
    sigma = 0.04
    for scale in [0.25, 0.5, 0.75, 1.0, 1.25, 1.5]:
        LL = scale * L
        ee = np.linalg.eigvalsh(LL)
        observed, se = run_rollout(LL, alpha, gamma, sigma, T=4000, burn=1200, d=d, reps=50, seed=500+int(scale*100))
        predicted = stationary_energy_formula(ee, alpha, gamma, sigma, d)
        rows_gap.append({'scale': scale, 'lambda2': float(ee[1]), 'observed': observed, 'predicted': predicted, 'se': se})
    gap_monotonic = all(rows_gap[i+1]['observed'] > rows_gap[i]['observed'] for i in range(len(rows_gap)-1))

    # Baseline comparison: same initial nonconstant state, with and without
    # persistent noise, at long horizon.
    baseline_long, _ = run_rollout(L, alpha, gamma, 0.0, T=100, burn=99, d=d, reps=1, seed=7, initial=initial)
    idea_energy, idea_se = run_rollout(L, alpha, gamma, 0.04, T=3500, burn=1000, d=d, reps=80, seed=777, initial=initial)

    result = {
        'setup': {'n': n, 'd': d, 'alpha': alpha, 'gamma': gamma, 'contraction_rho': float(rho), 'lambda2': float(eigs[1]), 'seed': SEED},
        'prediction_1_sigma_squared_scaling': {'rows': rows_sigma, 'fit_slope_energy_per_sigma2': slope, 'R2_through_origin': slope_r2},
        'prediction_2_deterministic_decay': {'q_lambda2': float(q2), 'predicted_half_life_steps': pred_half, 'observed_first_half_step': int(obs_half), 'initial_energy': float(energies[0]), 'energy_at_20': float(energies[20])},
        'prediction_3_fixed_topology_gap_sweep': {'rows': rows_gap, 'strictly_monotone_observed': gap_monotonic},
        'baseline_vs_persistent_noise': {'deterministic_energy_at_100': baseline_long, 'persistent_noise_stationary_energy': idea_energy, 'persistent_noise_se': idea_se},
    }
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()