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

SEED = 3082
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.set_num_threads(4)
try:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if device.type == "cuda":
        torch.cuda.set_device(0)
        _ = torch.zeros(1, device=device)
except Exception:
    device = torch.device("cpu")

# A small robust-control Bellman envelope.  The action is a in [-1,1],
# and the adversarial disturbance is w in {-d,+d}.
def q_exact(x, a, d=0.18):
    z1 = -(x + a - d) ** 2 - 0.1 * a ** 2
    z2 = -(x + a + d) ** 2 - 0.1 * a ** 2
    return torch.minimum(z1, z2)

def exact_target(x, n=4001):
    a = torch.linspace(-1., 1., n, device=x.device)
    return q_exact(x[:, None], a[None, :]).amax(1)

def math_check():
    x = torch.linspace(-1.2, 1.2, 401)
    a = torch.linspace(-1., 1., 10001)
    vals = q_exact(x[:, None], a[None, :])
    da = float(a[1] - a[0])
    L = float((vals[:, 1:] - vals[:, :-1]).abs().max() / da)
    rows = []
    for M in [5, 9, 17, 33, 65]:
        grid = torch.linspace(-1., 1., M)
        disc = q_exact(x[:, None], grid[None, :]).amax(1)
        err = float((vals.amax(1) - disc).abs().max())
        rows.append({"M": M, "delta": 2.0/(M-1), "max_error": err,
                     "L_delta": L * 2.0/(M-1)})
    # max is monotone in every branch: increasing one branch cannot decrease output
    branches = torch.randn(200, 11)
    base = branches.max(1).values
    bumped = branches.clone(); bumped[:, 4] += torch.rand(200) * 2.0
    monotone_violations = int((bumped.max(1).values < base - 1e-7).sum())
    slopes = np.array([r["max_error"] for r in rows])
    deltas = np.array([r["delta"] for r in rows])
    slope = float(np.polyfit(np.log(deltas), np.log(slopes + 1e-12), 1)[0])
    return {"action_lipschitz_estimate": L, "discretization": rows,
            "loglog_error_slope": slope,
            "max_monotonicity_violations": monotone_violations}

class Baseline(nn.Module):
    def __init__(self, width=64):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(1, width), nn.Tanh(),
                                 nn.Linear(width, width), nn.Tanh(),
                                 nn.Linear(width, 1))
    def forward(self, x):
        return self.net(x[:, None]).squeeze(1)

class Envelope(nn.Module):
    def __init__(self, M=17, width=32, tau=0.03):
        super().__init__()
        self.M, self.tau = M, tau
        self.actions = torch.linspace(-1., 1., M)
        self.branches = nn.ModuleList([
            nn.Sequential(nn.Linear(1, width), nn.Tanh(), nn.Linear(width, width),
                          nn.Tanh(), nn.Linear(width, 1)) for _ in range(M)])
        # Learned branch penalties, initialized near the known running cost.
        self.penalty = nn.Parameter(0.1 * self.actions.square())
    def forward(self, x, hard=False):
        ys = torch.stack([net(x[:, None]).squeeze(1) for net in self.branches], 1)
        q = ys - self.penalty[None, :]
        return q.max(1).values if hard else self.tau * torch.logsumexp(q / self.tau, 1)

def train(model, xtr, ytr, steps=1400):
    model.to(device)
    opt = torch.optim.Adam(model.parameters(), lr=2e-3)
    for step in range(steps):
        ix = torch.randint(0, len(xtr), (128,), device=device)
        pred = model(xtr[ix])
        loss = (pred-ytr[ix]).square().mean()
        opt.zero_grad(); loss.backward(); opt.step()
    return model

def experiment():
    gen = torch.Generator(device=device).manual_seed(SEED)
    xtr = (torch.rand(1024, generator=gen, device=device)*2.4-1.2)
    ytr = exact_target(xtr)
    xte = torch.linspace(-1.2, 1.2, 801, device=device)
    yte = exact_target(xte)
    results = {}
    for name, model in [("baseline", Baseline()), ("envelope", Envelope())]:
        train(model, xtr, ytr)
        with torch.no_grad():
            pred = model(xte, hard=True) if name == "envelope" else model(xte)
            mse = float((pred-yte).square().mean())
            mae = float((pred-yte).abs().mean())
            worst = float((pred-yte).abs().max())
        results[name] = {"mse": mse, "mae": mae, "max_abs_error": worst}
    # Refinement uses the exact finite action envelope, isolating the claimed grid effect.
    refine = {}
    for M in [5, 9, 17, 33, 65]:
        grid = torch.linspace(-1., 1., M, device=device)
        with torch.no_grad():
            p = q_exact(xte[:, None], grid[None, :]).amax(1)
            refine[str(M)] = float((p-yte).abs().max())
    results["refinement_max_error"] = refine
    return results

if __name__ == "__main__":
    out = {"device": str(device), "math": math_check(), "experiment": experiment()}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))