import math, json
import numpy as np

# E_{1/2}(-x): exact series for small x and erfcx asymptotic for large x.
def mlhalf_x(x):
    x=float(x)
    if x<3.5:
        term=total=1.0
        for n in range(1,500):
            term *= -x*math.gamma(.5*(n-1)+1)/math.gamma(.5*n+1)
            total += term
            if abs(term)<2e-15*max(1,abs(total)): break
        return max(0.,total)
    z=1/x**2
    return (1/math.sqrt(math.pi)/x)*(1-z/2+3*z*z/4-15*z**3/8+105*z**4/16)

def target(l,tau,alpha=.5): return mlhalf_x(((l+1)/tau)**alpha) if alpha==.5 else None

def simplex(v):
    u=np.sort(v)[::-1]; css=np.cumsum(u)-1; q=np.arange(1,len(v)+1)
    ok=np.where(u-css/q>0)[0]; j=int(ok[-1]) if len(ok) else 0
    return np.maximum(v-css[j]/(j+1),0)

def fit(tau=16, taus=(1,4,16,64,256,1024), L=4096):
    l=np.unique(np.r_[0,np.logspace(0,math.log10(L),320).astype(int)])
    r=np.exp(-1/np.asarray(taus,float)); A=np.array([[(1-x)*x**z for x in r] for z in l])
    y=np.array([target(z,tau) for z in l]); w=np.ones(len(r))/len(r)
    step=1/(np.linalg.norm(A,2)**2+1e-12)
    for _ in range(10000):
        nw=simplex(w-step*1.8*A.T@(A@w-y))
        if np.max(abs(nw-w))<1e-11: w=nw; break
        w=nw
    return l,r,w,y,A@w

def sparse_retention(gap, taus, w):
    r=np.exp(-1/np.asarray(taus,float)); v=np.zeros(len(r))
    v=(1-r)*1 # impulse at t=0
    for _ in range(gap): v=r*v
    return float(np.dot(w,v))

def run_opt(kind, seed=7, steps=500):
    rng=np.random.default_rng(seed); d=20; x=rng.normal(size=d); Q=np.linspace(1,30,d)
    m=np.zeros(d); taus=np.array([1,4,16,64,256,1024.]); r=np.exp(-1/taus); w=np.ones(6)/6
    vv=np.zeros(d) if kind=='adam' else np.zeros((6,d))
    if kind=='idea':
        _,r,w,_,_=fit(16,tuple(taus.astype(int)))
    for t in range(1,steps+1):
        g=Q*x + .03*rng.normal(size=d)
        m=.9*m+.1*g
        if kind=='adam':
            vv=.999*vv+.001*g*g; den=np.sqrt(vv/(1-.999**t))+1e-8
        else:
            vv=r[:,None]*vv+(1-r[:,None])*g[None,:]*g[None,:]
            den=np.sqrt(sum(w[k]*vv[k]/(1-r[k]**t) for k in range(len(r))))+1e-8
        x-=.03*m/(1-.9**t)/den
    return float(.5*np.sum(Q*x*x))

def main():
    taus=np.array([1,4,16,64,256,1024.]); l,r,w,y,hat=fit()
    # Prediction 1: asymptotic ML slope is -alpha=-.5.
    ll=np.arange(512,4096); slope=np.polyfit(np.log(ll),np.log([target(z,16) for z in ll]),1)[0]
    # Prediction 2: exponential bank has a less negative late slope than Adam beta=.999.
    late=np.arange(1000,4000); bank=np.array([sum(w[k]*(1-r[k])*r[k]**z for k in range(6)) for z in late])
    ok=bank>1e-300; bs=np.polyfit(np.log(late[ok]),np.log(bank[ok]),1)[0]
    adam_slope=math.log(.999)
    # Prediction 3: old impulse retention grows with longest available memory.
    gaps=[1,16,128,512]; ret=[sparse_retention(g,taus,w) for g in gaps]
    # Compare equal toy steps.
    base=[run_opt('adam',s) for s in range(3)]; idea=[run_opt('idea',s) for s in range(3)]
    out={'alpha':.5,'fit_weights':w.tolist(),'kernel_rel_rmse':float(np.linalg.norm(hat-y)/np.linalg.norm(y)),
      'predictions':{'tail_slope_pred':-.5,'tail_slope_obs':float(slope),'tol':.08,
      'late_bank_log_slope_obs':float(bs),'single_ema_log_slope_obs':adam_slope,
      'sparse_gaps':gaps,'sparse_retention':ret,'retention_monotone':bool(np.all(np.diff(ret)<0))},
      'toy_final_losses':{'adamw':base,'mittag_leffler':idea},
      'math_normalization_note':'The stated h_hat uses (1-rho)rho^l while sum_l h_hat=1, whereas h_0 is generally <1; fitting with sum(w)=1 cannot exactly match h_0.'}
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()