import sys, json, math, random
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, make_model, train_model, evaluate, sweep_baseline, make_report

SEEDS = tuple(range(8))
# Same union on both sides: baseline is evaluated at every idea LR.
GRID = [{'lr': 0.0015, 'weight_decay': 0.0},
        {'lr': 0.0030, 'weight_decay': 0.0},
        {'lr': 0.0060, 'weight_decay': 0.0}]
EPOCHS = 10
BATCH = 128
DEGREE = 8


def graph(n, degree, seed):
    """Randomly relabelled circulant, exactly degree regular (including self)."""
    if degree % 2 or degree >= n: raise ValueError('degree must be even and < n')
    rng = np.random.default_rng(seed)
    adj = [set() for _ in range(n)]
    for i in range(n):
        for z in range(1, degree//2 + 1):
            adj[i].add((i-z) % n); adj[i].add((i+z) % n)
    p = rng.permutation(n); out = [None] * n
    for i in range(n): out[int(p[i])] = sorted(int(p[j]) for j in adj[i])
    return out


class SparseSelfAttention(nn.Module):
    def __init__(self, d=64, heads=2, degree=8, n=32, seed=0):
        super().__init__(); assert d % heads == 0
        self.d, self.h, self.dk, self.degree = d, heads, d//heads, degree
        a = graph(n, degree, seed)
        self.register_buffer('idx', torch.tensor(a, dtype=torch.long), persistent=False)
        self.qkv = nn.Linear(d, 3*d); self.proj = nn.Linear(d, d)
        self.last_attention = None

    def forward(self, x):
        b,n,d = x.shape; qkv = self.qkv(x).view(b,n,3,self.h,self.dk)
        q,k,v = qkv[:,:,0],qkv[:,:,1],qkv[:,:,2]
        jj = self.idx[:n]
        kg = k[:,jj,:]; vg = v[:,jj,:]
        scores = (q.unsqueeze(2) * kg).sum(-1) / math.sqrt(self.dk)
        att = torch.softmax(scores, dim=2)
        self.last_attention = att.detach()
        y = (att.unsqueeze(-1) * vg).sum(2).transpose(1,2).reshape(b,n,d)
        return self.proj(y)


class SparseLayer(nn.Module):
    def __init__(self, seed):
        super().__init__(); self.attn=SparseSelfAttention(seed=seed, degree=DEGREE)
        self.n1=nn.LayerNorm(64); self.ff=nn.Sequential(nn.Linear(64,128),nn.ReLU(),nn.Linear(128,64)); self.n2=nn.LayerNorm(64)
    def forward(self,x):
        x=self.n1(x+self.attn(x)); return self.n2(x+self.ff(x))


class SparseTransformer(nn.Module):
    def __init__(self, n=32, out_dim=1, seed=0):
        super().__init__(); self.inp=nn.Linear(1,64); self.pos=nn.Parameter(torch.zeros(1,n,64)); nn.init.normal_(self.pos,std=.02)
        self.enc=nn.ModuleList([SparseLayer(seed+i*7919) for i in range(2)]); self.head=nn.Linear(n*64,out_dim)
    def forward(self,x):
        h=self.inp(x.unsqueeze(-1))+self.pos[:,:x.shape[1]]
        for layer in self.enc: h=layer(h)
        return self.head(h.reshape(x.shape[0],-1))


def baseline_fn(cfg):
    def run(seed):
        torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
        d=get_dataset('sequence',seed,n_train=400,n_test=200)
        m=make_model('transformer_tiny',d['input_shape'],d['out_dim'])
        _, metric, _=train_model(m,d,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=cfg['weight_decay'],log=lambda *_:None)
        return float(metric)
    return run


def idea_fn(cfg):
    def run(seed):
        torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
        d=get_dataset('sequence',seed,n_train=400,n_test=200)
        m=SparseTransformer(n=d['input_shape'][0],out_dim=d['out_dim'],seed=seed)
        _, metric, _=train_model(m,d,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=cfg['weight_decay'],log=lambda *_:None)
        return float(metric)
    return run


def signature(seed, cfg):
    """Measured on a trained sparse model: weighted support and boundary growth."""
    d=get_dataset('sequence',seed,n_train=400,n_test=200)
    m=SparseTransformer(n=32,out_dim=1,seed=seed)
    m,_,_=train_model(m,d,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=cfg['weight_decay'],log=lambda *_:None)
    m.eval()
    dev=next(m.parameters()).device
    with torch.no_grad(): m(d['xte'][:64].to(dev))
    # Average effective outgoing support (mass >= 1% of a row) and external boundary.
    U=set(range(4)); layer=m.enc[-1]; att=layer.attn.last_attention.mean((0,3)).cpu().numpy()
    support=[]
    for i in U: support.append({int(layer.attn.idx[i,j]) for j in range(DEGREE) if att[i,j] >= .01})
    ext=set().union(*support)-U
    observed=len(ext)/len(U)
    eps=1.0; predicted=eps/(math.log(3*len(U)/2)**2)
    return {'subset_size':4,'k':2,'epsilon':eps,'predicted_external_ratio':predicted,
            'observed_effective_external_ratio':float(observed),
            'predicted_growth_factor':1+predicted,'observed_growth_factor':1+observed,
            'confirmed': bool(observed+1e-9 >= predicted*.8)}


def main():
    # Baseline sweep uses 4 seeds; final best config is reevaluated on all 8 by harness.
    base=sweep_baseline(baseline_fn,GRID,seeds=(0,1,2,3))
    # Equal-sized idea sweep, selecting on the same four seeds and then full 8.
    tried=[]
    for cfg in GRID:
        r=evaluate(idea_fn(cfg),seeds=(0,1,2,3)); tried.append({'cfg':cfg,'mean':r['mean']})
    best=min(GRID,key=lambda c: next(x['mean'] for x in tried if x['cfg']==c))
    idea=evaluate(idea_fn(best),seeds=SEEDS)
    base['idea_union_sweep']=tried
    sig=signature(0,best)
    rep=make_report('sequence','transformer_tiny',base,idea,{
        'mechanism_signature':sig,
        'implementation':{'degree':DEGREE,'epochs':EPOCHS,'n_train':400,'n_test':200,'matched_lr_grid':GRID},
        'structural_match':'sequence-level multi-token forecast; sparse self-attention is the sole architectural intervention'
    })
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))

if __name__=='__main__': main()