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

SEED = 1729
np.random.seed(SEED)
random.seed(SEED)

# Contact splitting for f(x)=lambda*x^2/2, M=1.
def split_step(x, p, s, h, lam, gamma):
    # K(h/2)
    s += 0.5*h*0.5*p*p
    x += 0.5*h*p
    # V(h/2)
    f = 0.5*lam*x*x
    p -= 0.5*h*lam*x
    s -= 0.5*h*f
    # D(h)
    r = math.exp(-gamma*h)
    p *= r
    s *= r
    # V(h/2)
    f = 0.5*lam*x*x
    p -= 0.5*h*lam*x
    s -= 0.5*h*f
    # K(h/2)
    s += 0.5*h*0.5*p*p
    x += 0.5*h*p
    return x, p, s

def H(x, p, s, lam, gamma):
    return 0.5*p*p + 0.5*lam*x*x + gamma*s

def run_split(h, lam=1.0, gamma=0.2, n=2000, x=1.0, p=0.0, s=0.0):
    hs, xs = [], []
    for _ in range(n + 1):
        hs.append(H(x, p, s, lam, gamma))
        xs.append(x)
        x, p, s = split_step(x, p, s, h, lam, gamma)
    return np.asarray(hs), np.asarray(xs)

def map_matrix(h, lam, gamma):
    r = math.exp(-gamma*h)
    b = h/2
    a = h*lam/2
    def mp(x, p):
        x1 = x + b*p
        p1 = p - a*x1
        p1 = r*p1
        p1 = p1 - a*x1
        return np.array([x1 + b*p1, p1])
    return np.column_stack([mp(1, 0), mp(0, 1)])

def slope_log_abs(y, h, start=100):
    z = np.abs(y[start:])
    ok = z > 1e-12
    if ok.sum() < 20:
        return float('nan')
    t = np.arange(len(z))[ok] * h
    return float(np.polyfit(t, np.log(z[ok]), 1)[0])

def split_opt_step(x, p, s, h, lam, gamma):
    # Same contact splitting, interpreted as an optimizer update.
    return split_step(x, p, s, h, lam, gamma)

def momentum_step(x, v, lr, beta, lam):
    v = beta*v - lr*lam*x
    return x + v, v

def main():
    lam, gamma = 1.0, 0.2
    results = {"seed": SEED, "predictions": {}, "comparison": {}}

    # Prediction 1: exact contact trajectory has log-energy slope -gamma.
    hs, _ = run_split(1e-4, lam, gamma, n=1000)
    slope = slope_log_abs(hs, 1e-4, start=100)
    results["predictions"]["exact_contact_rate"] = {
        "predicted": -gamma, "observed": slope,
        "abs_error": abs(slope + gamma)
    }

    # Prediction 2: symmetric splitting global error is O(h^2): halving h
    # should reduce fixed-time state error by about 4.
    T = 8.0
    ref_h = 1e-4
    xr = run_split(ref_h, lam, gamma, n=round(T/ref_h))[1][-1]
    errs = []
    for h in [0.08, 0.04, 0.02, 0.01]:
        xh = run_split(h, lam, gamma, n=round(T/h))[1][-1]
        errs.append((h, abs(xh-xr)))
    ratios = [errs[i][1]/errs[i+1][1] for i in range(len(errs)-1)]
    results["predictions"]["second_order_state_error"] = {
        "predicted_halving_ratio": 4.0,
        "observed_errors": errs,
        "observed_halving_ratios": ratios,
        "median_ratio": float(np.median(ratios))
    }

    # Prediction 3: stability boundary is determined by spectral radius=1,
    # and for gamma=0 the undamped oscillator has the Verlet boundary h*sqrt(lam)=2.
    boundary_rows = []
    for g in [0.0, 0.2, 1.0]:
        grid = np.linspace(0.01, 4.0, 4000)
        rho = np.array([max(abs(np.linalg.eigvals(map_matrix(h, lam, g)))) for h in grid])
        stable = np.where(rho <= 1.000001)[0]
        observed = float(grid[stable[-1]]) if len(stable) else 0.0
        boundary_rows.append({"gamma": g, "predicted_gamma0_boundary": 2.0/math.sqrt(lam), "observed_boundary": observed})
    results["predictions"]["stability_boundary"] = boundary_rows

    # Prediction 2b: the symmetric map's local certificate-rate defect is O(h^2).
    # Ignore the final transient where H can approach floating-point zero.
    residual_rows = []
    for h in [0.08, 0.04, 0.02, 0.01]:
        eh, _ = run_split(h, lam, gamma, n=round(4.0/h), x=1.0, p=0.0, s=0.3)
        valid = (np.abs(eh[:-1]) > 1e-8) & (np.abs(eh[1:]) > 1e-8)
        q = np.log(np.abs(eh[1:][valid])/np.abs(eh[:-1][valid]))/h + gamma
        residual_rows.append((h, float(np.median(np.abs(q)))))
    residual_ratios = [residual_rows[i][1]/residual_rows[i+1][1] for i in range(3)]
    results["predictions"]["certificate_rate_residual"] = {
        "predicted_halving_ratio": 4.0,
        "observed_median_abs_residuals": residual_rows,
        "observed_halving_ratios": residual_ratios,
        "median_ratio": float(np.median(residual_ratios))
    }

    # Small optimization comparison: quadratic with curvature spectrum, equal
    # number of gradient evaluations (split uses two; momentum uses two updates).
    curv = np.array([0.2, 1.0, 4.0, 12.0])
    x0 = np.ones_like(curv)
    def f(x): return float(0.5*np.sum(curv*x*x))
    h = 0.28; gamma2 = 0.2; beta = math.exp(-gamma2*h)
    x, p, s = x0.copy(), np.zeros_like(x0), 0.0
    xb, v = x0.copy(), np.zeros_like(x0)
    idea_losses, base_losses = [], []
    for k in range(120):
        idea_losses.append(f(x)); base_losses.append(f(xb))
        # vectorized separable contact split
        s += 0.5*h*0.5*np.sum(p*p)
        x += 0.5*h*p
        p -= 0.5*h*curv*x; s -= 0.5*h*f(x)
        p *= math.exp(-gamma2*h); s *= math.exp(-gamma2*h)
        p -= 0.5*h*curv*x; s -= 0.5*h*f(x)
        s += 0.5*h*0.5*np.sum(p*p); x += 0.5*h*p
        # two standard momentum updates for equal gradient evaluations
        for _ in range(2): xb, v = momentum_step(xb, v, h/2, beta, curv)
    results["comparison"] = {
        "h": h, "steps": 120, "idea_final_loss": idea_losses[-1],
        "baseline_final_loss": base_losses[-1],
        "idea_min_loss": min(idea_losses), "baseline_min_loss": min(base_losses),
        "idea_losses_first_last": [idea_losses[0], idea_losses[-1]],
        "baseline_losses_first_last": [base_losses[0], base_losses[-1]]
    }
    Path("results.json").write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))

if __name__ == "__main__":
    main()