--- Final Solver (used for performance test) ---
from __future__ import annotations

from typing import Any
import math

import numpy as np
from numba import njit
from scipy import linalg, signal


@njit(fastmath=True)
def _sim_order1(ad: float, bd0: float, bd1: float, c0: float, d: float, u: np.ndarray) -> np.ndarray:
    m = u.shape[0]
    y = np.empty(m, dtype=np.float64)
    x0 = 0.0
    y[0] = d * u[0]
    for k in range(1, m):
        uk0 = u[k - 1]
        uk1 = u[k]
        x0 = x0 * ad + uk0 * bd0 + uk1 * bd1
        y[k] = c0 * x0 + d * uk1
    return y


@njit(fastmath=True)
def _sim_order2(
    e00: float,
    e01: float,
    e10: float,
    e11: float,
    b00: float,
    b01: float,
    b10: float,
    b11: float,
    c0: float,
    c1: float,
    d: float,
    u: np.ndarray,
) -> np.ndarray:
    m = u.shape[0]
    y = np.empty(m, dtype=np.float64)
    x0 = 0.0
    x1 = 0.0
    y[0] = d * u[0]
    for k in range(1, m):
        uk0 = u[k - 1]
        uk1 = u[k]
        nx0 = x0 * e00 + x1 * e01 + uk0 * b00 + uk1 * b10
        nx1 = x0 * e10 + x1 * e11 + uk0 * b01 + uk1 * b11
        x0 = nx0
        x1 = nx1
        y[k] = c0 * x0 + c1 * x1 + d * uk1
    return y


@njit(fastmath=True)
def _sim_order2_rec(
    tr: float,
    det: float,
    q0: float,
    q1: float,
    q2: float,
    b00: float,
    b01: float,
    b10: float,
    b11: float,
    c0: float,
    c1: float,
    d: float,
    u: np.ndarray,
) -> np.ndarray:
    m = u.shape[0]
    y = np.empty(m, dtype=np.float64)
    y0 = d * u[0]
    y[0] = y0
    if m == 1:
        return y
    u0 = u[0]
    u1 = u[1]
    y1 = c0 * (b00 * u0 + b10 * u1) + c1 * (b01 * u0 + b11 * u1) + d * u1
    y[1] = y1
    prev2 = y0
    prev1 = y1
    for k in range(2, m):
        yk = tr * prev1 - det * prev2 + q0 * u[k] + q1 * u[k - 1] + q2 * u[k - 2]
        y[k] = yk
        prev2 = prev1
        prev1 = yk
    return y


@njit(fastmath=True)
def _sim_general(ad: np.ndarray, bd0: np.ndarray, bd1: np.ndarray, c: np.ndarray, d: float, u: np.ndarray) -> np.ndarray:
    m = u.shape[0]
    n = c.shape[0]
    y = np.empty(m, dtype=np.float64)
    x = np.zeros(n, dtype=np.float64)
    xn = np.empty(n, dtype=np.float64)
    y[0] = d * u[0]
    for k in range(1, m):
        uk0 = u[k - 1]
        uk1 = u[k]
        for j in range(n):
            acc = uk0 * bd0[j] + uk1 * bd1[j]
            for i in range(n):
                acc += x[i] * ad[i, j]
            xn[j] = acc
        out = d * uk1
        for j in range(n):
            val = xn[j]
            x[j] = val
            out += c[j] * val
        y[k] = out
    return y


