import json, math, random
from pathlib import Path
import numpy as np

SEED = 7
np.random.seed(SEED); random.seed(SEED)


def toy_checks():
    lam = 1.0; beta = 0.9
    # Prediction 1: ||B^-1 g|| <= ||g||/lambda.
    ratios = []
    for _ in range(2000):
        v = np.random.exponential(1.0, 30)
        g = np.random.randn(30)
        ratios.append(np.linalg.norm(g/(lam+v))/(np.linalg.norm(g)/lam))
    bound_max = float(max(ratios))

    # Prediction 2: under constant excitation q, v converges to q and
    # quadratic-error contraction tends to |1-eta*q/(lambda+q)|.
    qs = [0.05, 0.2, 1.0, 4.0]
    eta = 0.8
    contraction = []
    for q in qs:
        v = 0.0
        for _ in range(600):
            v = beta*v + (1-beta)*q
        measured = abs(1 - eta*q/(lam+v))
        pred = abs(1 - eta*q/(lam+q))
        contraction.append({'q':q, 'predicted_factor':pred, 'measured_factor':float(measured), 'v_final':v})

    # Prediction 3: stability boundary for constant q is eta < 2(lambda+q)/q.
    # Find the largest eta on a grid that does not grow from e=1 in 300 steps.
    stability=[]
    for q in qs:
        pred_eta = 2*(lam+q)/q
        grid=np.linspace(.1, min(pred_eta*1.4, 45), 300)
        stable_eta=.0
        for x in grid:
            v=0.; e=1.
            for _ in range(300):
                v=beta*v+(1-beta)*q
                e *= 1-x*q/(lam+v)
            if abs(e) < 1e3: stable_eta=float(x)
        stability.append({'q':q, 'predicted_boundary':pred_eta, 'grid_observed_boundary':stable_eta})

    # Claimed low-observability protection: q=0 gives v=0, hence no extra damping.
    # Inject unbiased gradient noise and measure variance slope.
    slopes=[]
    for q in [0., 1., 4.]:
        v=0.; th=0.; ys=[]
        for t in range(3000):
            v=beta*v+(1-beta)*q
            th -= 0.5*np.random.randn()/(lam+v)
            if t>=1000: ys.append(th*th)
        slopes.append({'q':q, 'mean_theta2_last2000':float(np.mean(ys)), 'final_v':v,
                       'predicted_update_noise_scale':1/(lam+q)**2})
    return {'bound': {'max_ratio_observed':bound_max, 'predicted_upper_bound':1.0},
            'contraction_sweep':contraction, 'stability_sweep':stability,
            'noise_drift_sweep':slopes}


def mlp_experiment():
    # Numpy two-layer regression keeps the comparison transparent and fast.
    rng=np.random.RandomState(SEED)
    n,d,h=2400,20,32
    X=rng.randn(n,d); true_w=rng.randn(d,1)
    y=np.tanh(X@true_w*.35)+.08*rng.randn(n,1)
    Xt,Xv=X[:2000],X[2000:]; yt,yv=y[:2000],y[2000:]
    # same initialization/data order for SGD and diagonal observability method
    W1_0=rng.randn(d,h)*.15; W2_0=rng.randn(h,1)*.15
    def run(kind, steps=700, mask_p=.8):
        W1,W2=W1_0.copy(),W2_0.copy(); local_rng=np.random.RandomState(SEED+{'sgd':1,'adam':2,'mop':3}[kind]); v1=np.zeros_like(W1); v2=np.zeros_like(W2)
        lr=.035 if kind!='adam' else .008; b1=.9; b2=.999; m1=np.zeros_like(W1);m2=np.zeros_like(W2); z1=np.zeros_like(W1);z2=np.zeros_like(W2)
        drift=[]; losses=[]
        for t in range(steps):
            ix=local_rng.choice(len(Xt),64,False); xb=Xt[ix].copy(); yb=yt[ix]
            mask=(local_rng.rand(*xb.shape)>mask_p).astype(float); xb*=mask
            a=xb@W1; h1=np.tanh(a); pred=h1@W2; diff=pred-yb
            g2=h1.T@diff/len(ix); g1=xb.T@((diff@W2.T)*(1-h1*h1))/len(ix)
            # diagonal J^T J proxy: parameter sensitivity squared, averaged per parameter
            sens1=(xb[:,:,None]**2)*(W2.T[None,:,:]**2)*(1-h1[:,None,:]**2)**2
            sens2=h1*h1
            vv1=sens1.mean(axis=0); vv2=sens2.mean(axis=0)[:,None]
            if kind=='mop':
                v1=b1*v1+(1-b1)*vv1; v2=b1*v2+(1-b1)*vv2
                W1-=lr*g1/(1+v1); W2-=lr*g2/(1+v2)
            elif kind=='sgd': W1-=lr*g1; W2-=lr*g2
            else:
                m1=b1*m1+(1-b1)*g1; m2=b1*m2+(1-b1)*g2; z1=b2*z1+(1-b2)*g1*g1; z2=b2*z2+(1-b2)*g2*g2
                W1-=lr*(m1/(1-b1**(t+1)))/(np.sqrt(z1/(1-b2**(t+1)))+1e-8); W2-=lr*(m2/(1-b1**(t+1)))/(np.sqrt(z2/(1-b2**(t+1)))+1e-8)
            if t%50==0:
                hv=np.tanh(Xv@W1)@W2; losses.append(float(np.mean((hv-yv)**2)))
                drift.append(float(np.linalg.norm(W1-W1_0)+np.linalg.norm(W2-W2_0)))
        return {'final_val_mse':losses[-1], 'loss_trace':losses, 'parameter_drift':drift}
    out={k:run(k) for k in ['sgd','adam','mop']}
    return out

if __name__=='__main__':
    result={'toy':toy_checks(), 'mlp':mlp_experiment()}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))