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


def exact_amplitude(mu, eps, mu0, r0):
    """|z(mu)| for epsilon dz/dmu = (mu+i beta)z."""
    return r0 * math.exp((mu * mu - mu0 * mu0) / (2.0 * eps))


def predicted_exit(eps, mu0, r0, rmax):
    target = eps * math.log(rmax / r0)
    return math.sqrt(mu0 * mu0 + 2.0 * target)


def simulate_mode(eps, mu0=-0.5, mu_end=1.0, r0=1e-3, beta=2.0):
    """Directly propagate the linear complex Hopf mode on a ramp grid."""
    mu = mu0
    z = complex(r0, 0.0)
    rows = []
    while mu < mu_end - 1e-12:
        dmu = min(eps, mu_end - mu)
        mid = mu + 0.5 * dmu
        alpha = mid
        z *= np.exp((alpha + 1j * beta) * dmu / eps)
        mu += dmu
        rows.append((mu, abs(z), alpha, alpha * dmu))
    return np.asarray(rows)


def integrated_growth_scheduler(eps, mu0=-0.5, mu_end=1.0, r0=1e-3,
                                rmax=0.1, beta=2.0, noise=0.0, seed=0):
    """Permit the delayed passage and stop when signed B reaches epsilon log(rmax/r0).

    The monitor estimates alpha; the plant evolves with the true alpha. Noise is
    optional measurement noise on alpha, included to test robustness of the toy rule.
    """
    rng = np.random.default_rng(seed)
    target = eps * math.log(rmax / r0)
    mu, z, B = mu0, complex(r0, 0.0), 0.0
    first_cross = None
    max_r = r0
    while mu < mu_end - 1e-12:
        dmu = min(eps, mu_end - mu)
        mid = mu + 0.5 * dmu
        true_alpha = mid
        observed_alpha = true_alpha + (noise * rng.standard_normal() if noise else 0.0)
        if first_cross is None and observed_alpha >= 0:
            first_cross = mid
        # Signed accumulated growth is the quantity in the paper's formula.
        B_next = B + dmu * observed_alpha
        if B_next >= target:
            # interpolate only for reporting; this avoids an artificially large
            # one-step overshoot in the reported control parameter.
            frac = (target - B) / (B_next - B) if B_next != B else 1.0
            exit_mu = mu + frac * dmu
            return {"exit_mu": exit_mu, "first_cross": first_cross,
                    "budget": target, "target": target,
                    "predicted_amplitude": exact_amplitude(exit_mu, eps, mu0, r0),
                    "post_crossing_delay": exit_mu,
                    "steps": int(round((exit_mu - mu0) / eps))}
        B = B_next
        z *= np.exp((true_alpha + 1j * beta) * dmu / eps)
        mu += dmu
        max_r = max(max_r, abs(z))
    return {"exit_mu": mu_end, "first_cross": first_cross, "budget": B,
            "target": target, "predicted_amplitude": abs(z),
            "post_crossing_delay": mu_end, "steps": int(round((mu_end-mu0)/eps))}


def instantaneous_clipping(eps, mu0=-0.5):
    """Baseline controller: stop as soon as Re(lambda)=alpha crosses zero."""
    return {"exit_mu": 0.0, "post_crossing_delay": 0.0,
            "amplitude_at_exit": exact_amplitude(0.0, eps, mu0, 1e-3)}


def main():
    mu0, r0, rmax, beta = -0.5, 1e-3, 0.1, 2.0
    checks, comparisons = [], []
    for eps in (0.02, 0.01, 0.005):
        trajectory = simulate_mode(eps, mu0, 1.0, r0, beta)
        mus, radii = trajectory[:, 0], trajectory[:, 1]
        formula = np.array([exact_amplitude(m, eps, mu0, r0) for m in mus])
        rel_error = float(np.max(np.abs(radii - formula) / formula))
        crossings = np.flatnonzero(radii >= rmax)
        observed = float(mus[crossings[0]]) if len(crossings) else None
        pred = predicted_exit(eps, mu0, r0, rmax)
        idea = integrated_growth_scheduler(eps, mu0, 1.0, r0, rmax, beta)
        base = instantaneous_clipping(eps, mu0)
        checks.append({"eps": eps, "max_relative_amplitude_error": rel_error,
                       "predicted_exit_mu": pred, "grid_observed_exit_mu": observed,
                       "idea_exit_mu": idea["exit_mu"],
                       "idea_exit_amplitude": idea["predicted_amplitude"],
                       "budget_error": idea["budget"] - idea["target"]})
        comparisons.append({"eps": eps, "baseline": base, "idea": idea,
                            "extra_stable_ramp": idea["exit_mu"] - base["exit_mu"]})

    # Noisy monitor repeat: the plant remains bounded at the requested threshold
    # while the estimated exit varies according to monitor noise.
    noisy = [integrated_growth_scheduler(0.01, mu0, 1.0, r0, rmax, beta,
                                         noise=0.01, seed=s)
             for s in range(10)]
    out = {"parameters": {"mu0": mu0, "r0": r0, "rmax": rmax, "beta": beta},
           "math_checks": checks, "baseline_vs_idea": comparisons,
           "noisy_monitor_exits": noisy}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()