import json, math, random
import numpy as np
import torch
from torch import nn

SEED=467
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DT=torch.float64

# Finite latent space makes the partition function exact (no Monte Carlo ambiguity).
xvals=torch.tensor([-2.,-1.,0.,1.,2.],dtype=DT)
prior=torch.tensor([.08,.20,.44,.20,.08],dtype=DT); prior/=prior.sum()
sigma=.65

def make_data(n, seed):
    g=torch.Generator().manual_seed(seed)
    ix=torch.multinomial(prior,n,replacement=True,generator=g)
    x=xvals[ix]
    y=x+sigma*torch.randn(n,generator=g,dtype=DT)
    return x,y

def base_features(x,y):
    # frozen nonlinear feature map; feature dimension is deliberately small
    z=torch.stack([x, y, x*y, x*x, y*y, torch.sin(y), torch.cos(y)],-1)
    W=torch.tensor([[.7,-.2,.4,.1,-.3,.5,.2],[-.4,.8,.2,-.5,.3,.1,-.6],
                    [.3,.2,-.7,.6,.1,-.4,.5],[.5,.1,.3,.4,-.6,.2,.7],
                    [-.2,.6,.5,.2,.4,-.5,.1], [.4,-.3,.1,-.2,.7,.6,-.4]],dtype=DT)
    b=torch.tensor([.1,-.2,.05,.3,-.1,.15],dtype=DT)
    return torch.tanh(z@W.T+b)

def phi_particles(y):
    # [B,K,D], with each observation paired to all prior-supported x
    B=y.numel(); K=xvals.numel()
    return base_features(xvals[None,:].expand(B,K), y[:,None].expand(B,K)).reshape(B,K,-1)

def exact_loss(theta, x, y):
    ph=phi_particles(y); target=base_features(x,y)
    logits=-torch.einsum('bkd,d->bk',ph,theta)+torch.log(prior)[None,:]
    return (torch.einsum('d,bd->b',theta,target)+torch.logsumexp(logits,1)).mean()

def posterior(theta,y):
    ph=phi_particles(y)
    return torch.softmax(-torch.einsum('bkd,d->bk',ph,theta)+torch.log(prior)[None,:],1)

class NeuralScore(nn.Module):
    def __init__(self,d=6,seed=0):
        super().__init__(); torch.manual_seed(seed)
        self.net=nn.Sequential(nn.Linear(7,16),nn.Tanh(),nn.Linear(16,1)).double()
    def forward(self,x,y):
        z=torch.stack([x,y,x*y,x*x,y*y,torch.sin(y),torch.cos(y)],-1)
        return self.net(z).squeeze(-1)

def neural_loss(model,x,y):
    B=y.numel(); xx=xvals[None,:].expand(B,-1); yy=y[:,None].expand(B,xvals.numel())
    s=model(xx.reshape(-1),yy.reshape(-1)).reshape(B,-1)
    st=model(x,y)
    return (st+torch.logsumexp(-s+torch.log(prior)[None,:],1)).mean()

def true_post(x_unused,y):
    # generative posterior p(x|y) under the known Gaussian observation model
    ll=-.5*((y[:,None]-xvals[None,:])/sigma)**2
    return torch.softmax(ll+torch.log(prior)[None,:],1)

def posterior_kl(pred,y):
    t=true_post(None,y); return (t*(torch.log(t+1e-12)-torch.log(pred+1e-12))).sum(1).mean().item()

def coverage(pred,y,level=.9):
    # shortest contiguous x interval containing requested posterior mass, then test true x
    # This is a discrete, transparent calibration proxy.
    t=true_post(None,y); covered=[]
    for i in range(y.numel()):
        p=pred[i]; best=None
        for lo in range(5):
            for hi in range(lo,5):
                if p[lo:hi+1].sum()>=level and (best is None or hi-lo<best[0]): best=(hi-lo,lo,hi)
        # expected coverage under true posterior, not sampled latent, is lower variance
        covered.append(t[i,best[1]:best[2]+1].sum().item())
    return float(np.mean(covered))

