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

SEEDS = tuple(range(8))
# Union of all tried lrs is shared by both systems; baseline's decisive knob is attention type,
# and the standard Transformer has no additional exposed method knob in this fixed harness.
GRID = [{'lr': 1e-3, 'epochs': 18}, {'lr': 3e-3, 'epochs': 18}, {'lr': 6e-3, 'epochs': 18}]

class DyadicMixer(nn.Module):
    def __init__(self, d, nhead=2):
        super().__init__()
        self.d = d
        self.nhead = nhead
        self.qkv = nn.Linear(d, 3*d)
        self.out = nn.Linear(d, d)
        self.norm = nn.LayerNorm(d)
        self.ff = nn.Sequential(nn.Linear(d, 128), nn.ReLU(), nn.Linear(128, d))
        self.ffnorm = nn.LayerNorm(d)
        self.alpha = nn.Parameter(torch.tensor(0.1))
        self.beta = nn.Parameter(torch.tensor(1.0))

    def forward(self, x):
        # x is B,L,d. Boundary ranges are clipped and divided by valid counts.
        h = self.norm(x)
        v = self.qkv(h).chunk(3, dim=-1)[2]
        B,L,D = v.shape
        pref = torch.cat((torch.zeros(B,1,D,device=v.device,dtype=v.dtype), v.cumsum(1)), 1)
        out = torch.zeros_like(v)
        K = max(1, int(math.log2(L)))
        idx = torch.arange(L, device=v.device)
        for k in range(K):
            lo, hi = 2**k, min(2**(k+1), L)
            # left indices [i-hi, i-lo), right [i+lo, i+hi)
            la = (idx-hi).clamp(0,L); lb = (idx-lo).clamp(0,L)
            ra = (idx+lo).clamp(0,L); rb = (idx+hi).clamp(0,L)
            lc = (lb-la).clamp(min=1).to(v.dtype)[None,:,None]
            rc = (rb-ra).clamp(min=1).to(v.dtype)[None,:,None]
            left = (pref[:,lb]-pref[:,la])/lc
            right = (pref[:,rb]-pref[:,ra])/rc
            # k'=k+1 gives the stated 1/[k'(k'+1) log 2] mass.
            a = 1.0/((k+1)*(k+2)*math.log(2.0))
            out = out + a*(left+right)
        z = x + self.alpha * self.out(out)
        z = z + self.beta * self.ff(self.ffnorm(z))
        return z

class TailTransformer(nn.Module):
    def __init__(self, win=32, d=64, depth=2, idea=True):
        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)
        if idea:
            self.layers = nn.ModuleList([DyadicMixer(d) for _ in range(depth)])
        else:
            layer = nn.TransformerEncoderLayer(d, nhead=2, dim_feedforward=128,
                batch_first=True, dropout=0.0, activation='relu')
            self.enc = nn.TransformerEncoder(layer, depth)
        self.head = nn.Linear(win*d,1)
    def forward(self,x):
        h=self.inp(x.unsqueeze(-1))+self.pos[:,:x.shape[1]]
        if self.idea:
            for layer in self.layers: h=layer(h)
        else: h=self.enc(h)
        return self.head(h.reshape(x.shape[0],-1))

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

def run_one(kind, cfg, seed, capture=False):
    seed_all(seed)
    ds=get_dataset('sequence', seed, n_train=400, n_test=400)
    net=TailTransformer(win=ds['input_shape'][0], idea=(kind=='idea'))
    net, metric, hist=train_model(net, ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=128, log=lambda *_: None)
    if capture:
        # Trained-model signature: compare observed influence of distant versus near input
        # coordinates by finite perturbations on the held-out task model.
        dev=next(net.parameters()).device; x=ds['xte'][:128].to(dev)
        with torch.no_grad():
            base=net(x).squeeze(-1)
            near=x.clone(); near[:,1]+=0.1
            far=x.clone(); far[:,-1]+=0.1
            dn=(net(near).squeeze(-1)-base).abs().mean().item()
            df=(net(far).squeeze(-1)-base).abs().mean().item()
        return float(metric), {'near_influence':dn,'far_influence':df,'far_near_ratio':df/(dn+1e-12)}
    return float(metric)

def main():
    # Math sanity check is included before training: tail and dyadic-mass ratios.
    rs=np.arange(1,2000000,dtype=np.float64); kval=1/((rs+2)*np.log(rs+2)**2)
    math_rows=[]
    for R in [8,32,128,512]:
        tail=float(kval[R:].sum()+1/np.log(2000002)); math_rows.append({'R':R,'ratio':tail/(1/np.log(R+2))})
    def base_factory(cfg): return lambda s: run_one('baseline',cfg,s)
    baseline=sweep_baseline(base_factory, GRID)
    # Idea uses exactly the same 3-point union and is evaluated on all paired seeds.
    idea_trials=[]
    for cfg in GRID:
        r=evaluate(lambda s,cfg=cfg: run_one('idea',cfg,s), seeds=tuple(range(4)))
        idea_trials.append({'cfg':cfg,'mean':r['mean']})
    best=min(idea_trials,key=lambda z:z['mean'])['cfg']
    idea=evaluate(lambda s: run_one('idea',best,s), seeds=SEEDS)
    # Signature from two independently trained systems at their selected settings.
    _, bsig=run_one('baseline',baseline['best_cfg'],0,True)
    _, isig=run_one('idea',best,0,True)
    signature={'prediction':'dyadic tail should preserve more distant influence than standard local/attention-free mixing',
               'baseline_observed':bsig,'idea_observed':isig,
               'predicted_far_near_ratio_idea_gt_baseline':True,
               'confirmed': bool(isig['far_near_ratio'] > bsig['far_near_ratio'])}
    report=make_report('sequence','transformer_tiny',baseline,idea,{'mechanism_signature':signature,
        'math_sanity':math_rows,'idea_sweep':idea_trials,
        'track_justification':'Sequence forecast has multi-token correlations and is the mandated structural match for attention/SSM ideas.'})
    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()