--- Final Solver (used for performance test) ---
import numpy as np
from scipy.linalg import expm as _expm
from scipy.signal import tf2ss as _tf2ss


class Solver:
    def solve(self, problem, **kwargs):
        num = np.atleast_1d(np.asarray(problem["num"], dtype=np.float64))
        den = np.atleast_1d(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)

        # Transfer function to state-space
        A, B, C, D = _tf2ss(num, den)
        n = A.shape[0]
        B = B.ravel()
        C = C.ravel()
        d0 = float(D.ravel()[0])

        # Discretize using matrix exponential (zero-order hold for uniform grid)
        dt = t[1] - t[0]
        if n <= 2 and N > 2 and np.allclose(np.diff(t), dt):
            # Fast path for 1st/2nd order systems with uniform time step
            Ad = _expm(A * dt)
            # Bd = A^{-1} (Ad - I) B
            if n == 1:
                a = float(A[0, 0])
                if abs(a) > 1e-30:
                    bd = (float(Ad[0, 0]) - 1.0) / a * float(B[0])
                else:
                    bd = dt * float(B[0])
                ad = float(Ad[0, 0])
                c0 = float(C[0])
                x = 0.0
                yout = np.empty(N)
                for i in range(N):
                    yout[i] = c0 * x + d0 * u[i]
                    if i < N - 1:
                        x = ad * x + bd * u[i]
            else:  # n == 2
                # Solve A * Bd = (Ad - I) * B
                # A = [[a00,a01],[a10,a11]]
                a00, a01 = float(A[0, 0]), float(A[0, 1])
                a10, a11 = float(A[1, 0]), float(A[1, 1])
                det_A = a00 * a11 - a01 * a10
                em = Ad - np.eye(2)
                eb0 = float(em[0, 0]) * float(B[0]) + float(em[0, 1]) * float(B[1])
                eb1 = float(em[1, 0]) * float(B[0]) + float(em[1, 1]) * float(B[1])
                if abs(det_A) > 1e-30:
                    bd0 = (a11 * eb0 - a01 * eb1) / det_A
                    bd1 = (-a10 * eb0 + a00 * eb1) / det_A
                else:
                    bd0 = dt * float(B[0])
                    bd1 = dt * float(B[1])
                ad00, ad01 = float(Ad[0, 0]), float(Ad[0, 1])
                ad10, ad11 = float(Ad[1, 0]), float(Ad[1, 1])
                c0, c1 = float(C[0]), float(C[1])
                x0 = 0.0
                x1 = 0.0
                yout = np.empty(N)
                for i in range(N):
                    yout[i] = c0 * x0 + c1 * x1 + d0 * u[i]
                    if i < N - 1:
                        n0 = ad00 * x0 + ad01 * x1 + bd0 * u[i]
                        n1 = ad10 * x0 + ad11 * x1 + bd1 * u[i]
                        x0 = n0
                        x1 = n1
        else:
            # General case
            dt_arr = np.diff(t)
            unique_dts = np.unique(dt_arr)
            cache = {}
            for dt_i in unique_dts:
                Ad_i = _expm(A * dt_i)
                if n == 1:
                    a = float(A[0, 0])
                    if abs(a) > 1e-30:
                        bd_i = (float(Ad_i[0, 0]) - 1.0) / a * float(B[0])
                    else:
                        bd_i = dt_i * float(B[0])
                    cache[dt_i] = (float(Ad_i[0, 0]), bd_i)
                else:
                    em_i = Ad_i - np.eye(n)
                    if abs(np.linalg.det(A)) > 1e-30:
                        bd_i = np.linalg.solve(A, em_i @ B)
                    else:
                        bd_i = np.linalg.solve(A + 1e-12 * np.eye(n), em_i @ B)
                    cache[dt_i] = (Ad_i, bd_i)

            if n == 1:
                c0 = float(C[0])
                x = 0.0
                yout = np.empty(N)
                for i in range(N):
                    yout[i] = c0 * x + d0 * u[i]
                    if i < N - 1:
                        ad_i, bd_i = cache[dt_arr[i]]
                        x = ad_i * x + bd_i * u[i]
            elif n == 2:
                c0, c1 = float(C[0]), float(C[1])
                x0 = 0.0
                x1 = 0.0
                yout = np.empty(N)
                for i in range(N):
                    yout[i] = c0 * x0 + c1 * x1 + d0 * u[i]
                    if i < N - 1:
                        Ad_i, bd_i = cache[dt_arr[i]]
                        n0 = float(Ad_i[0, 0]) * x0 + float(Ad_i[0, 1]) * x1 + float(bd_i[0]) * u[i]
                        n1 = float(Ad_i[1, 0]) * x0 + float(Ad_i[1, 1]) * x1 + float(bd_i[1]) * u[i]
                        x0 = n0
                        x1 = n1
            else:
                Cv = C.copy()
                X = np.zeros(n)
                yout = np.empty(N)
                for i in range(N):
                    yout[i] = Cv @ X + d0 * u[i]
                    if i < N - 1:
                        Ad_i, bd_i = cache[dt_arr[i]]
                        X = Ad_i @ X + bd_i * u[i]

        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: False
Total Baseline Time: 0.0000s
Total Solver Time:   0.0000s
Raw Speedup:         0.0000 x
Final Reward (Score): 1.0000
---------------------------
FF

=================================== FAILURES ===================================
_____________________________ test_solver_validity _____________________________

performance_results = {'raw_speedup': 0.0, 'speedup': 1.0, 'validity': False}

    def test_solver_validity(performance_results):
        """Ensures the solver produces valid outputs."""
        if not performance_results["validity"]:
>           pytest.fail("Solver produced invalid solutions on the test set.")
E           Failed: Solver produced invalid solutions on the test set.

/tests/test_outputs.py:200: Failed
------------------------------ Captured log setup ------------------------------
ERROR    root:evaluator.py:187 Solution verification failed. Max absolute error: 1.12e-01, Max relative error: 1.65e+03 (rtol=1e-05, atol=1e-08)
ERROR    test_outputs:test_outputs.py:93 Solver produced an invalid solution during warmup.
WARNING  test_outputs:test_outputs.py:172 Invalid solution detected. Assigning mercy score of 1.0.
_____________________________ test_solver_speedup ______________________________

performance_results = {'raw_speedup': 0.0, 'speedup': 1.0, 'validity': False}

    def test_solver_speedup(performance_results):
        """Checks if the solver effectively optimized the code."""
>       assert performance_results["raw_speedup"] > 1.0, \
            f"Solver was not faster than baseline (Speedup: {performance_results['raw_speedup']:.2f}x)"
E       AssertionError: Solver was not faster than baseline (Speedup: 0.00x)
E       assert 0.0 > 1.0

/tests/test_outputs.py:205: AssertionError
==================================== PASSES ====================================
=========================== short test summary info ============================
PASSED ../tests/test_outputs.py::test_solver_exists
FAILED ../tests/test_outputs.py::test_solver_validity - Failed: Solver produc...
FAILED ../tests/test_outputs.py::test_solver_speedup - AssertionError: Solver...
========================= 2 failed, 1 passed in 1.00s ==========================