def convex_fit(x,y,init):
    # Newton with covariance Hessian; damping only for numerical stability.
    th=init.clone(); history=[]
    for it in range(40):
        th.requires_grad_(True)
        loss=exact_loss(th,x,y); g=torch.autograd.grad(loss,th,create_graph=False)[0]
        ph=phi_particles(y); r=posterior(th,y); mu=(r[:,:,None]*ph).sum(1)
        centered=ph-mu[:,None,:]
        H=torch.einsum('bk,bkd,bke->de',r,centered,centered)/y.numel()
        H=H+1e-7*torch.eye(th.numel(),dtype=DT)
        try: step=torch.linalg.solve(H,g)
        except RuntimeError: step=torch.linalg.lstsq(H,g[:,None]).solution[:,0]
        # backtracking preserves objective descent
        old=loss.item(); alpha=1.
        with torch.no_grad():
            while alpha>1e-5 and exact_loss(th-alpha*step,x,y).item()>old-1e-4*alpha*(g@step).item(): alpha*=.5
            th=(th-alpha*step).detach()
        history.append(old)
        if g.norm().item()<2e-8: break
    return th,history

def run():
    # Core calculus check at a random parameter and random direction.
    x,y=make_data(64,11); theta=torch.randn(6,dtype=DT)*.3; theta.requires_grad_()
    L=exact_loss(theta,x,y); grad=torch.autograd.grad(L,theta,create_graph=True)[0]
    H=torch.autograd.functional.hessian(lambda t: exact_loss(t,x,y),theta).detach().numpy()
    eig=np.linalg.eigvalsh((H+H.T)/2)
    # finite difference directional derivative agrees with Hessian-vector product
    v=torch.randn(6,dtype=DT); v/=v.norm(); eps=1e-4
    tp=(theta.detach()+eps*v).requires_grad_(True)
    tm=(theta.detach()-eps*v).requires_grad_(True)
    gp=torch.autograd.grad(exact_loss(tp,x,y),tp)[0]
    gm=torch.autograd.grad(exact_loss(tm,x,y),tm)[0]
    fd=((gp-gm)/(2*eps)).dot(v).item(); quad=(v@torch.tensor(H)@v).item()
    check={'hessian_min_eigenvalue':float(eig[0]),'gradient_hvp_fd':fd,'gradient_hvp_formula':quad,'hvp_abs_error':abs(fd-quad)}

    trainx,trainy=make_data(256,21); testx,testy=make_data(512,22)
    convex=[]; histories=[]
    for s in [0,1,2,3,4]:
        th,h=convex_fit(trainx,trainy,torch.randn(6,dtype=DT)*2*(s+1)); histories.append({'start':s,'iterations':len(h),'final_train_objective':h[-1]})
        pr=posterior(th,testy); convex.append({'objective':exact_loss(th,testx,testy).item(),'posterior_KL':posterior_kl(pr,testy),'coverage90':coverage(pr,testy),'theta':th.tolist()})
    neural=[]
    for s in [0,1,2,3,4]:
        model=NeuralScore(seed=100+s); opt=torch.optim.Adam(model.parameters(),lr=.025)
        hs=[]
        for it in range(500):
            opt.zero_grad(); loss=neural_loss(model,trainx,trainy); loss.backward(); opt.step(); hs.append(loss.item())
        B=testy.numel(); xx=xvals[None,:].expand(B,-1); yy=testy[:,None].expand(B,xvals.numel())
        with torch.no_grad(): pr=torch.softmax(-model(xx.reshape(-1),yy.reshape(-1)).reshape(B,-1)+torch.log(prior)[None,:],1)
        neural.append({'objective':neural_loss(model,testx,testy).item(),'posterior_KL':posterior_kl(pr,testy),'coverage90':coverage(pr,testy),'final_train_objective':hs[-1]})
    out={'math_check':check,'convex_runs':convex,'convex_histories':histories,'neural_runs':neural,
         'summary':{'convex_objective_mean_sd':[float(np.mean([z['objective'] for z in convex])),float(np.std([z['objective'] for z in convex]))],
                    'neural_objective_mean_sd':[float(np.mean([z['objective'] for z in neural])),float(np.std([z['objective'] for z in neural]))],
                    'convex_KL_mean_sd':[float(np.mean([z['posterior_KL'] for z in convex])),float(np.std([z['posterior_KL'] for z in convex]))],
                    'neural_KL_mean_sd':[float(np.mean([z['posterior_KL'] for z in neural])),float(np.std([z['posterior_KL'] for z in neural]))],
                    'convex_coverage_mean_sd':[float(np.mean([z['coverage90'] for z in convex])),float(np.std([z['coverage90'] for z in convex]))],
                    'neural_coverage_mean_sd':[float(np.mean([z['coverage90'] for z in neural])),float(np.std([z['coverage90'] for z in neural]))]}}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': run()