import json
import math
from pathlib import Path
import numpy as np

# Scalar center-manifold normal form:
# z_{k+1}=-(1+mu) z_k + a z_k^3.
# The fixed point has J=-(1+mu), and a symmetric 2-cycle satisfies
# f(z)=-z, giving |z|=sqrt(mu/a) for mu>0.

def iterate(mu, a=1.0, z0=0.03, steps=4000):
    z = float(z0)
    hist = []
    for _ in range(steps):
        z = -(1.0 + mu) * z + a * z**3
        hist.append(z)
        if not np.isfinite(z) or abs(z) > 1e6:
            return np.asarray(hist), False
    return np.asarray(hist), True

def cycle_amplitude(hist, tail=500):
    if len(hist) < tail or not np.all(np.isfinite(hist)):
        return float("nan")
    return float(np.mean(np.abs(hist[-tail:])))

def a_k(hist, tail=500):
    if len(hist) < tail + 2 or not np.all(np.isfinite(hist)):
        return float("nan")
    return float(np.mean(np.abs(hist[-tail:] - hist[-tail-2:-2])) /
                 (np.mean(np.abs(hist[-tail:])) + 1e-12))

def jacobian_eigenvalue(mu):
    return -(1.0 + mu)

def corrected_guard_mu(mu, delta=0.05):
    # lambda >= -1+delta; lambda=-1-mu, so mu <= -delta.
    return min(mu, -delta)

def run():
    a = 2.0
    delta = 0.05

    # Prediction 1: the fixed point loses stability at mu=0 (lambda=-1).
    mus_boundary = np.linspace(-0.20, 0.20, 81)
    stable = []
    for mu in mus_boundary:
        h, ok = iterate(mu, a=a, z0=1e-4, steps=1500)
        stable.append(ok and abs(h[-1]) < 1e-3 and abs(h[-1] - h[-3]) < 1e-3)
    stable_mus = mus_boundary[np.asarray(stable)]
    unstable_mus = mus_boundary[~np.asarray(stable)]
    observed_boundary = float((stable_mus[-1] + unstable_mus[0]) / 2)

    # Prediction 2: period-two amplitude scales as sqrt(mu/a).
    mu_amp = np.linspace(0.01, 0.18, 12)
    observed_amp, predicted_amp = [], []
    for mu in mu_amp:
        h, _ = iterate(mu, a=a, z0=0.01, steps=5000)
        observed_amp.append(cycle_amplitude(h))
        predicted_amp.append(math.sqrt(mu / a))
    observed_amp, predicted_amp = np.asarray(observed_amp), np.asarray(predicted_amp)
    rel_amp_error = float(np.mean(np.abs(observed_amp - predicted_amp) / predicted_amp))
    coeff = float(np.dot(np.sqrt(mu_amp), observed_amp) /
                  np.dot(np.sqrt(mu_amp), np.sqrt(mu_amp)))

    # Prediction 3: amplitude changes as 1/sqrt(a).
    mu_scale = 0.08
    a_values = np.asarray([0.5, 1.0, 2.0, 4.0])
    scale_obs, scale_pred = [], []
    for av in a_values:
        h, _ = iterate(mu_scale, a=av, z0=0.01, steps=5000)
        scale_obs.append(cycle_amplitude(h))
        scale_pred.append(math.sqrt(mu_scale / av))
    scale_obs, scale_pred = np.asarray(scale_obs), np.asarray(scale_pred)
    scale_error = float(np.mean(np.abs(scale_obs-scale_pred)/scale_pred))

    # Small same-task comparison: train scalar recurrent gain toward an
    # unstable target mu=+0.15. Baseline follows task loss; corrected guard
    # projects onto lambda >= -1+delta after each update.
    target, beta, lr = 0.15, 5.0, 0.08
    mu_base, mu_guard, mu_literal = -0.10, -0.10, -0.10
    for _ in range(100):
        mu_base -= lr * 2.0 * (mu_base-target)
        corrected_grad = 2.0 * beta * max(0.0, mu_guard + delta)
        mu_guard -= lr * (2.0*(mu_guard-target) + corrected_grad)
        mu_guard = corrected_guard_mu(mu_guard, delta)
        # Literal formula from the prompt: [max(0, lambda+1-delta)]^2.
        # Since lambda=-1-mu, this is max(0,-mu-delta)^2 and cannot stop
        # positive mu (the flip-unstable region).
        literal_grad = -2.0 * beta * max(0.0, -mu_literal-delta)
        mu_literal -= lr * (2.0*(mu_literal-target) + literal_grad)

    hb, _ = iterate(mu_base, a=a, z0=0.01, steps=3000)
    hg, _ = iterate(corrected_guard_mu(mu_guard, delta), a=a, z0=0.01, steps=3000)
    hl, _ = iterate(mu_literal, a=a, z0=0.01, steps=3000)
    test_mu = 0.10
    lam = jacobian_eigenvalue(test_mu)

    out = {
        "normal_form": {"a": a, "delta": delta},
        "prediction_1_boundary": {
            "predicted_mu": 0.0, "observed_mu_midpoint": observed_boundary,
            "last_stable_grid_mu": float(stable_mus[-1]),
            "first_unstable_grid_mu": float(unstable_mus[0]),
            "lambda_at_observed_midpoint": jacobian_eigenvalue(observed_boundary)},
        "prediction_2_sqrt_scaling": {
            "mu": mu_amp.tolist(), "observed_amplitude": observed_amp.tolist(),
            "predicted_amplitude": predicted_amp.tolist(),
            "mean_relative_error": rel_amp_error,
            "observed_sqrt_coefficient": coeff,
            "predicted_sqrt_coefficient": 1.0/math.sqrt(a)},
        "prediction_3_cubic_scaling": {
            "a": a_values.tolist(), "observed_amplitude": scale_obs.tolist(),
            "predicted_amplitude": scale_pred.tolist(), "mean_relative_error": scale_error},
        "guard_training": {
            "target_mu": target, "baseline_mu": float(mu_base),
            "corrected_guard_mu": float(mu_guard), "literal_formula_mu": float(mu_literal),
            "baseline_lambda": jacobian_eigenvalue(mu_base),
            "corrected_guard_lambda": jacobian_eigenvalue(mu_guard),
            "literal_formula_lambda": jacobian_eigenvalue(mu_literal),
            "baseline_A_K": a_k(hb), "corrected_guard_A_K": a_k(hg),
            "literal_formula_A_K": a_k(hl),
            "baseline_cycle_amplitude": cycle_amplitude(hb),
            "corrected_guard_cycle_amplitude": cycle_amplitude(hg),
            "literal_formula_cycle_amplitude": cycle_amplitude(hl)},
        "formula_sign_check": {
            "mu": test_mu, "lambda": lam,
            "displayed_penalty": max(0.0, lam + 1.0 - delta)**2,
            "corrected_lower_bound_penalty": max(0.0, (-1.0 + delta) - lam)**2,
            "note": "The displayed formula penalizes lambda above -1+delta, rather than lambda below it; corrected results are reported separately."}}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    run()