--- Final Solver (used for performance test) ---
"""
Fast LTI System Simulation Solver.

Uses exact discrete-time update with linear interpolation of input to match
scipy.signal.lsim output to machine precision, accelerated with specialized
Numba JIT kernels for small system orders.
"""
import numpy as np
from scipy import signal
from scipy.linalg import expm
import numba


# ----- Specialized kernels for small n -----

@numba.njit(cache=True, fastmath=True)
def _sim_n1(Ad, Bd0, Bd1, C, D, u, N):
    a00 = Ad[0, 0]
    b00 = Bd0[0]
    b10 = Bd1[0]
    c0 = C[0]
    yout = np.empty(N)
    x0 = 0.0
    for k in range(N - 1):
        y = c0 * x0 + D * u[k]
        yout[k] = y
        x0 = a00 * x0 + b00 * u[k] + b10 * u[k + 1]
    yout[N - 1] = c0 * x0 + D * u[N - 1]
    return yout


@numba.njit(cache=True, fastmath=True)
def _sim_n2(Ad, Bd0, Bd1, C, D, u, N):
    a00 = Ad[0, 0]; a01 = Ad[0, 1]
    a10 = Ad[1, 0]; a11 = Ad[1, 1]
    b00 = Bd0[0]; b01 = Bd0[1]
    b10 = Bd1[0]; b11 = Bd1[1]
    c0 = C[0]; c1 = C[1]
    yout = np.empty(N)
    x0 = 0.0
    x1 = 0.0
    for k in range(N - 1):
        y = c0 * x0 + c1 * x1 + D * u[k]
        yout[k] = y
        new_x0 = a00 * x0 + a01 * x1 + b00 * u[k] + b10 * u[k + 1]
        new_x1 = a10 * x0 + a11 * x1 + b01 * u[k] + b11 * u[k + 1]
        x0 = new_x0
        x1 = new_x1
    yout[N - 1] = c0 * x0 + c1 * x1 + D * u[N - 1]
    return yout


@numba.njit(cache=True, fastmath=True)
def _sim_n3(Ad, Bd0, Bd1, C, D, u, N):
    a00 = Ad[0, 0]; a01 = Ad[0, 1]; a02 = Ad[0, 2]
    a10 = Ad[1, 0]; a11 = Ad[1, 1]; a12 = Ad[1, 2]
    a20 = Ad[2, 0]; a21 = Ad[2, 1]; a22 = Ad[2, 2]
    b00 = Bd0[0]; b01 = Bd0[1]; b02 = Bd0[2]
    b10 = Bd1[0]; b11 = Bd1[1]; b12 = Bd1[2]
    c0 = C[0]; c1 = C[1]; c2 = C[2]
    yout = np.empty(N)
    x0 = x1 = x2 = 0.0
    for k in range(N - 1):
        y = c0 * x0 + c1 * x1 + c2 * x2 + D * u[k]
        yout[k] = y
        nx0 = a00 * x0 + a01 * x1 + a02 * x2 + b00 * u[k] + b10 * u[k + 1]
        nx1 = a10 * x0 + a11 * x1 + a12 * x2 + b01 * u[k] + b11 * u[k + 1]
        nx2 = a20 * x0 + a21 * x1 + a22 * x2 + b02 * u[k] + b12 * u[k + 1]
        x0 = nx0; x1 = nx1; x2 = nx2
    yout[N - 1] = c0 * x0 + c1 * x1 + c2 * x2 + D * u[N - 1]
    return yout


@numba.njit(cache=True, fastmath=True)
def _sim_n4(Ad, Bd0, Bd1, C, D, u, N):
    a00=Ad[0,0];a01=Ad[0,1];a02=Ad[0,2];a03=Ad[0,3]
    a10=Ad[1,0];a11=Ad[1,1];a12=Ad[1,2];a13=Ad[1,3]
    a20=Ad[2,0];a21=Ad[2,1];a22=Ad[2,2];a23=Ad[2,3]
    a30=Ad[3,0];a31=Ad[3,1];a32=Ad[3,2];a33=Ad[3,3]
    b00=Bd0[0];b01=Bd0[1];b02=Bd0[2];b03=Bd0[3]
    b10=Bd1[0];b11=Bd1[1];b12=Bd1[2];b13=Bd1[3]
    c0=C[0];c1=C[1];c2=C[2];c3=C[3]
    yout=np.empty(N)
    x0=x1=x2=x3=0.0
    for k in range(N-1):
        y=c0*x0+c1*x1+c2*x2+c3*x3+D*u[k]
        yout[k]=y
        nx0=a00*x0+a01*x1+a02*x2+a03*x3+b00*u[k]+b10*u[k+1]
        nx1=a10*x0+a11*x1+a12*x2+a13*x3+b01*u[k]+b11*u[k+1]
        nx2=a20*x0+a21*x1+a22*x2+a23*x3+b02*u[k]+b12*u[k+1]
        nx3=a30*x0+a31*x1+a32*x2+a33*x3+b03*u[k]+b13*u[k+1]
        x0=nx0;x1=nx1;x2=nx2;x3=nx3
    yout[N-1]=c0*x0+c1*x1+c2*x2+c3*x3+D*u[N-1]
    return yout


