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

TRACK='multitoken_diffusion'
MODEL='mlp_tiny'
DIM=8
BETA=8.0

def floor_estimate(noisy_t, bank, t):
    # Gaussian VP corruption: alpha=exp(-beta*t/2), sigma^2=1-alpha^2.
    alpha=torch.exp(-BETA*t/2).unsqueeze(1)
    sigma2=(1-alpha*alpha).clamp_min(1e-5)
    logits=-((noisy_t[:,None,:]-alpha[:,None,:]*bank[None,:,:])**2).sum(-1)/(2*sigma2)
    pi=torch.softmax(logits, dim=1)
    mu=(pi[:,:,None]*bank[None,:,:]).sum(1)
    var=(pi*((bank[None,:,:]-mu[:,None,:])**2).sum(-1)).sum(1)
    return (alpha[:,0]**2/(sigma2[:,0]**2))*var

def train(seed, cfg, corrected):
    torch.manual_seed(seed); np.random.seed(seed)
    ds=get_dataset(TRACK, seed, n_train=400, n_test=200)
    model=make_model(MODEL, ds['input_shape'], DIM)
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        model=model.to(device)
        opt=torch.optim.Adam(model.parameters(), lr=cfg['lr'])
        x=ds['xtr'].to(device); clean=ds['ytr'].reshape(-1,DIM).to(device)
        bank=clean.detach().clone()
        hist=[]
        for ep in range(cfg['epochs']):
            model.train(); perm=torch.randperm(len(x),device=device); total=0.
            for i in range(0,len(x),128):
                ix=perm[i:i+128]; inp=x[ix]; y=clean[ix]
                noisy=inp[:,:DIM]; t=inp[:,DIM]
                alpha=torch.exp(-BETA*t/2).unsqueeze(1)
                sigma2=(1-alpha*alpha).clamp_min(1e-5)
                target=(alpha*y-noisy)/sigma2
                raw=((model(inp)-target)**2).sum(1)
                if corrected:
                    floor=floor_estimate(noisy,bank,t).detach()
                    loss=(raw-floor).mean()
                else: loss=raw.mean()
                opt.zero_grad(); loss.backward(); opt.step(); total += float(loss.detach())*len(ix)
            hist.append(total/len(x))
        model.eval()
        with torch.no_grad():
            # Posterior-mean denoising readout from predicted score; standard task MSE.
            inp=ds['xte'].to(device); noisy=inp[:,:DIM]; t=inp[:,DIM]
            alpha=torch.exp(-BETA*t/2).unsqueeze(1); sigma2=(1-alpha*alpha).clamp_min(1e-5)
            score=model(inp); pred=(noisy+sigma2*score)/alpha.clamp_min(1e-4)
            metric=float(((pred-ds['yte'].reshape(-1,DIM).to(device))**2).mean())
            # behavior signature: gradient equality and output agreement versus a fresh batch
            raw=((score-((alpha*ds['yte'].reshape(-1,DIM).to(device)-noisy)/sigma2))**2).sum(1)
        return metric, {'final_train_loss':hist[-1], 'mean_raw_dsm':float(raw.mean()), 'pred_std':float(pred.std())}
    except RuntimeError:
        # CPU retry is explicit per bench requirement.
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        device='cpu'; model=make_model(MODEL, ds['input_shape'], DIM).to(device)
        opt=torch.optim.Adam(model.parameters(),lr=cfg['lr']); x=ds['xtr']; clean=ds['ytr'].reshape(-1,DIM); bank=clean.detach().clone()
        for _ in range(cfg['epochs']):
            for i in range(0,len(x),128):
                inp=x[i:i+128]; y=clean[i:i+128]; noisy=inp[:,:DIM]; t=inp[:,DIM]
                a=torch.exp(-BETA*t/2).unsqueeze(1); s=(1-a*a).clamp_min(1e-5); target=(a*y-noisy)/s
                raw=((model(inp)-target)**2).sum(1); loss=(raw-floor_estimate(noisy,bank,t).detach()).mean() if corrected else raw.mean()
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            inp=ds['xte']; noisy=inp[:,:DIM]; t=inp[:,DIM]; a=torch.exp(-BETA*t/2).unsqueeze(1); s=(1-a*a).clamp_min(1e-5)
            pred=(noisy+s*model(inp))/a.clamp_min(1e-4); metric=float(((pred-ds['yte'].reshape(-1,DIM))**2).mean())
        return metric, {'fallback':'cpu'}

def factory(cfg, corrected):
    return lambda seed: train(seed,cfg,corrected)[0]

def main():
    # Shared union: every lr tested by the idea is also swept for baseline.
    grid=[{'lr':1e-3,'epochs':15},{'lr':3e-3,'epochs':15},{'lr':6e-3,'epochs':15}]
    small=tuple(range(8))
    base=sweep_baseline(lambda c:factory(c,False),grid,seeds=small)
    bcfg=base['best_cfg']
    base_full=evaluate(factory(bcfg,False),seeds=small)
    idea_runs=[(c,evaluate(factory(c,True),seeds=small)) for c in grid]
    icfg, idea=min(idea_runs,key=lambda z:z[1]['mean'])
    # NN-scale mechanism signature: measure loss invariance and observed paired outputs.
    probe=[]
    for s in small:
        a,_=train(s,icfg,False); b,_=train(s,icfg,True); probe.append({'seed':s,'baseline_metric':a,'idea_metric':b,'delta':b-a})
    signature={'prediction':'detached floor subtraction leaves parameter gradients unchanged under identical weighting','predicted_vs_observed':{'predicted_gradient_change':0.0,'observed_metric_delta_mean':float(np.mean([p['delta'] for p in probe])),'observed_abs_delta_max':float(np.max(np.abs([p['delta'] for p in probe])))},'trained_model_observations':probe,'confirmed':False}
    report=make_report(TRACK,MODEL,{'best_cfg':bcfg,'sweep':base['sweep'],'full':base_full},idea,signature)
    report['idea_sweep']=[{'cfg':c,'result':r} for c,r in idea_runs]
    report['custom_track']={'name':TRACK,'file':'bench_exp.py','domain':'diffusion-sampling'}
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()