import json
import math
import numpy as np


def decimated(A, Q, d):
    n = A.shape[0]
    Ad = np.eye(n)
    Qd = np.zeros_like(Q)
    # A^j Q (A^j)^T, j=0,...,d-1
    Aj = np.eye(n)
    for j in range(d):
        Qd += Aj @ Q @ Aj.T
        Aj = A @ Aj
    return Aj, (Qd + Qd.T) * 0.5


def riccati(A, H, Q, R, d, max_iter=10000, tol=1e-11):
    Ad, Qd = decimated(A, Q, d)
    P = Qd.copy()
    for _ in range(max_iter):
        S = H @ P @ H.T + R
        # Joseph-equivalent form of the stated prior Riccati update
        K = np.linalg.solve(S, H @ P @ Ad.T).T
        Pn = Ad @ P @ Ad.T - K @ S @ K.T + Qd
        Pn = (Pn + Pn.T) * 0.5
        if np.max(np.abs(Pn - P)) < tol:
            P = Pn
            break
        P = Pn
    return P, Ad, Qd


def choose_stride(A, H, Q, R, tau, dmax):
    vals = []
    for d in range(1, dmax + 1):
        P, _, _ = riccati(A, H, Q, R, d)
        vals.append(float(np.linalg.eigvalsh(P).max()))
    feasible = [d + 1 for d, v in enumerate(vals) if v <= tau]
    return (max(feasible) if feasible else 1), vals


def math_check(seed=7):
    rng = np.random.default_rng(seed)
    # Stable, observable 2-D system; empirical covariance is measured over many
    # independent d-step prior rollouts initialized from the filtered covariance.
    A = np.array([[0.88, 0.12], [-0.04, 0.82]])
    H = np.eye(2)
    Q = np.diag([0.015, 0.008])
    R = np.diag([0.025, 0.025])
    P1, _, _ = riccati(A, H, Q, R, 1)
    rows = []
    for d in (1, 2, 4, 8):
        P, Ad, Qd = riccati(A, H, Q, R, d)
        # P is the steady-state prior at an observation interval of d.
        # The d-step prior starts immediately after an observation, hence
        # sample from the corresponding posterior covariance P_plus.
        S = H @ P @ H.T + R
        K = P @ H.T @ np.linalg.inv(S)
        P_plus = (P - K @ S @ K.T + (P - K @ S @ K.T).T) * 0.5
        x = rng.multivariate_normal(np.zeros(2), P_plus, size=60000)
        noise = rng.multivariate_normal(np.zeros(2), Qd, size=60000)
        e = (x @ Ad.T) + noise
        empirical = np.cov(e, rowvar=False, bias=True)
        pred = float(np.linalg.eigvalsh(P).max())
        emp = float(np.linalg.eigvalsh(empirical).max())
        rows.append({"d": d, "predicted_lambda_max": pred,
                     "empirical_lambda_max": emp,
                     "relative_error": abs(pred - emp) / max(pred, 1e-12)})
    # monotonic signal expected for this stable system
    monotone = all(rows[i]["predicted_lambda_max"] <= rows[i+1]["predicted_lambda_max"] + 1e-10
                   for i in range(len(rows)-1))
    tau = 0.13
    selected = choose_stride(A, H, Q, R, tau, 8)[0]
    return {"rows": rows, "predicted_monotone": monotone,
            "tau": tau, "selected_stride": selected,
            "max_relative_error": max(r["relative_error"] for r in rows)}


class ToyWorld:
    # Mild nonlinear latent dynamics; observation is the latent vector plus noise.
    def __init__(self, A, q=0.006, r=0.04, seed=0):
        self.A = A
        self.q = q
        self.r = r
        self.rng = np.random.default_rng(seed)

    def step(self, z):
        return self.A @ z + np.array([0.035 * np.tanh(z[1]),
                                      -0.025 * np.tanh(z[0])]) + self.rng.normal(0, math.sqrt(self.q), 2)

    def observe(self, z):
        return z + self.rng.normal(0, math.sqrt(self.r), 2)


def rollout(strategy, A, H, Q, R, tau, seed, T=100, d_fixed=4):
    # A scheduler decides when to read an expensive observation. Between reads,
    # the model rolls forward using its local linear dynamics (here the known toy model).
    world = ToyWorld(A, q=float(Q[0, 0]), r=float(R[0, 0]), seed=seed)
    true = np.zeros(2)
    estimate = np.zeros(2)
    P = np.eye(2) * 0.05
    calls = 0
    sqerr = []
    t = 0
    strides = []
    while t < T:
        if strategy == "adaptive":
            d, _ = choose_stride(A, H, Q, R, tau, 8)
        else:
            d = d_fixed
        d = min(d, T - t)
        # Observation update at current time (the expensive encoder call).
        y = world.observe(true)
        calls += 1
        S = H @ P @ H.T + R
        K = P @ H.T @ np.linalg.inv(S)
        estimate = estimate + K @ (y - H @ estimate)
        P = (np.eye(2) - K @ H) @ P
        # Predict d transitions, recording prediction error at every frame.
        for _ in range(d):
            true = world.step(true)
            estimate = A @ estimate + np.array([0.035 * np.tanh(estimate[1]),
                                                -0.025 * np.tanh(estimate[0])])
            P = A @ P @ A.T + Q
            sqerr.append(float(np.sum((estimate - true) ** 2)))
            t += 1
        strides.append(d)
    return {"rmse": float(np.sqrt(np.mean(sqerr))), "encoder_calls": calls,
            "mean_stride": float(np.mean(strides)), "strides": strides}


def experiment():
    A = np.array([[0.88, 0.12], [-0.04, 0.82]])
    H = np.eye(2)
    Q = np.diag([0.006, 0.006])
    R = np.diag([0.04, 0.04])
    # Chosen from the analytical covariance curve: d=4 is feasible, d=5 is not.
    tau = 0.105
    check = math_check()
    out = {"math_check": check, "runs": {}}
    for strategy, d in (("baseline_stride_1", 1), ("fixed_stride_4", 4), ("fixed_stride_8", 8), ("adaptive", 0)):
        vals = [rollout(strategy, A, H, Q, R, tau, seed=100+i, T=240, d_fixed=d)
                for i in range(12)]
        out["runs"][strategy] = {
            "rmse_mean": float(np.mean([v["rmse"] for v in vals])),
            "rmse_std": float(np.std([v["rmse"] for v in vals])),
            "calls_mean": float(np.mean([v["encoder_calls"] for v in vals])),
            "mean_stride": float(np.mean([v["mean_stride"] for v in vals])),
            "example_strides": vals[0]["strides"][:12]
        }
    # Claimed stability signal: increasing the dominant eigenvalue increases
    # predicted uncertainty and can force a shorter admissible interval.
    probe = []
    for radius in (0.90, 0.95, 0.98, 0.99, 0.995):
        Ap = np.diag([radius, 0.75])
        selected, curve = choose_stride(Ap, H, Q, R, tau, 16)
        probe.append({"dominant_radius": radius, "selected_stride": selected,
                      "lambda_max_by_d": curve})
    out["near_unit_radius_probe"] = probe
    return out


if __name__ == "__main__":
    print(json.dumps(experiment(), indent=2))