--- Final Solver (used for performance test) ---
"""
Euclidean projection onto the probability simplex.
O(n) linear-time algorithm (Condat et al.) in numba,
with parallel reduction for the scan step.
"""

from typing import Any

import numpy as np
from numba import njit, prange


@njit(cache=True, parallel=True, fastmath=True)
def _project_numba_parallel(v, u, x):
    """
    Project v onto the probability simplex in O(n) expected time.
    Uses parallel reduction for the sum/count scan step.
    """
    n = len(v)
    if n == 1:
        x[0] = 1.0
        return

    for i in range(n):
        u[i] = v[i]

    s = 0.0
    rho = 0
    lo = 0
    hi = n

    while lo < hi:
        size = hi - lo
        if size == 1:
            val = u[lo]
            if (s + val - 1.0) < val * (rho + 1):
                s += val
                rho += 1
            break

        mid = lo + (size >> 1)
        a, b, c = u[lo], u[mid], u[hi - 1]
        if a > b:
            pivot = c if b > c else (c if a > c else a)
        else:
            pivot = a if a > c else (c if b > c else b)

        # Parallel sum/count reduction
        sum_geq = 0.0
        count_geq = 0
        for i in prange(lo, hi):
            if u[i] >= pivot:
                sum_geq += u[i]
                count_geq += 1

        if (s + sum_geq - 1.0) < pivot * (rho + count_geq):
            s += sum_geq
            rho += count_geq
            w = lo
            for i in range(lo, hi):
                if u[i] < pivot:
                    u[w] = u[i]
                    w += 1
            hi = w
        else:
            w = lo
            for i in range(lo, hi):
                if u[i] > pivot:
                    u[w] = u[i]
                    w += 1
            hi = w

    theta = (s - 1.0) / rho
    for i in prange(n):
        val = v[i] - theta
        x[i] = val if val > 0.0 else 0.0


@njit(cache=True, fastmath=True)
def _project_numba_serial(v, u, x):
    """Serial version for small n where parallel overhead hurts."""
    n = len(v)
    if n == 1:
        x[0] = 1.0
        return

    for i in range(n):
        u[i] = v[i]

    s = 0.0
    rho = 0
    lo = 0
    hi = n

    while lo < hi:
        size = hi - lo
        if size == 1:
            val = u[lo]
            if (s + val - 1.0) < val * (rho + 1):
                s += val
                rho += 1
            break

        mid = lo + (size >> 1)
        a, b, c = u[lo], u[mid], u[hi - 1]
        if a > b:
            pivot = c if b > c else (c if a > c else a)
        else:
            pivot = a if a > c else (c if b > c else b)

        sum_geq = 0.0
        count_geq = 0
        for i in range(lo, hi):
            if u[i] >= pivot:
                sum_geq += u[i]
                count_geq += 1

        if (s + sum_geq - 1.0) < pivot * (rho + count_geq):
            s += sum_geq
            rho += count_geq
            w = lo
            for i in range(lo, hi):
                if u[i] < pivot:
                    u[w] = u[i]
                    w += 1
            hi = w
        else:
            w = lo
            for i in range(lo, hi):
                if u[i] > pivot:
                    u[w] = u[i]
                    w += 1
            hi = w

    theta = (s - 1.0) / rho
    for i in range(n):
        val = v[i] - theta
        x[i] = val if val > 0.0 else 0.0


# Warm up both versions with various sizes
for _n in [1, 2, 3, 5, 10, 50, 100, 1000, 10000]:
    _wy = np.random.RandomState(42).randn(_n)
    _wu = np.empty(_n, dtype=np.float64)
    _wx = np.empty(_n, dtype=np.float64)
    _project_numba_serial(_wy, _wu, _wx)

for _n in [10000, 100000]:
    _wy = np.random.RandomState(42).randn(_n)
    _wu = np.empty(_n, dtype=np.float64)
    _wx = np.empty(_n, dtype=np.float64)
    _project_numba_parallel(_wy, _wu, _wx)

# Pre-allocated workspace
_MAX_PREALLOC = 2_000_000
_ws_u = np.empty(_MAX_PREALLOC, dtype=np.float64)
_ws_x = np.empty(_MAX_PREALLOC, dtype=np.float64)

# Threshold for parallel vs serial
_PARALLEL_THRESHOLD = 5000


def _convert_to_array(y):
    """Convert input to numpy float64 array as fast as possible."""
    if isinstance(y, np.ndarray):
        if y.dtype == np.float64 and y.ndim == 1:
            return y
        return np.asarray(y, dtype=np.float64).ravel()
    return np.fromiter(y, dtype=np.float64, count=len(y))


class Solver:
    def __init__(self):
        pass

    def solve(self, problem, **kwargs) -> Any:
        y = _convert_to_array(problem["y"])
        n = len(y)

        if n <= _MAX_PREALLOC:
            u = _ws_u[:n]
            x = _ws_x[:n]
        else:
            u = np.empty(n, dtype=np.float64)
            x = np.empty(n, dtype=np.float64)

        if n >= _PARALLEL_THRESHOLD:
            _project_numba_parallel(y, u, x)
        else:
            _project_numba_serial(y, u, x)

        if n <= _MAX_PREALLOC:
            return {"solution": x.copy()}
        return {"solution": x}
--- 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.4818s
Total Solver Time:   83.8062s
Raw Speedup:         0.0773 x
Final Reward (Score): 1.0000
---------------------------
.F

=================================== FAILURES ===================================
_____________________________ test_solver_speedup ______________________________

performance_results = {'raw_speedup': 0.07734229814291786, 'speedup': 1.0, 'validity': True}

    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.08x)
E       assert 0.07734229814291786 > 1.0

/tests/test_outputs.py:205: AssertionError
==================================== PASSES ====================================
=========================== short test summary info ============================
PASSED ../tests/test_outputs.py::test_solver_exists
PASSED ../tests/test_outputs.py::test_solver_validity
FAILED ../tests/test_outputs.py::test_solver_speedup - AssertionError: Solver...
=================== 1 failed, 2 passed in 1808.45s (0:30:08) ===================
