import json, math, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

SEED=2433
random.seed(SEED); np.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.empty_cache()
except Exception:
    device='cpu'

a=0.7; sigma=0.35; beta=0.8; dt=0.04

def exact_V(x): return x*x
def drift(x): return -a*x
def exact_r(x): return (beta+2*a-sigma*sigma)*x*x-sigma*sigma

def generator_torch(v,x):
    g=torch.autograd.grad(v.sum(),x,create_graph=True)[0]
    h=torch.autograd.grad(g.sum(),x,create_graph=True)[0]
    return drift(x)*g+0.5*sigma*sigma*h

def analytic_generator(x):
    # For V=x^2: f V' + .5 sigma^2 V'' = -2a x^2 + sigma^2.
    return -2*a*x*x+sigma*sigma

def math_checks():
    xs=np.linspace(-.9,.9,101)
    x=torch.tensor(xs[:,None],dtype=torch.float64,requires_grad=True)
    ad=float(np.max(np.abs(generator_torch(x*x,x).detach().numpy().ravel()-analytic_generator(xs))))
    # Antithetic Monte Carlo makes the leading O(dt) Euler bias measurable.
    rng=np.random.default_rng(SEED); x0=.63; hs=[.04,.02,.01,.005]; errs=[]
    for h in hs:
        z=rng.normal(size=1000000); z=np.concatenate([z,-z])
        xn=x0+drift(x0)*h+sigma*np.sqrt(h)*z
        errs.append(float(abs(np.mean((xn*xn-x0*x0)/h)-analytic_generator(np.array([x0]))[0])))
    return {'autodiff_generator_max_error':ad,'weak_dt':hs,'weak_generator_abs_error':errs,
            'weak_error_ratio_last_first':errs[-1]/errs[0],
            'predicted_euler_bias_ratio':hs[-1]/hs[0]}

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

def make_data(n=512):
    rng=np.random.default_rng(SEED+4)
    x=rng.uniform(-1,1,(n,1)); z=rng.normal(size=(n,1))
    xp=x+drift(x)*dt+sigma*np.sqrt(dt)*z
    r=exact_r(x)*dt
    return tuple(torch.tensor(q,dtype=torch.float32,device=device) for q in (x,xp,r))

def train(mode, lam_final=10., steps=500, threshold=.0007, patience=5, seed=SEED):
    torch.manual_seed(seed); np.random.seed(seed)
    net=Net().to(device); opt=torch.optim.Adam(net.parameters(),lr=3e-3)
    x,xp,r=make_data(); coll=torch.linspace(-1,1,96,device=device).reshape(-1,1)
    lam=0.; ramp_steps=[]; ema=None; good=0; history=[]
    for k in range(steps):
        opt.zero_grad(); vx=net(x); td=(r+math.exp(-beta*dt)*net(xp)-vx).pow(2).mean()
        c=coll.detach().clone().requires_grad_(True); vc=net(c)
        H=generator_torch(vc,c)-beta*vc+exact_r(c); pde=H.pow(2).mean()
        b=net(torch.zeros(16,1,device=device)).pow(2).mean()
        with torch.no_grad():
            t=float(td.detach()); var=float(vx.var().detach())
            ema=t if ema is None else .95*ema+.05*t
        if mode=='td': use=0.
        elif mode=='fixed': use=lam_final
        else:
            if ema < threshold and var < .30: good+=1
            else: good=0
            if good>=patience and lam < lam_final:
                lam=.1 if lam==0 else min(lam_final,2*lam); ramp_steps.append(k+1); good=0
            use=lam
        loss=td+use*pde+2*b; loss.backward(); opt.step()
        if k in (0,4,9,24,49,99,199,299,399,499):
            with torch.no_grad():
                grid=torch.linspace(-1,1,401,device=device).reshape(-1,1)
                rmse=(net(grid)-grid.pow(2)).pow(2).mean().sqrt().item()
            history.append([k+1,float(td.detach()),float(pde.detach()),rmse,use])
    return {'final_td':history[-1][1],'final_pde':history[-1][2],'value_rmse':history[-1][3], 'ramp_steps':ramp_steps,'history':history,'final_lambda':lam}

def gate_sweep():
    # Prediction: stricter TD threshold cannot ramp earlier; threshold zero never ramps.
    out=[]
    for th in [0.,.0003,.0007,.002,.02]:
        z=train('continuation',steps=180,threshold=th,patience=5)
        out.append({'threshold':th,'first_ramp_step':z['ramp_steps'][0] if z['ramp_steps'] else None,'num_ramps':len(z['ramp_steps'])})
    return out

def lambda_sweep():
    # Prediction: lambda=0 is exactly TD-only; fixed strong PDE perturbs TD upward.
    out=[]
    for mode in ['td','fixed','continuation']:
        z=train(mode,steps=300,threshold=(.002 if mode=='continuation' else .0007))
        out.append({'mode':mode,'td':z['final_td'],'pde':z['final_pde'],'rmse':z['value_rmse']})
    return out

def main():
    out={'device':device,'math':math_checks(),'gate_sweep':gate_sweep(),'matched_training':lambda_sweep()}
    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__': main()