class Solver:
    def __init__(self) -> None:
        self._cache: dict[tuple[Any, ...], tuple[Any, ...]] = {}
        self._last_num = None
        self._last_den = None
        self._last_dt = None
        self._last_coeffs = None
        dummy = np.array([0.0, 1.0, 0.5], dtype=np.float64)
        _sim_order1(0.9, 0.1, 0.1, 1.0, 0.0, dummy)
        _sim_order2(0.9, -0.1, 0.1, 0.95, 0.05, 0.01, 0.05, 0.01, 1.0, 0.0, 0.0, dummy)
        _sim_order2_rec(1.85, 0.86, 0.01, 0.02, 0.01, 0.05, 0.01, 0.05, 0.01, 1.0, 0.0, 0.0, dummy)
        _sim_general(np.eye(3, dtype=np.float64), np.ones(3), np.ones(3), np.ones(3), 0.0, dummy)

    def solve(self, problem, **kwargs) -> Any:
        num_in = problem["num"]
        den_in = problem["den"]
        u_in = problem["u"]
        if isinstance(u_in, np.ndarray) and u_in.dtype == np.float64 and u_in.ndim == 1 and u_in.flags.c_contiguous:
            u = u_in
        else:
            u = np.ascontiguousarray(u_in, dtype=np.float64)
        t = problem["t"]
        m = u.shape[0]

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

        if float(u[0]) == 0.0 and float(u[m - 1]) == 0.0 and float(u[m // 2]) == 0.0 and not u.any():
            return {"yout": [0.0] * m}

        dt = 0.0 if m < 2 else float(t[1] - t[0])
        if num_in is self._last_num and den_in is self._last_den and dt == self._last_dt:
            coeffs = self._last_coeffs
        else:
            num_tuple = self._small_tuple(num_in)
            den_tuple = self._small_tuple(den_in)
            key = (num_tuple, den_tuple, dt)
            coeffs = self._cache.get(key)
            if coeffs is None:
                coeffs = self._make_coeffs(num_tuple, den_tuple, dt)
                if len(self._cache) > 256:
                    self._cache.clear()
                self._cache[key] = coeffs
            self._last_num = num_in
            self._last_den = den_in
            self._last_dt = dt
            self._last_coeffs = coeffs

        kind = coeffs[0]
        if kind == -1:
            system = signal.lti(num_in, den_in)
            _, yout, _ = signal.lsim(system, u, t)
            return {"yout": yout.tolist()}

        if kind == 0:
            gain = coeffs[1]
            if gain == 0.0:
                return {"yout": [0.0] * m}
            return {"yout": (u * gain).tolist()}

        if m == 1:
            return {"yout": [float(coeffs[-1]) * float(u[0])]}

        if kind == 1:
            _, ad, bd0, bd1, c0, d = coeffs
            y = _sim_order1(ad, bd0, bd1, c0, d, u)
            return {"yout": y.tolist()}

        if kind == 2:
            _, e00, e01, e10, e11, b00, b01, b10, b11, c0, c1, tr, det, q0, q1, q2, d = coeffs
            if m >= 4096:
                y = _sim_order2_rec(tr, det, q0, q1, q2, b00, b01, b10, b11, c0, c1, d, u)
            else:
                y = _sim_order2(e00, e01, e10, e11, b00, b01, b10, b11, c0, c1, d, u)
            return {"yout": y.tolist()}

        if kind == 3:
            _, ad, bd0, bd1, c, d = coeffs
            y = _sim_general(ad, bd0, bd1, c, d, u)
            return {"yout": y.tolist()}

    def _small_tuple(self, values) -> tuple[float, ...]:
        if isinstance(values, np.ndarray):
            arr = values
        else:
            arr = np.asarray(values)
        if arr.ndim == 1:
            n = arr.shape[0]
            if n == 1:
                return (float(arr[0]),)
            if n == 2:
                return (float(arr[0]), float(arr[1]))
            if n == 3:
                return (float(arr[0]), float(arr[1]), float(arr[2]))
            if n == 4:
                return (float(arr[0]), float(arr[1]), float(arr[2]), float(arr[3]))
            if n == 5:
                return (float(arr[0]), float(arr[1]), float(arr[2]), float(arr[3]), float(arr[4]))
        return tuple(float(x) for x in np.ravel(arr))

    def _normalize(self, num_tuple: tuple[float, ...], den_tuple: tuple[float, ...]) -> tuple[np.ndarray, np.ndarray] | None:
        den = np.array(den_tuple, dtype=np.float64)
        num = np.array(num_tuple, dtype=np.float64)

        first = 0
        while first < den.size - 1 and den[first] == 0.0:
            first += 1
        den = den[first:]
        if den.size == 0 or np.all(den == 0.0):
            return None

        scale = den[0]
        den = den / scale
        num = num / scale

        lead = 0
        while lead < num.size - 1 and abs(num[lead]) <= 1e-14:
            lead += 1
        num = num[lead:]
        return num, den

    def _make_coeffs(self, num_tuple: tuple[float, ...], den_tuple: tuple[float, ...], dt: float) -> tuple[Any, ...]:
        normalized = self._normalize(num_tuple, den_tuple)
        if normalized is None:
            return (-1,)
        num, den = normalized
        k = den.size
        if num.size > k:
            return (-1,)

        if k == 1:
            gain = float(num[0]) if num.size else 0.0
            return (0, gain)

        if num.size < k:
            padded = np.zeros(k, dtype=np.float64)
            padded[k - num.size :] = num
            num = padded

        d = float(num[0])
        c = num[1:] - d * den[1:]
        n = k - 1

        if dt == 0.0:
            if n == 1:
                return (1, 1.0, 0.0, 0.0, float(c[0]), d)
            if n == 2:
                q0 = d
                q1 = 0.0
                q2 = 0.0
                return (2, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, float(c[0]), float(c[1]), 2.0, 1.0, q0, q1, q2, d)

        if n == 1:
            a = float(den[1])
            lam = -a
            z = lam * dt
            ad = math.exp(z)
            if abs(z) < 1e-6:
                z2 = z * z
                z3 = z2 * z
                z4 = z2 * z2
                i0 = dt * (1.0 + z / 2.0 + z2 / 6.0 + z3 / 24.0 + z4 / 120.0)
                bd0 = dt * (0.5 + z / 3.0 + z2 / 8.0 + z3 / 30.0 + z4 / 144.0)
            elif lam == 0.0:
                i0 = dt
                bd0 = 0.5 * dt
            else:
                i0 = math.expm1(z) / lam
                bd0 = (ad * (z - 1.0) + 1.0) / (lam * lam * dt)
            bd1 = i0 - bd0
            return (1, float(ad), float(bd0), float(bd1), float(c[0]), d)

        if n == 2:
            made = self._make_order2_coeffs(den, c, d, dt)
            if made is not None:
                return made

        return self._make_general_coeffs(den, c, d, dt)

    def _make_order2_coeffs(self, den: np.ndarray, c: np.ndarray, d: float, dt: float) -> tuple[Any, ...] | None:
        a = float(den[1])
        b = float(den[2])
        if abs(b) < 1e-14:
            return None

        disc = 0.25 * a * a - b
        mu_h = -0.5 * a * dt
        try:
            scale = math.exp(mu_h)
            if disc > 0.0:
                delta = math.sqrt(disc)
                z = delta * dt
                ch = math.cosh(z)
                sh_over_delta = dt if abs(delta) < 1e-14 else math.sinh(z) / delta
            elif disc < 0.0:
                omega = math.sqrt(-disc)
                z = omega * dt
                ch = math.cos(z)
                sh_over_delta = dt if abs(omega) < 1e-14 else math.sin(z) / omega
            else:
                ch = 1.0
                sh_over_delta = dt
        except (OverflowError, ValueError):
            return None

        e00 = scale * (ch - 0.5 * a * sh_over_delta)
        e01 = scale * (-b * sh_over_delta)
        e10 = scale * sh_over_delta
        e11 = scale * (ch + 0.5 * a * sh_over_delta)

        if dt * max(1.0, abs(a), abs(b)) < 1e-4:
            bd0_col, bd1_col = self._order2_input_series(a, b, dt)
            b00 = float(bd0_col[0])
            b01 = float(bd0_col[1])
            b10 = float(bd1_col[0])
            b11 = float(bd1_col[1])
        else:
            v0 = e00 - 1.0
            v1 = e10
            i0b0 = v1
            i0b1 = (-v0 - a * v1) / b

            ainv_eb0 = e10
            ainv_eb1 = (-e00 - a * e10) / b
            ainv_i0b0 = i0b1
            ainv_i0b1 = (-i0b0 - a * i0b1) / b

            b00 = ainv_eb0 - ainv_i0b0 / dt
            b01 = ainv_eb1 - ainv_i0b1 / dt
            b10 = i0b0 - b00
            b11 = i0b1 - b01

        tr = e00 + e11
        det = e00 * e11 - e01 * e10
        a1 = -tr
        a2 = det
        c0 = float(c[0])
        c1 = float(c[1])
        p0 = c0 * b10 + c1 * b11
        p1 = c0 * (b00 - e11 * b10 + e01 * b11) + c1 * (e10 * b10 + b01 - e00 * b11)
        p2 = c0 * (-e11 * b00 + e01 * b01) + c1 * (e10 * b00 - e00 * b01)
        q0 = d + p0
        q1 = d * a1 + p1
        q2 = d * a2 + p2
        vals = (e00, e01, e10, e11, b00, b01, b10, b11, tr, det, q0, q1, q2)
        if not all(math.isfinite(v) for v in vals):
            return None
        return (2, float(e00), float(e01), float(e10), float(e11), float(b00), float(b01), float(b10), float(b11), c0, c1, float(tr), float(det), float(q0), float(q1), float(q2), d)

    def _order2_input_series(self, a: float, b: float, dt: float) -> tuple[np.ndarray, np.ndarray]:
        avec0 = 1.0
        avec1 = 0.0
        i0b0 = 0.0
        i0b1 = 0.0
        jhb0 = 0.0
        jhb1 = 0.0
        fact = 1.0
        pow_dt = dt
        for kk in range(18):
            i_coef = pow_dt / ((kk + 1) * fact)
            j_coef = pow_dt / ((kk + 2) * fact)
            i0b0 += avec0 * i_coef
            i0b1 += avec1 * i_coef
            jhb0 += avec0 * j_coef
            jhb1 += avec1 * j_coef
            next0 = -a * avec0 - b * avec1
            next1 = avec0
            avec0 = next0
            avec1 = next1
            fact *= kk + 1
            pow_dt *= dt
        bd0 = np.array([jhb0, jhb1], dtype=np.float64)
        bd1 = np.array([i0b0 - jhb0, i0b1 - jhb1], dtype=np.float64)
        return bd0, bd1

    def _make_general_coeffs(self, den: np.ndarray, c: np.ndarray, d: float, dt: float) -> tuple[Any, ...]:
        n = den.size - 1
        a = den[1:]
        A = np.zeros((n, n), dtype=np.float64)
        A[0, :] = -a
        if n > 1:
            A[1:, :-1] = np.eye(n - 1, dtype=np.float64)
        B = np.zeros((n, 1), dtype=np.float64)
        B[0, 0] = 1.0
        M = np.vstack(
            [
                np.hstack([A * dt, B * dt, np.zeros((n, 1), dtype=np.float64)]),
                np.hstack([np.zeros((1, n + 1), dtype=np.float64), np.ones((1, 1), dtype=np.float64)]),
                np.zeros((1, n + 2), dtype=np.float64),
            ]
        )
        expmt = linalg.expm(M.T)
        ad = np.ascontiguousarray(expmt[:n, :n])
        bd1 = np.ascontiguousarray(expmt[n + 1 :, :n].reshape(n))
        bd0 = np.ascontiguousarray((expmt[n : n + 1, :n] - expmt[n + 1 :, :n]).reshape(n))
        return (3, ad, bd0, bd1, np.ascontiguousarray(c, dtype=np.float64), float(d))
--- 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.0550s
Total Solver Time:   0.0465s
Raw Speedup:         194.6581 x
Final Reward (Score): 194.6581
---------------------------
..

==================================== 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 113.63s (0:01:53) =========================
