"""Certified tube wrapper MVP and numerical verification.

Run with: python3 tube_experiment.py
All experiments are deterministic and use only numpy.
"""
import json
import math
from pathlib import Path
import numpy as np

SEED = 7


def propagate(A_abs, e0, w):
    """Elementwise interval-radius recursion e' = |A|e+w."""
    e = np.asarray(e0, dtype=float).copy()
    hist = [e.copy()]
    for A, wi in zip(A_abs, w):
        e = np.asarray(A) @ e + np.asarray(wi)
        hist.append(e.copy())
    return np.asarray(hist)


def spectral_radius(A):
    return float(np.max(np.abs(np.linalg.eigvals(A))))


def verify_math():
    # Prediction 1: scalar recursion diverges iff gamma*Lambda >= 1.
    # At equality it grows linearly for positive disturbance; above it grows exponentially.
    gammas = np.linspace(0.60, 1.40, 161)
    lam = 1.0
    N = 80
    terminal = []
    for g in gammas:
        e = 0.0
        for _ in range(N):
            e = (g * lam) * e + 0.01
        terminal.append(e)
    terminal = np.asarray(terminal)
    # finite/unstable transition estimated by a large finite-horizon slope jump
    stable_med = np.median(terminal[gammas < .98])
    unstable_med = np.median(terminal[gammas > 1.02])
    boundary_obs = gammas[np.argmin(np.abs(np.log1p(terminal) -
                                           np.median(np.log1p(terminal))))]
    # More meaningful classification: ratio e_N/e_(N-10) exceeds 1.2 on unstable side.
    growth = []
    for g in gammas:
        e1 = e2 = 0.
        for i in range(N):
            e1 = g * e1 + .01
            if i == N - 11:
                old = e1
            if i == N - 1:
                e2 = e1
        growth.append(e2 / max(old, 1e-12))
    growth = np.asarray(growth)
    idx = np.where(growth > 1.20)[0]
    boundary_growth = float(gammas[idx[0]]) if len(idx) else float('nan')

    # Prediction 2: steady radius is w/(1-r), exactly for scalar constant A.
    rs = np.array([0.10, 0.25, 0.40, 0.55, 0.70, 0.80, 0.90])
    w0 = .02
    measured = []
    predicted = []
    for r in rs:
        e = 0.
        for _ in range(300):
            e = r * e + w0
        measured.append(e)
        predicted.append(w0 / (1-r))
    measured, predicted = np.asarray(measured), np.asarray(predicted)
    rel_err = float(np.max(np.abs(measured-predicted)/predicted))

    # Prediction 3: multiplicative conversion is linear in beta and |u|.
    B, u = 0.35, 0.8
    betas = np.array([0., .05, .10, .20, .40, .70])
    wm = np.abs(B) * betas * abs(u)
    slope = float(np.polyfit(betas, wm, 1)[0])
    predicted_slope = abs(B*u)

    # A 2D componentwise tube is also checked against sampled errors.
    A = np.array([[.55, .10], [.05, .35]])
    e0 = np.array([.01, .02])
    w = np.array([.015, .01])
    h = 30
    bound = propagate([np.abs(A)]*h, e0, [w]*h)[-1]
    rng = np.random.default_rng(SEED)
    x = e0.copy()
    max_ratio = 0.
    for _ in range(20000):
        x = A @ x + rng.uniform(-w, w)
        max_ratio = max(max_ratio, float(np.max(np.abs(x) / bound)))
    return {
        "boundary_prediction": 1.0,
        "boundary_observed_growth_test": boundary_growth,
        "boundary_terminal_below_0.98": float(stable_med),
        "boundary_terminal_above_1.02": float(unstable_med),
        "steady_gain_max_relative_error": rel_err,
        "steady_gain_predicted": (w0/(1-rs)).tolist(),
        "steady_gain_measured": measured.tolist(),
        "multiplicative_slope_predicted": predicted_slope,
        "multiplicative_slope_observed": slope,
        "sampled_2d_bound_max_ratio": max_ratio,
    }


def policy(x):
    # A small nonlinear learned-policy surrogate; its Jacobian is available analytically.
    return -0.52*x + 0.08*np.tanh(2*x)


def policy_jac(x):
    return -0.52 + 0.16/(np.cosh(2*x)**2)


def nominal_f(x, u):
    return 0.82*x + u


def nominal_fx(x, u):
    return .82


def nominal_fu(x, u):
    return 1.0


def tube_radii(xhat, horizon, beta, mode):
    """Propagate learned closed-loop Jacobian and additive+mult uncertainty."""
    e = 0.
    out = []
    # horizon-specific model Jacobian evaluated along rollout
    z = xhat
    for j in range(horizon):
        u = policy(z)
        acl = abs(nominal_fx(z,u) + nominal_fu(z,u)*policy_jac(z))
        # Calibrated additive residual plus parameter/gain uncertainty.
        wadd = .018
        wmult = .20 * beta * abs(u)  # |B| pbar |u|
        e = acl*e + wadd + wmult
        out.append(e)
        z = nominal_f(z,u)
    if mode == "constant":
        # Stationary worst-case radius from calibration, deliberately conservative.
        rmax = .46
        wmax = .018 + .20*beta*.65
        steady = wmax/(1-rmax)
        return np.full(horizon, steady)
    if mode == "nominal":
        return np.zeros(horizon)
    return np.asarray(out)


def run_control(mode, beta, seed=SEED, episodes=300):
    rng = np.random.default_rng(seed)
    horizon, limit = 8, 1.0
    violations = fallbacks = 0
    total_abs_x = total_u = 0.
    total_steps = episodes*35
    for _ in range(episodes):
        x = rng.uniform(-.72, .72)
        for t in range(35):
            # The nominal controller is the learned policy, with a tube-aware bounded correction.
            u_nom = policy(x)
            radii = tube_radii(x, horizon, beta, mode)
            # Use the first-step state tightening. If nominal state is outside it, fallback.
            if mode != "nominal" and abs(x) > limit-radii[0]:
                fallbacks += 1
                # conservative bounded fallback, observable in the log
                u = float(np.clip(-.30*x, -.45, .45))
            else:
                # nominal MPC-like correction: robust modes reduce action near boundary
                correction = 0.
                if mode == "horizon":
                    correction = -.10*x*min(1., radii[0]/.20)
                elif mode == "constant":
                    correction = -.06*x
                u = float(np.clip(u_nom + correction, -1., 1.))
            # True plant has randomized gain and additive bounded/noisy disturbance.
            gain = 1.0 + rng.uniform(-beta, beta)
            d = rng.uniform(-.018, .018) + rng.normal(0., .006)
            x = .82*x + gain*u + d
            violations += int(abs(x) > limit)
            total_abs_x += abs(x)
            total_u += abs(u)
    return {
        "violation_rate": violations/total_steps,
        "fallback_rate": fallbacks/total_steps,
        "mean_abs_state": total_abs_x/total_steps,
        "mean_abs_action": total_u/total_steps,
    }


def main():
    math_result = verify_math()
    mini = {}
    for beta in (0.10, 0.25, 0.40):
        mini[str(beta)] = {m: run_control(m, beta) for m in ("nominal", "constant", "horizon")}
    result = {"math_verification": math_result, "mini_experiment": mini}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()