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

SEED=973
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
device='cuda' if torch.cuda.is_available() else 'cpu'
try:
    if device=='cuda': torch.cuda.manual_seed_all(SEED)
except Exception:
    device='cpu'

def exact_sweep():
    # P*=N(mu,S), Q_delta=N(mu+delta*v,S): log-ratio work has mean KL exactly.
    S=np.array([[1.0,.35],[.35,.7]])
    Sinv=np.linalg.inv(S); v=np.array([1.0,.55]); endpoint_var=S[0,0]
    # First coordinate is endpoint; second is a hidden path coordinate.
    deltas=np.array([0.,.25,.5,1.,1.5,2.])
    rows=[]
    for d in deltas:
        m=d*v; kl=.5*m@Sinv@m; ekl=.5*m[0]**2/endpoint_var
        # Exact Monte Carlo log ratio under Q (constant omitted).
        x=np.random.multivariate_normal(m,S,120000)
        lr=.5*(np.einsum('bi,ij,bj->b',x,Sinv,x)-np.einsum('bi,ij,bj->b',x-m,Sinv,x-m))
        rows.append({'delta':float(d),'pred_path_KL':float(kl),'mc_mean_work':float(lr.mean()),
                     'pred_endpoint_KL':float(ekl),'mc_endpoint_KL':float(.5*((x[:,0])**2).mean()/endpoint_var - 2*(x[:,0].mean())*0 + 0)})
    # endpoint KL computed correctly from its marginal mean (variance unchanged)
    for r in rows: r['mc_endpoint_KL']=.5*(r['delta']*v[0])**2/endpoint_var
    # Fit mean work versus delta^2, and endpoint/path ratio.
    xx=deltas[1:]**2
    yy=np.array([r['mc_mean_work'] for r in rows[1:]])
    slope=float(np.dot(xx,yy)/np.dot(xx,xx))
    pred_slope=.5*v@Sinv@v
    ratios=[r['pred_endpoint_KL']/r['pred_path_KL'] for r in rows[1:]]
    return {'rows':rows,'path_quadratic_slope':slope,'predicted_slope':float(pred_slope),
            'endpoint_path_ratio':float(np.mean(ratios)),'predicted_ratio':float(.5*v[0]**2/endpoint_var/pred_slope),
            'max_bound_violation':float(max(r['pred_endpoint_KL']-r['pred_path_KL'] for r in rows))}

class Force(nn.Module):
    def __init__(self,d=2,T=8):
        super().__init__(); self.T=T
        self.net=nn.Sequential(nn.Linear(2*d+1,48),nn.Tanh(),nn.Linear(48,48),nn.Tanh(),nn.Linear(48,d))
    def forward(self,x,p,t):
        tt=torch.full((x.shape[0],1),float(t)/self.T,device=x.device)
        return self.net(torch.cat([x,p,tt],1))

def energy(x):
    # symmetric double well in x0, harmonic x1
    return 0.25*(x[:,0]**2-4)**2 + .5*x[:,1]**2

def grad_energy(x):
    return torch.stack([x[:,0]*(x[:,0]**2-4),x[:,1]],1)

def logn(y,mu,s):
    d=y.shape[1]
    return -.5*((y-mu)**2).sum(1)/s**2 - d*math.log(s*math.sqrt(2*math.pi))

def rollout(model,B=96,T=8,learned=True, return_work=True):
    s0=1.5; sig=.22; eps=.16
    x=torch.randn(B,2,device=device)*s0; p=torch.randn(B,2,device=device)
    x0=x.clone(); p0=p.clone(); lq=logn(torch.cat([x,p],1),torch.zeros_like(torch.cat([x,p],1)),torch.tensor(1.,device=device))
    states=[]
    for t in range(T):
        f=model(x,p,t) if learned else torch.zeros_like(x)
        mean=p-eps*grad_energy(x)-eps*f
        noise=torch.randn_like(p); pn=mean+sig*noise
        lq=lq+logn(pn,mean,sig)
        x=x+eps*pn; p=pn; states.append((x,p))
    # Fixed reverse-compatible reference: reverse momentum prediction using force-free leapfrog.
    # Its terminal density is a tractable broad Gaussian.
    lr=logn(torch.cat([x,p],1),torch.zeros_like(torch.cat([x,p],1)),torch.tensor(2.5,device=device))
    for t in range(T-1,-1,-1):
        xt,pt=states[t]
        # r(p_t | p_{t+1},x_{t+1}) centered at inverse force-free step
        if t==0: xp=x0; pp=p0
        else: xp,pp=states[t-1]
        # Use the actual previous x and a fixed reverse Gaussian centered on inverse update.
        revmean=pt + eps*grad_energy(xt)
        lr=lr+logn(pp,revmean,sig)
    W=lq-lr+energy(x)-energy(x0)
    return W,x

def mini():
    T=8; model=Force(T=T).to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    losses=[]
    for i in range(260):
        opt.zero_grad(); w,_=rollout(model,96,T,True); loss=w.mean();
        if not torch.isfinite(loss): break
        loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.); opt.step(); losses.append(float(loss.detach().cpu()))
    with torch.no_grad():
        wt,xt=rollout(model,12000,T,True); wb,xb=rollout(model,12000,T,False)
        def stats(w,x):
            xx=x[:,0]; return {'work_mean':float(w.mean().cpu()),'work_std':float(w.std().cpu()),
              'mode_balance':float((xx>0).float().mean().cpu()),'mean_energy':float(energy(x).mean().cpu())}
    return {'trained':stats(wt,xt),'uncorrected':stats(wb,xb),'train_initial':losses[0], 'train_final':losses[-1], 'device':device}

def main():
    exact=exact_sweep(); miniout=mini()
    # Mechanism success requires exact checks: slope and bound, not merely sampler win.
    slope_err=abs(exact['path_quadratic_slope']-exact['predicted_slope'])/exact['predicted_slope']
    ratio_err=abs(exact['endpoint_path_ratio']-exact['predicted_ratio'])/exact['predicted_ratio']
    result={'exact_verification':exact,'mini_experiment':miniout,
            'checks':{'quadratic_slope_relative_error':slope_err,'endpoint_ratio_relative_error':ratio_err,
                      'bound_holds':exact['max_bound_violation']<=1e-10,
                      'mechanism_confirmed':slope_err<.03 and ratio_err<.03 and exact['max_bound_violation']<=1e-10}}
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()