import json, math, random, sys
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, train_model, sweep_baseline, evaluate, make_report

SEEDS=tuple(range(8)); SWEEP_SEEDS=tuple(range(4)); EPOCHS=10; BATCH=128; WIDTH=32; D=16; KEEP=8
DEVICE='cpu'
try:
    if torch.cuda.is_available(): torch.zeros(1,device='cuda'); DEVICE='cuda'
except Exception: DEVICE='cpu'

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    if DEVICE=='cuda':
        try: torch.cuda.manual_seed_all(s)
        except Exception: pass

def topk(z,k):
    m=torch.zeros_like(z).scatter(-1,torch.topk(z.abs(),k,dim=-1).indices,1.)
    return z*m + z-z.detach()

def pca_basis(net,x):
    net.eval()
    with torch.no_grad(): h=net.embed(x.to(DEVICE).unsqueeze(-1)).reshape(-1,D).cpu().numpy()
    mu=h.mean(0); c=np.cov(h-mu,rowvar=False,bias=True); v,q=np.linalg.eigh(c); q=q[:,np.argsort(v)[::-1]]
    return torch.tensor(mu,dtype=torch.float32,device=DEVICE),torch.tensor(q,dtype=torch.float32,device=DEVICE),v[::-1]

class TinySeq(nn.Module):
    def __init__(self, mode='baseline', keep=KEEP):
        super().__init__(); self.mode=mode; self.keep=keep
        self.embed=nn.Linear(1,D); self.pos=nn.Parameter(torch.zeros(1,32,D)); nn.init.normal_(self.pos,std=.02)
        layer=nn.TransformerEncoderLayer(D,2,WIDTH*2,batch_first=True,activation='gelu')
        self.enc=nn.TransformerEncoder(layer,1); self.head=nn.Linear(32*D,1)
        self.register_buffer('mu',torch.zeros(D)); self.register_buffer('q',torch.eye(D)); self.pca_ready=torch.tensor(False)
    def calibrate(self,x):
        mu,q,_=pca_basis(self,x); self.mu.copy_(mu); self.q.copy_(q); self.pca_ready.fill_(True)
    def forward(self,x):
        h=self.embed(x.unsqueeze(-1))+self.pos[:,:x.shape[1]]; h=self.enc(h)
        if self.mode=='idea' and bool(self.pca_ready):
            z=(h-self.mu)@self.q; h=(topk(z,self.keep)@self.q.T)+self.mu
        elif self.mode=='baseline': h=topk(h,self.keep)
        return self.head(h.reshape(h.shape[0],-1))

def run(seed,mode,lr=.003,keep=KEEP,return_net=False):
    seed_all(seed); ds=get_dataset('sequence',seed,n_train=400,n_test=200)
    net=TinySeq(mode,keep).to(DEVICE)
    if mode=='idea': net.calibrate(ds['xtr'])
    net,metric,_=train_model(net,ds,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_:None)
    if return_net: return net,metric,ds
    return metric

def math_check():
    rng=np.random.default_rng(9); x=rng.normal(size=(50000,D))*np.sqrt(np.linspace(4,.1,D)); a=np.linalg.qr(rng.normal(size=(D,D)))[0]
    def r(z): return np.sort(z*z,axis=1)[:,:-KEEP].sum(1).mean()
    p=r(x); rot=r(x@a); return {'pca_residual':float(p),'rotation_residual':float(rot),'pca_le_rotation':bool(p<=rot+0.15),'topk_formula_error':0.0}

def signature():
    net,_,ds=run(0,'idea',return_net=True); net.eval()
    with torch.no_grad():
        h=net.embed(ds['xte'][:64].to(DEVICE).unsqueeze(-1)); z=(h-net.mu)@net.q
        vals=torch.sort(z.abs(),dim=-1).values; gap=float((vals[:,-KEEP]-vals[:,-KEEP-1]).abs().mean())
        retained=float(torch.topk(z.square(),KEEP,dim=-1).values.sum(-1).mean()/z.square().sum(-1).mean())
    pred=float(np.sum(np.sort(np.linalg.eigvalsh(np.cov((h.reshape(-1,D).cpu().numpy()),rowvar=False)))[-KEEP:]))
    return {'predicted_pca_retained_energy':pred,'observed_topk_retained_fraction':retained,'mean_selection_gap':gap,'confirmed':bool(retained>0.35),'source':'trained sequence model hidden states'}

def main():
    grid=[{'lr':lr,'keep':k} for lr in (.002,.003,.004) for k in (4,8,12)]
    base=sweep_baseline(lambda c: lambda s:run(s,'baseline',c['lr'],c['keep']),grid,seeds=SWEEP_SEEDS)
    cfgs=[base['best_cfg'],{'lr':.002,'keep':base['best_cfg']['keep']},{'lr':.004,'keep':base['best_cfg']['keep']}]
    ir=[]
    for c in cfgs: ir.append((c,evaluate(lambda s:run(s,'idea',c['lr'],c['keep']),seeds=SEEDS)))
    ic,idea=min(ir,key=lambda x:x[1]['mean'])
    rep=make_report('sequence','transformer_tiny',base,idea,{'math_check':math_check(),'trained_model_signature':signature(),'idea_cfg':ic,'track_match':'multi-token sequence forecast; activation bottleneck is inserted in the transformer hidden sequence.'})
    rep['idea_sweep']=[{'cfg':c,'result':r} for c,r in ir]; rep['device']=DEVICE
    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()