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

DEVICE='cuda' if torch.cuda.is_available() else 'cpu'
try: torch.set_num_threads(4)
except Exception: pass

# Registered bench track; its target is a bounded 2-D energy surface.
TRACK='bounded_energy_regression'
K=64
side=8
axis=torch.linspace(-1.5,1.5,side)
grid=torch.cartesian_prod(axis,axis)
LOGPRIOR=torch.full((K,),-math.log(K))

def seed_all(s):
    np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)

def data(seed):
    return bench_get_dataset(TRACK,seed,n_train=400,n_test=400)

def phi(latent, obs):
    # Frozen neural-style features, shared by both separately trained systems.
    a,b=latent[...,0],latent[...,1]
    o=obs[...,0] if obs.shape[-1] else obs
    z=torch.stack((a,b,o,a*o,b*o,a*a,b*b,o*o,torch.sin(a),torch.cos(b)),-1)
    W=torch.tensor([[.31,-.22,.18,.11,-.14,.25,.07,-.09,.13,.21],
                    [-.12,.28,.24,-.19,.17,.08,-.16,.12,.22,-.11],
                    [.21,.07,-.31,.26,.09,-.12,.19,.15,-.18,.17],
                    [.16,.19,.05,.22,-.27,.13,.11,-.21,.14,.09],
                    [-.24,.11,.16,.08,.21,-.17,.23,.18,.07,-.15],
                    [.09,-.18,.27,.14,.12,.22,-.08,.16,-.11,.24]],device=latent.device)
    bias=torch.tensor([.1,-.2,.05,.3,-.1,.15],device=latent.device)
    return torch.tanh(z@W.t()+bias)

class NeuralPotential(nn.Module):
    def __init__(self,hidden=32):
        super().__init__(); self.net=nn.Sequential(nn.Linear(10,hidden),nn.Tanh(),nn.Linear(hidden,1))
    def forward(self,latent,obs):
        a,b=latent[...,0],latent[...,1]; o=obs[...,0] if obs.shape[-1] else obs
        z=torch.stack((a,b,o,a*o,b*o,a*a,b*b,o*o,torch.sin(a),torch.cos(b)),-1)
        return self.net(z).squeeze(-1)

def convex_loss(theta,latent_obs,latent_true,obs,ret=False):
    # observation is scalar target energy; latent_true is the coordinate pair.
    b=obs.shape[0]; particles=grid.to(obs.device)[None].expand(b,-1,-1); oo=obs[:,None].expand(b,K,1)
    p=phi(particles,oo); target=phi(latent_true,obs)
    logits=-torch.einsum('bkd,d->bk',p,theta)+LOGPRIOR.to(obs.device)
    loss=(torch.einsum('bd,d->b',target,theta)+torch.logsumexp(logits,1)).mean()
    return (loss,p,torch.softmax(logits,1)) if ret else loss

def neural_loss(net,latent_true,obs,ret=False):
    b=obs.shape[0]; particles=grid.to(obs.device)[None].expand(b,-1,-1); oo=obs[:,None].expand(b,K,1)
    logits=-net(particles.reshape(-1,2),oo.reshape(-1,1)).reshape(b,K)+LOGPRIOR.to(obs.device)
    loss=(net(latent_true,obs)+torch.logsumexp(logits,1)).mean()
    return (loss,logits,torch.softmax(logits,1)) if ret else loss

def posterior_mean(model,obs,idea):
    b=obs.shape[0]; p=grid.to(obs.device)[None].expand(b,-1,-1); oo=obs[:,None].expand(b,K,1)
    if idea: logits=-torch.einsum('bkd,d->bk',phi(p,oo),model)+LOGPRIOR.to(obs.device)
    else: logits=-model(p.reshape(-1,2),oo.reshape(-1,1)).reshape(b,K)+LOGPRIOR.to(obs.device)
    return (torch.softmax(logits,1)[...,None]*p).sum(1)

def train_one(seed,lr,idea,epochs=28,return_model=False):
    seed_all(seed); d=data(seed)
    # Bench x is latent coordinate and y is observed energy.
    latent=torch.as_tensor(d['xtr'],device=DEVICE); obs=torch.as_tensor(d['ytr'],device=DEVICE)
    if idea:
        model=nn.Parameter(torch.randn(6,device=DEVICE)*.15); opt=torch.optim.Adam([model],lr=lr)
    else:
        model=NeuralPotential().to(DEVICE); opt=torch.optim.Adam(model.parameters(),lr=lr)
    n=len(obs)
    for _ in range(epochs):
        order=torch.randperm(n,device=DEVICE)
        for j in range(0,n,64):
            ix=order[j:j+64]; opt.zero_grad(set_to_none=True)
            loss=convex_loss(model,None,latent[ix],obs[ix]) if idea else neural_loss(model,latent[ix],obs[ix])
            loss.backward(); opt.step()
    lt=torch.as_tensor(d['xte'],device=DEVICE); ot=torch.as_tensor(d['yte'],device=DEVICE)
    with torch.no_grad(): metric=((posterior_mean(model,ot,idea)-lt)**2).mean().item()
    return (metric,model,d) if return_model else metric

GRID_CFG=[{'lr':.01},{'lr':.03},{'lr':.08}]
def base_fn(cfg): return lambda s: train_one(s,cfg['lr'],False)
def idea_fn(cfg): return lambda s: train_one(s,cfg['lr'],True)

def signature(seed,lr):
    _,_,d=train_one(seed,lr,False,return_model=True); _,theta,_=train_one(seed,lr,True,return_model=True)
    obs=torch.as_tensor(d['yte'][:96],device=DEVICE); true=torch.as_tensor(d['xte'][:96],device=DEVICE)
    with torch.no_grad():
        _,p,r=convex_loss(theta,None,true,obs,True)
        target=phi(true,obs); grad=(target-(r[...,None]*p).sum(1)).mean(0)
        mu=(r[...,None]*p).sum(1); c=p-mu[:,None,:]
        H=torch.einsum('bk,bkd,bke->de',r,c,c)/len(obs)
        mineig=float(torch.linalg.eigvalsh((H+H.T)/2)[0])
    return {'prediction':'trained convex posterior head has covariance Hessian PSD',
            'observed_gradient_norm':float(grad.norm()),
            'observed_covariance_min_eigenvalue':mineig,
            'confirmed':bool(mineig>=-1e-5)}

def main():
    base=sweep_baseline(base_fn,GRID_CFG)
    trials=[{'cfg':c,'result':evaluate(idea_fn(c))} for c in GRID_CFG]
    best=min(trials,key=lambda z:z['result']['mean'])
    rep=make_report(TRACK,'energy_potential_mlp',base,best['result'],{
        'idea_sweep':trials,
        'structural_match':'Registered bounded_energy_regression: 2-D energy surface; latent coordinates are normalized particles conditioned on observed energy.',
        'mechanism_signature':signature(0,best['cfg']['lr'])})
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()