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

# Same architecture as bench.transformer_tiny, with a fixed structured token mask.
class MaskedTransformer(nn.Module):
    def __init__(self, win=32, d=64, depth=2, keep=None):
        super().__init__(); self.win=win
        self.inp=nn.Linear(1,d); self.pos=nn.Parameter(torch.zeros(1,win,d)); nn.init.normal_(self.pos,std=.02)
        layer=nn.TransformerEncoderLayer(d,nhead=2,dim_feedforward=128,batch_first=True,dropout=0.0)
        self.enc=nn.TransformerEncoder(layer,depth); self.head=nn.Linear(win*d,1)
        self.register_buffer('mask', torch.ones(win) if keep is None else torch.tensor(keep,dtype=torch.float32))
    def forward(self,x):
        x=x*self.mask[:x.shape[1]]
        h=self.inp(x.unsqueeze(-1))+self.pos[:,:x.shape[1]]
        return self.head(self.enc(h).reshape(x.shape[0],-1))

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

def interactions(x):
    # Local interaction graph from calibration covariance; diagonal is marginal importance.
    a=x-x.mean(0,keepdims=True); cov=(a.T@a)/max(1,len(a)-1)
    sd=np.sqrt(np.maximum(np.diag(cov),1e-8)); corr=cov/(sd[:,None]*sd[None,:])
    np.fill_diagonal(corr,0); return corr

def magnitude_keep(x,k):
    # Standard activation-magnitude token pruning.
    score=np.mean(np.abs(x),axis=0); ix=np.argsort(score)[-k:]
    m=np.zeros(x.shape[1],dtype=np.float32); m[ix]=1; return m

def boundary_keep(x,k,eta=2,threshold=0.12):
    # Boundary-compressed DP on a chain ordering. Q support contains local, strong
    # interactions; state key is (count,last-eta bits), and distant interactions are truncated.
    q=np.abs(interactions(x)); n=q.shape[0]; adj=(q>=threshold).astype(np.int8); np.fill_diagonal(adj,0)
    # retain only nearest graph-neighbor interactions, making the intended sparse graph explicit
    for i in range(n):
        for j in range(n):
            if abs(i-j)>eta: adj[i,j]=0
    states={(0,0):(0.0,[])}; peaks=[1]
    for i in range(n):
        nxt={}
        for (used,hist),(cost,prefix) in states.items():
            for bit in (0,1):
                if used+bit>k: continue
                # reward important tokens, penalize selected pair interactions locally
                val=cost-bit*float(np.mean(np.abs(x[:,i])))
                for dist in range(1,min(eta,i)+1):
                    if bit and ((hist>>(dist-1))&1) and adj[i-dist,i]:
                        val += 0.20*float(q[i-dist,i])
                nh=((hist<<1)|bit)&((1<<eta)-1)
                key=(used+bit,nh)
                if key not in nxt or val<nxt[key][0]: nxt[key]=(val,prefix+[bit])
        states=nxt; peaks.append(len(states))
    candidates=[v for (u,_),v in states.items() if u==k]
    _,bits=min(candidates,key=lambda z:z[0]); return np.asarray(bits,dtype=np.float32),max(peaks),int(adj.sum()/2)

def run_one(seed,lr,method,epochs=10,n=1200):
    seed_all(seed); ds=get_dataset('sequence',seed,n_train=n,n_test=400); x=ds['xtr'].numpy()
    if method=='baseline': keep=magnitude_keep(x,16); peak=0; edges=0
    else: keep,peak,edges=boundary_keep(x,16,eta=2,threshold=.12)
    # train_model is the canonical benchmark path; only the fixed structured mask differs.
    seed_all(seed)
    net=MaskedTransformer(32,64,2,keep)
    net,metric,hist=train_model(net,ds,epochs=epochs,lr=lr,batch=128,log=lambda *a:None)
    return {'seed':seed,'metric':float(metric),'keep':int(keep.sum()),'peak_states':peak,'graph_edges':edges,
            'retained_fraction':float(keep.mean()),'history_last':float(hist[-1])}

def baseline_sweep(lrs,epochs=10):
    out=[]
    for lr in lrs:
        vals=[run_one(s,lr,'baseline',epochs)['metric'] for s in range(4)]
        out.append({'config':{'lr':lr,'epochs':epochs,'keep':16},'mean':float(np.mean(vals)),'per_seed':vals})
    best=min(out,key=lambda z:z['mean'])
    return {'grid':out,'best':best['config'],'full':{'per_seed':[run_one(s,best['config']['lr'],'baseline',epochs)['metric'] for s in range(8)]}}

def main():
    # Union parity: all idea lrs are also baseline sweep values.
    lrs=[0.0015,0.003,0.006]; epochs=10; t=time.time()
    base=baseline_sweep(lrs,epochs)
    idea_runs=[]
    for lr in lrs:
        rr=[run_one(s,lr,'idea',epochs) for s in range(8)]
        idea_runs.append({'config':{'lr':lr,'epochs':epochs,'eta':2,'keep':16},'mean':float(np.mean([r['metric'] for r in rr])),'runs':rr})
    best=min(idea_runs,key=lambda z:z['mean']); idea={'config':best['config'],'per_seed':[r['metric'] for r in best['runs']], 'details':best['runs']}
    # Behavioral signature: prediction from omitted empirical interaction mass vs observed
    sig=[]
    for r in best['runs']:
        s=r['seed']; ds=get_dataset('sequence',s,n_train=1200,n_test=400); x=ds['xte'].numpy(); q=np.abs(interactions(x));
        # Reconstruct selected masks deterministically and measure trained-model masked sensitivity.
        keep,_,_=boundary_keep(get_dataset('sequence',s,n_train=1200,n_test=400)['xtr'].numpy(),16,2,.12)
        pred=float(q[np.outer(1-keep,1-keep).astype(bool)].mean()) if np.any(keep==0) else 0.0
        seed_all(s); net=MaskedTransformer(32,64,2,keep); net,_,_=train_model(net,ds,epochs=epochs,lr=best['config']['lr'],batch=128,log=lambda *a:None)
        dev=next(net.parameters()).device
        xt=ds['xte'].to(dev)
        with torch.no_grad():
            trained_mask=net.mask.detach().clone()
            net.mask.fill_(1.0)
            full=net(xt)
            net.mask.copy_(trained_mask)
            masked=net(xt)
        obs=float(((full-masked)**2).mean()); sig.append((pred,obs))
    pairs=np.asarray(sig); corr=float(np.corrcoef(pairs.T)[0,1]) if np.std(pairs[:,0])>0 and np.std(pairs[:,1])>0 else 0.0
    extra={'prediction':'omitted-token interaction mass predicts output sensitivity','predicted_mean':float(pairs[:,0].mean()),'observed_mean':float(pairs[:,1].mean()),'correlation':corr,'confirmed':bool(corr>0.3)}
    rep=make_report('sequence','transformer_tiny',base,idea,extra)
    rep['idea_sweep']=[{'config':z['config'],'mean':z['mean'],'per_seed':[r['metric'] for r in z['runs']]} for z in idea_runs]
    rep['runtime_sec']=time.time()-t; rep['structural_match']='multi-token sequence forecasting with token interaction graph'; rep['state_counts']=[r['peak_states'] for r in best['runs']]
    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()