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

SEED = 2951
EPOCHS = 8
NTR, NTE = 400, 200
LRS = [1e-3, 3e-3, 1e-2]
TAUS = [0.5, 1.0, 2.0]

class DenseLayer(nn.Module):
    def __init__(self, d=64, tau=1.0):
        super().__init__(); self.tau=tau
        self.q=nn.Linear(d,d); self.k=nn.Linear(d,d); self.v=nn.Linear(d,d)
        self.o=nn.Linear(d,d); self.ff=nn.Sequential(nn.Linear(d,128),nn.ReLU(),nn.Linear(128,d))
        self.n1=nn.LayerNorm(d); self.n2=nn.LayerNorm(d)
    def forward(self,x):
        q,k,v=self.q(x),self.k(x),self.v(x)
        a=F.softmax(torch.matmul(q,k.transpose(-1,-2))/(math.sqrt(x.shape[-1])*self.tau),-1)
        z=self.o(torch.matmul(a,v)); return self.n2(self.n1(x+z)+self.ff(self.n1(x+z))),a

class RoleFillerLayer(nn.Module):
    """Role-filler attention: positions are orthonormal roles; projected token states are fillers."""
    def __init__(self, win=32, d=64, tau=1.0):
        super().__init__(); self.win=win; self.tau=tau
        self.role=nn.Parameter(torch.randn(win,d)*0.02)
        self.qf=nn.Linear(d,d); self.ffill=nn.Linear(d,d); self.out=nn.Linear(d,d)
        self.ff=nn.Sequential(nn.Linear(d,128),nn.ReLU(),nn.Linear(128,d))
        self.n1=nn.LayerNorm(d); self.n2=nn.LayerNorm(d)
    def forward(self,x):
        # O[source, role, filler]; query at each target role specifies target role and filler.
        h=x + self.role[:x.shape[1]].unsqueeze(0)
        fillers=self.ffill(h)
        q=self.qf(h)
        scores=torch.einsum('bid,bjd->bij',q,fillers)/(math.sqrt(x.shape[-1])*self.tau)
        a=F.softmax(scores,-1)
        # Extract the target role filler from each retrieved object and rebind to target.
        z=torch.einsum('bij,bjd->bid',a,fillers)
        z=self.out(z)
        u=self.n1(x+z); return self.n2(u+self.ff(u)),a

class Net(nn.Module):
    def __init__(self, kind='dense', win=32, d=64, tau=1.0):
        super().__init__(); self.kind=kind; self.win=win
        self.inp=nn.Linear(1,d); self.pos=nn.Parameter(torch.randn(1,win,d)*.02)
        Layer=DenseLayer if kind=='dense' else RoleFillerLayer
        self.layers=nn.ModuleList([Layer(d=d,tau=tau) if kind=='dense' else Layer(win=win,d=d,tau=tau) for _ in range(2)])
        self.head=nn.Linear(win*d,1); self.last_attention=None
    def forward(self,x):
        h=self.inp(x.unsqueeze(-1))+self.pos[:,:x.shape[1]]; ats=[]
        for layer in self.layers: h,a=layer(h); ats.append(a)
        self.last_attention=ats[-1].detach(); return self.head(h.reshape(x.shape[0],-1))

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

def run(kind, cfg, seed, return_model=False):
    seed_all(seed)
    ds=get_dataset('sequence',seed,n_train=NTR,n_test=NTE)
    net=Net(kind=kind,win=ds['input_shape'][0],tau=cfg['tau'])
    trained, metric, hist=train_model(net,ds,epochs=EPOCHS,lr=cfg['lr'],batch=128,log=lambda *_:None)
    if return_model: return float(metric), trained, ds
    return float(metric)

def main():
    # Baseline sweep includes every idea-side lr and every attention temperature.
    grid=[{'lr':lr,'tau':tau} for lr in LRS for tau in TAUS]
    base=sweep_baseline(lambda c: lambda s: run('dense',c,s),grid)
    # Idea at the best baseline setting and two nearby settings; all are in baseline grid.
    idea_cfgs=[base['best_cfg'], {'lr':1e-3,'tau':1.0}, {'lr':1e-2,'tau':1.0}]
    idea_runs=[]
    for cfg in idea_cfgs:
        r=evaluate(lambda s,c=cfg: run('role',c,s))
        idea_runs.append({'cfg':cfg,'result':r})
    best=min(idea_runs,key=lambda z:z['result']['mean'])
    # NN-scale mechanism signature: measured attention entropy for learned exact-role
    # (diagonal target/source) versus mismatched-role (off-diagonal) queries.
    ent_exact=[]; ent_wrong=[]
    for s in range(8):
        metric,net,ds=run('role',best['cfg'],s,True)
        x=torch.as_tensor(ds['xte'][:64],dtype=torch.float32)
        try:
            device=next(net.parameters()).device
            with torch.no_grad(): net(x.to(device))
            a=net.last_attention.cpu().numpy()
            diag=np.arange(a.shape[1]); exact=a[:,diag,diag]
            wrong=a[:,diag,(diag+1)%a.shape[2]]
            # NN-scale signature: exact role binding should receive more mass
            # than a mismatched adjacent role.
            ent_exact.append(float(exact.mean()))
            ent_wrong.append(float(wrong.mean()))
        except Exception as e:
            print('signature probe failed',repr(e))
    sig={'prediction':'role-selective exact queries should be more concentrated than mismatched queries',
         'predicted_exact_minus_mismatched_mass':0.2,'observed_exact_role_mass':float(np.mean(ent_exact)),
         'observed_mismatched_role_mass':float(np.mean(ent_wrong)),
         'observed_gap':float(np.mean(ent_exact)-np.mean(ent_wrong)),
         'confirmed':bool(ent_exact and np.mean(ent_exact)-np.mean(ent_wrong) >= 0.2),
         'note':'probability mass is measured from attention of trained role-filler models on held-out sequence windows'}
    rep=make_report('sequence','transformer_tiny',base,best['result'],{'idea_sweep':idea_runs,'mechanism_signature':sig})
    # Preserve explicit mechanism_signature at top-level as required by stage-2 runner.
    rep['mechanism_signature']=sig; rep['idea_sweep']=idea_runs
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()