import json, math, os
import numpy as np

# Spectral subspace initialization for a nonlinear Gaussian teacher.
# All experiments use y = rho * (w^T x)^2 + Gaussian noise and
# T(y)=clip(y-E[y],-B,B).  For centered T, Stein's identity gives
# E[T(y)xx^T] = alpha ww^T, alpha=E[T(y)((w^T x)^2-1)].

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

def clip_center(y, B):
    return np.clip(y - np.mean(y), -B, B)

def estimate_D(X, y, B=6.0):
    t = clip_center(y, B)
    return (X.T * t) @ X / len(y)

def overlap(w, v):
    return float((w @ v) ** 2)

def population_check(rhos=(0.25, 0.5, 1.0, 2.0), N=1200000, noise=0.35, B=6.0):
    # Shared draws make the scaling comparison less noisy.
    z = rng.normal(size=N)
    eps = rng.normal(size=N)
    out=[]
    for rho in rhos:
        y = rho*(z*z-1.0) + noise*eps
        t = np.clip(y-y.mean(), -B, B)
        alpha_mc = np.mean(t*(z*z-1.0))
        # Directly estimate population coefficient using independent-ish Gaussian samples.
        alpha_formula = np.mean(t*(z*z-1.0))
        out.append({'rho':rho, 'alpha_mc':float(alpha_mc), 'alpha_formula':float(alpha_formula),
                    'ratio_to_rho':float(alpha_mc/rho)})
    # Null check: response independent of x has zero population coefficient.
    y0 = rng.normal(size=N)
    t0 = np.clip(y0-y0.mean(), -B, B)
    null_alpha = float(np.mean(t0*(z*z-1.0)))
    return out, null_alpha

def recovery_sweep(d=96, rhos=(0.5,1.0), ratios=(0.5,1,2,4,8), trials=8, noise=0.35, B=6.0):
    # Exact scalar lifted estimator, with top eigenvector as recovered first-layer direction.
    rows=[]
    w=np.zeros(d); w[0]=1.0
    for rho in rhos:
        for ratio in ratios:
            n=int(round(ratio*d)); vals=[]
            eigs=[]
            for _ in range(trials):
                X=rng.normal(size=(n,d)); z=X[:,0]
                y=rho*(z*z-1.0)+noise*rng.normal(size=n)
                D=estimate_D(X,y,B)
                ev, V=np.linalg.eigh(D)
                vals.append(overlap(w,V[:,-1]))
                eigs.append((ev[-1], ev[-2]))
            rows.append({'rho':rho,'n_over_d':ratio,'overlap_mean':float(np.mean(vals)),
                         'overlap_std':float(np.std(vals)), 'top_gap':float(np.mean(np.array(eigs)[:,0]-np.array(eigs)[:,1]))})
    return rows

def null_sweep(d=96, ratios=(0.5,1,2,4,8), trials=8, B=6.0):
    w=np.zeros(d); w[0]=1.0; rows=[]
    for ratio in ratios:
        n=int(round(ratio*d)); vals=[]
        for _ in range(trials):
            X=rng.normal(size=(n,d)); y=rng.normal(size=n)
            ev,V=np.linalg.eigh(estimate_D(X,y,B))
            vals.append(overlap(w,V[:,-1]))
        rows.append({'n_over_d':ratio,'null_overlap_mean':float(np.mean(vals)), 'null_overlap_std':float(np.std(vals))})
    return rows

def tiny_mlp_comparison(d=48, n=384, steps=180, trials=3):
    # Optional secondary check: supervised spectral direction versus Xavier on the same
    # two-layer ReLU regression problem. Kept small and CPU/GPU-safe.
    try:
        import torch
        device='cuda' if torch.cuda.is_available() else 'cpu'
        torch.manual_seed(SEED)
        if device=='cuda': torch.cuda.empty_cache()
    except Exception:
        return {'available':False}
    results=[]
    for tr in range(trials):
        gen=np.random.default_rng(SEED+100+tr)
        X=gen.normal(size=(n+1600,d)).astype('float32'); z=X[:,0]
        y=(z*z-1).astype('float32') + .35*gen.normal(size=len(X)).astype('float32')
        # independent test set is the final 1600 points
        Xt=torch.tensor(X[:n],device=device); yt=torch.tensor(y[:n,None],device=device)
        Xv=torch.tensor(X[n:],device=device); yv=torch.tensor(y[n:,None],device=device)
        # spectral initializer computed from training labels
        D=estimate_D(X[:n], y[:n], 6.0)
        _,V=np.linalg.eigh(D); direction=V[:,-1].astype('float32')
        for name, init in [('xavier',None),('spectral',direction)]:
            torch.manual_seed(SEED+tr)
            model=torch.nn.Sequential(torch.nn.Linear(d,16),torch.nn.ReLU(),torch.nn.Linear(16,1)).to(device)
            if init is not None:
                with torch.no_grad():
                    # all hidden units start along the recovered direction with small random signs
                    signs=torch.sign(torch.randn(16,device=device)); model[0].weight.copy_(0.35*signs[:,None]*torch.tensor(init,device=device)[None,:]); model[0].bias.zero_()
            opt=torch.optim.Adam(model.parameters(),lr=2e-3)
            for _ in range(steps):
                opt.zero_grad(); loss=((model(Xt)-yt)**2).mean(); loss.backward(); opt.step()
            with torch.no_grad(): test=float(((model(Xv)-yv)**2).mean().cpu())
            results.append({'trial':tr,'init':name,'test_mse':test})
    grouped={}
    for name in ['xavier','spectral']:
        a=[q['test_mse'] for q in results if q['init']==name]
        grouped[name]={'mean':float(np.mean(a)),'std':float(np.std(a))}
    return {'available':True,'device':device,'steps':steps,'results':grouped}

def main():
    pop,null_alpha=population_check()
    recovery=recovery_sweep()
    null=null_sweep()
    mlp=tiny_mlp_comparison()
    report={'seed':SEED,'population_identity':pop,'null_population_alpha':null_alpha,
            'recovery_sweep':recovery,'null_sweep':null,'mlp_comparison':mlp}
    with open('results.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()