import os, sys, json, math
import numpy as np
import torch
import torch.nn as nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import evaluate, sweep_baseline, make_report, permutation_pvalue, get_dataset as bench_get_dataset
from custom_diffusion_track import META

SEEDS = tuple(range(8))
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

class ScoreNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(4+8+1,64), nn.SiLU(), nn.Linear(64,64), nn.SiLU(), nn.Linear(64,8))
    def forward(self, c, x, t):
        return self.net(torch.cat([c,x,t],1))

def fit(seed, lr, epochs=12):
    torch.manual_seed(seed); np.random.seed(seed)
    d = bench_get_dataset('conditional_sequence_diffusion', seed, 400, 100)
    c = torch.tensor(d['xtr'], dtype=torch.float32)
    y = torch.tensor(d['ytr'], dtype=torch.float32).reshape(-1, 8)
    net = ScoreNet()
    try:
        dev = torch.device(DEVICE); net.to(dev); c,y=c.to(dev),y.to(dev)
        opt = torch.optim.Adam(net.parameters(), lr=lr)
        gen = torch.Generator(device=dev).manual_seed(seed+991)
        for ep in range(epochs):
            perm = torch.randperm(len(c), generator=gen, device=dev)
            for ii in range(0,len(c),128):
                ix=perm[ii:ii+128]; t=torch.rand((len(ix),1),generator=gen,device=dev)
                noise=torch.randn((len(ix),8),generator=gen,device=dev)
                xt=torch.sqrt(1-t)*y[ix] + torch.sqrt(t)*noise
                pred=net(c[ix],xt,t)
                loss=((pred-y[ix])**2).mean()
                opt.zero_grad(); loss.backward(); opt.step()
        return net, d
    except RuntimeError:
        net=ScoreNet(); opt=torch.optim.Adam(net.parameters(),lr=lr)
        for ep in range(epochs):
            ix=torch.randperm(len(c), device='cpu')
            for ii in range(0,len(c),128):
                j=ix[ii:ii+128]; t=torch.rand(len(j),1); noise=torch.randn(len(j),8)
                xt=torch.sqrt(1-t)*y[j]+torch.sqrt(t)*noise
                loss=((net(c[j],xt,t)-y[j])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
        return net,d

def density_grid(net,d,n,seed):
    dev=next(net.parameters()).device
    c=torch.tensor(d['xte'][:64],dtype=torch.float32,device=dev); y=torch.tensor(d['yte'],dtype=torch.float32,device=dev).reshape(-1,8)[:64]
    rng=torch.Generator(device=dev).manual_seed(seed+2222); noise=torch.randn(y.shape,generator=rng,device=dev)
    ts=torch.linspace(.01,.99,33,device=dev); vals=[]
    with torch.no_grad():
        for t in ts:
            x=torch.sqrt(1-t)*y+torch.sqrt(t)*noise
            h=1e-3; t2=torch.clamp(t+h,max=.999)
            p1=net(c,x,t.expand(len(c),1)); p2=net(c,x,t2.expand(len(c),1))
            vals.append(float(((p2-p1)/h).pow(2).mean().cpu()))
    a=np.maximum(np.asarray(vals),1e-8); w=np.sqrt(a); cum=np.r_[0,np.cumsum((w[1:]+w[:-1])*np.diff(ts.cpu().numpy())/2)]
    targets=np.linspace(cum[0],cum[-1],n+1); return np.interp(targets,cum,ts.cpu().numpy()), a, ts.cpu().numpy()

def sample_metric(net,d,n,adaptive,seed, collect=False):
    dev=next(net.parameters()).device; c=torch.tensor(d['xte'],dtype=torch.float32,device=dev)
    gen=torch.Generator(device=dev).manual_seed(seed+7788); x=torch.randn((len(c),8),generator=gen,device=dev)
    if adaptive: grid,a,ts=density_grid(net,d,n,seed)
    else: grid=np.linspace(.01,.99,n+1); a=None; ts=None
    net.eval()
    with torch.no_grad():
        for hi,lo in zip(grid[::-1][:-1],grid[::-1][1:]):
            t=torch.full((len(c),1),float(hi),device=dev)
            x0=net(c,x,t); step=(hi-lo)/max(hi,1e-6); x=(1-step)*x+step*x0
    mse=float(((x-torch.tensor(d['yte'],dtype=torch.float32,device=dev).reshape(-1, 8))**2).mean().cpu())
    if collect: return mse,grid,a,ts
    return mse

def train_eval(seed,cfg,adaptive):
    net,d=fit(seed,float(cfg['lr']),int(cfg['epochs']))
    return sample_metric(net,d,int(cfg['steps']),adaptive,seed)

def make_fn(cfg,adaptive):
    return lambda seed: train_eval(seed,cfg,adaptive)

def main():
    # Same union of method hyperparameters is tested on both sides.
    grid=[{'lr':x,'epochs':12,'steps':n} for x in (0.001,0.003,0.01) for n in (8,16)]
    base=sweep_baseline(lambda cfg: make_fn(cfg,False),grid,seeds=SEEDS)
    # idea is evaluated at every baseline candidate, then best selected fairly
    idea_trials=[]
    for cfg in grid:
        r=evaluate(make_fn(cfg,True),SEEDS); idea_trials.append({'cfg':cfg,'mean':r['mean'],'full':r})
    best=min(idea_trials,key=lambda z:z['mean']); idea=best['full']; idea_cfg=best['cfg']
    rep=make_report(META['name'],'score_mlp',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']},idea,extra={})
    # measured signature from one trained benchmark model: allocation vs temporal density,
    # and predicted equalization vs observed one-step proxy on the same learned net.
    net,d=fit(0,float(idea_cfg['lr']),int(idea_cfg['epochs']))
    _,g,a,ts=sample_metric(net,d,int(idea_cfg['steps']),True,0,True)
    widths=np.diff(g); amid=np.interp((g[:-1]+g[1:])/2,ts,a)
    corr=float(np.corrcoef(np.log(widths),np.log(1/np.sqrt(amid)))[0,1])
    uniform=np.linspace(.01,.99,int(idea_cfg['steps'])+1)
    # local model-change proxy: squared temporal change times interval^2
    costs=np.interp((g[:-1]+g[1:])/2,ts,a)*widths**2
    rep['mechanism_signature']={'density_mean':float(a.mean()),'density_max':float(a.max()),'width_inverse_sqrt_corr':corr,'observed_cost_cv':float(np.std(costs)/np.mean(costs)),'predicted_cost_equalization':'low CV under sqrt allocation','confirmed':bool(corr>.8 and np.std(costs)/np.mean(costs)<.5)}
    rep['idea']['trials']=[{'cfg': z['cfg'], 'mean': z['mean']} for z in idea_trials]
    rep['custom_track']={'name':META['name'],'file':'custom_diffusion_track.py','domain':META['domain']}
    rep['comparison']['permutation_pvalue']=permutation_pvalue([i-j for i,j in zip(rep['idea']['per_seed'],rep['baseline']['full']['per_seed'])])
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()