import os, sys, json, math, importlib.util
import numpy as np
import torch
import torch.nn as nn

sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import make_model, make_report
from bench.protocol import permutation_pvalue

HERE = os.path.dirname(os.path.abspath(__file__))
seeds = list(range(8))

spec = importlib.util.spec_from_file_location('pointset_fusion_track', os.path.join(HERE, 'pointset_fusion_track.py'))
track = importlib.util.module_from_spec(spec); spec.loader.exec_module(track)


def kappa(m, beta):
    if m * beta <= 1: raise ValueError('m beta must exceed 1')
    return beta * beta / (8.0 * (m * beta - 1.0) * (2 * m + 1.0))


def prior_score(x, beta=2.0, delta=0.08, correction=True, clip=20.0):
    # x: [batch,m], centered and locally scaled before applying theorem prior.
    m = x.shape[1]
    scale = x.std(dim=1, keepdim=True).clamp_min(0.25)
    a = (x - x.mean(dim=1, keepdim=True)) / scale
    d = a[:, :, None] - a[:, None, :]
    mask = 1.0 - torch.eye(m, device=x.device)[None]
    rep = beta * (d / (d*d + delta*delta) * mask).sum(dim=2)
    corr = -2.0 * kappa(m, beta) * (d * mask).sum(dim=2) if correction else 0.0
    q = rep + corr
    return q.clamp(-clip, clip)


def train_one(seed, cfg, idea):
    torch.manual_seed(seed); np.random.seed(seed)
    d0 = track.get_dataset(seed, n_train=400, n_test=160)
    xtr = torch.tensor(d0['xtr']); ytr = torch.tensor(d0['ytr'])
    xte = torch.tensor(d0['xte']); yte = torch.tensor(d0['yte'])
    net = make_model('mlp_tiny', (6,), 6)
    dev = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net = net.to(dev); xtr=xtr.to(dev); ytr=ytr.to(dev); xte=xte.to(dev); yte=yte.to(dev)
        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
        gen = torch.Generator(device=dev).manual_seed(seed + 991)
        bs=128
        for _ in range(18):
            net.train(); perm=torch.randperm(len(xtr), generator=gen, device=dev)
            for start in range(0, len(xtr), bs):
                ix=perm[start:start+bs]; clean=ytr[ix]; noisy=xtr[ix]
                t=torch.rand((len(ix),1), generator=gen, device=dev)
                sigma=0.08 + 0.34*t; alpha=torch.sqrt(1.0-sigma*sigma)
                xt=alpha*noisy + sigma*torch.randn(xt_shape := noisy.shape, generator=gen, device=dev)
                # Network predicts the diffusion score; denoising estimate is the
                # standard Tweedie reconstruction used by both systems.
                pred=net(xt)
                target=(clean-xt)/(sigma*sigma)
                loss=((pred-target)**2).mean()
                if idea:
                    # Local score prior evaluated on the model's own denoised
                    # point-set estimate, not on an oracle target.
                    xhat=xt + sigma*sigma*pred
                    q=prior_score(xhat, beta=cfg['beta'], delta=cfg['delta'], correction=True)
                    loss=loss + cfg['lam'] * ((pred-q)**2).mean() * (sigma < 0.55).float().mean()
                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(), 10.0); opt.step()
        net.eval()
        with torch.no_grad():
            # Fixed moderate-noise test score, standard denoising MSE.
            sigma=torch.full((len(xte),1), 0.25, device=dev)
            xt=xte + sigma*torch.randn(xte.shape, generator=gen, device=dev)
            pred=net(xt); xhat=xt + sigma*sigma*pred
            metric=((xhat-yte)**2).mean().item()
            # Behavioral signature: normalized short-gap collision rate of the
            # trained system's predictions, measured identically for both sides.
            ss=torch.sort(xhat, dim=1).values
            gaps=ss[:,1:]-ss[:,:-1]; gaps=gaps/(gaps.mean(dim=1,keepdim=True)+1e-6)
            collision=(gaps < 0.10).float().mean().item()
            mean_gap=gaps.mean().item()
        return {'metric': float(metric), 'collision': float(collision), 'mean_gap': float(mean_gap)}
    except RuntimeError:
        # Robust CPU fallback for shared/unsupported CUDA environments.
        if dev == 'cuda':
            torch.cuda.empty_cache()
            old=torch.cuda.is_available
            torch.cuda.is_available=lambda: False
            try: return train_one(seed, cfg, idea)
            finally: torch.cuda.is_available=old
        raise


