--- Final Solver (used for performance test) ---
import numpy as np
from scipy.signal import tf2ss
from scipy.linalg import expm
from scipy.signal import ss2tf as _ss2tf
from scipy.signal import lfilter as _lfilter
from typing import Any

# Attempt JIT compilation with numba for a very fast inner loop.
_sim_loop_jit = None
try:
    from numba import njit

    @njit(cache=True)
    def _sim_loop_numba(E11, e12, e13, c, d, u, slopes):
        n = E11.shape[0]
        N = len(u)
        x = np.zeros(n)
        yout = np.empty(N)
        yout[0] = d * u[0]
        for k in range(N - 1):
            x = E11 @ x + e12 * u[k] + e13 * slopes[k]
            yout[k + 1] = c @ x + d * u[k + 1]
        return yout

    # Warm up JIT compilation (does not count toward runtime).
    _d_E = np.eye(1)
    _d_v = np.zeros(1)
    _d_u = np.array([0.0, 0.0])
    _d_s = np.array([0.0])
    _sim_loop_numba(_d_E, _d_v, _d_v, _d_v, 0.0, _d_u, _d_s)
    _sim_loop_jit = _sim_loop_numba
except Exception:
    pass


class Solver:
    def __init__(self):
        pass

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

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

        # Convert transfer function H(s) = num(s)/den(s) to state space.
        A, B, C, D = tf2ss(num, den)
        n = A.shape[0]
        b = B.ravel()            # (n,)
        c = C.ravel()            # (n,)
        d_val = float(D.flat[0]) if D.size > 0 else 0.0

        if N == 1:
            return {"yout": [d_val * float(u[0])]}

        # Build augmented matrix M (constant across all intervals).
        na = n + 2
        M = np.zeros((na, na))
        if n > 0:
            M[:n, :n] = A
            M[:n, n] = b
        M[n, n + 1] = 1.0

        dt = np.diff(t)
        h0 = dt[0]
        tol = max(abs(h0) * 1e-9, 1e-12)
        uniform = (len(dt) == 1) or bool(np.all(np.abs(dt - h0) <= tol))

        if uniform:
            h = h0
            E = expm(M * h)
            E11 = E[:n, :n]
            e12 = E[:n, n]
            e13 = E[:n, n + 1]
            slopes = (u[1:] - u[:-1]) / h

            if n == 0:
                # Pure gain system
                yout = d_val * u
            else:
                # Use lfilter (C-implemented) for the discrete-time simulation.
                # Discrete system: x_{k+1} = A_d @ x_k + b1*u_k + b2*slope_k
                #                  y_k     = c @ x_k  + d*u_k
                # This is a 2-input SISO discrete system.
                Ad = np.ascontiguousarray(E11)
                Cd = c.reshape(1, n)

                # Build B column for input 1 (u)
                Bd1 = e12.reshape(n, 1)
                num_d1, den_d = _ss2tf(Ad, Bd1, Cd, np.array([[d_val]]))

                # Build B column for input 2 (slope)
                Bd2 = e13.reshape(n, 1)
                num_d2, _ = _ss2tf(Ad, Bd2, Cd, np.array([[0.0]]))

                # Pad slopes to length N (last slope is unused since n_states
                # at N-1 doesn't need to propagate further)
                slopes_padded = np.empty(N)
                slopes_padded[:N - 1] = slopes
                slopes_padded[N - 1] = 0.0

                # Apply IIR filters (C implementation, very fast)
                yout = _lfilter(num_d1.ravel(), den_d.ravel(), u)
                yout = yout + _lfilter(num_d2.ravel(), den_d.ravel(), slopes_padded)

            return {"yout": yout.tolist()}
        else:
            # Non-uniform time steps: compute expm once per unique step size.
            x = np.zeros(n)
            yout = np.empty(N)
            yout[0] = d_val * float(u[0])
            cache: dict[float, tuple] = {}
            for k in range(N - 1):
                hk = float(dt[k])
                if hk not in cache:
                    Ek = expm(M * hk)
                    cache[hk] = (
                        np.ascontiguousarray(Ek[:n, :n]),
                        np.ascontiguousarray(Ek[:n, n]),
                        np.ascontiguousarray(Ek[:n, n + 1]),
                    )
                E11_k, e12_k, e13_k = cache[hk]
                slope = (u[k + 1] - u[k]) / hk
                x = E11_k @ x + e12_k * u[k] + e13_k * slope
                yout[k + 1] = c @ x + d_val * float(u[k + 1])

            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: 8.9268s
Total Solver Time:   0.1084s
Raw Speedup:         82.3679 x
Final Reward (Score): 82.3679
---------------------------
..

==================================== 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.13s (0:01:52) =========================
