"""A saturated one-dimensional mass, PD/PID, and discrete LQR; NumPy only."""
import json
import numpy as np


DT, MASS, DRAG = 0.01, 1.0, 0.2
A = np.array([[1.0, DT - DRAG * DT * DT / (2 * MASS)],
              [0.0, 1 - DRAG * DT / MASS]])
B = np.array([[DT * DT / (2 * MASS)], [DT / MASS]])


def lqr():
    q, r = np.diag([8.0, 1.0]), np.array([[0.5]])
    p = q.copy()
    for _ in range(10000):
        gain = np.linalg.solve(r + B.T @ p @ B, B.T @ p @ A)
        nxt = q + A.T @ p @ A - A.T @ p @ B @ gain
        if np.max(np.abs(nxt - p)) < 1e-10:
            return gain
        p = nxt
    raise RuntimeError("Riccati iteration did not converge")


def simulate(mode, gain=None):
    state, integral, squared_error = np.zeros(2), 0.0, 0.0
    saturated, peak = 0, 0.0
    for step in range(2000):
        error = 1.0 - state[0]
        if mode == "lqr":
            raw = float((-gain @ (state - np.array([1.0, 0.0]))).item())
        else:
            trial = integral + error * DT
            raw = 6.0 * error - 3.5 * state[1] + (1.5 * trial if mode == "pid" else 0.0)
            # Conditional integration: do not integrate further into saturation.
            if abs(raw) <= 2.0 or raw * error < 0:
                integral = trial
        effort = float(np.clip(raw, -2.0, 2.0))
        saturated += int(abs(raw) > 2.0)
        # A constant unknown external force after t=5 s tests disturbance rejection.
        disturbance = 0.3 if step >= 500 else 0.0
        state = A @ state + B[:, 0] * (effort + disturbance)
        squared_error += error * error
        peak = max(peak, float(state[0]))
    return {"final_position_m": float(state[0]), "final_velocity_m_s": float(state[1]),
            "rmse_m": float(np.sqrt(squared_error / 2000)),
            "overshoot_m": max(0.0, peak - 1.0), "saturated_steps": saturated}


if __name__ == "__main__":
    gain = lqr()
    report = {"dt_s": DT, "numpy": np.__version__, "lqr_gain": gain.tolist(),
              "controllers": {mode: simulate(mode, gain) for mode in ("pd", "pid", "lqr")}}
    assert abs(report["controllers"]["pid"]["final_position_m"] - 1) < 0.01
    assert all(np.isfinite(item["rmse_m"]) for item in report["controllers"].values())
    print(json.dumps(report, indent=2))