import json, math, itertools
from pathlib import Path
import numpy as np
import torch
from torch import nn

SEEDS = list(range(8))
# Shared search space: every idea setting is evaluated by baseline too.
CONFIGS = [(lr, noise) for lr in (0.01, 0.03, 0.05) for noise in (0.005, 0.02)]
EPOCHS, BATCH = 24, 64

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(10, 32), nn.Tanh(), nn.Linear(32, 1))
    def forward(self, x): return self.net(x)

def data(seed, n=400):
    r = np.random.default_rng(seed)
    x = r.normal(size=(n, 10)).astype('float32')
    y = (np.sin(x[:,0]) + .5*x[:,1]**2 - .3*x[:,2] + .15*x[:,3]*x[:,4] + .1*r.normal(size=n)).astype('float32')
    ix = r.permutation(n); tr, te = ix[:300], ix[300:]
    return torch.tensor(x[tr]), torch.tensor(y[tr,None]), torch.tensor(x[te]), torch.tensor(y[te,None])

def risk_stats(model, center, threshold, pi_floor=1e-3):
    # A is a deliberately predefined unsafe parameter region: distance from
    # the initial basin exceeds threshold. Empirical stationary risk is the
    # tail of a local Gaussian fit to parameter perturbation samples.
    with torch.no_grad():
        d = torch.cat([(p-center[i]).flatten() for i,p in enumerate(model.parameters())])
    radius = float(torch.linalg.vector_norm(d))
    # local quadratic proxy: stationary radius scale estimated from current weights
    scale = max(float(torch.linalg.vector_norm(torch.cat([p.flatten() for p in model.parameters()])))/12., 1e-3)
    z = max((threshold-radius)/scale, -8.)
    pi = float(.5*math.erfc(z/math.sqrt(2)))
    return min(max(pi, pi_floor), 1.-pi_floor), radius

def train(seed, lr, noise, controlled):
    torch.manual_seed(seed); np.random.seed(seed)
    x,y,xe,ye = data(seed)
    model=MLP(); center=[p.detach().clone() for p in model.parameters()]
    # Threshold makes excursions measurable but not trivially impossible.
    threshold=1.8
    # Conservative certificate parameters (m is discrete relaxation proxy).
    m=.12; delta=.18; chi2=9.0; stop=None; unsafe=[]
    gen=torch.Generator().manual_seed(seed+1000)
    for ep in range(EPOCHS):
        perm=torch.randperm(len(x), generator=gen)
        for st in range(0,len(x),BATCH):
            pred=model(x[perm[st:st+BATCH]]); loss=((pred-y[perm[st:st+BATCH]])**2).mean()
            model.zero_grad(); loss.backward()
            use_noise = noise if (not controlled or stop is None) else 0.
            with torch.no_grad():
                for p in model.parameters():
                    if p.grad is not None:
                        p -= lr*p.grad + math.sqrt(2*lr*use_noise)*torch.randn_like(p)
        pi, radius=risk_stats(model, center, threshold)
        bound=min(1., pi+math.sqrt(pi*chi2)*math.exp(-m*(ep+1)))
        unsafe.append(float(radius>threshold))
        if controlled and stop is None and bound <= delta: stop=ep+1
    with torch.no_grad():
        mse=float(((model(xe)-ye)**2).mean())
    return {'mse':mse,'max_unsafe':max(unsafe),'mean_unsafe':float(np.mean(unsafe)),
            'post_stop_unsafe':float(np.mean(unsafe[stop-1:])) if stop else float(np.mean(unsafe)),
            'stop_epoch':stop,'final_radius':risk_stats(model,center,threshold)[1],
            'predicted_bound_final':bound}

def mean_metric(rows,k): return float(np.mean([r[k] for r in rows]))
def perm_p(a,b):
    d=np.array([b[i]-a[i] for i in range(len(a))]); obs=abs(d.mean()); rng=np.random.default_rng(991)
    cnt=0; n=20000
    for _ in range(n):
        if abs((d*rng.choice([-1,1],len(d))).mean())>=obs: cnt+=1
    return float((cnt+1)/(n+1))

def main():
    allres={}
    for cfg in CONFIGS:
        lr,noise=cfg
        base=[train(s,lr,noise,False) for s in SEEDS]
        idea=[train(s,lr,noise,True) for s in SEEDS]
        allres[f'{lr}_{noise}']={'lr':lr,'noise':noise,'baseline':base,'idea':idea,
          'baseline_mse':mean_metric(base,'mse'),'idea_mse':mean_metric(idea,'mse')}
    # baseline sweep chooses lowest mean MSE; idea reports its best setting,
    # both on the same union of configurations.
    best_base=min(allres.values(),key=lambda z:z['baseline_mse'])
    best_idea=min(allres.values(),key=lambda z:z['idea_mse'])
    a=best_base['baseline']; b=best_idea['idea']
    delta=mean_metric(b,'mse')-mean_metric(a,'mse')
    # Signature is behavior measured on trained models, not an identity.
    pred=float(np.mean([r['predicted_bound_final'] for r in b]))
    obs=float(np.mean([r['mean_unsafe'] for r in b]))
    sig={'predicted_final_bound':pred,'observed_mean_unsafe_frequency':obs,
         'ratio_observed_to_bound':obs/max(pred,1e-9),
         'confirmed': bool(obs <= pred*1.25 + .02)}
    report={'track':'local_friedman_mlp_optimizer','task':'regression','metric':'test_mse',
      'baseline_sweep':{k:{'mean_mse':v['baseline_mse'],'lr':v['lr'],'noise':v['noise']} for k,v in allres.items()},
      'best_baseline':{'lr':best_base['lr'],'noise':best_base['noise'],'per_seed':a},
      'idea_best':{'lr':best_idea['lr'],'noise':best_idea['noise'],'per_seed':b},
      'paired_delta_mean_idea_minus_baseline':delta,'permutation_p_value':perm_p([r['mse'] for r in a],[r['mse'] for r in b]),
      'mechanism_signature':sig,'all_configs':allres,
      'custom_track':{'name':'local_friedman_mlp_optimizer','file':'local_nn_bench.py','domain':'tabular/optimizer'},
      'note':'Official bench path /home/maxwelhelp/all/math2nn/bench and README were absent; this is a fallback, not official bench evidence.'}
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps({'delta':delta,'p':report['permutation_p_value'],'signature':sig,'best_base':best_base['baseline_mse'],'best_idea':best_idea['idea_mse']},indent=2))
if __name__=='__main__': main()