@numba.njit(cache=True, fastmath=True)
def _sim_n5(Ad, Bd0, Bd1, C, D, u, N):
    a00=Ad[0,0];a01=Ad[0,1];a02=Ad[0,2];a03=Ad[0,3];a04=Ad[0,4]
    a10=Ad[1,0];a11=Ad[1,1];a12=Ad[1,2];a13=Ad[1,3];a14=Ad[1,4]
    a20=Ad[2,0];a21=Ad[2,1];a22=Ad[2,2];a23=Ad[2,3];a24=Ad[2,4]
    a30=Ad[3,0];a31=Ad[3,1];a32=Ad[3,2];a33=Ad[3,3];a34=Ad[3,4]
    a40=Ad[4,0];a41=Ad[4,1];a42=Ad[4,2];a43=Ad[4,3];a44=Ad[4,4]
    b00=Bd0[0];b01=Bd0[1];b02=Bd0[2];b03=Bd0[3];b04=Bd0[4]
    b10=Bd1[0];b11=Bd1[1];b12=Bd1[2];b13=Bd1[3];b14=Bd1[4]
    c0=C[0];c1=C[1];c2=C[2];c3=C[3];c4=C[4]
    yout=np.empty(N)
    x0=x1=x2=x3=x4=0.0
    for k in range(N-1):
        y=c0*x0+c1*x1+c2*x2+c3*x3+c4*x4+D*u[k]
        yout[k]=y
        nx0=a00*x0+a01*x1+a02*x2+a03*x3+a04*x4+b00*u[k]+b10*u[k+1]
        nx1=a10*x0+a11*x1+a12*x2+a13*x3+a14*x4+b01*u[k]+b11*u[k+1]
        nx2=a20*x0+a21*x1+a22*x2+a23*x3+a24*x4+b02*u[k]+b12*u[k+1]
        nx3=a30*x0+a31*x1+a32*x2+a33*x3+a34*x4+b03*u[k]+b13*u[k+1]
        nx4=a40*x0+a41*x1+a42*x2+a43*x3+a44*x4+b04*u[k]+b14*u[k+1]
        x0=nx0;x1=nx1;x2=nx2;x3=nx3;x4=nx4
    yout[N-1]=c0*x0+c1*x1+c2*x2+c3*x3+c4*x4+D*u[N-1]
    return yout


@numba.njit(cache=True, fastmath=True)
def _sim_n6(Ad, Bd0, Bd1, C, D, u, N):
    a00=Ad[0,0];a01=Ad[0,1];a02=Ad[0,2];a03=Ad[0,3];a04=Ad[0,4];a05=Ad[0,5]
    a10=Ad[1,0];a11=Ad[1,1];a12=Ad[1,2];a13=Ad[1,3];a14=Ad[1,4];a15=Ad[1,5]
    a20=Ad[2,0];a21=Ad[2,1];a22=Ad[2,2];a23=Ad[2,3];a24=Ad[2,4];a25=Ad[2,5]
    a30=Ad[3,0];a31=Ad[3,1];a32=Ad[3,2];a33=Ad[3,3];a34=Ad[3,4];a35=Ad[3,5]
    a40=Ad[4,0];a41=Ad[4,1];a42=Ad[4,2];a43=Ad[4,3];a44=Ad[4,4];a45=Ad[4,5]
    a50=Ad[5,0];a51=Ad[5,1];a52=Ad[5,2];a53=Ad[5,3];a54=Ad[5,4];a55=Ad[5,5]
    b00=Bd0[0];b01=Bd0[1];b02=Bd0[2];b03=Bd0[3];b04=Bd0[4];b05=Bd0[5]
    b10=Bd1[0];b11=Bd1[1];b12=Bd1[2];b13=Bd1[3];b14=Bd1[4];b15=Bd1[5]
    c0=C[0];c1=C[1];c2=C[2];c3=C[3];c4=C[4];c5=C[5]
    yout=np.empty(N)
    x0=x1=x2=x3=x4=x5=0.0
    for k in range(N-1):
        y=c0*x0+c1*x1+c2*x2+c3*x3+c4*x4+c5*x5+D*u[k]
        yout[k]=y
        nx0=a00*x0+a01*x1+a02*x2+a03*x3+a04*x4+a05*x5+b00*u[k]+b10*u[k+1]
        nx1=a10*x0+a11*x1+a12*x2+a13*x3+a14*x4+a15*x5+b01*u[k]+b11*u[k+1]
        nx2=a20*x0+a21*x1+a22*x2+a23*x3+a24*x4+a25*x5+b02*u[k]+b12*u[k+1]
        nx3=a30*x0+a31*x1+a32*x2+a33*x3+a34*x4+a35*x5+b03*u[k]+b13*u[k+1]
        nx4=a40*x0+a41*x1+a42*x2+a43*x3+a44*x4+a45*x5+b04*u[k]+b14*u[k+1]
        nx5=a50*x0+a51*x1+a52*x2+a53*x3+a54*x4+a55*x5+b05*u[k]+b15*u[k+1]
        x0=nx0;x1=nx1;x2=nx2;x3=nx3;x4=nx4;x5=nx5
    yout[N-1]=c0*x0+c1*x1+c2*x2+c3*x3+c4*x4+c5*x5+D*u[N-1]
    return yout


