"""Constant-velocity Kalman filter, seeded noise and periodic missing measurements."""
import json
import numpy as np


def run(seed=7):
    rng = np.random.default_rng(seed)
    dt = 0.1
    a = np.array([[1.0, dt], [0.0, 1.0]])
    h = np.array([[1.0, 0.0]])
    acceleration = np.array([[0.5 * dt * dt], [dt]])
    q = acceleration @ acceleration.T * 0.2 ** 2
    r = np.array([[0.3 ** 2]])
    truth = np.array([0.0, 0.8])
    mean, cov = np.zeros(2), np.diag([1.0, 1.0])
    raw_errors, filtered_errors, missing_errors = [], [], []
    for step in range(300):
        truth = a @ truth + acceleration[:, 0] * rng.normal(0, 0.2)
        mean, cov = a @ mean, a @ cov @ a.T + q
        if step % 10 < 8:
            z = truth[0] + rng.normal(0, 0.3)
            innovation = z - float((h @ mean).item())
            s = h @ cov @ h.T + r
            gain = np.linalg.solve(s, h @ cov).T
            mean = mean + gain[:, 0] * innovation
            # Joseph form reduces numerical asymmetry/PSD loss.
            ikh = np.eye(2) - gain @ h
            cov = ikh @ cov @ ikh.T + gain @ r @ gain.T
            raw_errors.append((z - truth[0]) ** 2)
            filtered_errors.append((mean[0] - truth[0]) ** 2)
        else:
            missing_errors.append((mean[0] - truth[0]) ** 2)
        assert np.linalg.eigvalsh(cov).min() >= -1e-12
    return {"seed": seed, "numpy": np.__version__, "measurements": len(raw_errors),
            "raw_position_rmse_m": float(np.sqrt(np.mean(raw_errors))),
            "filtered_position_rmse_same_frames_m": float(np.sqrt(np.mean(filtered_errors))),
            "prediction_only_rmse_m": float(np.sqrt(np.mean(missing_errors)))}


if __name__ == "__main__":
    report = run()
    assert report["filtered_position_rmse_same_frames_m"] < report["raw_position_rmse_m"]
    print(json.dumps(report, indent=2))