import json, math, time
from dataclasses import dataclass
import numpy as np

@dataclass
class Controller:
    eta0: float
    B0: int
    rho0: float
    kind: str = "polyak"
    safety: float = 0.5
    beta: float = 1.5

    def eta(self, B, rho):
        if self.kind == "nesterov":
            scale = self.safety * (B / self.B0) ** self.beta * (1-rho)/(1-self.rho0)
        else:
            scale = self.safety * (B / self.B0) * (1-rho)/(1-self.rho0)
        return self.eta0 * min(1.0, scale)

def spectral_radius(kind, eta, rho, lam=1.0):
    if kind == "polyak":
        A = np.array([[1-eta*lam, -eta*rho], [lam, rho]], float)
    elif kind == "nesterov":
        A = np.array([[1-eta*lam, -eta*rho*(1-eta*lam)],
                      [lam, rho*(1-eta*lam)]], float)
    else:
        return abs(1-eta*lam)
    return float(np.max(np.abs(np.linalg.eigvals(A))))

def critical_eta(kind, rho, lam=1.0):
    lo, hi = 0.0, 1.0 / lam
    while spectral_radius(kind, hi, rho, lam) < 1.0:
        hi *= 2
    for _ in range(70):
        mid = (lo + hi) / 2
        if spectral_radius(kind, mid, rho, lam) < 1.0:
            lo = mid
        else:
            hi = mid
    return lo

def stochastic_run(kind, B, eta, rho, seed, steps=400, d=8, noise=1.0):
    rng = np.random.default_rng(seed)
    x = rng.normal(size=d)
    v = np.zeros(d)
    losses = []
    updates = []
    for _ in range(steps):
        if kind == "nesterov":
            grad = x - eta * rho * v + rng.normal(size=d) * noise / math.sqrt(B)
        else:
            grad = x + rng.normal(size=d) * noise / math.sqrt(B)
        old = x.copy()
        v = rho * v + grad
        x = x - eta * v
        losses.append(0.5 * float(np.dot(x, x)))
        updates.append(float(np.linalg.norm(x-old)))
        if not np.isfinite(losses[-1]) or losses[-1] > 1e12:
            return {"stable": False, "final_loss": float("inf"), "median_loss": float("inf"),
                    "update": float("inf"), "steps": len(losses)}
    return {"stable": bool(np.isfinite(losses).all()),
            "final_loss": losses[-1], "median_loss": float(np.median(losses[-50:])),
            "update": float(np.median(updates[-50:])), "steps": steps}

def scan_stability(kind, B, rho, eta_grid, repeats=3):
    out=[]
    for eta in eta_grid:
        runs=[stochastic_run(kind,B,eta,rho,100+r,steps=300) for r in range(repeats)]
        out.append((eta, all(r["stable"] for r in runs)))
    stable=[e for e, ok in out if ok]
    return max(stable) if stable else 0.0

def main():
    rho0=0.9; B0=1; eta0=0.05; beta=1.5
    # Core math: compare critical-rate ratios to the predicted capped scaling.
    rho_values=[0.0,0.5,0.9,0.95]
    math_rows=[]
    for kind in ["polyak","nesterov"]:
        base=critical_eta(kind,rho0)
        for rho in rho_values:
            c=critical_eta(kind,rho)
            predicted=min(1.0, (1-rho)/(1-rho0))
            math_rows.append({"kind":kind,"rho":rho,"critical":c,
                              "ratio":c/base,"predicted_uncapped_ratio":predicted})
    # Batch controller experiment with fixed number of samples (steps shrink as B grows).
    rows=[]
    for kind in ["polyak","nesterov"]:
        ctrl=Controller(eta0,B0,rho0,kind,0.5,beta)
        for B in [1,2,4,8,16]:
            eta=ctrl.eta(B,rho0)
            fixed=stochastic_run(kind,B,eta0,rho0,123,steps=max(20,800//B))
            controlled=stochastic_run(kind,B,eta,rho0,123,steps=max(20,800//B))
            rows.append({"kind":kind,"B":B,"eta_controller":eta,
                         "fixed_stable":fixed["stable"],"controller_stable":controlled["stable"],
                         "fixed_final_loss":fixed["final_loss"],"controller_final_loss":controlled["final_loss"],
                         "controller_update":controlled["update"]})
    # Empirical maximum stable rates at each batch, plus theoretical deterministic boundary.
    scan=[]
    for kind in ["polyak","nesterov"]:
        for B in [1,2,4,8,16]:
            grid=np.geomspace(0.002,2.0,45)
            empirical=scan_stability(kind,B,rho0,grid)
            scan.append({"kind":kind,"B":B,"empirical_max_eta":empirical,
                         "deterministic_quadratic_eta":critical_eta(kind,rho0)})
    result={"math_check":math_rows,"controller_runs":rows,"stability_scan":scan,
            "config":{"eta0":eta0,"rho0":rho0,"beta":beta,"safety":0.5}}
    with open("results.json","w") as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__ == "__main__":
    main()