import json
import math
import numpy as np
from scipy.optimize import brentq

SEED = 2029
np.random.seed(SEED)


def ring_adj(n):
    A = np.zeros((n, n))
    for i in range(n):
        A[i, (i - 1) % n] = 1.0
        A[i, (i + 1) % n] = 1.0
    return A


def complete_adj(n):
    A = np.ones((n, n)) - np.eye(n)
    return A


def delayed_linear(A, eps, k, delay_steps, h, steps, x0=None, targets=None):
    n = len(A)
    d = A.sum(1)
    x = np.ones(n) if x0 is None else np.asarray(x0, float).copy()
    targets = np.zeros(n) if targets is None else np.asarray(targets, float)
    # Constant prehistory is the standard well-posed history for this test.
    hist = [x.copy() for _ in range(delay_steps + 1)]
    xs = np.empty((steps + 1, n)); xs[0] = x
    for t in range(steps):
        stale = hist[0]
        v = -eps * (x - targets) + k * (A @ stale - d * x)
        x = x + h * v
        hist.pop(0); hist.append(x.copy())
        xs[t + 1] = x
    return xs


def rate_from_trace(trace, h, lo, hi):
    y = np.abs(trace[lo:hi])
    t = np.arange(lo, hi) * h
    slope = np.polyfit(t, np.log(np.maximum(y, 1e-14)), 1)[0]
    return float(-slope)


def exact_slow_rate(eps, kd, tau):
    if tau == 0: return eps
    # q is the positive decay rate of q=eps+kd*(1-exp(-q*tau)).
    # For x(t)=exp(-q t), q=eps+kd*(1-exp(q*tau)); the
    # principal real decay root lies between 0 and eps.
    f = lambda q: q - eps - kd * (1.0 - math.exp(q * tau))
    return brentq(f, 0.0, eps)


def main():
    h = 0.001
    n = 4
    A = ring_adj(n); degree = 2.0
    eps = 0.2
    # Prediction 1: slow collective rate scales as 1/(1+k*d*tau).
    speed_rows = []
    for k in [0.25, 0.5, 1.0, 2.0]:
        for tau in [0.0, 0.25, 0.5, 1.0]:
            delay_steps = round(tau / h)
            tr = delayed_linear(A, eps, k, delay_steps, h, 12000)
            # discard transient, and use mean (consensus) mode
            q = rate_from_trace(tr.mean(1), h, 4000, 11000)
            approx = eps / (1 + k * degree * tau)
            exact = exact_slow_rate(eps, k * degree, tau)
            speed_rows.append(dict(k=k, tau=tau, observed=q, approx=approx, exact=exact,
                                   approx_rel_error=abs(q-approx)/approx))

    # Prediction 2: the relevant dimensionless synchronization ratio is
    # rho=eps*||J||/(k*lambda_2); disagreement should rise around rho ~ 1.
    # Heterogeneous quadratic minima create a persistent disagreement signal.
    targets = np.array([-1., 1., -1., 1.])
    lam2 = 2.0  # ring of four
    sync_rows = []
    for k in [0.05, 0.1, 0.2, 0.4, 0.8]:
        tr = delayed_linear(A, eps, k, 0, h, 10000, x0=np.zeros(n), targets=targets)
        tail = tr[-1000:]
        disagreement = np.mean(np.sum((tail - tail.mean(1, keepdims=True))**2, axis=1))
        rho = eps / (k * lam2)
        sync_rows.append(dict(k=k, rho=rho, disagreement=disagreement))

    # Prediction 3: with no local field, disagreement decay is set by k*lambda_2.
    gap_rows = []
    for name, B, gap in [('ring', A, 2.0), ('complete', complete_adj(n), 4.0)]:
        for k in [0.1, 0.2, 0.4]:
            x0 = np.array([1., -1., 0.5, 0.2])
            tr = delayed_linear(B, 0.0, k, 0, h, 5000, x0=x0)
            dis = np.sqrt(np.mean((tr - tr.mean(1, keepdims=True))**2, axis=1))
            q = rate_from_trace(dis, h, 500, 4000)
            gap_rows.append(dict(graph=name, k=k, lambda2=gap, observed_rate=q,
                                 predicted_rate=k*gap))

    # Small optimizer-facing comparison: four workers minimize local quadratics.
    # Independent SGD is the standard local baseline; delayed diffusion is the idea.
    opt = {}
    for name, delay in [('local', None), ('delayed_consensus', 5)]:
        x = np.zeros(n); hist = [x.copy()] * (delay + 1 if delay is not None else 1)
        for _ in range(3000):
            if delay is None:
                x += 0.01 * (-eps * (x - targets))  # gradient descent on 0.5*(x-target)^2
            else:
                stale = hist[0]
                x += 0.01 * (-0.2 * (x-targets) + 0.8 * (A@stale-degree*x))
                hist.pop(0); hist.append(x.copy())
        obj = float(np.mean(0.5*(x-targets)**2))
        dis = float(np.mean((x-x.mean())**2))
        opt[name] = {'final_mean_local_objective': obj, 'final_disagreement': dis}

    # Aggregate checks deliberately distinguish the first-order approximation
    # from the exact characteristic-root prediction.
    mean_approx_err = float(np.mean([r['approx_rel_error'] for r in speed_rows if r['tau'] > 0]))
    mean_exact_err = float(np.mean([abs(r['observed']-r['exact'])/r['exact'] for r in speed_rows if r['tau'] > 0]))
    # A useful mechanism check: rate decreases monotonically with k*tau.
    monotone = all(speed_rows[i]['observed'] >= speed_rows[i+1]['observed'] - 1e-4
                   for i in range(len(speed_rows)-1) if speed_rows[i]['k'] == speed_rows[i+1]['k'])
    out = {'seed': SEED, 'step_h': h, 'speed_scaling': speed_rows,
           'synchronization_ratio_sweep': sync_rows, 'spectral_gap_sweep': gap_rows,
           'optimizer_comparison': opt,
           'summary': {'mean_first_order_relative_error': mean_approx_err,
                       'mean_exact_root_relative_error': mean_exact_err,
                       'speed_monotone_in_delay': monotone,
                       'rho_one_disagreement_at': min(sync_rows, key=lambda r: abs(r['rho']-1))}}
    with open('results.json', 'w') as f: json.dump(out, f, indent=2)
    print(json.dumps(out['summary'], indent=2))

if __name__ == '__main__': main()