import sys, json, math, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
from scipy.stats import qmc, norm
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, sweep_baseline, evaluate, make_report

SEED=2108
TRACK='conditional_multitoken_diffusion'
MODEL='mlp_tiny'
LRS=[1e-3,3e-3,1e-2]
EPOCHS=12
BATCH=64

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 gaussian_corruption(n,d,seed,kind):
    if kind=='iid':
        z=np.random.RandomState(seed).standard_normal((n,d)).astype('float32')
    else:
        # Independent Owen scrambles per epoch/seed; inverse-CDF maps cube to N(0,1).
        p=qmc.Sobol(d=d, scramble=True, seed=int(seed)).random(n)
        z=norm.ppf(np.clip(p,1e-6,1-1e-6)).astype('float32')
    return torch.from_numpy(z)

def run_one(seed,lr,kind,return_model=False):
    seed_all(seed)
    d=get_dataset(TRACK,seed,n_train=400,n_test=200)
    # The custom track exposes an 8-token target; adapt the harness' generic
    # regression reshape back to the declared multi-token output dimension.
    d['ytr']=d['ytr'].reshape(-1,8); d['yte']=d['yte'].reshape(-1,8)
    dev='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net=make_model(MODEL,d['input_shape'],d['out_dim']).to(dev)
        opt=torch.optim.Adam(net.parameters(),lr=lr)
        loss=nn.MSELoss()
        x=d['xtr'].to(dev); y=d['ytr'].to(dev)
        n=len(x)
        for ep in range(EPOCHS):
            # Same examples and targets; only diffusion perturbation differs.
            perm=torch.randperm(n,device=dev)
            xep=x[perm]; yep=y[perm]
            z=gaussian_corruption(n,8,seed*1000+ep,kind).to(dev)
            # Mild additional terminal-noise transport, with observed sigma retained.
            xin=xep.clone(); sig=xin[:,9:10].clamp(.04,1.)
            xin[:,:8]=xin[:,:8] + .18*sig*z
            for j in range(0,n,BATCH):
                opt.zero_grad(set_to_none=True)
                pred=net(xin[j:j+BATCH]); l=loss(pred,yep[j:j+BATCH])
                l.backward(); opt.step()
        with torch.no_grad():
            pred=net(d['xte'].to(dev)); metric=float(loss(pred,d['yte'].to(dev)).cpu())
        if return_model: return metric, net, d
        return metric
    except RuntimeError:
        # Explicit CPU fallback for constrained shared GPU slots.
        torch.cuda.empty_cache()
        osdev=torch.device('cpu')
        seed_all(seed)
        d=get_dataset(TRACK,seed,n_train=400,n_test=200)
        d['ytr']=d['ytr'].reshape(-1,8); d['yte']=d['yte'].reshape(-1,8)
        net=make_model(MODEL,d['input_shape'],d['out_dim']).to(osdev)
        opt=torch.optim.Adam(net.parameters(),lr=lr); loss=nn.MSELoss()
        x=d['xtr']; y=d['ytr']; n=len(x)
        for ep in range(EPOCHS):
            perm=torch.randperm(n); xin=x[perm].clone(); yep=y[perm]
            z=gaussian_corruption(n,8,seed*1000+ep,kind); xin[:,:8]+=0.18*xin[:,9:10].clamp(.04,1.)*z
            for j in range(0,n,BATCH):
                opt.zero_grad(); l=loss(net(xin[j:j+BATCH]),yep[j:j+BATCH]); l.backward(); opt.step()
        with torch.no_grad(): return float(loss(net(d['xte']),d['yte']))

def fn(kind,lr): return lambda s: run_one(s,lr,kind)

def mechanism_signature():
    # Re-test the stage-1 prediction on outputs of trained systems: batch-mean
    # variance should decay faster for scrambled Sobol than IID perturbations.
    seed=77; lr=3e-3; metric,net,d=run_one(seed,lr,'sobol',True)
    dev=next(net.parameters()).device; x=d['xte'][:128].to(dev)
    vals={}
    for kind in ('iid','sobol'):
        means=[]
        for r in range(16):
            xx=x.clone(); z=gaussian_corruption(len(x),8,90000+r,kind).to(dev)
            xx[:,:8]+=0.18*xx[:,9:10].clamp(.04,1.)*z
            with torch.no_grad(): means.append(float(net(xx).mean().cpu()))
        vals[kind]=float(np.std(means,ddof=1))
    # A second size check uses prefixes of power-of-two Sobol sets and IID sets.
    slopes={}
    for kind in ('iid','sobol'):
        ns=[16,32,64,128]; sds=[]
        for n in ns:
            mm=[]
            for r in range(12):
                xx=x[:n].clone(); z=gaussian_corruption(n,8,120000+r,kind).to(dev)
                xx[:,:8]+=0.18*xx[:,9:10].clamp(.04,1.)*z
                with torch.no_grad(): mm.append(float(net(xx).mean().cpu()))
            sds.append(np.std(mm,ddof=1))
        slopes[kind]=float(np.polyfit(np.log2(ns),np.log2(np.maximum(sds,1e-12)),1)[0])
    confirmed=vals['sobol'] < vals['iid'] and slopes['sobol'] < slopes['iid']
    return {'prediction':'scrambled Sobol batch-mean variance lower and slope closer to -1 than IID',
            'trained_model_batch_mean_sd':vals,'observed_log2N_slopes':slopes,
            'variance_ratio_sobol_over_iid':vals['sobol']/vals['iid'], 'confirmed':bool(confirmed)}

def main():
    grid=[{'lr':v} for v in LRS]
    base=sweep_baseline(lambda c: fn('iid',c['lr']),grid)
    # Same union of learning rates; select best idea setting on the same 4-seed sweep.
    idea_trials=[]
    for cfg in grid:
        r=evaluate(fn('sobol',cfg['lr']))
        idea_trials.append({'cfg':cfg,'mean':r['mean'],'std':r['std']})
    best=min(idea_trials,key=lambda x:x['mean'])['cfg']
    idea=evaluate(fn('sobol',best['lr']))
    rep=make_report(TRACK,MODEL,base,idea,{'idea_sweep':idea_trials,'selected_idea_cfg':best,
                                           'mechanism_signature':mechanism_signature(),
                                           'custom_track':{'name':TRACK,'file':'bench/custom_tracks/conditional_multitoken_diffusion.py','domain':'diffusion-sampling'}})
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()