import json, math, time
import numpy as np

SEED = 2693
rng = np.random.default_rng(SEED)


def relax_impulse(beta, n):
    # X_0=1, X_t=0 afterwards; state at integer n is beta**n.
    return beta ** n


def relax_constant(beta, n):
    # X_t=1 from t=0, M_-1=0.
    return 1.0 - beta ** (n + 1)


def polar_ns(a, iters=5):
    # Newton-Schulz approximation to the polar factor, normalized for stability.
    x = a.astype(np.float64, copy=True)
    norm = np.linalg.norm(x, 2)
    if norm == 0:
        return np.zeros_like(x)
    x /= norm
    for _ in range(iters):
        x = 0.5 * x @ (3.0 * np.eye(x.shape[1]) - x.T @ x)
    return x


def update_mode(m, g, beta):
    return beta * m + (1.0 - beta) * g


def optimizer_run(kind, x, y, steps=350, lr=0.025, beta=.95, bf=.9, bs=.99, w=.5):
    W = np.zeros((x.shape[1], y.shape[1]), dtype=np.float64)
    m = np.zeros_like(W); mf = np.zeros_like(W); ms = np.zeros_like(W)
    losses=[]; cosines=[]
    t0=time.perf_counter()
    for _ in range(steps):
        pred=x@W; g=x.T@(pred-y)/len(x)
        if kind == 'adamw':
            # Included as a standard reference, with fixed hyperparameters.
            m = .9*m + .1*g
            v = .999*(v if 'v' in locals() else np.zeros_like(W)) + .001*g*g
            W -= lr * m/(np.sqrt(v)+1e-8)
            upd = m/(np.sqrt(v)+1e-8)
        elif kind == 'muon':
            m = update_mode(m,g,beta); upd=polar_ns(m)
            W -= lr*upd
        else:
            mf=update_mode(mf,g,bf); ms=update_mode(ms,g,bs)
            mix=w*mf+(1-w)*ms; upd=polar_ns(mix)
            W -= lr*upd
        losses.append(float(np.mean((pred-y)**2)/2))
        gn=np.linalg.norm(g); un=np.linalg.norm(upd)
        cosines.append(float(np.sum(g*upd)/(gn*un+1e-12)))
    return {'final_loss':losses[-1], 'best_loss':min(losses), 'loss_100':losses[99],
            'mean_update_cos':float(np.mean(cosines)), 'seconds':time.perf_counter()-t0}


def main():
    # Stage-1 prediction A: tau ordering implies beta ordering and measured relaxation ordering.
    eta=1.0; tauf=2.0; taus=20.0
    bf=math.exp(-eta/tauf); bs=math.exp(-eta/taus)
    ngrid=np.arange(0,1000)
    hf=int(ngrid[np.argmin(np.abs(np.array([relax_constant(bf,n) for n in ngrid])-.5))])
    hs=int(ngrid[np.argmin(np.abs(np.array([relax_constant(bs,n) for n in ngrid])-.5))])
    pred_hf=math.log(.5)/math.log(bf)-1
    pred_hs=math.log(.5)/math.log(bs)-1

    # Prediction B: impulse residual is exactly beta^n; sweep beta and compare fitted log slope.
    beta_sweep=[.5,.8,.9,.95,.99]
    impulse=[]
    for b in beta_sweep:
        n=np.arange(1,51)
        slope=np.polyfit(n,np.log(np.array([relax_impulse(b,int(k)) for k in n])),1)[0]
        impulse.append({'beta':b,'predicted_log_slope':math.log(b),'observed_log_slope':float(slope),
                        'residual_at_20':relax_impulse(b,20),'predicted_residual_at_20':b**20})

    # Prediction C: sinusoidal response. For input cos(omega t), |H| for EMA is analytic.
    def gain(b,om):
        return (1-b)/math.sqrt(1+b*b-2*b*math.cos(om))
    freq_rows=[]
    for om in [0.1,0.5,1.5,math.pi]:
        N=4000; t=np.arange(N); signal=np.cos(om*t)
        mf=ms=0.; out=[]
        for z in signal:
            mf=bf*mf+(1-bf)*z; ms=bs*ms+(1-bs)*z; out.append(.5*mf+.5*ms)
        out=np.asarray(out); tail=slice(1000,None)
        # projection amplitude avoids phase sensitivity
        obs=2*abs(np.mean(out[tail]*np.exp(-1j*om*t[tail])))
        pred=.5*gain(bf,om)+.5*gain(bs,om)
        freq_rows.append({'omega':om,'predicted_gain':pred,'observed_gain':float(obs)})

    # Deterministic low-dimensional matrix regression; same polar routine for Muon variants.
    local=np.random.default_rng(SEED)
    X=local.normal(size=(256,16)); trueW=local.normal(size=(16,8)); Y=X@trueW + .05*local.normal(size=(256,8))
    runs={
      'Muon':optimizer_run('muon',X,Y,beta=.95),
      'BiMaxwell_(.90,.99,.5)':optimizer_run('bi',X,Y,bf=.90,bs=.99,w=.5),
      'BiMaxwell_(.95,.995,.5)':optimizer_run('bi',X,Y,bf=.95,bs=.995,w=.5),
      'AdamW':optimizer_run('adamw',X,Y),
    }
    result={'seed':SEED,'predictions':{
      'A_half_life':{'beta_fast':bf,'beta_slow':bs,'predicted_steps_fast':pred_hf,'observed_steps_fast':hf,'predicted_steps_slow':pred_hs,'observed_steps_slow':hs},
      'B_impulse_decay':impulse,'C_frequency_gain':freq_rows},'mini_experiment':runs}
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()