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

SEED=2049
np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type=='cuda':
        torch.zeros(1,device=device).sum().item()
except Exception:
    device=torch.device('cpu')

def directions(n,d,rng):
    q=rng.normal(size=(n,d)); return q/np.linalg.norm(q,axis=1,keepdims=True)

def wos_ball(x, K, M, d, lam=0., delta=1e-4, corrected=False, seed=0):
    rng=np.random.default_rng(seed); x=np.asarray(x); B=len(x)
    z=np.repeat(x[:,None,:],K,axis=1); t=np.zeros((B,K)); y=np.zeros((B,K))
    for j in range(M):
        r=1.-np.linalg.norm(z,axis=2)
        active=r>delta
        dt=r*r/((d if corrected else 2*d))
        y += np.exp(-lam*t)*dt # h1=1, h0=0
        u=directions(B*K,d,rng).reshape(B,K,d)
        z += r[...,None]*u; t += dt
        if not active.any(): break
    # zero boundary data, approximate truncation by zero
    return y.mean(axis=1), y, t

def wos_square(x,K,M,delta=2e-3,seed=0):
    rng=np.random.default_rng(seed); x=np.asarray(x); B=len(x); d=2
    z=np.repeat(x[:,None,:],K,axis=1); t=np.zeros((B,K)); y=np.zeros((B,K))
    for j in range(M):
        r=np.minimum(z,1-z).min(axis=2); dt=r*r/(2*d)
        f=2*(z[:,:,0]*(1-z[:,:,0])+z[:,:,1]*(1-z[:,:,1]))
        y += f*dt
        u=directions(B*K,d,rng).reshape(B,K,d); z += r[...,None]*u; t += dt
        if np.max(r)<delta: break
    return y.mean(axis=1)

def analytic_square(x):
    return x[:,0]*(1-x[:,0])*x[:,1]*(1-x[:,1])

def analytic_ball(x): return (1-(x*x).sum(axis=1))/(2*x.shape[1])

class MLP(nn.Module):
    def __init__(self,d):
        super().__init__(); self.net=nn.Sequential(nn.Linear(d,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1))
    def forward(self,x): return self.net(x).squeeze(-1)

def train_wos(idea=True, steps=350, d=2, K=8):
    torch.manual_seed(SEED); rng=np.random.default_rng(SEED+3)
    model=MLP(d).to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    start=time.time()
    for s in range(steps):
        x=rng.uniform(.03,.97,size=(64,d))
        if idea: y=wos_square(x,K,25,seed=SEED+s)
        else: y=analytic_square(x) # ideal exact target is a favorable standard supervised reference
        xt=torch.tensor(x,dtype=torch.float32,device=device); yt=torch.tensor(y,dtype=torch.float32,device=device)
        loss=((model(xt)-yt)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
    grid=np.stack(np.meshgrid(np.linspace(.03,.97,25),np.linspace(.03,.97,25),indexing='ij'),-1).reshape(-1,d)
    with torch.no_grad(): pred=model(torch.tensor(grid,dtype=torch.float32,device=device)).cpu().numpy()
    rel=np.sqrt(np.mean((pred-analytic_square(grid))**2))/np.sqrt(np.mean(analytic_square(grid)**2))
    return float(rel),time.time()-start

def main():
    # Prediction 1: exact Brownian ball exit time is r^2/d; supplied formula predicts r^2/(2d).
    dvals=[1,2,4,8]; r=.37; exit_rows=[]
    for d in dvals:
        supplied=r*r/(2*d); exact=r*r/d
        exit_rows.append({'d':d,'predicted_supplied':supplied,'exact_MC_theory':exact,'ratio_exact/supplied':exact/supplied})
    # Direct one-step isotropic identity verifies E|z+rU|^2=|z|^2+r^2 and Poisson exit-time equation.
    rng=np.random.default_rng(SEED); z=np.array([.2,-.1]); rr=.4; u=directions(500000,2,rng)
    radial_increment=float(np.mean(np.sum((z+rr*u)**2,axis=1)-np.sum(z*z)))
    # Prediction 2: with h1=1 and zero boundary, infinite-depth WOS converges to half the exact ball solution under supplied dt.
    x=np.array([[0.,0.],[.3,0.],[.0,.6]])
    exact=analytic_ball(x); rows=[]
    for M in [2,4,8,16,32,64]:
        est,_,_=wos_ball(x,30000,M,2,seed=SEED+M)
        rows.append({'M':M,'mean_relative_to_exact':float(np.mean(est/exact)),'abs_bias_center':float(abs(est[0]-exact[0]))})
    # Prediction 3: MC variance of the mean scales as 1/K.
    x0=np.array([[.2,.15]])
    _, paths,_=wos_ball(x0,20000,32,2,seed=77)
    var1=float(np.var(paths[0],ddof=1)); varrows=[]
    for K in [1,2,4,8,16,32,64,128]:
        # grouped means from common samples; observed K*Var(mean) should be constant
        m=paths[0][:len(paths[0])//K*K].reshape(-1,K).mean(1)
        var=float(np.var(m,ddof=1)); varrows.append({'K':K,'variance':var,'K_times_variance':K*var,'ratio_to_K1':var/var1})
    # small comparison: WOS stochastic target versus exact-target MLP (same model/steps).
    idea_rel,idea_sec=train_wos(True); base_rel,base_sec=train_wos(False)
    out={'device':str(device),'exit_time_check':exit_rows,'sphere_second_moment_increment':radial_increment,'predicted_increment':rr*rr,'finite_depth_source_check':rows,'variance_scaling':{'single_path_variance':var1,'rows':varrows},'training_comparison':{'wos_target_relative_rmse':idea_rel,'exact_target_relative_rmse':base_rel,'wos_seconds':idea_sec,'exact_seconds':base_sec},'note':'The authoritative formula uses r^2/(2d); standard Brownian motion uses r^2/d, so the source estimator is systematically mis-scaled by approximately 1/2.'}
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()