import json
import numpy as np


def interpolation_residual(xi, gi, fi, xj, gj, fj, mu, L):
    """H_mu,L(i,j), in the paper's s.q + tr(GQ) form, evaluated from samples."""
    dx = xi - xj
    dg = gi - gj
    if L <= mu:
        raise ValueError("need L > mu")
    trace_q = (-np.dot(gj, dx)
               - np.dot(dg - mu * dx, dg - mu * dx) / (2.0 * (L - mu))
               - mu * np.dot(dx, dx) / 2.0)
    return (fi - fj) + trace_q


def estimate_curvature(xs, gs, eps=1e-12):
    vals_L, vals_mu = [], []
    for x0, x1, g0, g1 in zip(xs[:-1], xs[1:], gs[:-1], gs[1:]):
        dx, dg = x1-x0, g1-g0
        n = max(np.dot(dx, dx), eps)
        vals_L.append(np.linalg.norm(dg) / np.sqrt(n))
        vals_mu.append(max(0.0, np.dot(dg, dx) / n))
    return max(vals_L), min(vals_mu)


def rho_scalar(eta, beta, lam):
    # z_{k+1}=(1+beta-eta*lam)z_k-beta*z_{k-1}
    roots = np.roots([1.0, -(1.0 + beta - eta*lam), beta])
    return float(np.max(np.abs(roots)))


def search_certificate(mu, L, beta_grid=None, eta_grid=None):
    """Small robust Lyapunov proxy: minimize worst endpoint spectral radius.
    The endpoint check is exact for scalar quadratics and is conservative for an
    interval. The returned support is the active curvature endpoints."""
    mu = max(float(mu), 1e-5)
    L = max(float(L), mu * 1.001)
    if beta_grid is None: beta_grid = np.linspace(0, .95, 40)
    if eta_grid is None: eta_grid = np.linspace(.001, 1.99 / L, 160)
    best = None
    for beta in beta_grid:
        for eta in eta_grid:
            r = max(rho_scalar(eta, beta, mu), rho_scalar(eta, beta, L))
            if best is None or r < best[0]:
                best = (r, eta, beta)
    return {"rho": best[0], "eta": best[1], "beta": best[2],
            "support": ["mu_endpoint", "L_endpoint"]}


def verify_math(seed=0):
    rng = np.random.default_rng(seed)
    mu, L = .7, 4.0
    # Random diagonal quadratics have exact [mu,L] smooth/strong convex bounds.
    A = np.diag([mu, 1.5, L])
    b = rng.normal(size=3)
    f = lambda x: .5*x@A@x + b@x
    g = lambda x: A@x+b
    xs = [rng.normal(size=3) for _ in range(10)]
    gs = [g(x) for x in xs]
    fs = [f(x) for x in xs]
    residuals = [interpolation_residual(xs[i],gs[i],fs[i],xs[j],gs[j],fs[j],mu,L)
                 for i in range(len(xs)) for j in range(len(xs)) if i != j]
    # The same expression must be nonnegative; deliberately underspecifying L breaks it.
    bad = [interpolation_residual(xs[0],gs[0],fs[0],xs[1],gs[1],fs[1],mu,1.0)
           if 1.0 > mu else 0.0]
    return {"min_valid_residual": float(min(residuals)),
            "max_abs_pair_asymmetry": float(max(abs(residuals[k] - residuals[k]) for k in range(len(residuals)))),
            "underspecified_L_residual": float(bad[0])}


def run_experiment(seed=7, steps=180):
    rng = np.random.default_rng(seed)
    # Two parameter blocks: a global step is dictated by the sharp block.
    curvatures = [np.array([1., 3., 10.]), np.array([.08, .2, .5])]
    xs0 = [rng.normal(size=len(c)) for c in curvatures]
    baseline_x = [x.copy() for x in xs0]
    idea_x = [x.copy() for x in xs0]
    base_eta = 1.0 / max(c.max() for c in curvatures)
    base_beta = .9
    certs = []
    # Warmup secants are obtained at a small stable step, as in the proposal.
    for c, x in zip(curvatures, xs0):
        histx, histg = [x.copy()], [c*x]
        for _ in range(8):
            x = x - .05 * (c*x) + rng.normal(0, .002, size=x.shape)
            histx.append(x.copy()); histg.append(c*x)
        Lhat, mhat = estimate_curvature(histx, histg)
        certs.append(search_certificate(mhat, Lhat))
    base_losses, idea_losses = [], []
    def loss(blocks): return float(sum(.5*np.sum(c*x*x) for c,x in zip(curvatures,blocks)))
    spikes_b = spikes_i = 0
    for t in range(steps):
        oldb = loss(baseline_x); oldi = loss(idea_x)
        # standard global momentum SGD
        if t == 0: prevb = [x.copy() for x in baseline_x]
        newb=[]
        for c,x,p in zip(curvatures,baseline_x,prevb):
            newb.append(x - base_eta*(c*x) + base_beta*(x-p))
        prevb, baseline_x = baseline_x, newb
        # certified per-block momentum update
        if t == 0: previ = [x.copy() for x in idea_x]
        newi=[]
        for c,x,p,cert in zip(curvatures,idea_x,previ,certs):
            newi.append(x - cert['eta']*(c*x) + cert['beta']*(x-p))
        previ, idea_x = idea_x, newi
        lb, li = loss(baseline_x), loss(idea_x)
        spikes_b += int(lb > oldb * 1.05); spikes_i += int(li > oldi * 1.05)
        base_losses.append(lb); idea_losses.append(li)
    return {"baseline": {"initial_loss": loss(xs0), "final_loss": base_losses[-1],
                          "loss_20pct": base_losses[int(steps*.2)], "eta": base_eta,
                          "beta": base_beta, "spikes": spikes_b},
            "idea": {"initial_loss": loss(xs0), "final_loss": idea_losses[-1],
                     "loss_20pct": idea_losses[int(steps*.2)], "spikes": spikes_i,
                     "certificates": certs},
            "curvatures": [c.tolist() for c in curvatures]}


if __name__ == '__main__':
    out = {"math": verify_math(), "experiment": run_experiment()}
    print(json.dumps(out, indent=2))