import sys, json, importlib.util
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 train_model, evaluate, sweep_baseline, make_report, count_params
TRACK='multitoken_diffusion'; DIM=8
_spec=importlib.util.spec_from_file_location('registered_multitoken_diffusion','/home/maxwelhelp/all/math2nn/bench/custom_tracks/multitoken_diffusion.py')
_mod=importlib.util.module_from_spec(_spec); _spec.loader.exec_module(_mod)
class LocalBaseline(nn.Module):
    def __init__(self, hidden=32, radius=1):
        super().__init__(); k=2*radius+1
        self.body=nn.Sequential(nn.Conv1d(2,hidden,k,padding=radius),nn.ReLU(),nn.Conv1d(hidden,hidden,k,padding=radius),nn.ReLU(),nn.Conv1d(hidden,1,1))
    def forward(self,x):
        z=x[:,:DIM].unsqueeze(1); t=x[:,DIM:].unsqueeze(1).expand(-1,1,DIM)
        return self.body(torch.cat([z,t],1)).squeeze(1)
class RGPyramid(nn.Module):
    def __init__(self, hidden=32, radius=1):
        super().__init__(); k=2*radius+1
        self.fine=nn.Sequential(nn.Conv1d(2,hidden,k,padding=radius),nn.ReLU(),nn.Conv1d(hidden,hidden,k,padding=radius),nn.ReLU())
        self.coarse=nn.Sequential(nn.Conv1d(2,hidden,k,padding=radius),nn.ReLU(),nn.Conv1d(hidden,hidden,1),nn.ReLU())
        self.out=nn.Conv1d(2*hidden,1,1)
    def forward(self,x):
        z=x[:,:DIM].unsqueeze(1); t=x[:,DIM:].unsqueeze(1).expand(-1,1,DIM)
        coarse=F.avg_pool1d(torch.cat([z,t],1),2,2); c=F.interpolate(self.coarse(coarse),size=DIM,mode='linear',align_corners=False)
        f=self.fine(torch.cat([z,t],1)); return self.out(torch.cat([f,c],1)).squeeze(1)
def dataset(seed): return _mod.get_dataset(seed,n_train=400,n_test=200)
def run(kind,seed,cfg,detail=False):
    torch.manual_seed(seed); np.random.seed(seed); ds=dataset(seed)
    model=LocalBaseline(cfg['hidden'],1) if kind=='baseline' else RGPyramid(cfg['hidden'],1)
    model,train_metric,_=train_model(model,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,log=lambda *a,**k:None)
    if model is None: raise RuntimeError('bench training returned None')
    # Evaluation on CPU avoids a shared-GPU Conv1d engine failure; both systems use this identically.
    model=model.cpu().eval()
    with torch.no_grad(): pred=model(ds['xte']).numpy()
    metric=float(np.mean((pred-ds['yte'].numpy())**2))
    if not detail:return metric
    x=ds['xte'][:128].clone(); x2=x.clone(); x2[:,0]+=0.5
    with torch.no_grad(): delta=np.abs(model(x2).numpy()-model(x).numpy())
    return metric,{'params':count_params(model),'mean_output_sensitivity':float(delta.mean()),'distant_sensitivity_0_to_7':float(delta[:,7].mean()),'near_sensitivity_0_to_1':float(delta[:,1].mean()),'train_metric':float(train_metric)}
def factory(kind,cfg): return lambda seed: run(kind,seed,cfg)
def main():
    grid=[{'lr':1e-3,'epochs':12,'hidden':32},{'lr':3e-3,'epochs':12,'hidden':32},{'lr':6e-3,'epochs':12,'hidden':32}]
    base=sweep_baseline(lambda c:factory('baseline',c),grid,seeds=(0,1,2,3)); bcfg=base['best_cfg']
    base_full=evaluate(factory('baseline',bcfg),seeds=tuple(range(8)))
    idea_runs=[(c,evaluate(factory('idea',c),seeds=tuple(range(8)))) for c in grid]
    best_cfg,best=min(idea_runs,key=lambda q:q[1]['mean'])
    isig=[run('idea',s,best_cfg,True)[1] for s in range(8)]; bsig=[run('baseline',s,bcfg,True)[1] for s in range(8)]
    im=float(np.mean([q['distant_sensitivity_0_to_7'] for q in isig])); bm=float(np.mean([q['distant_sensitivity_0_to_7'] for q in bsig]))
    signature={'prediction':'coarse pooled conditioning transmits distant-token influence beyond a radius-1 local field','predicted_vs_observed':{'predicted':'idea distant sensitivity > baseline','observed_idea_mean_distant_sensitivity':im,'observed_baseline_mean_distant_sensitivity':bm},'trained_model_observations':isig,'confirmed':bool(im>bm)}
    report=make_report(TRACK,'custom_local_vs_rg_pyramid',{'best_cfg':bcfg,'sweep':base['sweep'],'full':base_full},best,signature)
    report['idea_sweep']=[{'cfg':c,'result':r} for c,r in idea_runs]
    report['custom_track']={'name':TRACK,'file':'/home/maxwelhelp/all/math2nn/bench/custom_tracks/multitoken_diffusion.py','domain':'diffusion-sampling'}
    report['architecture']={'baseline':'local radius-1 convolutional velocity field','idea':'dyadic pooled coarse field plus local fine residual conditioning','baseline_params':count_params(LocalBaseline()),'idea_params':count_params(RGPyramid())}
    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
if __name__=='__main__':main()