import json, math, random, sys
import numpy as np
import torch
from torch import nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
import bench

SEEDS = tuple(range(8))
TRACK='sequence'; MODEL='transformer_tiny'; NTR=400; NTE=200
# Shared search space: every idea learning rate is also a baseline configuration.
LR_GRID=[0.001,0.003,0.009]
EPOCHS=12; BATCH=64

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)

def make(seed):
    seed_all(seed); ds=bench.get_dataset(TRACK, seed, NTR, NTE)
    return ds, bench.make_model(MODEL, ds['input_shape'], ds['out_dim'])

def target_cov(x, model):
    # NNGP for the first shared linear projection, pooled over sequence positions.
    # PyTorch Linear(1,d) has init variance 1/3; positional vectors are shared.
    c=1.0/3.0
    z=x.mean(1)
    K=c*(z[:,None]*z[None,:])
    p=model.pos[:, :x.shape[1], :].mean(1)
    K=K + (p @ p.T)/p.shape[-1]
    return K.detach()

def pooled_cov(h):
    q=h.mean(1)
    return q @ q.T / q.shape[-1]

def train_idea(seed, lr, lam, return_model=False):
    ds, net=make(seed)
    x,y=ds['xtr'],ds['ytr']; lossf=nn.MSELoss()
    # Robust CUDA fallback, matching the benchmark's allowed device policy.
    devices=['cuda','cpu'] if torch.cuda.is_available() else ['cpu']
    last=None
    for dev in devices:
        try:
            net=net.to(dev); xx,yy=x.to(dev),y.to(dev)
            opt=torch.optim.Adam(net.parameters(),lr=lr)
            for ep in range(EPOCHS):
                net.train(); perm=torch.randperm(len(xx),device=dev)
                for i in range(0,len(xx),BATCH):
                    ix=perm[i:i+BATCH]; xb,yb=xx[ix],yy[ix]
                    h=net.inp(xb.unsqueeze(-1))+net.pos[:,:xb.shape[1]]
                    pred=net.head(net.enc(h).reshape(xb.shape[0],-1))
                    task=lossf(pred,yb)
                    cov=((pooled_cov(h)-target_cov(xb,net))**2).mean()
                    loss=task+lam*cov
                    opt.zero_grad(); loss.backward(); opt.step()
            net.eval()
            with torch.no_grad():
                pred=net(ds['xte'].to(dev)); metric=float(((pred-ds['yte'].to(dev))**2).mean())
            if return_model: return net, metric, ds, dev
            return metric
        except RuntimeError as e:
            last=e
            if dev=='cuda':
                net=net.to('cpu'); continue
            raise
    raise last

def train_base(seed, lr, return_model=False):
    ds,net=make(seed)
    out=bench.train_model(net,ds,epochs=EPOCHS,lr=lr,batch=BATCH,weight_decay=0.0,log=lambda _:None)
    if return_model: return out[0],float(out[1]),ds
    return float(out[1])

def signature(base_cfg, idea_cfg):
    bd=[]; idd=[]
    for s in SEEDS:
        bm,bmse,ds=train_base(s,base_cfg['lr'],True)
        im,imse,ids,dev=train_idea(s,idea_cfg['lr'],idea_cfg['lambda_cov'],True)
        with torch.no_grad():
            xb=ds['xte'][:64].to(next(bm.parameters()).device)
            hb=bm.inp(xb.unsqueeze(-1))+bm.pos[:,:xb.shape[1]]
            kb=target_cov(xb,bm).to(hb.device)
            xi=ids['xte'][:64].to(dev); hi=im.inp(xi.unsqueeze(-1))+im.pos[:,:xi.shape[1]]
            ki=target_cov(xi,im).to(hi.device)
            bd.append(float(((pooled_cov(hb)-kb)**2).mean()))
            idd.append(float(((pooled_cov(hi)-ki)**2).mean()))
    predicted=1/math.sqrt(64)
    observed_ratio=float(np.mean(idd)/np.mean(bd))
    return {'quantity':'pooled input-projection covariance deviation on held-out benchmark windows',
            'predicted_finite_width_scale_at_width_64':predicted,
            'observed_baseline_mean':float(np.mean(bd)),
            'observed_idea_mean':float(np.mean(idd)),
            'observed_idea_over_baseline':observed_ratio,
            'confirmed': bool(observed_ratio < 0.9),
            'note':'The O(n^-1/2) prediction is a scale law, not an absolute covariance value; this signature tests reduction on trained systems.'}

def main():
    base_grid=[{'lr':v} for v in LR_GRID]
    base=bench.sweep_baseline(lambda c: lambda s: train_base(s,c['lr']),base_grid,seeds=SEEDS)
    idea_trials=[]
    for cfg in [{'lr':base['best_cfg']['lr'],'lambda_cov':1e-4}, {'lr':0.001,'lambda_cov':1e-3}, {'lr':0.009,'lambda_cov':1e-2}]:
        r=bench.evaluate(lambda s,c=cfg: train_idea(s,c['lr'],c['lambda_cov']),SEEDS)
        idea_trials.append({'cfg':cfg,'result':r})
    best=min(idea_trials,key=lambda z:z['result']['mean'])
    rep=bench.make_report(TRACK,MODEL,base,best['result'],{'mechanism_signature':signature(base['best_cfg'],best['cfg']), 'idea_sweep':idea_trials, 'epochs':EPOCHS,'batch':BATCH})
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()