import math, json
import numpy as np

EPS = 1e-12

def huber(x, eps, delta):
    if x >= delta: return delta*abs(x) - delta*delta/2
    if x <= -eps: return eps*abs(x) - eps*eps/2
    return .5*x*x

def huber_grad(x, eps, delta):
    if x >= delta: return delta
    if x <= -eps: return -eps
    return x

def trajectory(etas, m, criterion):
    """The paper's asymmetric-Huber construction, with 1-indexed m."""
    before = float(sum(etas[:m-1])); after = float(sum(etas[m:]))
    delta = 1.0/(1.0 + before)
    overshoot = (float(etas[m-1])-1.0)*delta
    eps = overshoot/(1.0 + (2.0*after if criterion == 'R' else after))
    x = 1.0; grads=[]
    for eta in etas:
        g = huber_grad(x, eps, delta)
        grads.append(g); x -= float(eta)*g
    initial = huber(1.0, eps, delta)
    R = huber(x, eps, delta) / .5
    G = (.5 * huber_grad(x, eps, delta)**2) / initial
    rhsR = (float(etas[m-1])-1)**2 / ((1+before)**2 * (1+2*after))
    rhsG = (float(etas[m-1])-1)**2 / ((1+2*before) * (1+after)**2)
    return R, G, rhsR, rhsG

def verify_math(seed=7, trials=2000):
    rng=np.random.default_rng(seed); worstR=1e9; worstG=1e9; failures=0
    for _ in range(trials):
        n=int(rng.integers(2,12)); etas=rng.uniform(.03,2.5,n)
        m=int(rng.integers(1,n+1))
        if etas[m-1] <= 1: etas[m-1]=1.01+rng.random()*1.5
        RR,_,rR,_=trajectory(etas,m,'R'); _,GG,_,rG=trajectory(etas,m,'G')
        worstR=min(worstR,RR/rR); worstG=min(worstG,GG/rG)
        failures += int(RR+1e-9 < rR or GG+1e-9 < rG)
    a=np.array([.15,.25,1.8,.4,.1]); b=np.array([1.8,.15,.25,.4,.1])
    return {'trials':int(trials),'failures':int(failures),
      'min_R_ratio':float(worstR),'min_G_ratio':float(worstG),
      'ordering_R_rhs':[float(trajectory(a,3,'R')[2]),float(trajectory(b,1,'R')[2])],
      'ordering_G_rhs':[float(trajectory(a,3,'G')[3]),float(trajectory(b,1,'G')[3])]}

def run_controller(proposals, rho=.02, beta=.9, mode='controller', grad_clip=1.0):
    x=1.0; S=0.; prevq=None; v=1.0; records=[]
    for proposal in proposals:
        raw_g=x; q=raw_g*raw_g
        ratio=1.0 if prevq is None else (q+EPS)/(prevq+EPS)
        v=beta*v+(1-beta)*ratio
        cap=1+math.sqrt(1+2*S)*math.sqrt(rho*v+EPS)
        eta=min(float(proposal),cap) if mode=='controller' else float(proposal)
        g=max(-grad_clip,min(grad_clip,raw_g)) if mode=='grad_clip' else raw_g
        x=x-eta*g; S+=eta; prevq=q
        records.append((x,eta,cap,ratio,g))
    return x, records

def experiment():
    # Late, isolated proposals are deliberately beyond the unit-quadratic stability limit.
    n=64; proposals=np.full(n,.2); proposals[[7,23,39,55]]=20.0
    out={}
    for mode in ['sgd','grad_clip','controller']:
        x, rec=run_controller(proposals,mode=mode)
        losses=[.5*r[0]*r[0] for r in rec]
        out[mode]={'final_loss':float(losses[-1]) if np.isfinite(losses[-1]) else 'nonfinite',
          'max_abs_x':float(max(abs(r[0]) for r in rec)) if all(np.isfinite(r[0]) for r in rec) else 'nonfinite',
          'max_loss':float(max(losses)) if all(np.isfinite(z) for z in losses) else 'nonfinite',
          'clipped_steps':int(sum(r[1]<p-1e-10 for r,p in zip(rec,proposals))),
          'spike_etas':[float(rec[i][1]) for i in [7,23,39,55]],
          'loss_after_spikes':[float(losses[i+1]) for i in [7,23,39,55]]}
    _,r=run_controller(np.where(np.arange(n)==55,20.0,.2),mode='controller')
    out['single_late_spike']={'proposal':20.0,'accepted':float(r[55][1]),'cap':float(r[55][2])}
    return out

if __name__=='__main__':
    print(json.dumps({'math_check':verify_math(),'experiment':experiment()},indent=2))