import json
import math
import numpy as np

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


def companion(a, alpha, delay):
    """Companion matrix for x[t+1] = a*x[t] + alpha*x[t-delay]."""
    n = delay + 1
    C = np.zeros((n, n), dtype=complex if np.iscomplexobj(alpha) else float)
    C[0, 0] = a
    C[0, -1] = alpha
    C[1:, :-1] = np.eye(delay)
    return C


def rho(a, alpha, delay):
    return float(np.max(np.abs(np.linalg.eigvals(companion(a, alpha, delay)))))


def empirical_exponent(a, alpha, delay, steps=5000):
    C = companion(a, alpha, delay)
    v = rng.normal(size=delay + 1) + 1j * rng.normal(size=delay + 1)
    v /= np.linalg.norm(v)
    logs = []
    for _ in range(steps):
        v = C @ v
        n = np.linalg.norm(v)
        logs.append(math.log(max(n, 1e-300)))
        v /= max(n, 1e-300)
    # discard transient; renormalization makes this a Lyapunov estimate
    return float(np.mean(logs[500:]))


def nontrivial_eigenvalues(W):
    vals = np.linalg.eigvals(W)
    # For row-stochastic W, remove the eigenvalue nearest one.
    j = int(np.argmin(np.abs(vals - 1)))
    return np.delete(vals, j)


def disagreement_rollout(W, a, sigma, delay, steps=180):
    n = W.shape[0]
    # History is a transverse random perturbation around a nonzero synchronized value.
    hist = np.zeros((delay + 1, n))
    hist[:] = 1.0 + 0.03 * rng.normal(size=(delay + 1, n))
    hist -= hist.mean(axis=1, keepdims=True)
    ds = []
    for t in range(steps):
        x = hist[-1]
        delayed = hist[0]
        y = a * x + sigma * W @ delayed
        ds.append(float(np.mean((x - x.mean()) ** 2)))
        hist = np.concatenate([hist[1:], y[None]], axis=0)
    tail = np.array(ds[20:])
    slope = float(np.polyfit(np.arange(len(tail)), np.log(tail + 1e-30), 1)[0])
    return ds, slope


def make_symmetric_ring(n):
    W = np.zeros((n, n))
    for i in range(n):
        W[i, (i - 1) % n] = 0.5
        W[i, (i + 1) % n] = 0.5
    return W


def make_directed_heterogeneous(n, trials=3000):
    # Directed, nonreciprocal, heterogeneous indegrees; retain row normalization and W1=1.
    best, best_score = None, float('inf')
    for _ in range(trials):
        W = rng.dirichlet(0.22 * np.ones(n), size=n)
        vals = nontrivial_eigenvalues(W)
        score = float(np.max(np.abs(vals)))
        if score < best_score:
            best, best_score = W, score
    return best


def main():
    a = 0.72
    # Prediction 1: exact transition at rho=1 for every delay.
    boundary_rows = []
    for d in [0, 1, 2, 4, 7]:
        # bisection on positive real alpha
        lo, hi = 0.0, 2.0
        for _ in range(55):
            mid = (lo + hi) / 2
            if rho(a, mid, d) < 1:
                lo = mid
            else:
                hi = mid
        alpha_star = (lo + hi) / 2
        boundary_rows.append({"delay": d, "predicted_alpha_rho1": alpha_star,
                              "rho_at_boundary": rho(a, alpha_star, d)})

    # Prediction 2: measured finite-rollout exponent equals log spectral radius.
    exponent_rows = []
    for d in [1, 2, 4]:
        for alpha in [0.05, 0.15, 0.30, 0.50]:
            r = rho(a, alpha, d)
            e = empirical_exponent(a, alpha, d)
            exponent_rows.append({"delay": d, "alpha": alpha, "log_rho": math.log(r),
                                  "measured_exponent": e, "abs_error": abs(e-math.log(r))})

    # Prediction 3: increasing delay shrinks the stable positive-alpha interval.
    delay_rows = []
    for d in range(0, 9):
        lo, hi = 0.0, 2.0
        for _ in range(50):
            mid = (lo + hi) / 2
            if rho(a, mid, d) < 1: lo = mid
            else: hi = mid
        delay_rows.append({"delay": d, "stable_alpha_max": (lo+hi)/2})

    n = 8
    W_base = make_symmetric_ring(n)
    W_idea = make_directed_heterogeneous(n)
    vals_base = nontrivial_eigenvalues(W_base)
    vals_idea = nontrivial_eigenvalues(W_idea)
    sigma = 0.42
    delay = 2
    ds_b, slope_b = disagreement_rollout(W_base, a, sigma, delay)
    ds_i, slope_i = disagreement_rollout(W_idea, a, sigma, delay)
    network = {
        "a": a, "sigma": sigma, "delay": delay,
        "baseline_max_transverse_abs_eigenvalue": float(np.max(np.abs(vals_base))),
        "idea_max_transverse_abs_eigenvalue": float(np.max(np.abs(vals_idea))),
        "baseline_predicted_worst_log_rho": max(math.log(rho(a, sigma*v, delay)) for v in vals_base),
        "idea_predicted_worst_log_rho": max(math.log(rho(a, sigma*v, delay)) for v in vals_idea),
        "baseline_disagreement_log_slope": slope_b,
        "idea_disagreement_log_slope": slope_i,
        "baseline_final_disagreement": ds_b[-1],
        "idea_final_disagreement": ds_i[-1],
        "baseline_W": W_base.tolist(), "idea_W": W_idea.tolist()
    }
    result = {"seed": SEED, "predictions": {
        "boundary_rho_equals_one": boundary_rows,
        "exponent_log_rho": exponent_rows,
        "delay_stability_boundary": delay_rows}, "network_comparison": network}
    with open("results.json", "w") as f: json.dump(result, f, indent=2)
    print(json.dumps({"boundary": boundary_rows, "max_exponent_error": max(x["abs_error"] for x in exponent_rows),
                      "delay_boundary": delay_rows, "network": {k:v for k,v in network.items() if not k.endswith("_W")}}, indent=2))

if __name__ == "__main__": main()