import json
import numpy as np


def simulate(lambdas, eta, kappa, x0, y0, steps):
    H = np.diag(np.asarray(lambdas, dtype=float))
    x, y = np.array(x0, dtype=float).copy(), np.array(y0, dtype=float).copy()
    xs, ys, losses = [x.copy()], [y.copy()], []
    for _ in range(steps):
        # Exact discrete Euler version of the proposed coupled optimizer.
        # Simultaneous Euler update: both right-hand sides use the old state.
        xo, yo = x.copy(), y.copy()
        gx, gy = H @ xo, H @ yo
        x = xo - eta * (gx + kappa * (xo - yo))
        y = yo - eta * (gy + kappa * (yo - xo))
        xs.append(x.copy()); ys.append(y.copy())
        losses.append(0.5 * (x @ H @ x + y @ H @ y))
    return np.asarray(xs), np.asarray(ys), np.asarray(losses)


def modal_rate(series, eta, start=0, end=None):
    z = np.abs(np.asarray(series)[start:end])
    z = z[z > 1e-14]
    if len(z) < 2:
        return float("nan")
    # Geometric per-step rate, converted to continuous-time equivalent.
    return -np.polyfit(np.arange(len(z)), np.log(z), 1)[0] / eta


def run():
    # Diagonal quadratic is a local Hessian model, with an intentionally separated
    # slow mode so the predicted transition is visible.
    lam = np.array([0.10, 0.70, 2.00])
    eta = 0.05
    x0 = np.array([1.0, 0.8, 0.4])
    results = {"lambdas": lam.tolist(), "eta": eta}

    # Prediction 1: prepared packet cancels c_1 exactly, while ordinary packet
    # retains c_1=x_1. Also test a tenfold cancellation perturbation.
    y_unprepared = x0.copy()
    y_prepared = x0.copy(); y_prepared[0] = -x0[0]
    y_perturbed = y_prepared.copy(); y_perturbed[0] += 0.10 * x0[0]
    rows = []
    for name, y0 in [("unprepared", y_unprepared), ("prepared", y_prepared),
                     ("10pct_residual", y_perturbed)]:
        xs, ys, _ = simulate(lam, eta, 0.35, x0, y0, 30)
        c1 = 0.5 * (xs[:, 0] + ys[:, 0])
        rows.append({"case": name, "initial_abs_c1_over_abs_a1": float(abs(c1[0]) / abs(x0[0])),
                     "abs_c1_after_30": float(abs(c1[-1]))})
    results["cancellation"] = rows

    # Prediction 2: after cancellation, rate is min(lambda_2, lambda_1+2*kappa).
    # Rates are measured from the exact modal trajectories and compared to theory.
    rate_rows = []
    for k in [0.0, 0.05, 0.15, 0.30, 0.60]:
        xs, ys, _ = simulate(lam, eta, k, x0, y_prepared, 180)
        c2 = 0.5 * (xs[:, 1] + ys[:, 1])
        d1 = 0.5 * (xs[:, 0] - ys[:, 0])
        # The slowest surviving term is whichever of c2 and d1 decays slower.
        r2 = -np.log(abs(1 - eta * lam[1])) / eta
        r1d = -np.log(abs(1 - eta * (lam[0] + 2*k))) / eta
        predicted = min(r2, r1d)
        def fit_rate(z):
            zz = np.abs(z[:80])
            idx = zz > 1e-12
            return float(-np.polyfit(np.arange(len(zz))[idx], np.log(zz[idx]), 1)[0] / eta)
        observed = min(fit_rate(c2), fit_rate(d1))
        rate_rows.append({"kappa": k, "predicted_rate": float(predicted),
                          "observed_modal_rate": float(observed),
                          "lambda1_plus_2k": float(lam[0] + 2*k),
                          "lambda2": float(lam[1])})
    results["rate_sweep"] = rate_rows

    # Prediction 3: Euler stability boundary eta*(lambda_max+2*kappa)=2.
    # Sweep around predicted kappa critical value.
    kcrit = (2.0 / eta - lam.max()) / 2.0
    stability_rows = []
    for k in [kcrit - 0.5, kcrit - 0.05, kcrit, kcrit + 0.05, kcrit + 0.5]:
        xs, ys, loss = simulate(lam, eta, k, x0, -x0, 200)
        rho = max(abs(1 - eta * lam.max()), abs(1 - eta * (lam.max() + 2*k)))
        stability_rows.append({"kappa": float(k), "predicted_spectral_radius": float(rho),
                               "predicted_stable": bool(rho <= 1 + 1e-12),
                               "observed_unstable_200_steps": bool((not np.all(np.isfinite(loss))) or loss[-1] > loss[0] * 1.01),
                               "loss_final": float(loss[-1])})
    results["stability_sweep"] = {"predicted_kappa_critical": float(kcrit), "rows": stability_rows}

    # Baseline comparison at equal coupled update steps: prepared packet has a
    # substantially lower late loss once the slow mode is removed.
    xb, yb, lb = simulate(lam, eta, 0.35, x0, y_unprepared, 100)
    xp, yp, lp = simulate(lam, eta, 0.35, x0, y_prepared, 100)
    results["baseline_vs_idea"] = {"unprepared_final_loss": float(lb[-1]),
                                    "prepared_final_loss": float(lp[-1]),
                                    "loss_ratio_idea_over_baseline": float(lp[-1] / lb[-1])}
    with open("results.json", "w") as f:
        json.dump(results, f, indent=2)
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    run()