--- Final Solver (used for performance test) ---
from typing import Any
import os
import sys
import numpy as np

_DIR = os.path.dirname(os.path.abspath(__file__))
if _DIR not in sys.path:
    sys.path.insert(0, _DIR)

try:
    import _simplex
    _csolve = _simplex.solve
except Exception:
    _csolve = None

from numba import njit


@njit(cache=True)
def _condat_simplex(y):
    n = y.shape[0]
    x = np.empty(n, dtype=np.float64)
    if n == 0:
        return x
    if n == 1:
        x[0] = 1.0
        return x
    v = np.empty(n, dtype=np.float64)
    vt = np.empty(n, dtype=np.float64)
    v[0] = y[0]
    vlen = 1
    vtlen = 0
    rho = y[0] - 1.0
    for i in range(1, n):
        yi = y[i]
        if yi > rho:
            rho += (yi - rho) / (vlen + 1)
            if rho > yi - 1.0:
                v[vlen] = yi
                vlen += 1
            else:
                for j in range(vlen):
                    vt[vtlen] = v[j]
                    vtlen += 1
                v[0] = yi
                vlen = 1
                rho = yi - 1.0
    for i in range(vtlen):
        yi = vt[i]
        if yi > rho:
            v[vlen] = yi
            vlen += 1
            rho += (yi - rho) / vlen
    while True:
        removed = False
        cur = vlen
        write = 0
        for read in range(vlen):
            val = v[read]
            if val <= rho:
                cur -= 1
                rho += (rho - val) / cur
                removed = True
            else:
                v[write] = val
                write += 1
        vlen = write
        if not removed:
            break
    theta = rho
    for i in range(n):
        d = y[i] - theta
        x[i] = d if d > 0.0 else 0.0
    return x


class Solver:
    def __init__(self):
        if _csolve is None:
            _condat_simplex(np.array([0.5, 0.3, 0.2, -0.1, 1.2], dtype=np.float64))
            _condat_simplex(np.array([1.0], dtype=np.float64))
            _condat_simplex(np.array([0.4, 0.4], dtype=np.float64))

    def solve(self, problem, **kwargs) -> Any:
        y = problem["y"]
        if _csolve is not None:
            return {"solution": _csolve(y)}
        arr = np.asarray(y, dtype=np.float64).ravel()
        return {"solution": _condat_simplex(arr)}
--- 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: 6.4719s
Total Solver Time:   1.8754s
Raw Speedup:         3.4509 x
Final Reward (Score): 3.4509
---------------------------
..

==================================== 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 248.12s (0:04:08) =========================
