import sys,json,random
from pathlib import Path
import numpy as np, torch
from torch import nn
import torch.nn.functional as F
sys.path.insert(0,'/home/maxwelhelp/all/math2nn')
from bench import get_dataset,make_model,train_model,sweep_baseline,make_report
EPOCHS=12; NTRAIN=1200; NTEST=400; LRS=[1e-3,3e-3,6e-3]
def seed_all(s):
 random.seed(s);np.random.seed(s);torch.manual_seed(s)
 if torch.cuda.is_available():
  try: torch.cuda.manual_seed_all(s)
  except Exception: pass
class SubsetAttn(nn.Module):
 def __init__(self,d=64,n=4):
  super().__init__();self.n=n
  self.score=nn.Linear(2*d,1);self.pair=nn.Sequential(nn.Linear(2*d,32),nn.Tanh(),nn.Linear(32,1));self.dec=nn.Sequential(nn.Linear(2*d,64),nn.Tanh(),nn.Linear(64,d))
  self.register_buffer('masks',torch.tensor([[(m>>i)&1 for i in range(n)] for m in range(1,1<<n)],dtype=torch.float32))
 def posterior(self,x,z,boost=0.):
  B,n,d=x.shape;M=self.masks.to(x);zz=z[:,None,:].expand(-1,n,-1)
  s=self.score(torch.cat((x,zz),-1)).squeeze(-1);xa=x[:,:,None,:].expand(-1,-1,n,-1);xb=x[:,None,:,:].expand(-1,n,-1,-1)
  q=self.pair(torch.cat((xa,xb),-1)).squeeze(-1);q=(q+q.transpose(1,2))/2;logq=F.logsigmoid(q)
  agg=torch.einsum('kn,bnd->bkd',M,x);dec=self.dec(torch.cat((agg,z[:,None,:].expand(-1, M.shape[0],-1)),-1))
  pm=torch.einsum('ki,kj->kij',M,M)*(1-torch.eye(n,device=x.device));lw=torch.einsum('bn,kn->bk',s,M)+.5*torch.einsum('bij,kij->bk',logq,pm)
  if boost: lw=lw+boost*M[:,0]*M[:,1]
  pi=torch.softmax(lw,-1);r=torch.einsum('bk,kn->bn',pi,M);u=torch.einsum('bk,bkd->bd',pi,dec)
  return r,u,(pi*M[:,0]*M[:,1]).sum(-1)
 def forward(self,x):
  B,T,D=x.shape;parts=[]
  for st in range(0,T,self.n):
   c=x[:,st:st+self.n];actual=c.shape[1]
   if actual<self.n:c=torch.cat((c,c[:,-1:].expand(-1,self.n-actual,-1)),1)
   z=c.mean(1);r,u,_=self.posterior(c,z);parts.append(c+r.unsqueeze(-1)*u.unsqueeze(1))
  return torch.cat(parts,1)[:,:T]
class MatchedTransformer(nn.Module):
 def __init__(self,win,out_dim,idea=False,d=64,depth=2):
  super().__init__();self.idea=idea;self.inp=nn.Linear(1,d);self.pos=nn.Parameter(torch.zeros(1,win,d));nn.init.normal_(self.pos,std=.02)
  self.norms=nn.ModuleList([nn.LayerNorm(d) for _ in range(depth)]);self.ff=nn.ModuleList([nn.Sequential(nn.Linear(d,128),nn.GELU(),nn.Linear(128,d)) for _ in range(depth)])
  if idea:self.attn=nn.ModuleList([SubsetAttn(d,4) for _ in range(depth)])
  else:self.attn=nn.ModuleList([nn.MultiheadAttention(d,2,dropout=0.,batch_first=True) for _ in range(depth)])
  self.head=nn.Linear(win*d,out_dim);self.last_pair=None
 def forward(self,x):
  h=self.inp(x.unsqueeze(-1))+self.pos[:,:x.shape[1]]
  for i in range(len(self.ff)):
   q=self.norms[i](h)
   if self.idea:
    a=self.attn[i](q);h=h+a
   else:
    a,_=self.attn[i](q,q,q,need_weights=False);h=h+a
   h=h+self.ff[i](self.norms[i](h))
  return self.head(h.reshape(x.shape[0],-1))
def train_one(kind,cfg,seed,keep=False):
 seed_all(seed);ds=get_dataset('sequence',seed,n_train=NTRAIN,n_test=NTEST)
 net=MatchedTransformer(ds['input_shape'][0],ds['out_dim'],kind=='idea')
 out=train_model(net,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,weight_decay=0.,log=lambda *_:None)
 return (float(out[1]),net,ds) if keep else float(out[1])
def main():
 grid=[{'lr':x,'epochs':EPOCHS} for x in LRS]
 base=sweep_baseline(lambda c:lambda s:train_one('baseline',c,s),grid)
 trials=[]
 for c in grid:
  vals=[train_one('idea',c,s) for s in range(8)];trials.append({'cfg':c,'mean':float(np.mean(vals)),'per_seed':vals})
 best=min(trials,key=lambda x:x['mean']);idea={'mean':float(np.mean(best['per_seed'])),'std':float(np.std(best['per_seed'])),'per_seed':best['per_seed'],'n':8,'best_cfg':best['cfg'],'sweep':trials}
 _,net,ds=train_one('idea',best['cfg'],0,True);dev=next(net.parameters()).device
 with torch.no_grad():
  h=net.inp(ds['xte'][:128].to(dev).unsqueeze(-1))+net.pos[:,:32];q=net.norms[0](h);z=q[:,:4].mean(1);_,_,b=net.attn[0].posterior(q[:,:4],z);_,_,a=net.attn[0].posterior(q[:,:4],z,2.)
 sig={'quantity':'trained-model pair inclusion mass','predicted_direction':'compatibility boost increases joint inclusion','boost':2.,'observed_before':float(b.mean()),'observed_after':float(a.mean()),'observed_change':float((a-b).mean()),'confirmed':bool(float(a.mean())>float(b.mean()))}
 rep=make_report('sequence','transformer_tiny',base,idea,{'mechanism_signature':sig,'selection':{'baseline_best_cfg':base['best_cfg'],'idea_best_cfg':best['cfg'],'shared_lr_union':LRS},'structural_match':'multi-token correlated sequence window'})
 rep['protocol_note']='Eight paired seeds; three shared learning rates; reduced 1200/400 samples and 12 epochs for runtime.'
 Path('bench_report.json').write_text(json.dumps(rep,indent=2));print(json.dumps(rep,indent=2))
if __name__=='__main__':main()