import json, math
import numpy as np

SEED = 2807

def boundary(lam, r, sig, B):
    return 2*lam/(r*(lam*lam + sig*sig/B))

def toy():
    rng=np.random.default_rng(SEED); lam,r,sig=1.,.8,1.5
    contraction=[]
    for alpha in (.2,.5,1.0,1.4):
        a=1-alpha*r*lam; x=1.; ratios=[]
        for _ in range(30):
            y=a*x; ratios.append(abs(y/x)); x=y
        contraction.append({'alpha':alpha,'predicted':abs(a),'observed':float(np.mean(ratios))})
    # Directly estimate E[(x_{k+1}/x_k)^2], avoiding rare-event trajectory bias.
    rows=[]; Bs=[1,2,4,8,16,32]
    for B in Bs:
        pred=boundary(lam,r,sig,B)
        grid=np.linspace(pred*.72,pred*1.28,41)
        vals=[]
        for alpha in grid:
            xi=rng.normal(0,sig/math.sqrt(B),size=500000)
            factor=1-alpha*r*(lam+xi)
            m2=float(np.mean(factor*factor))
            vals.append((alpha,m2-1))
        cross=min(vals,key=lambda z:abs(z[1]))[0]
        for (a,s),(b,t) in zip(vals[:-1],vals[1:]):
            if s*t<=0:
                cross=a+(0-s)*(b-a)/(t-s); break
        rows.append({'B':B,'predicted_alpha_boundary':pred,'observed_alpha_boundary':float(cross),'relative_error':float(abs(cross-pred)/pred)})
    # Additive response-noise floor: stationary E[x^2] = alpha^2*sigma^2/B/(1-a^2).
    alpha=.25; add_sig=1.2; scale=[]; a=1-alpha*r*lam
    predicted_per_B=(alpha*alpha*add_sig*add_sig)/(1-a*a)
    for B in (1,4,16):
        x=np.zeros(50000); vals=[]
        for k in range(2200):
            x=a*x + alpha*rng.normal(0,add_sig/math.sqrt(B),size=x.size)
            if k>=1000: vals.append(np.mean(x*x))
        obs=float(np.mean(vals))
        scale.append({'B':B,'predicted_second_moment':predicted_per_B/B,'observed_second_moment':obs,'B_times_moment':B*obs})
    return {'contraction':contraction,'boundary':rows,'variance_scaling':scale}

def optimizer_demo():
    rng=np.random.default_rng(SEED+1); n=1600; d=12
    X=rng.normal(size=(n,d)); true=rng.normal(size=d); y=X@true+rng.normal(0,.25,n)
    def run(kind, seed):
        rr=np.random.default_rng(seed); w=np.zeros(d); m=np.zeros(d); v=np.zeros(d); losses=[]; alpha=.35; B=32
        for k in range(300):
            ix=rr.integers(0,n,B); xb=X[ix]; yb=y[ix]; g=xb.T@(xb@w-yb)/B
            if kind=='adam':
                m=.9*m+.1*g; v=.999*v+.001*g*g; step=.06*m/(np.sqrt(v)+1e-8)
            elif kind=='prox': step=.35*g
            else: step=alpha*.35*g
            w-=step; losses.append(float(np.mean((X@w-y)**2)))
        return losses
    out={k:run(k,SEED+2) for k in ('adam','prox','relaxed_prox')}
    return {k:{'final_loss':v[-1],'best_loss':min(v)} for k,v in out.items()}

if __name__=='__main__':
    result={'toy':toy(),'optimizer_demo':optimizer_demo()}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))