import json
import numpy as np
from scipy.integrate import quad
from nonlinear_drift import scalar_pairwise_rate

LAM, D = 1.0, 0.5


def exact_variance(beta):
    def unnorm(x):
        return np.exp(-(LAM*x*x/2.0 + beta*x**4/4.0)/D)
    z = quad(unnorm, -np.inf, np.inf, epsabs=1e-12, epsrel=1e-12)[0]
    m2 = quad(lambda x: x*x*unnorm(x), -np.inf, np.inf,
              epsabs=1e-12, epsrel=1e-12)[0] / z
    return m2


def simulate(beta, seed=123, n=4000, burn=3000, sample_steps=7000, dt=.005):
    rng = np.random.default_rng(seed)
    x = np.zeros(n)
    for _ in range(burn):
        x += (-LAM*x - beta*x**3)*dt + np.sqrt(2*D*dt)*rng.normal(size=n)
    vals = []
    for k in range(sample_steps):
        x += (-LAM*x - beta*x**3)*dt + np.sqrt(2*D*dt)*rng.normal(size=n)
        if k % 10 == 0:
            vals.append(np.mean(x*x))
    return float(np.mean(vals)), float(np.std(vals) / np.sqrt(len(vals)))


def main():
    wide_betas = np.array([0., .02, .05, .1, .2, .5, 1.0])
    exact_wide = np.array([exact_variance(b) for b in wide_betas])
    baseline = D / LAM
    # First-order perturbation: Var/baseline = 1 - 3 beta D/lambda^2 + O(beta^2).
    tiny_betas = np.array([0., .0001, .0005, .001, .002, .005])
    exact_tiny = np.array([exact_variance(b) for b in tiny_betas])
    relative = exact_tiny / baseline - 1.0
    slope = float(np.polyfit(tiny_betas, relative, 1)[0])
    predicted_slope = -3 * D / LAM**2

    # Prediction 1: pairwise rate is exactly >= lambda, equality only at x=y=0.
    rng = np.random.default_rng(7)
    x, y = rng.normal(size=100000), rng.normal(size=100000)
    rates = scalar_pairwise_rate(x, y, LAM, .5)
    contraction = {
        "predicted_min_rate": LAM,
        "observed_min_rate": float(rates.min()),
        "observed_fraction_below_lambda": float(np.mean(rates < LAM - 1e-12)),
        "rate_at_equal_radius_1": float(scalar_pairwise_rate(1., 1., LAM, .5)),
        "predicted_rate_at_equal_radius_1": LAM + 3*.5
    }
    sims = {str(b): simulate(float(b), seed=100 + i) for i, b in enumerate([0., .5])}
    out = {
        "parameters": {"lambda": LAM, "D": D, "baseline_variance": baseline},
        "prediction_1_contraction": contraction,
        "prediction_2_small_beta_slope": {
            "prediction": "relative variance slope = -3D/lambda^2 + O(beta)",
            "predicted_relative_slope": predicted_slope,
            "observed_relative_slope": slope,
            "relative_error": abs(slope - predicted_slope) / abs(predicted_slope),
            "betas": tiny_betas.tolist(),
            "exact_variances": exact_tiny.tolist()
        },
        "prediction_3_monotonic_variance": {
            "prediction": "variance strictly decreases as beta increases",
            "betas": wide_betas.tolist(), "exact_variances": exact_wide.tolist(),
            "strictly_decreasing": bool(np.all(np.diff(exact_wide) < 0)),
            "relative_gap_at_beta_0.5": float(1 - exact_wide[5] / baseline)
        },
        "matched_euler_maruyama": {
            "linear_beta_0": sims["0.0"], "cubic_beta_0.5": sims["0.5"],
            "simulated_relative_gap": float(1 - sims["0.5"][0] / sims["0.0"][0]),
            "dt": .005, "burn_steps": 3000, "sample_steps": 7000, "n_parallel": 4000
        }
    }
    with open("results.json", "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()