import sys,json,random,math
import numpy as np
import torch
from torch import nn
sys.path.insert(0,'/home/maxwelhelp/all/math2nn')
from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
from bench.protocol import DEFAULT_SEEDS

TRACK='poisson_boundary'; MODEL='scalar_potential_backbone'; EPOCHS=35; NTRAIN=400; NTEST=400
LRS=[0.003,0.01,0.03]

def math_check():
    rng=np.random.default_rng(1286); w=rng.normal(size=(7,2)); w/=np.linalg.norm(w,axis=1,keepdims=True)
    nil=0.0
    for wi in w:
        k0=wi[:,None]; k1=np.array([[-wi[1],wi[0]]]); nil=max(nil,float(np.max(np.abs(k1@k0))))
    x=torch.tensor([[.37,-.21]],dtype=torch.float64,requires_grad=True); wt=torch.tensor(w[:1],dtype=torch.float64); b=torch.tensor([.41],dtype=torch.float64)
    s=x@wt.T+b; y=torch.relu(s)**3/6; g=torch.autograd.grad(y.sum(),x)[0].detach().numpy()[0]
    rhs=(torch.relu(s.detach())**2/2).numpy()[0,0]*w[0]
    return {'nilpotence_max_abs':nil,'derivative_identity_max_abs':float(np.max(np.abs(g-rhs)))}

class ExactPotential(nn.Module):
    def __init__(self,w,b,k=3):
        super().__init__(); self.register_buffer('w',torch.tensor(w)); self.register_buffer('b',torch.tensor(b)); self.coef=nn.Parameter(torch.zeros(len(w))); self.k=k
    def forward(self,x):
        s=x@self.w.T+self.b
        return (torch.relu(s)**self.k/math.factorial(self.k) @ self.coef)[:,None]

class StandardMLP(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(2,16),nn.Tanh(),nn.Linear(16,16),nn.Tanh(),nn.Linear(16,1))
    def forward(self,x): return self.net(x)

def run(kind,lr,seed,signature=False):
    np.random.seed(seed); random.seed(seed); torch.manual_seed(seed)
    d=get_dataset(TRACK,seed,NTRAIN,NTEST); ds=dict(d)
    for q in ('xtr','ytr','xte','yte'):
        ds[q]=d[q] if torch.is_tensor(d[q]) else torch.from_numpy(d[q])
    rng=np.random.default_rng(9001+seed*7919); w=rng.normal(size=(32,2)).astype('float32'); w/=np.linalg.norm(w,axis=1,keepdims=True); b=rng.uniform(-1,1,32).astype('float32')
    net=ExactPotential(w,b) if kind=='idea' else StandardMLP()
    net,metric,_=train_model(net,ds,epochs=EPOCHS,lr=lr,batch=128,log=lambda *_:None)
    if net is None:return float('nan'),None
    sig=None
    if signature:
        dev=next(net.parameters()).device; x=torch.as_tensor(d['xte'][:128],device=dev).requires_grad_(True); out=net(x)
        grad=torch.autograd.grad(out.sum(),x,create_graph=True)[0]
        c0=torch.autograd.grad(grad[:,0].sum(),x,retain_graph=True)[0]; c1=torch.autograd.grad(grad[:,1].sum(),x)[0]
        curl=c1[:,0]-c0[:,1]; val=float(torch.sqrt(torch.mean(curl.detach()**2)))
        sig={'predicted_gradient_curl_rms':0.0,'observed_gradient_curl_rms':val,'n_probe':128,'confirmed':val<1e-5}
    return float(metric),sig

def factory(kind):
    return lambda cfg: (lambda seed: run(kind,float(cfg['lr']),seed)[0])

if __name__=='__main__':
    mc=math_check()
    base=sweep_baseline(factory('baseline'),[{'lr':x} for x in LRS])
    trials=[]
    for x in LRS:
        r=evaluate(lambda seed,x=x:run('idea',x,seed)[0],DEFAULT_SEEDS); trials.append({'cfg':{'lr':x},'result':r})
    best=min(trials,key=lambda z:z['result']['mean']); idea=best['result']
    sig=run('idea',best['cfg']['lr'],0,True)[1]
    bsig=run('baseline',base['best_cfg']['lr'],0,True)[1]
    rep=make_report(TRACK,MODEL,base,idea,{'stage1_prediction':'d squared equals zero for the learned potential gradient','baseline_gradient_curl_rms':bsig['observed_gradient_curl_rms'],**sig,'idea_sweep':trials,'math_check':mc})
    rep['math_check']=mc; rep['idea_selected_cfg']=best['cfg']
    with open('bench_report_registered.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))