import json
import math
import numpy as np

SEED = 2771
UC = 2.0 / (3.0 * math.sqrt(3.0))

def F(x, u):
    return x - x**3 + u

def jac(x):
    return 1.0 - 3.0*x*x

def equilibria(u):
    roots = np.roots([1.0, 0.0, -1.0, -u])
    return sorted(float(r.real) for r in roots if abs(r.imag) < 1e-8)

def integrate_many(x0, u, T=35.0, damping=1.0):
    """Vectorized fixed-step RK4 integration of the scalar flow."""
    x = np.asarray(x0, dtype=float).copy()
    n = max(1, int(math.ceil(T / 0.04)))
    dt = T / n
    for _ in range(n):
        k1 = damping * F(x, u)
        k2 = damping * F(x + 0.5*dt*k1, u)
        k3 = damping * F(x + 0.5*dt*k2, u)
        k4 = damping * F(x + dt*k3, u)
        x += dt*(k1 + 2*k2 + 2*k3 + k4)/6.0
    return x

def integrate(x0, u, T=35.0, damping=1.0):
    return float(integrate_many(np.array([x0]), u, T, damping)[0])

def desired_eq(u):
    return max(equilibria(u))

def estimate_saddle_node():
    us = np.linspace(0.0, 0.45, 9001)
    counts = np.array([len(equilibria(u)) for u in us])
    return float(us[np.where(counts >= 3)[0][-1]])

def main():
    numeric_uc = estimate_saddle_node()
    u_vals = [0.0, 0.10, 0.20, 0.35, 0.38]
    stability = []
    for u in u_vals:
        roots = equilibria(u)
        xp = max(roots)
        stability.append({"u": u, "roots": roots, "desired_jacobian": jac(xp),
                          "locally_stable": jac(xp) < 0, "coexisting": len(roots) == 3})

    # At u=0, x_ref=1 and separatrix=0. Thus recovery is exactly the event
    # 1 + eps*z > 0 (up to numerical integration tolerance).
    eps_values = [0.10, 0.30, 0.50, 0.70, 0.90, 1.10]
    basin = []
    for eps in eps_values:
        rng = np.random.default_rng(SEED + int(eps * 1000))
        z = rng.normal(size=4000)
        xref = desired_eq(0.0)
        final = integrate_many(xref + eps*z, 0.0)
        basin.append({"epsilon": eps,
                      "B_hat": float(np.mean(np.abs(final-xref) < .08)),
                      "B_predicted_gaussian": float(0.5*(1+math.erf(1/(eps*math.sqrt(2))))),
                      "separatrix": 0.0, "reference": xref})

    # Finite-horizon comparison. The intervention is deliberately tested after
    # local stability is restored. The guard estimates basin recovery with probes
    # and retains damping when recovery is not demonstrated.
    def trial(mode, n=300):
        rng = np.random.default_rng(SEED + (1 if mode == 'guard' else 2))
        xref = desired_eq(0.0)
        failures = releases = 0
        for _ in range(n):
            x = xref + 0.95 * rng.normal()
            if mode == 'threshold':
                c = 1.0  # local Jacobian is safely negative, so threshold-only releases
                releases += 1
            else:
                probes = xref + 0.95*rng.normal(size=80)
                probe_final = integrate_many(probes, 0.0, T=6.0, damping=0.35)
                B = np.mean(np.abs(probe_final-xref) < .15)
                c = 1.0 if jac(xref) <= -0.5 and B >= .95 else 0.35
                releases += int(c == 1.0)
            y = integrate(x, 0.0, T=6.0, damping=c)
            failures += int(abs(y-xref) >= .15)
        return {"failure_rate": failures/n, "release_rate": releases/n}

    result = {
        "seed": SEED,
        "model": "dx/dt = x - x^3 + u (desired positive equilibrium)",
        "predictions": {
            "saddle_node": {"analytic_uc": UC, "numeric_last_three_root_u": numeric_uc,
                            "absolute_error": abs(UC-numeric_uc)},
            "local_vs_global": "desired Jacobian remains negative throughout the three-root bistable interval; local stability does not imply global recovery",
            "basin_transition": "at u=0, separatrix is x=0 and desired equilibrium is x=1; recovery probability is Phi(1/epsilon)"
        },
        "stability_sweep": stability,
        "basin_sweep": basin,
        "finite_horizon_comparison": {"threshold_only": trial("threshold"), "basin_guard": trial("guard")}
    }
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()