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

SEEDS=tuple(range(8)); E=8; K=4; EPS=1.0

class RoutedForecaster(nn.Module):
    def __init__(self, win, use_expand, lam, temperature):
        super().__init__(); self.use_expand=use_expand; self.lam=lam
        self.inp=nn.Linear(1,16); self.pos=nn.Parameter(torch.zeros(1,win,16)); nn.init.normal_(self.pos,.02)
        self.enc=nn.TransformerEncoder(nn.TransformerEncoderLayer(16,2,32,batch_first=True,dropout=0.),1)
        self.router=nn.Linear(16,E)
        self.experts=nn.ModuleList([nn.Sequential(nn.Linear(16,16),nn.ReLU(),nn.Linear(16,16)) for _ in range(E)])
        self.head=nn.Linear(16,1); self.temperature=temperature; self.last_sig={}
    def expansion_loss(self,p):
        n=p.shape[0]; rng=np.random.default_rng(1771); vals=[]; hard=[]; bounds=[]
        for m in [K,2*K,4*K,8*K]:
            if m<=n//2:
                for _ in range(1):
                    ix=torch.as_tensor(rng.choice(n,m,False),device=p.device); c=1-torch.prod(1-p[ix],0); C=c.sum()
                    b=EPS*m/(math.log(3*m/K)**2); vals.append(F.relu(torch.as_tensor(b,device=p.device)-C)**2)
                    hard.append(int(p[ix].argmax(1).unique().numel())); bounds.append(b)
        return torch.stack(vals).mean(),hard,bounds
    def forward(self,x):
        h=self.enc(self.inp(x.unsqueeze(-1))+self.pos[:,:x.shape[1]])
        p=F.softmax(self.router(h)/self.temperature,-1); ys=torch.stack([e(h) for e in self.experts],2)
        out=self.head((ys*p.unsqueeze(-1)).sum(2)[:,-1]); load=p.mean((0,1)); lb=E*(load**2).sum()
        ex,hard,bounds=self.expansion_loss(p.reshape(-1,E))
        counts=p.argmax(-1).reshape(-1).bincount(minlength=E).float()
        self.last_sig={'soft_C':float((1-torch.prod(1-p.reshape(-1,E),0)).sum().detach()),'mean_hard_neighborhood':float(np.mean(hard)),'violation_fraction':float(np.mean([a<b for a,b in zip(hard,bounds)])),'load_std':float(counts.std().detach())}
        return out,lb,ex

def train_one(seed,cfg,use_expand,return_sig=False):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    d=get_dataset('sequence',seed,n_train=200,n_test=100)
    # Small CPU path is deterministic and avoids shared-GPU contention; this is the harness fallback path.
    dev=torch.device('cpu')
    net=RoutedForecaster(d['input_shape'][0],use_expand,cfg['lambda'],cfg['temperature']).to(dev)
    opt=torch.optim.Adam(net.parameters(),lr=cfg['lr']); x=d['xtr'].to(dev); y=d['ytr'].reshape(-1,1).to(dev)
    for ep in range(cfg['epochs']):
        for ix in torch.randperm(len(x)).split(128):
            pred,lb,ex=net(x[ix]); warm=min(1.,(ep+1)/max(1,cfg['epochs']//10))
            loss=F.mse_loss(pred,y[ix])+.01*(lb-1)**2+(cfg['lambda']*warm*ex if use_expand else 0.)
            opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad(): pred,_,_=net(d['xte'].to(dev)); metric=float(F.mse_loss(pred,d['yte'].reshape(-1,1).to(dev)))
    return (metric,net.last_sig) if return_sig else metric

def main():
    base_grid=[{'lr':lr,'temperature':t,'lambda':0.,'epochs':3} for lr in [.001,.003] for t in [.7,1.0]]
    base=sweep_baseline(lambda c: lambda s: train_one(s,c,False),base_grid,seeds=(0,1,2,3))
    best=base['best_cfg']; idea_grid=[dict(best, **{'lambda': v}) for v in [.01,.03,.08]]
    # Three-config idea sweep on the same lr/temp selected by the baseline sweep.
    idea_per=[]; idea_sigs=[]
    for s in SEEDS:
        rr=[train_one(s,c,True,True) for c in idea_grid]; j=int(np.argmin([q[0] for q in rr])); idea_per.append(rr[j][0]); idea_sigs.append(rr[j][1])
    idea={'mean':float(np.mean(idea_per)),'std':float(np.std(idea_per)),'per_seed':idea_per,'n':8,'sweep':[{'cfg':c,'mean':float(np.mean([train_one(s,c,True) for s in (0,1,2,3)]))} for c in idea_grid]}
    base_sigs=[train_one(s,best,False,True)[1] for s in SEEDS]
    sig={'prediction':'expansion penalty increases token-group expert neighborhood and reduces load dispersion','baseline_observed':{k:float(np.mean([a[k] for a in base_sigs])) for k in base_sigs[0]},'idea_observed':{k:float(np.mean([a[k] for a in idea_sigs])) for k in idea_sigs[0]},'confirmed':bool(np.mean([a['mean_hard_neighborhood'] for a in idea_sigs])>np.mean([a['mean_hard_neighborhood'] for a in base_sigs]) and np.mean([a['load_std'] for a in idea_sigs])<np.mean([a['load_std'] for a in base_sigs]))}
    rep=make_report('sequence','transformer_tiny',base,idea,sig); rep['idea_hyperparameter_grid']=idea_grid; rep['notes']='Matched compact sequence forecaster; same 8-expert dense MoE and training budget; CPU fallback used to avoid shared GPU contention.'
    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()