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

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

def scm(n, sigma_x, seed):
    r = np.random.RandomState(seed)
    x = r.randn(n) * sigma_x
    u = r.randn(n); v = r.randn(n)
    m = x + u
    y = m + x + v
    return x.astype('float32'), m.astype('float32'), y.astype('float32')

# Stage 1: exact Gaussian calculations and sampling verify the claimed selection effect.
def math_check():
    # For M=X+U, Y=M+X+V: E[X|M=m] = s2/(s2+1)m and E[Y|M=m] changes with intervention variance.
    m0 = 1.0
    analytic = {}
    for s in (1.0, 2.0):
        analytic[str(s)] = (1.0 + s*s/(s*s+1.0))*m0
    x,m,y = scm(1500000, 1.0, 11)
    selected = np.abs(m-m0) < .015
    empirical_train = float(y[selected].mean())
    x,m,y = scm(1500000, 2.0, 12)
    selected = np.abs(m-m0) < .015
    empirical_shift = float(y[selected].mean())
    # Joint marginal variance is Var(Y)=4 Var(X)+2, while mediator-only cannot encode this shift.
    var_formula = {str(s): 4*s*s+2 for s in (1.0,2.0)}
    return {'analytic_EY_given_M1': analytic, 'empirical_EY_given_M1': {'sigma1': empirical_train, 'sigma2': empirical_shift}, 'marginal_Y_variance_formula': var_formula}

class Joint(nn.Module):
    def __init__(self):
        super().__init__()
        self.q = nn.Sequential(nn.Linear(1,32), nn.Tanh(), nn.Linear(32,2)) # mu, log sd of X
        self.m = nn.Sequential(nn.Linear(1,32), nn.Tanh(), nn.Linear(32,2)) # conditioned on x
        self.y = nn.Sequential(nn.Linear(2,32), nn.Tanh(), nn.Linear(32,2)) # x,m
    def forward(self,x,m):
        q=self.q(torch.zeros_like(x[:,None])); mp=self.m(x[:,None]); yp=self.y(torch.cat([x[:,None],m[:,None]],1))
        return q,mp,yp

def normal_nll(z, pars):
    mu, logsd = pars[:,0], pars[:,1].clamp(-5,3)
    return .5*((z-mu)/logsd.exp())**2 + logsd + .5*math.log(2*math.pi)

class Base(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(1,32),nn.Tanh(),nn.Linear(32,2))
    def forward(self,m): return self.net(m[:,None])

def train_eval():
    x,m,y=scm(24000,1.0,20)
    xt,mt,yt=scm(120000,2.0,21)
    X=torch.tensor(x,device=device); M=torch.tensor(m,device=device); Y=torch.tensor(y,device=device)
    # Baseline deliberately follows the common mediator-only predictor p(y|m).
    base=Base().to(device); opt=torch.optim.Adam(base.parameters(),lr=.01)
    for _ in range(350):
        ix=torch.randint(0,len(X),(128,),device=device)
        loss=normal_nll(Y[ix],base(M[ix])).mean(); opt.zero_grad(); loss.backward(); opt.step()
    # Joint likelihood q(x) + p(m|x) + p(y|x,m).
    joint=Joint().to(device); opt=torch.optim.Adam(joint.parameters(),lr=.01)
    for _ in range(450):
        ix=torch.randint(0,len(X),(128,),device=device)
        q,mp,yp=joint(X[ix],M[ix])
        loss=(normal_nll(X[ix],q)+normal_nll(M[ix],mp)+normal_nll(Y[ix],yp)).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        MT=torch.tensor(mt,device=device); XT=torch.tensor(xt,device=device); YT=torch.tensor(yt,device=device)
        base_nll=float(normal_nll(YT,base(MT)).mean().cpu())
        _,_,yp=joint(XT,MT)
        joint_nll=float(normal_nll(YT,yp).mean().cpu())
        # Estimated conditional mean at M=1 under the learned conditional model, versus baseline.
        one=torch.ones(2000,device=device); xx=torch.randn(2000,device=device)*2
        # draw mediator noise, then retain a narrow mediator selection
        mm=xx+torch.randn_like(xx)
        keep=(mm-1).abs()<.02
        _,_,pred=joint(xx[keep],mm[keep])
        joint_sel=float(pred[:,0].mean().cpu())
        bsel=float(base(torch.ones(2000,device=device))[:,0].mean().cpu())
        # Marginal predictive variance under target X~N(0,2), via structural heads.
        xs=torch.randn(20000,device=device)*2
        ms=joint.m(xs,)[0] if False else None
        mp=joint.m(xs[:,None]); mdraw=mp[:,0]+torch.randn_like(xs)*mp[:,1].clamp(-5,3).exp()
        yp=joint.y(torch.stack([xs,mdraw],1)); ydraw=yp[:,0]+torch.randn_like(xs)*yp[:,1].clamp(-5,3).exp()
        pred_var=float(ydraw.var().cpu())
    return {'test_shift_sigma2_nll':{'baseline_m_only':base_nll,'joint':joint_nll,'joint_minus_baseline':joint_nll-base_nll},'selected_M1_predicted_mean':{'baseline':bsel,'joint_conditional':joint_sel,'true_sigma2':1.8},'joint_marginal_predicted_Y_variance_sigma2':pred_var,'device':str(device)}

if __name__=='__main__':
    out={'math_check':math_check()}
    try: out['neural_experiment']=train_eval()
    except Exception as e:
        # Explicit CPU fallback as required for shared GPU environments.
        device=torch.device('cpu'); out['neural_experiment']=train_eval(); out['cuda_error']=repr(e)
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))