"""Cycle-aware heavy-ball safeguard: math check and tiny reproducible experiments."""
import json, math, random
import numpy as np


def spectral_radius(lam, lr, beta):
    A = np.array([[1 + beta - lr * lam, -beta], [1., 0.]])
    return float(np.max(np.abs(np.linalg.eigvals(A))))


def diagnostics(d_t, d_prev, d_lag2, eps=1e-12):
    a, b, c = (np.asarray(x, float) for x in (d_t, d_prev, d_lag2))
    corr = float(np.dot(a, c) / (np.linalg.norm(a)*np.linalg.norm(c) + eps))
    ratio = float(np.linalg.norm(a) / (np.linalg.norm(b) + eps))
    return corr, ratio


class CycleAwareHB:
    """Numpy HB optimizer with fixed random projection and periodic safeguard."""
    def __init__(self, shape, lr, beta, projection_dim=32, interval=20, hold=100, seed=0):
        self.lr0, self.beta0 = lr, beta
        self.lr, self.beta = lr, beta
        self.interval, self.hold = interval, hold
        self.rng = np.random.default_rng(seed)
        n = int(np.prod(shape))
        self.P = self.rng.normal(0, 1/math.sqrt(projection_dim), (projection_dim, n))
        self.v = np.zeros(n); self.prev_d = None; self.prev_prev_d = None
        self.cooldown = 0; self.triggers = 0; self.history = []

    def step(self, x, grad, curvature):
        old = x.copy(); self.v = self.beta*self.v + grad
        x = x - self.lr*self.v
        d = self.P @ (x-old)
        # Preserve d_{t-1} and d_{t-2} before installing current d_t.
        d_prev, d_lag2 = self.prev_d, self.prev_prev_d
        self.prev_prev_d, self.prev_d = d_prev, d
        step_no = len(self.history) + 1
        trig = False; c = float('nan'); r = float('nan'); rho = spectral_radius(curvature, self.lr, self.beta)
        if step_no % self.interval == 0 and d_lag2 is not None and d_prev is not None:
            c, r = diagnostics(d, d_prev, d_lag2)
            trig = c > .8 and abs(r-1) < .15 and rho >= .98
            if trig:
                self.beta = min(.5, .7*self.beta); self.lr *= .8
                self.v.fill(0); self.cooldown = self.hold; self.triggers += 1
        if self.cooldown > 0: self.cooldown -= 1
        elif not trig and (step_no % self.interval == 0) and spectral_radius(curvature, self.lr0, self.beta0) < .95:
            # Never recover into the originally diagnosed unstable regime.
            self.lr += .05*(self.lr0-self.lr); self.beta += .05*(self.beta0-self.beta)
        self.history.append((c,r,rho,trig,self.lr,self.beta))
        return x


def math_check():
    # For constant curvature the measured linear transition must equal the formula.
    lr, beta, lam = .039, .95, 100.
    A = np.array([[1+beta-lr*lam, -beta], [1., 0.]])
    z = np.array([.37, -.21]); observed = A @ z
    # direct recurrence gives the same next state
    x, old = z
    direct = np.array([(1+beta)*x - beta*old - lr*lam*x, x])
    cyc = np.array([[1.,0.],[-.5,.8660254],[-.5,-.8660254], [1.,0.]])
    cycle_corr, cycle_ratio = diagnostics(cyc[3]-cyc[2], cyc[2]-cyc[1], cyc[1]-cyc[0])
    two_corr, two_ratio = diagnostics(np.array([1.,0.]), np.array([-1.,0.]), np.array([1.,0.]))
    return {"matrix_error": float(np.max(np.abs(observed-direct))),
            "period3_lag2_corr": cycle_corr, "period3_ratio": cycle_ratio,
            "period2_lag2_corr": two_corr, "period2_ratio": two_ratio,
            "rho_near_unit": spectral_radius(lam,lr,beta),
            "rho_after_reduction": spectral_radius(lam,.8*lr,.7*beta)}


def quadratic_run(guarded, seed=7, steps=400):
    rng=np.random.default_rng(seed); n=24
    eig=np.geomspace(1,100,n); x=rng.normal(size=n); loss=[]
    if guarded: opt=CycleAwareHB((n,),.039,.95,seed=seed)
    v=np.zeros(n)
    for t in range(steps):
        g=eig*x
        if guarded: x=opt.step(x,g,float(eig.max()))
        else:
            v=.95*v+g; x=x-.039*v
        loss.append(float(.5*np.dot(eig*x,x)))
        if not np.isfinite(loss[-1]): break
    return {"final_loss":loss[-1], "best_loss":min(loss), "triggers":opt.triggers if guarded else 0,
            "loss_at_100":loss[min(99,len(loss)-1)], "loss_at_400":loss[-1]}


def main():
    np.random.seed(0); random.seed(0)
    out={"math":math_check(), "baseline":quadratic_run(False), "safeguard":quadratic_run(True)}
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()