--- Final Solver (used for performance test) ---
import numpy as np
from scipy.linalg import expm as scipy_expm
from numba import njit
from typing import Any


@njit(cache=True)
def _sim_even(Phi_xx, Phi_xu, Phi_xs, C, D_scalar, u, t, dt):
    n = len(t)
    nx = Phi_xx.shape[0]
    yout = np.empty(n)
    x = np.zeros(nx)
    yout[0] = D_scalar * u[0]
    inv_dt = 1.0 / dt
    for i in range(n - 1):
        slope = (u[i + 1] - u[i]) * inv_dt
        x = Phi_xx @ x + Phi_xu * u[i] + Phi_xs * slope
        yout[i + 1] = np.dot(C, x) + D_scalar * u[i + 1]
    return yout


@njit(cache=True)
def _sim_rk4(A, B, C, D_scalar, u, t, n_substeps):
    n = len(t)
    nx = A.shape[0]
    yout = np.empty(n)
    x = np.zeros(nx)
    yout[0] = D_scalar * u[0]

    for i in range(n - 1):
        dt = t[i + 1] - t[i]
        if dt <= 0.0:
            yout[i + 1] = np.dot(C, x) + D_scalar * u[i + 1]
            continue

        h = dt / n_substeps
        u0 = u[i]
        u1 = u[i + 1]
        slope = (u1 - u0) / dt
        half_h = 0.5 * h
        h6 = h / 6.0

        for step in range(n_substeps):
            t_local = step * h

            u_k1 = u0 + slope * t_local
            u_k23 = u0 + slope * (t_local + half_h)
            u_k4 = u0 + slope * (t_local + h)

            k1 = A @ x + B * u_k1
            k2 = A @ (x + half_h * k1) + B * u_k23
            k3 = A @ (x + half_h * k2) + B * u_k23
            k4 = A @ (x + h * k3) + B * u_k4

            x = x + h6 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)

        yout[i + 1] = np.dot(C, x) + D_scalar * u1
    return yout


@njit(cache=True)
def _sim_static(D_scalar, u, t):
    n = len(t)
    yout = np.empty(n)
    for i in range(n):
        yout[i] = D_scalar * u[i]
    return yout


def tf2ss_fast(num, den):
    """Fast transfer function to state-space matching scipy.signal.tf2ss convention."""
    num = np.asarray(num, dtype=np.float64).ravel()
    den = np.asarray(den, dtype=np.float64).ravel()
    
    # Strip leading zeros from den
    while len(den) > 1 and den[0] == 0.0:
        den = den[1:]
    
    nn = len(num)
    nd = len(den)
    
    if nn > nd:
        raise ValueError("System is improper")
    
    # Pad numerator
    if nn < nd:
        num = np.concatenate([np.zeros(nd - nn), num])
    
    nx = nd - 1
    
    if nx == 0:
        return np.zeros((0, 0)), np.zeros(0), np.zeros(0), float(num[0] / den[0])
    
    # Normalize denominator
    den = den / den[0]
    
    D = float(num[0])
    num_adj = num.copy()
    num_adj[1:] = num[1:] - D * den[1:]
    
    # Companion form matching scipy
    A = np.zeros((nx, nx))
    A[0, :] = -den[1:]
    if nx > 1:
        np.fill_diagonal(A[1:, :], 0.0)  # already zero
        for i in range(1, nx):
            A[i, i - 1] = 1.0
    
    B = np.zeros(nx)
    B[0] = 1.0
    
    C = num_adj[1:].copy()
    
    return np.ascontiguousarray(A), np.ascontiguousarray(B), np.ascontiguousarray(C), D


def expm_small_fast(M, t):
    """Fast matrix exponential for small matrices."""
    n = M.shape[0]
    if n == 0:
        return np.zeros((0, 0))
    if n == 1:
        return np.array([[np.exp(M[0, 0] * t)]])
    
    A = M * t
    norm_A = np.max(np.abs(A))
    
    s = max(0, int(np.ceil(np.log2(norm_A + 1e-16))))
    if s > 5:
        s = 5
    
    A_scaled = A / (2 ** s)
    
    result = np.eye(n)
    term = np.eye(n)
    for i in range(1, 20):
        term = term @ A_scaled / i
        result = result + term
        if np.max(np.abs(term)) < 1e-17:
            break
    
    for _ in range(s):
        result = result @ result
    
    return result


class Solver:
    def __init__(self):
        A = np.array([[0.0, 1.0], [-1.0, -0.2]])
        B = np.array([0.0, 1.0])
        C = np.array([1.0, 0.0])
        D = 0.0
        u = np.array([0.0, 0.1, 0.2, 0.1, 0.0])
        t = np.array([0.0, 0.25, 0.5, 0.75, 1.0])
        Phi_xx = np.eye(2)
        Phi_xu = np.zeros(2)
        Phi_xs = np.zeros(2)
        _sim_even(Phi_xx, Phi_xu, Phi_xs, C, D, u, t, 0.25)
        _sim_rk4(A, B, C, D, u, t, 8)
        _sim_static(D, u, t)
        
        M = np.array([[0.0, 1.0], [-1.0, -0.2]])
        expm_small_fast(M, 0.1)
        
        print("Solver initialized")

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

        n = len(t)

        try:
            A, B, C, D_scalar = tf2ss_fast(num, den)
        except ValueError:
            from scipy import signal
            A_ss, B_ss, C_ss, D_ss = signal.tf2ss(num, den)
            A = np.ascontiguousarray(A_ss, dtype=np.float64)
            B = np.ascontiguousarray(B_ss.ravel(), dtype=np.float64)
            C = np.ascontiguousarray(C_ss.ravel(), dtype=np.float64)
            D_scalar = float(D_ss.ravel()[0])
        
        nx = A.shape[0]

        if n <= 1 or nx == 0:
            if n == 0:
                return {"yout": []}
            yout = _sim_static(D_scalar, u, t)
            return {"yout": yout.tolist()}

        dts = np.diff(t)
        dt_mean = dts.mean()
        evenly_spaced = np.allclose(dts, dt_mean, rtol=1e-10)

        if evenly_spaced:
            dt = dt_mean
            aug_n = nx + 2
            M = np.zeros((aug_n, aug_n), dtype=np.float64)
            M[:nx, :nx] = A
            M[:nx, nx] = B
            M[nx, nx + 1] = 1.0

            if aug_n <= 10:
                E = expm_small_fast(M, dt)
            else:
                E = scipy_expm(M * dt)
            
            Phi_xx = np.ascontiguousarray(E[:nx, :nx], dtype=np.float64)
            Phi_xu = np.ascontiguousarray(E[:nx, nx], dtype=np.float64)
            Phi_xs = np.ascontiguousarray(E[:nx, nx + 1], dtype=np.float64)

            yout = _sim_even(Phi_xx, Phi_xu, Phi_xs, C, D_scalar, u, t, dt)
        else:
            eig_vals = np.linalg.eigvals(A)
            max_freq = max(np.max(np.abs(eig_vals)), 1.0)
            min_dt = dts.min()
            n_sub = max(16, int(np.ceil(min_dt * max_freq * 4)))
            n_sub = min(n_sub, 128)

            yout = _sim_rk4(A, B, C, D_scalar, u, t, n_sub)

        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 .Solver initialized

--- Performance Summary ---
Validity: True
Total Baseline Time: 9.0605s
Total Solver Time:   0.2741s
Raw Speedup:         33.0611 x
Final Reward (Score): 33.0611
---------------------------
..

==================================== 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 116.44s (0:01:56) =========================
