import json
import math
import numpy as np

# Continuation Maps toy MVP.
# Quadratic objective: L(w)=1/2 w^T A w, with A=diag(1,4).
# Decoupled weight decay changes the update multiplier to
# q_i = 1 - eta*(eigenvalue_i + lambda).  Thus stability requires
# max_i |q_i| < 1, and the first boundary is eta=2/(4+lambda).

SEED = 2208
A = np.array([1.0, 4.0])
W0 = np.array([1.0, 1.0])
STEPS = 160
WINDOW = 30

def run(eta, lam, steps=STEPS):
    w = W0.copy()
    norms = [np.linalg.norm(w)]
    losses = [0.5 * np.sum(A * w * w)]
    for _ in range(steps):
        g = A * w
        w = w - eta * (g + lam * w)
        norms.append(np.linalg.norm(w))
        losses.append(0.5 * np.sum(A * w * w))
    return np.asarray(losses), np.asarray(norms)

def feature(eta, lam):
    # Late-time normalized temporal growth, analogous to a drift/oscillation
    # feature. For this linear system F=0 exactly at unit spectral radius.
    _, n = run(eta, lam)
    r = n[-WINDOW:] / np.maximum(n[-WINDOW-1:-1], 1e-300)
    return float(np.mean(np.log(np.maximum(r, 1e-300))))

def exact_boundary(lam):
    return 2.0 / (4.0 + lam)

def bisect_eta(lam, lo, hi, tol=2e-7, maxit=60):
    flo, fhi = feature(lo, lam), feature(hi, lam)
    # Expand only if a caller supplied a non-bracketing interval.
    for _ in range(20):
        if flo <= 0 <= fhi: break
        if flo > 0:
            lo *= 0.5; flo = feature(lo, lam)
        if fhi < 0:
            hi *= 1.25; fhi = feature(hi, lam)
    for _ in range(maxit):
        mid = (lo + hi) / 2
        fm = feature(mid, lam)
        if abs(fm) < 1e-9 or hi - lo < tol: return mid
        if fm > 0: hi = mid
        else: lo = mid
    return (lo + hi) / 2

def main():
    np.random.seed(SEED)
    # Prediction 1: divergence starts at eta*(4+lambda)=2.
    lam_sweep = np.linspace(0, 4, 9)
    boundary_rows = []
    for lam in lam_sweep:
        pred = exact_boundary(lam)
        # Direct sign sweep, with a fine grid, independently of continuation.
        etas = np.linspace(0.05, 0.55, 1001)
        fs = np.array([feature(e, lam) for e in etas])
        ix = np.flatnonzero(fs >= 0)
        measured = float(etas[ix[0]]) if len(ix) else float('nan')
        boundary_rows.append({"lambda": float(lam), "predicted_eta": pred,
                              "measured_eta": measured,
                              "relative_error": abs(measured-pred)/pred})

    # Prediction 2: eta_boundary*(4+lambda) is constant (=2).
    scaled = [r["measured_eta"] * (4 + r["lambda"]) for r in boundary_rows]
    scaling_error = abs(np.mean(scaled) - 2.0) / 2.0

    # Prediction 3: secant predictor + one-dimensional correction follows the
    # same curve. Start with two measured boundary points and advance lambda.
    cont_lams = np.linspace(0, 4, 9)
    p0 = bisect_eta(cont_lams[0], 0.2, 0.7)
    p1 = bisect_eta(cont_lams[1], 0.2, 0.7)
    points = [p0, p1]
    for lam in cont_lams[2:]:
        pred = points[-1] + (points[-1] - points[-2])
        # A bracket around secant prediction; the correction is a scalar sweep.
        lo, hi = max(0.03, pred * 0.65), pred * 1.35
        corrected = bisect_eta(float(lam), lo, hi)
        points.append(corrected)
    continuation_err = [abs(e-exact_boundary(l))/exact_boundary(l)
                        for e,l in zip(points, cont_lams)]

    # Local normal-direction sign switch at a regular crossing.
    lam0 = 2.0
    e0 = exact_boundary(lam0)
    eps = 0.01 * e0
    local = {"eta_below": e0-eps, "F_below": feature(e0-eps, lam0),
             "eta_above": e0+eps, "F_above": feature(e0+eps, lam0),
             "sign_switch": feature(e0-eps,lam0) < 0 < feature(e0+eps,lam0)}

    # Standard baseline: evaluate a uniform 20x20 parameter grid and infer
    # each lambda-column boundary from the first unstable eta.
    grid_etas = np.linspace(0.05, 0.55, 20)
    grid_lams = np.linspace(0.0, 4.0, 20)
    grid_F = np.array([[feature(e, l) for e in grid_etas] for l in grid_lams])
    grid_boundary = []
    for j, l in enumerate(grid_lams):
        unstable = np.flatnonzero(grid_F[j] >= 0)
        grid_boundary.append(float(grid_etas[unstable[0]]) if len(unstable) else float('nan'))
    grid_boundary = np.asarray(grid_boundary)
    grid_valid = np.isfinite(grid_boundary)
    exact_grid = np.array([exact_boundary(l) for l in grid_lams])
    grid_rel_error = (np.abs(grid_boundary[grid_valid] - exact_grid[grid_valid]) /
                      exact_grid[grid_valid])

    result = {
      "seed": SEED,
      "predictions": {
        "stability": "divergence when eta*(4+lambda)>2",
        "scaling": "eta_boundary*(4+lambda)=2",
        "continuation": "secant plus scalar correction tracks boundary"
      },
      "direct_sweep": boundary_rows,
      "mean_relative_boundary_error": float(np.mean([r["relative_error"] for r in boundary_rows])),
      "scaling_mean": float(np.mean(scaled)),
      "scaling_relative_error": float(scaling_error),
      "continuation": {"lambdas": cont_lams.tolist(), "etas": points,
                        "relative_errors": continuation_err,
                        "mean_relative_error": float(np.mean(continuation_err)),
                        "runs": 2 + 7 * 2, "dense_grid_20x20_runs": 400,
                        "grid_mean_relative_error": float(np.mean(grid_rel_error)),
                        "grid_boundary_points": int(np.sum(grid_valid))},
      "local_normal_check": local
    }
    with open("results.json", "w") as f: json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))

if __name__ == "__main__": main()