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

SEED = 2446
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)


def certificate(a, b, k, r, eps, safety=1.0):
    """Scalar system x'=a*x+b*u+w, policy u=k*x, g(x)=safety-|x|."""
    delta = (abs(a) + abs(b) * abs(k)) * r + eps
    return safety - delta, delta


def empirical_deviation(a, b, k, r, eps, n=50000):
    x = np.random.uniform(-r, r, n)
    w = np.random.uniform(-eps, eps, n)
    return float(np.max(np.abs((a + b * k) * x + w)))


def exact_bound_sweep():
    rows = []
    for k in [0.0, 0.2, 0.5, 1.0, 1.5]:
        for eps in [0.01, 0.08, 0.20]:
            pred = certificate(0.7, 0.8, k, 0.12, eps)[1]
            obs = empirical_deviation(0.7, 0.8, k, 0.12, eps)
            rows.append({"k": k, "eps": eps, "predicted_delta": pred,
                         "observed_max": obs, "ratio_observed_to_bound": obs / pred})
    return rows


def boundary_sweep():
    # For g=1-|x|, the predicted disturbance boundary is exactly 1-(|a|+|b||k|)r.
    rows = []
    a, b, r = 0.7, 0.8, 0.12
    for k in [0.0, 0.25, 0.6, 1.0]:
        predicted = 1.0 - (abs(a) + abs(b) * abs(k)) * r
        grid = np.linspace(0, 1.05, 211)
        observed = None
        for eps in grid:
            # use exact worst-case envelope for a noiseless numerical boundary
            worst = (abs(a) + abs(b) * abs(k)) * r + eps
            if worst >= 1.0:
                observed = float(eps)
                break
        rows.append({"k": k, "predicted_eps_boundary": predicted,
                     "observed_eps_boundary": observed,
                     "relative_error": abs(observed-predicted)/predicted})
    return rows


def gain_scaling_sweep():
    rows = []
    a, b, r, eps = 0.7, 0.8, 0.12, 0.05
    base = certificate(a, b, 0.0, r, eps)[0]
    for k in [0.0, 0.25, 0.5, 0.75, 1.0]:
        m = certificate(a, b, k, r, eps)[0]
        predicted_drop = abs(b) * r * abs(k)
        rows.append({"k": k, "margin": m, "margin_drop": base-m,
                     "predicted_linear_drop": predicted_drop})
    return rows


class Policy(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.net = torch.nn.Sequential(torch.nn.Linear(1, 16), torch.nn.Tanh(),
                                       torch.nn.Linear(16, 1))
    def forward(self, x):
        return self.net(x)


def spectral_product(model):
    prod = 1.0
    for layer in model.modules():
        if isinstance(layer, torch.nn.Linear):
            prod *= float(torch.linalg.matrix_norm(layer.weight, ord=2).detach())
    return prod


def train(use_cert, seed=2446, steps=600):
    torch.manual_seed(seed)
    p = Policy()
    opt = torch.optim.Adam(p.parameters(), lr=3e-3)
    a, b, r, eps, safety = 0.7, 0.8, 0.12, 0.05, 1.0
    for _ in range(steps):
        x = torch.linspace(-1, 1, 64).unsqueeze(1)
        u = p(x)
        # reward prefers a nonzero action while penalizing next-state magnitude
        xn = a*x + b*u
        reward_loss = ((u-0.65*x)**2).mean() + 0.08*(xn**2).mean()
        lp = spectral_product(p)
        delta = (a + b*lp)*r + eps
        margin = safety - delta
        cert_loss = 2.0*torch.relu(torch.tensor(0.15-margin))**2 if use_cert else 0.0
        loss = reward_loss + cert_loss
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        xx=torch.linspace(-1,1,1001).unsqueeze(1)
        uu=p(xx); xn=a*xx+b*uu
        nominal_violation=float(torch.relu(torch.abs(xn)-safety).mean())
        action_sensitivity=float((p(xx[1:])-p(xx[:-1])).abs().max()/(xx[1]-xx[0]))
    lp=spectral_product(p)
    m=certificate(a,b,lp,r,eps,safety)[0]
    return {"lipschitz_upper_bound":lp, "certified_margin":m,
            "nominal_violation_mean":nominal_violation,
            "empirical_action_sensitivity":action_sensitivity}


def main():
    result = {"seed": SEED, "bound_sweep": exact_bound_sweep(),
              "boundary_sweep": boundary_sweep(), "gain_scaling": gain_scaling_sweep(),
              "training": {"baseline": train(False), "certified": train(True)}}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()