import json
import math
import random
from dataclasses import dataclass

import numpy as np
import torch

SEED = 2686
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.set_num_threads(4)


@dataclass
class HybridZonotope:
    """Small hybrid-zonotope container: c + G xi + H beta, with optional coupling."""
    c: np.ndarray
    G: np.ndarray
    H: np.ndarray
    E: np.ndarray | None = None
    F: np.ndarray | None = None
    b: np.ndarray | None = None

    def affine(self, A, B=None, offset=None):
        A = np.asarray(A)
        c = A @ self.c
        G = A @ self.G
        H = A @ self.H
        if B is not None:
            c = c + np.asarray(B) @ np.asarray(offset)
        return HybridZonotope(c, G, H, self.E, self.F, self.b)

    def interval(self):
        rad = np.abs(self.G).sum(axis=1) + np.abs(self.H).sum(axis=1)
        return self.c - rad, self.c + rad


def relu_interval_zonotope(z: HybridZonotope):
    """Sound interval relaxation of ReLU, retaining a zonotope approximation."""
    lo, hi = z.interval()
    out_lo, out_hi = np.maximum(lo, 0), np.maximum(hi, 0)
    c = (out_lo + out_hi) / 2
    G = np.diag((out_hi - out_lo) / 2)
    H = np.zeros((len(c), 0))
    return HybridZonotope(c, G, H)


def scalar_interval_policy(lo, hi, w1, b1, w2, b2):
    """Interval bound for a one-hidden-layer scalar ReLU policy."""
    # Each hidden preactivation is affine in scalar x.
    p0, p1 = w1[:, 0] * lo + b1, w1[:, 0] * hi + b1
    pl, ph = np.minimum(p0, p1), np.maximum(p0, p1)
    rl, rh = np.maximum(pl, 0), np.maximum(ph, 0)
    terms_l = np.where(w2[0] >= 0, w2[0] * rl, w2[0] * rh)
    terms_h = np.where(w2[0] >= 0, w2[0] * rh, w2[0] * rl)
    return float(terms_l.sum() + b2[0]), float(terms_h.sum() + b2[0])


def exact_radius(lam, r0, gamma, n):
    """Exact reachable radius for x+=lam*x+w, |w|<=gamma."""
    if abs(lam - 1.0) < 1e-12:
        return r0 + n * gamma
    return abs(lam) ** n * r0 + gamma * sum(abs(lam) ** i for i in range(n))


def mechanism_sweeps():
    r0, gamma, n = 0.2, 0.03, 8
    lam_values = [0.0, 0.2, 0.5, 0.8, 1.0, 1.2]
    rows = []
    for lam in lam_values:
        observed = exact_radius(lam, r0, gamma, n)
        predicted = (abs(lam) ** n * r0 + gamma * sum(abs(lam) ** i for i in range(n)))
        rows.append({"lambda": lam, "observed_rN": observed, "predicted_rN": predicted,
                     "abs_error": abs(observed - predicted)})

    # At fixed lambda<1, the terminal radius is affine in gamma with known slope.
    lam, r0, n = 0.6, 0.2, 7
    gammas = [0.0, 0.01, 0.02, 0.04, 0.08]
    slope = sum(lam ** i for i in range(n))
    scaling = [{"gamma": g, "observed_rN": exact_radius(lam, r0, g, n),
                "predicted_rN": lam ** n * r0 + g * slope} for g in gammas]

    # Constraint boundary: r_N <= R gives gamma* exactly.
    lam, r0, n, R = 0.7, 0.2, 6, 0.30
    geom = sum(lam ** i for i in range(n))
    gamma_star = (R - lam ** n * r0) / geom
    boundary = [{"gamma": g, "violation": max(0.0, exact_radius(lam, r0, g, n) - R)}
                for g in [0.0, gamma_star * 0.99, gamma_star, gamma_star * 1.01]]

    # Contraction boundary: disturbance-free radius decays iff |lambda|<1.
    contraction = [{"lambda": lam, "ratio_rN_r0": exact_radius(lam, 0.2, 0.0, 10) / 0.2,
                    "predicted_contracts": abs(lam) < 1.0}
                   for lam in [0.8, 0.99, 1.0, 1.01, 1.2]]
    return {"radius_recurrence": rows, "disturbance_linear_scaling": scaling,
            "constraint_boundary": {"predicted_gamma_star": gamma_star, "samples": boundary},
            "contraction_boundary": contraction}


def rollout_loss(k, disturbances, x0=0.8, horizon=8):
    x = torch.tensor(x0, dtype=torch.float32)
    loss = 0.0
    for t in range(horizon):
        loss = loss + x * x
        x = (1.0 - k) * x + disturbances[t]
    return loss + 2.0 * x * x


def reach_loss(k, gamma, x0=0.8, radius0=0.08, horizon=8, limit=0.45):
    """Differentiable interval certificate for x+=x+u+w, u=-k*x."""
    center = torch.tensor(x0, dtype=torch.float32)
    radius = torch.tensor(radius0, dtype=torch.float32)
    total = 0.0
    for _ in range(horizon):
        total = total + torch.relu(torch.abs(center) + radius - limit) ** 2
        center = (1.0 - k) * center
        radius = torch.abs(1.0 - k) * radius + gamma
        total = total + 0.5 * torch.relu(radius - limit) ** 2
    total = total + 5.0 * torch.relu(torch.abs(center) + radius - 0.25) ** 2
    # contraction deficit, active when the closed-loop gain is not below one
    total = total + 0.2 * torch.relu(torch.abs(1.0 - k) - 0.9) ** 2
    return total


def train_controllers(gamma, steps=500):
    # Same scalar controller and initialization; baseline sees nominal trajectories only.
    torch.manual_seed(SEED)
    raw_b = torch.nn.Parameter(torch.tensor(-0.2))
    opt_b = torch.optim.Adam([raw_b], lr=0.035)
    nominal = torch.zeros(8)
    for _ in range(steps):
        opt_b.zero_grad()
        k = torch.sigmoid(raw_b) * 2.0
        loss = rollout_loss(k, nominal)
        loss.backward(); opt_b.step()
    kb = float((torch.sigmoid(raw_b) * 2).detach())

    torch.manual_seed(SEED)
    raw_r = torch.nn.Parameter(torch.tensor(-0.2))
    opt_r = torch.optim.Adam([raw_r], lr=0.035)
    for _ in range(steps):
        opt_r.zero_grad()
        k = torch.sigmoid(raw_r) * 2.0
        loss = reach_loss(k, gamma)
        loss.backward(); opt_r.step()
    kr = float((torch.sigmoid(raw_r) * 2).detach())

    def cert(k):
        radii = [exact_radius(abs(1-k), 0.08, gamma, t) for t in range(9)]
        return max(radii), radii[-1], max(0.0, radii[-1] - 0.45)
    return {"gamma": gamma, "baseline_k": kb, "reach_k": kr,
            "baseline_certificate": cert(kb), "reach_certificate": cert(kr),
            "baseline_nominal_loss": float(rollout_loss(torch.tensor(kb), nominal)),
            "reach_training_loss": float(reach_loss(torch.tensor(kr), gamma))}


def main():
    results = {"seed": SEED, "mechanism": mechanism_sweeps(),
               "training": [train_controllers(g) for g in [0.0, 0.12, 0.30, 0.60]]}
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()