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

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

def roots(g, tau_r=1.0, tau_d=0.0, branches=range(-30,31)):
    """Roots of tau_r*lambda=-1+g exp(-lambda*tau_d), scalar real mode."""
    if tau_d == 0:
        return np.array([(-1.0 + g) / tau_r], dtype=complex)
    q = 1.0 / tau_r
    z = g * tau_d / tau_r * np.exp(tau_d / tau_r)
    return np.array([lambertw(z, k) / tau_d - q for k in branches])

def dominant(g, tau_r=1.0, tau_d=0.0):
    r = roots(g, tau_r, tau_d)
    return r[np.argmax(r.real)]

def discrete_growth(g, tau_r=1.0, tau_d=0.5, dt=0.002, duration=18.0):
    """Euler growth test initialized with the analytical dominant mode."""
    delay = int(round(tau_d / dt)); n = int(duration / dt)
    lam = dominant(g, tau_r, tau_d)
    times = (np.arange(n + delay + 1) - delay) * dt
    x = np.exp(lam * times).astype(complex)
    for i in range(delay, delay+n):
        x[i+1] = x[i] + dt/tau_r * (-x[i] + g*x[i-delay])
    t = np.arange(n) * dt
    sel = t > duration*0.25
    return float(np.polyfit(t[sel], np.log(np.abs(x[delay:delay+n][sel])+1e-30), 1)[0])

def hopf_prediction(tau_r=1.0, tau_d=0.5):
    # Negative feedback first Hopf crossing: omega*tau_d + atan(omega*tau_r)=pi.
    f = lambda om: om*tau_d + math.atan(om*tau_r) - math.pi
    om = brentq(f, 1e-9, 30.0)
    return om, math.sqrt(1.0 + (om*tau_r)**2)

def toy_verification():
    # Prediction 1: at zero delay the boundary is exactly g=1.
    gs = np.array([0.85, 0.98, 1.00, 1.02, 1.15])
    observed = [float(dominant(g).real) for g in gs]
    boundary = brentq(lambda g: dominant(g).real, .5, 1.5)

    # Prediction 2: for finite delay, instability begins at the Hopf gain.
    om, gcrit = hopf_prediction()
    sweep = np.linspace(-gcrit-0.12, -gcrit+0.12, 13)
    reals = np.array([dominant(g, tau_d=.5).real for g in sweep])
    cross = brentq(lambda g: dominant(g, tau_d=.5).real, -gcrit-.2, -gcrit+.2)

    # Prediction 3: dominant growth changes with gain according to the root;
    # validate against a direct delayed Euler rollout at three gains.
    gs3 = np.array([-gcrit-.04, -gcrit-.16, -gcrit-.30])
    predicted = np.array([dominant(g, tau_d=.5).real for g in gs3])
    measured = np.array([discrete_growth(g) for g in gs3])
    return {
        "zero_delay": {"predicted_boundary": 1.0, "observed_boundary": float(boundary),
                       "gains": gs.tolist(), "root_real_parts": observed},
        "delayed_hopf": {"predicted_omega": om, "predicted_gain": gcrit,
                          "observed_crossing": float(cross), "sweep_gains": sweep.tolist(),
                          "sweep_root_real": reals.tolist()},
        "growth_scaling": {"gains": gs3.tolist(), "predicted_real_root": predicted.tolist(),
                            "measured_euler_log_growth": measured.tolist(),
                            "mean_abs_error": float(np.mean(np.abs(predicted-measured)))}
    }

def stabilized_ring_demo(N=64, tau_r=1.0, tau_d=.5, dt=.002, duration=12.0):
    # Build a ring with selected Fourier gains 1.35 (unstable for this delay),
    # and an idea version that clips those gains to 0.82. This changes no parameter count.
    selected = [1,2,3,N-3,N-2,N-1]
    wh_base = np.zeros(N, dtype=complex)
    for k in selected: wh_base[k] = 1.35
    # small harmless modes make the kernel nontrivial
    wh_base[0] = .15
    wh_base[4] = wh_base[N-4] = .12
    wh_idea = wh_base.copy()
    for k in selected: wh_idea[k] *= .82/1.35
    # Fourier-domain DDE: each mode is independent; initialize broadband perturbation.
    def rollout(wh):
        delay = int(round(tau_d/dt)); n=int(duration/dt)
        z=np.zeros((n+delay+1,N), dtype=complex)
        z[:delay+1] = rng.normal(size=(delay+1,N)) + 1j*rng.normal(size=(delay+1,N))
        z[:delay+1] *= 1e-3 / np.sqrt(np.mean(np.abs(z[:delay+1])**2))
        for i in range(delay, delay+n):
            z[i+1] = z[i] + dt/tau_r*(-z[i] + wh*z[i-delay])
        norms=np.sqrt(np.mean(np.abs(z[delay:delay+n])**2,axis=1))
        return norms
    b=rollout(wh_base); s=rollout(wh_idea)
    def fit(norms):
        t=np.arange(len(norms))*dt
        sel=t>duration*.35
        return float(np.polyfit(t[sel], np.log(norms[sel]+1e-30),1)[0])
    return {"selected_modes": selected, "baseline_gain":1.35, "idea_gain":.82,
            "baseline_predicted_max_root":float(dominant(1.35,tau_d=tau_d).real),
            "idea_predicted_max_root":float(dominant(.82,tau_d=tau_d).real),
            "baseline_rollout_growth":fit(b), "idea_rollout_growth":fit(s),
            "baseline_final_over_initial":float(b[-1]/b[0]),
            "idea_final_over_initial":float(s[-1]/s[0]),
            "unstable_baseline":bool(b[-1] > 10*b[0]), "stable_idea":bool(s[-1] < 10*b[0])}

def main():
    result={"seed":SEED, "math_verification":toy_verification(), "ring_comparison":stabilized_ring_demo()}
    with open("results.json","w") as f: json.dump(result,f,indent=2)
    print(json.dumps(result, indent=2))

if __name__ == "__main__": main()