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

SEED=1286
EPOCHS=35
NTRAIN=400
NTEST=400
NEXACT=32
NBASE=16
LRS=[0.003,0.01,0.03]

def math_check():
    rng=np.random.default_rng(SEED); w=rng.normal(size=(7,2)); w/=np.linalg.norm(w,axis=1,keepdims=True)
    # 2D Koszul matrices K0: R -> R2 and K1: R2 -> R
    nil=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)))}

def features(x,w,b,m):
    s=x@w.T+b
    return torch.relu(s)**m/math.factorial(m)

class Exact(nn.Module):
    def __init__(self,w,b):
        super().__init__(); self.register_buffer('w',torch.tensor(w)); self.register_buffer('b',torch.tensor(b)); self.coef=nn.Parameter(torch.zeros(len(w)))
    def forward(self,x):
        f=features(x,self.w,self.b,1)
        return f @ (self.coef[:,None]*self.w)

class Independent(nn.Module):
    def __init__(self,w,b):
        super().__init__(); self.register_buffer('w',torch.tensor(w)); self.register_buffer('b',torch.tensor(b)); self.coef=nn.Parameter(torch.zeros(len(w),2))
    def forward(self,x):
        return features(x,self.w,self.b,1) @ self.coef

def make_system(kind,seed):
    rng=np.random.default_rng(SEED+int(seed)*7919)
    if kind=='idea':
        n=NEXACT
    else: n=NBASE
    w=rng.normal(size=(n,2)).astype(np.float32); w/=np.linalg.norm(w,axis=1,keepdims=True); b=rng.uniform(-1,1,size=n).astype(np.float32)
    return Exact(w,b) if kind=='idea' else Independent(w,b)

def run(kind,lr,seed, want_sig=False):
    np.random.seed(seed); random.seed(seed); torch.manual_seed(seed)
    raw=track.get_dataset(seed,NTRAIN,NTEST)
    ds=dict(raw)
    for key in ('xtr','ytr','xte','yte'):
        ds[key]=torch.from_numpy(raw[key])
    net=make_system(kind,seed)
    net,metric,hist=train_model(net,ds,epochs=EPOCHS,lr=lr,batch=128,log=lambda *_:None)
    if net is None: return float('nan'), None
    sig=None
    if want_sig:
        dev=next(net.parameters()).device
        x=torch.from_numpy(raw['xte'][:128]).to(dev).requires_grad_(True)
        out=net(x)
        g0=torch.autograd.grad(out[:,0].sum(),x,retain_graph=True)[0]
        g1=torch.autograd.grad(out[:,1].sum(),x)[0]
        curl=g1[:,0]-g0[:,1]
        sig={'observed_curl_rms':float(torch.sqrt(torch.mean(curl.detach()**2))), 'n_probe':128,
             'predicted_curl_rms':0.0, 'confirmed':bool(float(torch.sqrt(torch.mean(curl.detach()**2))) < 1e-6)}
    return float(metric),sig

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

if __name__=='__main__':
    mathres=math_check()
    # Baseline sweep includes every idea learning rate; same epochs and model budget.
    base=sweep_baseline(factory('baseline'),[{'lr':lr} for lr in LRS])
    # Explicit idea sweep at all same settings on full paired seeds.
    idea_trials=[]
    for lr in LRS:
        r=evaluate(lambda seed,lr=lr: run('idea',lr,seed)[0],DEFAULT_SEEDS)
        idea_trials.append({'cfg':{'lr':lr},'result':r})
    best=min(idea_trials,key=lambda z:z['result']['mean'])
    # Signature is measured on the selected trained systems, not an algebra-only toy.
    idea_sig=run('idea',best['cfg']['lr'],0,True)[1]
    rep=make_report('exact_gradient_harmonic_field','fixed_feature_mlp',base,best['result'],{
        'stage1_prediction':'gradient output has zero curl', **idea_sig,
        'baseline_observed_curl_rms':run('baseline',base['best_cfg']['lr'],0,True)[1]['observed_curl_rms'],
        'idea_sweep':idea_trials,'math_check':mathres})
    rep['custom_track']={'name':track.META['name'],'file':'custom_exact_gradient_track.py','domain':track.META['domain']}
    rep['math_check']=mathres
    rep['idea_selected_cfg']=best['cfg']
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))