--- Final Solver (used for performance test) ---
from typing import Any

import numpy as np
from numba import njit
from scipy.linalg import expm
from scipy.signal import tf2ss


@njit(cache=True, fastmath=True)
def _simulate(AdT, b0, b1, c, d, u):
    """First-order-hold state-space simulation, matching scipy.signal.lsim.

    AdT is the transpose of the discrete state-transition matrix (so the
    matrix-vector product walks contiguous rows). State starts at zero.
    """
    n_steps = u.shape[0]
    n = AdT.shape[0]
    yout = np.empty(n_steps, dtype=np.float64)
    x = np.zeros(n, dtype=np.float64)
    xnew = np.zeros(n, dtype=np.float64)

    yout[0] = u[0] * d  # xout[0] = 0
    for i in range(1, n_steps):
        ui1 = u[i - 1]
        ui = u[i]
        yi = 0.0
        for k in range(n):
            row = AdT[k]
            s = 0.0
            for j in range(n):
                s += x[j] * row[j]
            s += ui1 * b0[k] + ui * b1[k]
            xnew[k] = s
            yi += s * c[k]
        x, xnew = xnew, x
        yout[i] = yi + ui * d
    return yout


class Solver:
    def __init__(self):
        # Warm up / compile the JIT function (not counted toward runtime).
        AdT = np.eye(2, dtype=np.float64)
        b = np.ones(2, dtype=np.float64)
        c = np.ones(2, dtype=np.float64)
        u = np.zeros(4, dtype=np.float64)
        _simulate(AdT, b, b, c, 0.0, u)

    def solve(self, problem: dict, **kwargs) -> Any:
        num = np.atleast_1d(np.asarray(problem["num"], dtype=np.float64))
        den = np.atleast_1d(np.asarray(problem["den"], dtype=np.float64))
        u = np.ascontiguousarray(problem["u"], dtype=np.float64)
        t = problem["t"]
        n_steps = len(t)

        if n_steps == 0:
            return {"yout": []}

        A, B, C, D = tf2ss(num, den)
        n = A.shape[0]
        d = float(D.reshape(-1)[0]) if D.size else 0.0

        if n_steps == 1 or n == 0:
            yout = u * d
            return {"yout": yout.tolist()}

        dt = t[1] - t[0]

        # First-order-hold discretization (identical to scipy.signal.lsim).
        M = np.zeros((n + 2, n + 2), dtype=np.float64)
        M[:n, :n] = A * dt
        M[:n, n] = B[:, 0] * dt
        M[n, n + 1] = 1.0
        expMT = expm(M.T)
        Ad = expMT[:n, :n]
        Bd1 = expMT[n + 1, :n]
        Bd0 = expMT[n, :n] - Bd1

        AdT = np.ascontiguousarray(Ad.T)
        b0 = np.ascontiguousarray(Bd0)
        b1 = np.ascontiguousarray(Bd1)
        c = np.ascontiguousarray(C[0])

        yout = _simulate(AdT, b0, b1, c, d, u)
        return {"yout": yout.tolist()}
--- End Solver ---
Running performance test...
============================= test session starts ==============================
platform linux -- Python 3.12.13, pytest-9.0.2, pluggy-1.6.0
rootdir: /tests
plugins: anyio-4.13.0, jaxtyping-0.3.9
collected 3 items

../tests/test_outputs.py .
--- Performance Summary ---
Validity: True
Total Baseline Time: 9.0112s
Total Solver Time:   0.0714s
Raw Speedup:         126.2840 x
Final Reward (Score): 126.2840
---------------------------
..

==================================== PASSES ====================================
=========================== short test summary info ============================
PASSED ../tests/test_outputs.py::test_solver_exists
PASSED ../tests/test_outputs.py::test_solver_validity
PASSED ../tests/test_outputs.py::test_solver_speedup
======================== 3 passed in 112.92s (0:01:52) =========================