# General kernel for arbitrary n
@numba.njit(cache=True, fastmath=True)
def _sim_general(Ad, Bd0, Bd1, C, D, u, n, N):
    yout = np.empty(N)
    x = np.zeros(n)
    x_new = np.zeros(n)
    for k in range(N - 1):
        y = D * u[k]
        for i in range(n):
            y += C[i] * x[i]
        yout[k] = y
        for i in range(n):
            val = Bd0[i] * u[k] + Bd1[i] * u[k + 1]
            for j in range(n):
                val += Ad[i, j] * x[j]
            x_new[i] = val
        for i in range(n):
            x[i] = x_new[i]
    y = D * u[N - 1]
    for i in range(n):
        y += C[i] * x[i]
    yout[N - 1] = y
    return yout


_KERNELS = {
    1: _sim_n1,
    2: _sim_n2,
    3: _sim_n3,
    4: _sim_n4,
    5: _sim_n5,
    6: _sim_n6,
}


class Solver:
    def __init__(self):
        # Trigger Numba compilation for all specialized kernels (in __init__, not counted)
        for n in [1, 2, 3, 4, 5, 6]:
            Ad = np.eye(n, dtype=np.float64)
            Bd0 = np.zeros(n, dtype=np.float64)
            Bd1 = np.zeros(n, dtype=np.float64)
            C = np.zeros(n, dtype=np.float64)
            u = np.zeros(10, dtype=np.float64)
            kernel = _KERNELS[n]
            kernel(Ad, Bd0, Bd1, C, 0.0, u, 10)
        # General kernel
        Ad = np.eye(2, dtype=np.float64)
        Bd0 = np.zeros(2, dtype=np.float64)
        Bd1 = np.zeros(2, dtype=np.float64)
        C = np.zeros(2, dtype=np.float64)
        u = np.zeros(10, dtype=np.float64)
        _sim_general(Ad, Bd0, Bd1, C, 0.0, u, 2, 10)

    def _build_matrices(self, A, B, dt):
        """Build discrete-time matrices using scipy.signal.lsim's exact formulation."""
        n = A.shape[0]
        # M = [[A*dt, B*dt, 0], [0, 0, I], [0, 0, 0]] (column-vector augmented form)
        M = np.zeros((n + 2, n + 2))
        M[:n, :n] = A * dt
        M[:n, n] = B * dt
        M[n, n + 1] = 1.0
        expmM = expm(M)
        Ad = expmM[:n, :n]
        F = expmM[:n, n]
        G = expmM[:n, n + 1]
        Bd0 = F - G
        Bd1 = G
        return Ad, Bd0, Bd1

    def solve(self, problem):
        num = problem["num"]
        den = problem["den"]
        u_list = problem["u"]
        t_list = problem["t"]

        N = len(t_list)
        if N == 0:
            return {"yout": []}

        n = len(den) - 1

        # State-space via tf2ss (passing lists - scipy converts internally)
        A, B, C, D_arr = signal.tf2ss(num, den)
        A = np.asarray(A, dtype=np.float64)
        B = np.asarray(B, dtype=np.float64).flatten()
        C = np.asarray(C, dtype=np.float64).flatten()
        D = float(np.asarray(D_arr).item())

        # Direct gain (n=0)
        if n == 0:
            gain = num[0] / den[0]
            u_arr = np.fromiter(u_list, dtype=np.float64, count=N)
            yout = gain * u_arr
            return {"yout": yout.tolist()}

        # Single time point
        if N == 1:
            u0 = float(u_list[0])
            yout = np.array([D * u0])
            return {"yout": yout.tolist()}

        # Compute dt from list (avoid converting t)
        dt = float(t_list[1] - t_list[0])

        Ad, Bd0, Bd1 = self._build_matrices(A, B, dt)

        # Convert u to numpy array (fastest is fromiter)
        u_arr = np.fromiter(u_list, dtype=np.float64, count=N)

        if n in _KERNELS:
            kernel = _KERNELS[n]
            yout = kernel(Ad, Bd0, Bd1, C, D, u_arr, N)
        else:
            yout = _sim_general(Ad, Bd0, Bd1, C, D, u_arr, n, N)

        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.1733s
Total Solver Time:   0.1338s
Raw Speedup:         68.5801 x
Final Reward (Score): 68.5801
---------------------------
..

==================================== 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 117.15s (0:01:57) =========================
