import json, math, time
import numpy as np
from bi_maxwell_muon import BiMaxwellMuon, polar_newton_schulz

SEED = 2693

def first_half_step(beta):
    n = 0
    while 1-beta**(n+1) < .5:
        n += 1
    return n

def run(kind, X, Y, steps=350, lr=.025, beta=.95, bf=.9, bs=.99, w=.5):
    W=np.zeros((16,8)); m=np.zeros_like(W); mf=np.zeros_like(W); ms=np.zeros_like(W)
    v=np.zeros_like(W); losses=[]; cs=[]; start=time.perf_counter()
    for _ in range(steps):
        pred=X@W; g=X.T@(pred-Y)/len(X)
        if kind=='adamw':
            m=.9*m+.1*g; v=.999*v+.001*g*g; upd=m/(np.sqrt(v)+1e-8)
        elif kind=='muon':
            m=beta*m+(1-beta)*g; upd=polar_newton_schulz(m)
        else:
            mf=bf*mf+(1-bf)*g; ms=bs*ms+(1-bs)*g
            upd=polar_newton_schulz(w*mf+(1-w)*ms)
        W -= lr*upd
        losses.append(float(np.mean((pred-Y)**2)/2))
        cs.append(float((g*upd).sum()/(np.linalg.norm(g)*np.linalg.norm(upd)+1e-12)))
    return {'final_loss':losses[-1], 'loss_step_100':losses[99],
            'mean_gradient_update_cosine':float(np.mean(cs)),
            'seconds':time.perf_counter()-start}

def main():
    eta=1.; tf=2.; ts=20.; bf=math.exp(-eta/tf); bs=math.exp(-eta/ts)
    half=[]
    for b in [bf,bs,.8,.9,.99]:
        pred=math.log(.5)/math.log(b)-1
        half.append({'beta':b, 'predicted_continuous_index':pred,
                     'predicted_first_integer':math.ceil(pred),
                     'observed_first_integer':first_half_step(b)})
    slopes=[]
    for b in [.5,.8,.9,.95,.99]:
        n=np.arange(1,51); y=b**n
        obs=float(np.polyfit(n,np.log(y),1)[0])
        slopes.append({'beta':b,'predicted_log_slope':math.log(b),'observed_log_slope':obs})
    def gain(b,o): return (1-b)/math.sqrt(1+b*b-2*b*math.cos(o))
    freqs=[]
    for o in [.1,.5,1.5,2.5]:
        n=6000; t=np.arange(n); z=np.cos(o*t); a=c=0.; out=[]
        for q in z:
            a=bf*a+(1-bf)*q; c=bs*c+(1-bs)*q; out.append(.5*a+.5*c)
        tail=slice(1000,None); tt=t[tail]; yy=np.asarray(out)[tail]
        observed=2*abs(np.mean(yy*np.exp(-1j*o*tt)))
        predicted=.5*gain(bf,o)+.5*gain(bs,o)
        freqs.append({'omega':o,'predicted_gain':predicted,'observed_gain':float(observed),
                      'relative_error':float(abs(observed-predicted)/predicted)})
    r=np.random.default_rng(SEED); X=r.normal(size=(256,16)); true=r.normal(size=(16,8)); Y=X@true+.05*r.normal(size=(256,8))
    runs={'Muon':run('muon',X,Y), 'BiMaxwell_(.90,.99,.5)':run('bi',X,Y,bf=.9,bs=.99),
          'BiMaxwell_(.95,.995,.5)':run('bi',X,Y,bf=.95,bs=.995), 'AdamW':run('adamw',X,Y)}
    print(json.dumps({'predictions':{'half_life':half,'impulse_log_slope':slopes,'frequency_gain':freqs},'mini_experiment':runs},indent=2))
if __name__=='__main__': main()