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

OUT=Path('bench_report.json'); SEEDS=tuple(range(8)); SWEEP=(0,1,2,3)
LR_GRID=[1e-3,3e-3,1e-2]; EPOCHS=12; NTR=400; NTE=200
spec=importlib.util.spec_from_file_location('local_track','/home/maxwelhelp/all/math2nn/bench/custom_tracks/conditional_multitoken_diffusion.py')
track=importlib.util.module_from_spec(spec); spec.loader.exec_module(track)

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 ds(seed):
    d0=track.get_dataset(seed,NTR,NTE)
    return {'xtr':torch.as_tensor(d0['xtr'],dtype=torch.float32),'ytr':torch.as_tensor(d0['ytr'],dtype=torch.float32),
            'xte':torch.as_tensor(d0['xte'],dtype=torch.float32),'yte':torch.as_tensor(d0['yte'],dtype=torch.float32),
            'task':'regression','metric':'mse','input_shape':tuple(d0['xtr'].shape[1:]),'out_dim':8}

class MixtureNet(nn.Module):
    def __init__(self, shape, k=3):
        super().__init__(); self.base=make_model('transformer_tiny',shape,16*k); self.k=k
    def forward(self,x):
        q=self.base(x); return q[:,:8*self.k].reshape(-1,8,self.k), q[:,8*self.k:].reshape(-1,8,self.k)

def idea_train(seed,lr,return_sig=False):
    seed_all(seed); d=ds(seed); device='cuda' if torch.cuda.is_available() else 'cpu'
    try: net=MixtureNet(d['input_shape']).to(device)
    except Exception: device='cpu'; net=MixtureNet(d['input_shape']).to(device)
    opt=torch.optim.Adam(net.parameters(),lr=lr); x,y=d['xtr'].to(device),d['ytr'].to(device)
    for _ in range(EPOCHS):
        for ix in torch.randperm(len(x),device=device).split(128):
            mu,lg=net(x[ix]); yy=y[ix].unsqueeze(-1); sd=.12
            lp=torch.log_softmax(lg,2)-.5*((yy-mu)/sd)**2-math.log(sd)-.5*math.log(2*math.pi)
            loss=-torch.logsumexp(lp,2).mean(); opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5); opt.step()
    net.eval()
    with torch.no_grad():
        mu,lg=net(d['xte'].to(device)); pred=(torch.softmax(lg,2)*mu).sum(2); yy=d['yte'].to(device)
        mse=((pred-yy)**2).mean().item(); resid=(yy-pred).cpu().numpy(); v=resid.var()
        kurt=float(np.mean((resid-resid.mean())**4)/(v*v+1e-12)-3); spread=float((mu.max(2).values-mu.min(2).values).mean())
    return (mse,{'mean_component_spread':spread,'residual_excess_kurtosis':kurt}) if return_sig else mse

def base_train(seed,lr):
    seed_all(seed); d=ds(seed); net=make_model('transformer_tiny',d['input_shape'],d['out_dim'])
    return train_model(net,d,epochs=EPOCHS,lr=lr,batch=128)[1]

def main():
    # Normal PIT sanity check for the inverse-CDF transport construction.
    rng=np.random.default_rng(1342); z=rng.normal(size=20000); u=0.5*(1+np.vectorize(math.erf)(z/math.sqrt(2)))
    pit=float(np.max(np.abs(np.sort(u)-(np.arange(len(u))+.5)/len(u))))
    grid=[{'lr':v} for v in LR_GRID]
    base=sweep_baseline(lambda c:lambda s:base_train(s,c['lr']),grid,seeds=SWEEP)
    tried=[]
    for lr in LR_GRID: tried.append(evaluate(lambda s,lr=lr:idea_train(s,lr),seeds=SEEDS)|{'cfg':{'lr':lr,'components':3}})
    best=min(tried,key=lambda r:r['mean']); sigs=[idea_train(s,best['cfg']['lr'],True)[1] for s in SEEDS]
    sig={k:float(np.mean([z[k] for z in sigs])) for k in sigs[0]}; sig.update({'prediction':'multimodal conditional transport has separated learned components','confirmed':sig['mean_component_spread']>0.05})
    rep=make_report('conditional_multitoken_diffusion','transformer_tiny',base,best,extra=sig)
    rep['idea_sweep']=tried; rep['math_check']={'normal_inverse_cdf_pit_sup_deviation':pit,'expected_below':0.02}; rep['custom_track']={'name':'conditional_multitoken_diffusion','file':'/home/maxwelhelp/all/math2nn/bench/custom_tracks/conditional_multitoken_diffusion.py','domain':'diffusion-sampling'}
    rep['protocol_notes']='Matched eight-token diffusion task; independently trained shared transformer backbone; identical test MSE readout; baseline and idea share lr union.'
    OUT.write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()