import json
import math
from pathlib import Path

import numpy as np
from scipy.special import lambertw

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


def dominant_root(k, tau):
    """Dominant characteristic root of lambda + k exp(-lambda tau)=0."""
    # lambda*tau = W_j(-k*tau); enumerate branches near the dominant pair.
    roots = [lambertw(-k * tau, j) / tau for j in range(-8, 9)]
    return max(roots, key=lambda z: z.real)


def simulate_linear(k, tau, dt=0.001, duration=80.0, z0=1e-6):
    """Euler integration of dz/dt=-k z(t-tau), with constant prehistory."""
    n = int(duration / dt)
    m = max(1, int(round(tau / dt)))
    z = np.full(n + m + 1, z0, dtype=float)
    for i in range(m, n + m):
        z[i + 1] = z[i] - dt * k * z[i - m]
        if abs(z[i + 1]) > 1e100:
            z[i + 1:] = z[i + 1]
            break
    t = np.arange(n + 1) * dt
    return t, z[m : m + n + 1]


def measured_growth(t, z, unstable):
    """Fit log magnitude after a short transient; ignore numerical zeros."""
    if not unstable:
        return float(np.polyfit(t[-20_000:], np.log(np.maximum(np.abs(z[-20_000:]), 1e-300)), 1)[0])
    # A broad fit before saturation/overflow gives the envelope rate.
    mask = (np.abs(z) > 1e-8) & (np.abs(z) < 1e50) & (t > 5.0)
    if mask.sum() < 20:
        return float("nan")
    return float(np.polyfit(t[mask], np.log(np.abs(z[mask])), 1)[0])


def escape_times(k, tau, Ds, R=0.1, dt=0.001, duration=220.0):
    """Direct deterministic delayed evolution from fluctuation amplitude sqrt(D)."""
    r = max(0.0, dominant_root(k, tau).real)
    out = []
    for D in Ds:
        t, z = simulate_linear(k, tau, dt=dt, duration=duration, z0=math.sqrt(D))
        hit = np.flatnonzero(np.abs(z) >= R)
        out.append(float(t[hit[0]]) if len(hit) else float("inf"))
    return r, out


def double_well_grad(x):
    return x * (x * x - 1.0)


def run_double_well(eta, steps, delayed_m=0, noise=0.018, seed=0):
    """Same noisy SGD setup; delayed_m=0 is current-gradient SGD."""
    rg = np.random.default_rng(seed)
    x = 0.95
    history = [x]
    gradients = []
    for _ in range(steps):
        gradients.append(double_well_grad(x))
        g = gradients[-1] if delayed_m == 0 or len(gradients) <= delayed_m else gradients[-1 - delayed_m]
        x -= eta * g + math.sqrt(eta) * noise * rg.normal()
        history.append(x)
    h = np.asarray(history)
    # A crossing into the opposite basin is a clear toy escape event.
    crossings = np.flatnonzero(h < -0.5)
    return int(crossings[0]) if len(crossings) else None, float(np.mean((h[-1000:] ** 2 - 1.0) ** 2) / 4.0)


def main():
    k = 1.0
    tau_c = math.pi / (2 * k)
    multiples = [0.5, 0.9, 1.05, 1.3, 2.0]
    boundary_rows = []
    for mult in multiples:
        tau = mult * tau_c
        root = dominant_root(k, tau)
        unstable = root.real > 1e-10
        t, z = simulate_linear(k, tau)
        fit = measured_growth(t, z, unstable)
        boundary_rows.append({
            "multiple_tau_c": mult, "tau": tau,
            "predicted_real_lambda": float(root.real),
            "predicted_unstable": bool(unstable),
            "observed_growth_fit": fit,
            "observed_unstable": bool(fit > 1e-4),
        })

    # Prediction 1: transition occurs at k*tau=pi/2.
    # Prediction 2: above threshold, measured growth follows Re(lambda+).
    growth_rows = [r for r in boundary_rows if r["predicted_unstable"]]
    growth_abs_errors = [abs(r["observed_growth_fit"] - r["predicted_real_lambda"]) for r in growth_rows]

    # Prediction 3: escape time is affine in log(R/sqrt(D)), with slope 1/r.
    tau_escape = 1.3 * tau_c
    Ds = np.logspace(-12, -5, 8)
    r, times = escape_times(k, tau_escape, Ds)
    x = np.log(0.1 / np.sqrt(Ds))
    slope, intercept = np.polyfit(x, np.asarray(times), 1)
    escape_rows = [{"D": float(D), "predicted_T": float(T)} for D, T in zip(Ds, times)]

    # Secondary mini-comparison: fixed stale gradients versus calibrated burst delay.
    # At the well curvature k=2, tau_c=pi/4; m=ceil(1.1*tau_c/eta).
    eta = 0.02
    burst_m = int(math.ceil(1.1 * (math.pi / 4) / eta))
    baseline = [run_double_well(eta, 30000, delayed_m=0, seed=s) for s in range(12)]
    fixed_stale = [run_double_well(eta, 30000, delayed_m=burst_m, seed=s) for s in range(12)]
    def summarize(a):
        events = [x[0] for x in a if x[0] is not None]
        return {"escape_fraction": len(events) / len(a), "median_escape_step": float(np.median(events)) if events else None,
                "final_loss": float(np.mean([x[1] for x in a]))}

    result = {
        "seed": SEED, "tau_c_k1": tau_c,
        "boundary_sweep": boundary_rows,
        "boundary_prediction": {"predicted_multiple": 1.0,
                                 "observed_first_unstable_multiple": next((r["multiple_tau_c"] for r in boundary_rows if r["observed_unstable"]), None)},
        "growth_prediction": {"mean_absolute_rate_error": float(np.mean(growth_abs_errors)) if growth_abs_errors else None,
                               "predicted_rates": [r["predicted_real_lambda"] for r in growth_rows],
                               "observed_rates": [r["observed_growth_fit"] for r in growth_rows]},
        "escape_prediction": {"tau": tau_escape, "r": r, "predicted_slope_1_over_r": 1.0 / r,
                               "observed_slope": float(slope), "slope_relative_error": float(abs(slope - 1.0 / r) / (1.0 / r)),
                               "intercept": float(intercept), "points": escape_rows},
        "double_well": {"eta": eta, "calibrated_delay_steps": burst_m,
                         "current_gradient_sgd": summarize(baseline),
                         "fixed_delayed_gradient_sgd": summarize(fixed_stale)},
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()