import json
import time
import numpy as np
from scipy.integrate import solve_ivp

SEED = 2040
np.random.seed(SEED)

# Periodic contractive linear ODE:
#   h' = A(s) h + b, A(s) = -alpha I + q sin(2 pi s) K,
# with K=diag(1,-1).  Its phase average is Abar=-alpha I.
# Since K is symmetric, mu_2(A(s)) = -alpha + q sin(2 pi s)*lambda_max(K)
# and therefore mu_2 <= -(alpha-q), for q < alpha.
ALPHA = 1.2
K = np.diag([1.0, -1.0])
B = np.array([0.7, -0.35])
H0 = np.array([0.4, -0.8])
T_HORIZON = 4.0


def A_of_phase(s, q):
    return -ALPHA * np.eye(2) + q * np.sin(2.0 * np.pi * s) * K


def rhs_fast(t, h, eps, q):
    return A_of_phase(t / eps, q).dot(h) + B


def rhs_avg(t, h, q):
    # Exact phase average: average(sin)=0. q is retained only for API symmetry.
    return -ALPHA * h + B


def solve_fast(eps, q, h0=H0, horizon=T_HORIZON, dense=False):
    # Resolving each oscillation with at least 20 adaptive solver steps keeps
    # integration error well below the averaging error in the sweep.
    sol = solve_ivp(lambda t, h: rhs_fast(t, h, eps, q), (0, horizon), h0,
                    rtol=2e-10, atol=2e-12, max_step=eps / 25.0,
                    dense_output=dense)
    if not sol.success:
        raise RuntimeError(sol.message)
    return sol


def solve_avg(h0=H0, horizon=T_HORIZON, dense=False):
    sol = solve_ivp(lambda t, h: rhs_avg(t, h, 0.0), (0, horizon), h0,
                    rtol=2e-11, atol=2e-13, max_step=0.03, dense_output=dense)
    if not sol.success:
        raise RuntimeError(sol.message)
    return sol


def trajectory_error(eps, q, n=2001):
    ts = np.linspace(0, T_HORIZON, n)
    fast = solve_fast(eps, q, dense=True).sol(ts).T
    avg = solve_avg(dense=True).sol(ts).T
    err = np.linalg.norm(fast - avg, axis=1)
    return float(np.max(err)), float(np.sqrt(np.mean(err * err)))


def contraction_sweep(qs):
    rows = []
    for q in qs:
        # Direct grid verification of the matrix-measure certificate.
        phases = np.linspace(0, 1, 10001)
        mus = np.array([np.linalg.eigvalsh((A_of_phase(s, q) + A_of_phase(s, q).T) / 2).max()
                        for s in phases])
        predicted = -ALPHA + abs(q)
        rows.append({"q": float(q), "predicted_max_mu": float(predicted),
                     "observed_max_mu": float(mus.max()),
                     "certificate_holds": bool(mus.max() < 0)})
    return rows


def perturbation_decay(q, horizon=12.0, n=2401):
    # Same input, two initial states: their difference has an exact contraction
    # envelope. Measure empirical slope after transients, and test the pointwise
    # certificate log ratio <= -(alpha-q)t.
    d0 = np.array([1.0, -0.7])
    sol1 = solve_fast(0.08, q, h0=np.zeros(2), horizon=horizon, dense=True)
    sol2 = solve_fast(0.08, q, h0=d0, horizon=horizon, dense=True)
    ts = np.linspace(0, horizon, n)
    d = np.linalg.norm(sol2.sol(ts).T - sol1.sol(ts).T, axis=1)
    log_ratio = np.log(np.maximum(d, 1e-300) / np.linalg.norm(d0))
    # Fit the late-time slope, where periodic modulation averages out.
    mask = (ts >= 4.0) & (ts <= horizon)
    slope = float(np.polyfit(ts[mask], log_ratio[mask], 1)[0])
    cert_rate = ALPHA - abs(q)
    envelope_margin = float(np.max(log_ratio + cert_rate * ts))
    return {"q": float(q), "predicted_long_run_slope": -ALPHA,
            "observed_late_log_slope": slope,
            "predicted_certificate_rate": cert_rate,
            "max_log_envelope_margin": envelope_margin}


def rollout_speed(eps=0.0625, q=0.8, horizon=4.0, repeats=3):
    # Compare expensive phase-resolved solve with the averaged rollout. Both
    # solve the same initial-value problem and are evaluated at 2001 points.
    ts = np.linspace(0, horizon, 2001)
    fast_times, avg_times = [], []
    for _ in range(repeats):
        t0 = time.perf_counter()
        sf = solve_fast(eps, q, horizon=horizon, dense=True)
        _ = sf.sol(ts)
        fast_times.append(time.perf_counter() - t0)
        t0 = time.perf_counter()
        sa = solve_avg(horizon=horizon, dense=True)
        _ = sa.sol(ts)
        avg_times.append(time.perf_counter() - t0)
    return {"eps": eps, "fast_seconds_median": float(np.median(fast_times)),
            "averaged_seconds_median": float(np.median(avg_times)),
            "speedup": float(np.median(fast_times) / np.median(avg_times))}


def main():
    q = 0.8
    eps_list = [0.5, 0.25, 0.125, 0.0625, 0.03125]
    errors = []
    for eps in eps_list:
        mx, rms = trajectory_error(eps, q)
        errors.append({"eps": eps, "max_error": mx, "rms_error": rms})
    loge = np.log([r["max_error"] for r in errors])
    logeps = np.log(eps_list)
    scaling = float(np.polyfit(logeps, loge, 1)[0])
    result = {
        "seed": SEED,
        "system": {"alpha": ALPHA, "q": q, "K": K.tolist(), "B": B.tolist()},
        "predictions": {
            "averaging_error_power": "O(eps), predicted log-log slope 1",
            "contraction_certificate": "max_s mu_2(A(s)) = -alpha+q; boundary q=alpha",
            "perturbation_decay": "late log-distance slope equals -alpha, and is bounded by -(alpha-q)"
        },
        "averaging_sweep": {"rows": errors, "observed_loglog_slope": scaling},
        "contraction_boundary_sweep": contraction_sweep([0.0, 0.6, 1.19, 1.2, 1.4]),
        "perturbation_decay": [perturbation_decay(v) for v in [0.0, 0.4, 0.8, 1.1]],
        "rollout_speed": rollout_speed()
    }
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()