def eval_cfg(cfg, idea):
    vals=[]; cols=[]; gaps=[]
    for s in seeds:
        r=train_one(s,cfg,idea); vals.append(r['metric']); cols.append(r['collision']); gaps.append(r['mean_gap'])
    return {'mean':float(np.mean(vals)), 'std':float(np.std(vals)), 'per_seed':vals, 'n':len(vals),
            'collision_per_seed':cols, 'collision_mean':float(np.mean(cols)), 'mean_gap':float(np.mean(gaps)), 'cfg':cfg}


def main():
    # Union of all learning rates appears on both sides; weight decay is the
    # baseline's central optimizer knob and is swept equally for every lr.
    lrs=[1e-3, 2e-3, 3e-3]; wds=[0.0, 1e-4]
    base_grid=[{'lr':lr,'weight_decay':wd} for lr in lrs for wd in wds]
    base_sweep=[]
    for cfg in base_grid:
        r=eval_cfg(cfg,False); base_sweep.append({'cfg':cfg,'mean':r['mean']})
    best=min(base_sweep,key=lambda z:z['mean'])['cfg']
    base_full=eval_cfg(best,False)
    # Three idea settings: best baseline lr plus two nearby parity lrs.
    idea_grid=[{'lr':lr,'weight_decay':best['weight_decay'],'lam':lam,'beta':2.0,'delta':0.08}
               for lr,lam in [(best['lr'],0.10),(2e-3,0.10),(1e-3,0.10)]]
    idea_results=[eval_cfg(c,True) for c in idea_grid]
    idea=min(idea_results,key=lambda z:z['mean'])
    diffs=[a-b for a,b in zip(idea['per_seed'],base_full['per_seed'])]
    cmp={'delta_mean':float(np.mean(diffs)), 'idea_wins':sum(x<0 for x in diffs), 'n_pairs':8,
         'per_seed_diffs':[float(x) for x in diffs], 'p_value':float(permutation_pvalue(diffs))}
    if cmp['delta_mean']<0 and cmp['p_value']<0.05: cmp['verdict']='idea better (significant)'; cmp['system_worked']=True
    elif cmp['delta_mean']>0 and cmp['p_value']<0.05: cmp['verdict']='idea worse (significant)'; cmp['system_worked']=False
    else: cmp['verdict']='no significant win'; cmp['system_worked']=False
    # Retest stage-1 prediction at NN scale: prior should reduce predicted
    # short-gap collisions relative to baseline. This is not the primary metric.
    pred_delta=idea['collision_mean']-base_full['collision_mean']
    observed_delta=float(np.mean(np.array(idea['per_seed'])-np.array(base_full['per_seed'])))
    sig={'quantity':'normalized predicted short-gap collision rate',
         'baseline_predicted_collision':base_full['collision_mean'],
         'idea_predicted_collision':idea['collision_mean'],
         'predicted_delta':float(pred_delta),
         'observed_test_mse_delta':observed_delta,
         'confirmed':bool(pred_delta < 0)}
    base_block={'best_cfg':best,'sweep':base_sweep,'full':base_full}
    report=make_report('unordered_pointset_denoising','mlp_tiny',base_block,idea,
                       {'custom_track':{'name':'unordered_pointset_denoising','file':'pointset_fusion_track.py','domain':'point-set-diffusion'},
                        'idea_settings':idea_results,'mechanism_signature':sig})
    report['comparison']=cmp
    with open(os.path.join(HERE,'bench_report.json'),'w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()