import json, math
import numpy as np

# OU Langevin certificate and a certificate-controlled noise schedule.
# dX=-m X dt+sqrt(2T)dW, with stationary pi=N(0,T/m).
SEED = 551
m, T, eta = 1.0, 0.5, 0.02
x0_mean = 3.0
x0_var = T / m                 # shifted stationary Gaussian initial law
H = 1.5                         # unsafe region A={x>H}
steps = 500
n_particles = 20000
delta = 0.10


def normal_tail(z):
    return 0.5 * math.erfc(z / math.sqrt(2.0))


def stationary_risk(h=H):
    return normal_tail(h / math.sqrt(T / m))


def exact_risk(t, h=H):
    mu = x0_mean * math.exp(-m * t)
    var = (T / m) + (x0_var - T / m) * math.exp(-2 * m * t)
    return normal_tail((h - mu) / math.sqrt(var))


def exact_chi2():
    # chi^2(N(mu,s2)||N(0,s2)) = exp(mu^2/s2)-1.
    return math.exp(x0_mean ** 2 / x0_var) - 1.0


def certificate(t):
    p = stationary_risk()
    return p + math.sqrt(p * exact_chi2()) * math.exp(-m * t)


def sufficient_burn_in():
    p = stationary_risk()
    numerator = math.sqrt(p * exact_chi2())
    if delta <= p:
        return math.inf
    return max(0.0, math.log(numerator / (delta - p)) / m)


def verify_envelope():
    ts = np.linspace(0, 12, 241)
    risks = np.array([exact_risk(float(t)) for t in ts])
    bounds = np.array([certificate(float(t)) for t in ts])
    return {
        "max_exact_minus_bound": float(np.max(risks - bounds)),
        "min_bound_minus_exact": float(np.min(bounds - risks)),
        "exact_risk_t0": exact_risk(0.0),
        "exact_risk_t1": exact_risk(1.0),
        "exact_risk_t4": exact_risk(4.0),
        "stationary_pi_A": stationary_risk(),
        "chi0_squared": exact_chi2(),
        "bound_t0": certificate(0.0),
        "bound_t1": certificate(1.0),
        "bound_t4": certificate(4.0),
        "analytic_burn_in_for_delta": sufficient_burn_in(),
    }


def run_schedule(certificate_stop=False):
    # Same initial law and same random seed for both schedules. The controlled
    # schedule consumes no further noise after the certificate crossing.
    r = np.random.default_rng(SEED)
    x = r.normal(x0_mean, math.sqrt(x0_var), size=n_particles)
    unsafe_fractions, losses = [], []
    stop_step = None
    for k in range(steps):
        t = k * eta
        if certificate_stop and stop_step is None and certificate(t) <= delta:
            stop_step = k
        noise = 0.0 if (certificate_stop and stop_step is not None) else math.sqrt(2 * T * eta) * r.normal(size=n_particles)
        x = x - eta * m * x + noise
        unsafe_fractions.append(float(np.mean(x > H)))
        losses.append(float(np.mean(0.5 * m * x * x)))
    post = unsafe_fractions[stop_step:] if stop_step is not None else unsafe_fractions
    return {
        "max_unsafe_fraction": float(np.max(unsafe_fractions)),
        "mean_unsafe_fraction": float(np.mean(unsafe_fractions)),
        "final_unsafe_fraction": unsafe_fractions[-1],
        "mean_post_stop_unsafe_fraction": float(np.mean(post)),
        "final_loss": losses[-1],
        "mean_last_50_loss": float(np.mean(losses[-50:])),
        "stop_step": stop_step,
        "stop_time": None if stop_step is None else stop_step * eta,
    }


def main():
    out = {
        "setup": {"m": m, "T": T, "eta": eta, "H": H, "steps": steps,
                  "particles": n_particles, "delta": delta, "seed": SEED},
        "verification": verify_envelope(),
        "baseline": run_schedule(False),
        "idea": run_schedule(True),
    }
    with open("results.json", "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()