import json
import math
from pathlib import Path
import numpy as np

# Jacobian-Frozen Stable Rollouts: small, reproducible numerical study.
SEED = 7
rng = np.random.default_rng(SEED)


def g(x, u):
    """A small nonlinear neural-state-space-like transition."""
    x = np.asarray(x, dtype=float)
    u = np.asarray(u, dtype=float)
    # Stable rotation plus smooth nonlinear state and input interactions.
    y0 = 0.93 * x[0] + 0.16 * x[1] + 0.08 * np.tanh(x[0] * x[1]) + 0.10 * u[0]
    y1 = -0.16 * x[0] + 0.93 * x[1] + 0.06 * np.tanh(x[0] ** 2) - 0.07 * u[0]
    return np.array([y0, y1])


def jacobian_xu(x, u):
    """Analytic A,B for g, avoiding any hidden finite-difference claim."""
    x0, x1 = x
    q = 1.0 - np.tanh(x0 * x1) ** 2
    A = np.array([
        [0.93 + 0.08 * x1 * q, 0.16 + 0.08 * x0 * q],
        [-0.16 + 0.12 * x0 * (1.0 - np.tanh(x0 ** 2) ** 2), 0.93],
    ])
    B = np.array([[0.10], [-0.07]])
    return A, B


def frozen_parameters(x, u):
    A, B = jacobian_xu(x, u)
    c = g(x, u) - A @ x - B @ u
    return A, B, c


def rollout_nonlinear(x0, us):
    xs = [np.asarray(x0, dtype=float).copy()]
    for u in us:
        xs.append(g(xs[-1], u))
    return np.asarray(xs)


def rollout_frozen(x0, us, clip_radius=None):
    x0 = np.asarray(x0, dtype=float)
    A, B, c = frozen_parameters(x0, us[0])
    rho = max(abs(np.linalg.eigvals(A)))
    if clip_radius is not None and rho > clip_radius:
        # Radial matrix scaling preserves the locally affine direction while
        # imposing the requested spectral-radius monitor threshold.
        A = A * (clip_radius / rho)
        rho = max(abs(np.linalg.eigvals(A)))
        c = g(x0, us[0]) - A @ x0 - B @ us[0]
    xs = [x0.copy()]
    for u in us:
        xs.append(A @ xs[-1] + B @ u + c)
    return np.asarray(xs), rho


def remainder_scaling():
    """Check ||g(x+d)-g(x)-Jd|| = O(||d||^2) around a nonzero point."""
    x = np.array([0.42, -0.31]); u = np.array([0.13])
    A, B = jacobian_xu(x, u)
    direction = np.array([0.8, -0.6]); direction /= np.linalg.norm(direction)
    eps = np.array([0.2, 0.1, 0.05, 0.025, 0.0125])
    errors = []
    for e in eps:
        d = e * direction
        errors.append(np.linalg.norm(g(x + d, u) - g(x, u) - A @ d))
    slope = np.polyfit(np.log(eps), np.log(errors), 1)[0]
    ratios = (np.asarray(errors[:-1]) / np.asarray(errors[1:])).tolist()
    return {"eps": eps.tolist(), "errors": np.asarray(errors).tolist(),
            "loglog_slope": float(slope), "halve_error_ratios": ratios}


def stability_boundary():
    """Linear scalar control check: x_{t+1}=a*x_t has boundary |rho|=1."""
    rows = []
    for a in [0.80, 0.95, 0.99, 1.00, 1.01, 1.10]:
        x = 1e-3
        for _ in range(40):
            x = a * x
        rows.append({"a": a, "rho": abs(a), "growth_factor": abs(x / 1e-3),
                     "decays": bool(abs(x) < 1e-3)})
    return rows


def mini_experiment():
    # Constant but nonzero input makes this a genuine state-space rollout.
    H = 30
    us = np.tile(np.array([[0.12]]), (H, 1))
    records = []
    for radius in [0.03, 0.10, 0.25, 0.50, 0.90]:
        x0 = radius * np.array([0.8, -0.6])
        truth = rollout_nonlinear(x0, us)
        frozen, rho = rollout_frozen(x0, us)
        clipped, clipped_rho = rollout_frozen(x0, us, clip_radius=0.98)
        err = float(np.sqrt(np.mean((frozen - truth) ** 2)))
        clip_err = float(np.sqrt(np.mean((clipped - truth) ** 2)))
        records.append({"radius": radius, "rho_at_start": float(rho),
                        "nonlinear_final_norm": float(np.linalg.norm(truth[-1])),
                        "frozen_final_norm": float(np.linalg.norm(frozen[-1])),
                        "frozen_rmse": err, "clipped_frozen_rmse": clip_err,
                        "frozen_max_norm": float(np.max(np.linalg.norm(frozen, axis=1))),
                        "nonlinear_max_norm": float(np.max(np.linalg.norm(truth, axis=1)))})
    # A rough operation count illustrates why freezing can be cheaper: one
    # nonlinear evaluation per step versus one Jacobian plus affine matvecs.
    return records


def main():
    out = {"seed": SEED, "remainder_scaling": remainder_scaling(),
           "stability_boundary": stability_boundary(),
           "mini_experiment": mini_experiment(),
           "notes": "Frozen rollout uses exact local analytic Jacobian; clipping rescales A and recomputes c."}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()