--- Final Solver (used for performance test) ---
import os
import ctypes
import numpy as np
from scipy.signal import tf2ss
from scipy.linalg import expm
from typing import Any

# Shared C code
C_CODE = '''
#include <string.h>

#define MAX_STATES 512

void simulate_lsim(int n_steps, int n_states,
                   const double *restrict Ad,
                   const double *restrict Bd0,
                   const double *restrict Bd1,
                   const double *restrict C,
                   double D,
                   const double *restrict u,
                   double *restrict yout) {
    double x[MAX_STATES];
    double new_x[MAX_STATES];
    memset(x, 0, sizeof(double) * n_states);
    for (int i = 0; i < n_steps; i++) {
        if (i > 0) {
            for (int j = 0; j < n_states; j++) {
                double s = 0.0;
                for (int k = 0; k < n_states; k++) {
                    s += x[k] * Ad[k * n_states + j];
                }
                s += u[i-1] * Bd0[j];
                s += u[i] * Bd1[j];
                new_x[j] = s;
            }
            memcpy(x, new_x, sizeof(double) * n_states);
        }
        double y = D * u[i];
        for (int j = 0; j < n_states; j++) {
            y += C[j] * x[j];
        }
        yout[i] = y;
    }
}
'''

SO_PATH = '/app/simulate.so'
C_PATH = '/app/simulate.c'

class Solver:
    def __init__(self):
        # ensure C library is compiled
        if not os.path.exists(SO_PATH) or (os.path.exists(C_PATH) and os.path.getmtime(SO_PATH) < os.path.getmtime(C_PATH)):
            with open(C_PATH, 'w') as f:
                f.write(C_CODE)
            ret = os.system(
                'gcc -O3 -march=native -ffast-math -funroll-loops -shared -fPIC -o {} {}'.format(SO_PATH, C_PATH)
            )
            if ret != 0:
                raise RuntimeError('Failed to compile simulate.so')
        self.lib = ctypes.CDLL(SO_PATH)
        self.lib.simulate_lsim.argtypes = [
            ctypes.c_int, ctypes.c_int,
            ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double),
            ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double),
            ctypes.c_double, ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double)
        ]
        self.lib.simulate_lsim.restype = None
        # cache discrete system matrices
        self._cache = {}

    def solve(self, problem, **kwargs) -> Any:
        num = np.atleast_1d(problem["num"])
        den = np.atleast_1d(problem["den"])
        u = np.asarray(problem["u"], dtype=float)
        t = np.asarray(problem["t"], dtype=float)
        n_steps = t.size
        if n_steps == 0:
            return {"yout": []}
        dt = float(t[1] - t[0])
        cache_key = (tuple(num.ravel()), tuple(den.ravel()), dt)
        cached = self._cache.get(cache_key)
        if cached is not None:
            Ad, Bd0, Bd1, C, D = cached
        else:
            A, B, C_arr, D = tf2ss(num, den)
            A = np.asarray(A, dtype=float)
            B = np.asarray(B, dtype=float)
            C = np.asarray(C_arr, dtype=float).ravel()
            D = float(np.squeeze(D))
            n_states = A.shape[0]
            M = np.zeros((n_states + 2, n_states + 2), dtype=float)
            M[:n_states, :n_states] = A * dt
            M[:n_states, n_states:n_states + 1] = B * dt
            M[n_states, n_states + 1] = 1.0
            expMT = expm(M.T)
            Ad = np.ascontiguousarray(expMT[:n_states, :n_states], dtype=float)
            Bd1 = np.ascontiguousarray(expMT[n_states + 1:, :n_states].ravel(), dtype=float)
            Bd0 = np.ascontiguousarray((expMT[n_states:n_states + 1, :n_states].ravel() - Bd1), dtype=float)
            self._cache[cache_key] = (Ad, Bd0, Bd1, C, D)
        Ad, Bd0, Bd1, C, D = self._cache[cache_key]
        n_states = Ad.shape[0]
        yout = np.empty(n_steps, dtype=float)
        self.lib.simulate_lsim(
            n_steps, n_states,
            Ad.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            Bd0.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            Bd1.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            C.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            D,
            u.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            yout.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
        )
        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.2046s
Total Solver Time:   0.0732s
Raw Speedup:         125.8102 x
Final Reward (Score): 125.8102
---------------------------
..

==================================== 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 118.30s (0:01:58) =========================
