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

SEED = 2237
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_default_dtype(torch.float64)
DEVICE = 'cpu'
T, D, EPS = 18, 5, 1e-3
LAMBDA = 0.015

# Bounded open-loop signal u_k=tanh(theta_k); this is a minimal policy
# parameterization that remains differentiable when environment gradients exist.
def rollout(theta, a):
    u = torch.tanh(theta)
    x = torch.zeros((), dtype=theta.dtype, device=theta.device)
    rows = []
    for k in range(T + 1):
        rows.append(torch.stack((u[k], x, u[k]**2, x**2, u[k] * x)))
        if k < T:
            # candidate nonlinear world models
            x = 0.78 * x + a * torch.tanh(x)**3 + u[k]
    return torch.stack(rows)

def phi(theta, a):
    W = rollout(theta, a)
    G = EPS * torch.eye(D) + W.T @ W / (T + 1)
    sign, ld = torch.linalg.slogdet(G)
    return ld

def objective(theta, models, lam=LAMBDA):
    vals = torch.stack([phi(theta, float(a)) for a in models])
    # mild action penalty; the robust mechanism is the expectation over models
    return vals.mean() - lam * torch.mean(torch.tanh(theta)**2)

INIT = 0.35 * torch.sin(torch.arange(T + 1, dtype=torch.float64) * 1.37 + 0.21)

def optimize(models, steps=700, lr=0.035):
    theta = INIT.clone().requires_grad_()
    opt = torch.optim.Adam([theta], lr=lr)
    for _ in range(steps):
        opt.zero_grad(); loss = -objective(theta, models); loss.backward(); opt.step()
    with torch.no_grad():
        vals = np.array([phi(theta, float(a)).item() for a in models])
        return theta.detach(), vals, objective(theta, models).item()

def finite_difference_check():
    theta = torch.linspace(-.4, .4, T + 1, requires_grad=True)
    a = 0.08
    y = phi(theta, a); g = torch.autograd.grad(y, theta)[0].detach().numpy()
    h = 1e-5; fd=[]
    for i in range(T + 1):
        tp=theta.detach().clone(); tm=theta.detach().clone(); tp[i]+=h; tm[i]-=h
        fd.append((phi(tp,a).item()-phi(tm,a).item())/(2*h))
    fd=np.array(fd)
    return float(np.max(np.abs(g-fd))), float(np.linalg.norm(g-fd)/max(np.linalg.norm(fd),1e-12))

def spread_sweep():
    # Nominal design optimizes only center model; robust design optimizes expectation.
    out=[]
    for s in [0.0, .03, .06, .10, .14, .18]:
        models=np.array([-.10-s, -.10, -.10+s])
        tn, _vn_nom, _ = optimize([-.10])
        vn = np.array([phi(tn, float(a)).item() for a in models])
        tr, vr, _ = optimize(models)
        out.append({'spread':s, 'single_min':float(vn.min()),
                    'robust_min':float(vr.min()),
                    'gain_pct':float(100*(vr.min()-vn.min())/max(abs(vn.min()),1e-9)),
                    'single_mean':float(vn.mean()), 'robust_mean':float(vr.mean())})
    return out

def variance_sweep():
    # At a fixed theta, independently sample model indices and estimate grad Phi.
    theta,_v,_=optimize([-.10,-.16,-.04])
    models=np.array([-.16,-.10,-.04]); probs=np.ones(3)/3
    grads=[]
    for a in models:
        t=theta.clone().requires_grad_(True); g=torch.autograd.grad(phi(t,float(a)),t)[0]
        grads.append(g.detach().numpy())
    grads=np.array(grads); true=grads.mean(0)
    single_var=float(np.mean((grads-true)**2))
    rows=[]
    rng=np.random.default_rng(SEED+9)
    for B in [1,2,4,8,16,32,64]:
        errs=[]
        for _ in range(3000):
            ix=rng.integers(0,3,size=B); est=grads[ix].mean(0)
            errs.append(np.mean((est-true)**2))
        v=float(np.mean(errs)); rows.append({'B':B,'mse':v,'B_times_mse':B*v,'ratio_to_B1':v/max(single_var, 1e-30)})
    return rows

def main():
    fd_abs, fd_rel=finite_difference_check()
    spread=spread_sweep(); variance=variance_sweep()
    result={'seed':SEED,'device':DEVICE,'T':T,'epsilon':EPS,
            'finite_difference_max_abs':fd_abs,'finite_difference_relative_l2':fd_rel,
            'spread_sweep':spread,'variance_sweep':variance}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__=='__main__': main()