--- Final Solver (used for performance test) ---
from typing import Any

import hashlib
import importlib.util
import os
import subprocess
import sysconfig
import tempfile

import numpy as np
from numba import njit


_C_SOURCE = r'''
#include <Python.h>
#include <math.h>
#include <stdlib.h>

static inline int read_double(PyObject *obj, double *out) {
    if (PyFloat_CheckExact(obj)) {
        *out = PyFloat_AS_DOUBLE(obj);
        return 0;
    }
    if (PyLong_CheckExact(obj)) {
        *out = PyLong_AsDouble(obj);
        if (*out == -1.0 && PyErr_Occurred()) return -1;
        return 0;
    }
    *out = PyFloat_AsDouble(obj);
    if (*out == -1.0 && PyErr_Occurred()) return -1;
    return 0;
}

static inline double median3(double a, double b, double c) {
    if (a < b) {
        if (b < c) return b;
        if (a < c) return c;
        return a;
    }
    if (a < c) return a;
    if (b < c) return c;
    return b;
}

static double select_theta(double *work, Py_ssize_t n) {
    Py_ssize_t m = n;
    double total = 0.0;
    Py_ssize_t count = 0;

    while (m > 0) {
        double pivot = median3(work[0], work[m >> 1], work[m - 1]);
        Py_ssize_t left = 0;
        Py_ssize_t i = 0;
        Py_ssize_t right = m - 1;
        double sum_greater = 0.0;

        while (i <= right) {
            double value = work[i];
            if (value > pivot) {
                work[i] = work[left];
                work[left] = value;
                sum_greater += value;
                left++;
                i++;
            } else if (value < pivot) {
                work[i] = work[right];
                work[right] = value;
                right--;
            } else {
                i++;
            }
        }

        Py_ssize_t equal_count = right - left + 1;
        Py_ssize_t ge_count = left + equal_count;
        double ge_sum = sum_greater + pivot * (double)equal_count;

        if ((total + ge_sum) - (double)(count + ge_count) * pivot < 1.0) {
            total += ge_sum;
            count += ge_count;
            Py_ssize_t lower_count = m - ge_count;
            Py_ssize_t lower_start = ge_count;
            for (Py_ssize_t j = 0; j < lower_count; j++) {
                work[j] = work[lower_start + j];
            }
            m = lower_count;
        } else {
            m = left;
        }
    }

    return (total - 1.0) / (double)count;
}

static int project_core(PyObject *seq, double *out, Py_ssize_t n, int k) {
    if (n <= 0) {
        PyErr_SetString(PyExc_ValueError, "empty input");
        return -1;
    }
    if (n == 1) {
        out[0] = 1.0;
        return 0;
    }

    PyObject *fast = NULL;
    PyObject **items;
    if (PyList_Check(seq) || PyTuple_Check(seq)) {
        if (Py_SIZE(seq) != n) {
            PyErr_SetString(PyExc_ValueError, "length changed");
            return -1;
        }
        items = PySequence_Fast_ITEMS(seq);
    } else {
        fast = PySequence_Fast(seq, "expected a sequence");
        if (fast == NULL) return -1;
        if (PySequence_Fast_GET_SIZE(fast) != n) {
            Py_DECREF(fast);
            PyErr_SetString(PyExc_ValueError, "length changed");
            return -1;
        }
        items = PySequence_Fast_ITEMS(fast);
    }

    if (k > 96) k = 96;
    if (k > n) k = (int)n;
    double top[96];
    int size = 0;

    for (Py_ssize_t idx = 0; idx < n; idx++) {
        double value;
        if (read_double(items[idx], &value) < 0) {
            Py_XDECREF(fast);
            return -1;
        }
        if (size < k) {
            int j = size - 1;
            while (j >= 0 && top[j] < value) {
                top[j + 1] = top[j];
                j--;
            }
            top[j + 1] = value;
            size++;
        } else if (value > top[k - 1]) {
            int j = k - 2;
            while (j >= 0 && top[j] < value) {
                top[j + 1] = top[j];
                j--;
            }
            top[j + 1] = value;
        }
    }

    double cssv = 0.0;
    double theta = 0.0;
    int active = 0;
    for (int j = 0; j < k; j++) {
        double value = top[j];
        cssv += value;
        double candidate = (cssv - 1.0) / (double)(j + 1);
        if (value > candidate) {
            theta = candidate;
            active = j + 1;
        }
    }

    int success = 0;
    if (active < k) {
        success = top[active] <= theta + 1e-12;
    } else if (k == n) {
        success = 1;
    }

    if (success) {
        for (Py_ssize_t idx = 0; idx < n; idx++) {
            double value;
            if (read_double(items[idx], &value) < 0) {
                Py_XDECREF(fast);
                return -1;
            }
            value -= theta;
            out[idx] = value > 0.0 ? value : 0.0;
        }
        Py_XDECREF(fast);
        return 0;
    }

    double *work = (double *)malloc((size_t)n * sizeof(double));
    if (work == NULL) {
        Py_XDECREF(fast);
        PyErr_NoMemory();
        return -1;
    }
    for (Py_ssize_t idx = 0; idx < n; idx++) {
        double value;
        if (read_double(items[idx], &value) < 0) {
            free(work);
            Py_XDECREF(fast);
            return -1;
        }
        work[idx] = value;
        out[idx] = value;
    }
    Py_XDECREF(fast);

    theta = select_theta(work, n);
    free(work);
    for (Py_ssize_t idx = 0; idx < n; idx++) {
        double value = out[idx] - theta;
        out[idx] = value > 0.0 ? value : 0.0;
    }
    return 0;
}

static PyObject *project(PyObject *self, PyObject *args) {
    PyObject *seq;
    PyObject *out_obj;
    int k = 64;
    if (!PyArg_ParseTuple(args, "OO|i", &seq, &out_obj, &k)) return NULL;
    Py_ssize_t n = PyObject_Length(seq);
    if (n < 0) return NULL;

    Py_buffer view;
    if (PyObject_GetBuffer(out_obj, &view, PyBUF_WRITABLE) < 0) return NULL;
    if (view.len < (Py_ssize_t)(n * (Py_ssize_t)sizeof(double))) {
        PyBuffer_Release(&view);
        PyErr_SetString(PyExc_ValueError, "output buffer too small");
        return NULL;
    }

    int status = project_core(seq, (double *)view.buf, n, k);
    PyBuffer_Release(&view);
    if (status < 0) return NULL;
    Py_RETURN_NONE;
}

static PyMethodDef methods[] = {
    {"project", project, METH_VARARGS, "Project onto the probability simplex."},
    {NULL, NULL, 0, NULL}
};

static struct PyModuleDef moduledef = {
    PyModuleDef_HEAD_INIT,
    "_simplex_fast",
    NULL,
    -1,
    methods
};

PyMODINIT_FUNC PyInit__simplex_fast(void) {
    return PyModule_Create(&moduledef);
}
'''


def _load_c_projector():
    digest = hashlib.sha256(_C_SOURCE.encode()).hexdigest()[:16]
    tmpdir = tempfile.gettempdir()
    suffix = sysconfig.get_config_var("EXT_SUFFIX") or ".so"
    so_path = os.path.join(tmpdir, f"_simplex_fast_{digest}{suffix}")
    c_path = os.path.join(tmpdir, f"_simplex_fast_{digest}.c")
    if not os.path.exists(so_path):
        with open(c_path, "w", encoding="utf-8") as handle:
            handle.write(_C_SOURCE)
        include = sysconfig.get_paths()["include"]
        cmd = ["gcc", "-O3", "-shared", "-fPIC", f"-I{include}", c_path, "-o", so_path]
        subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    spec = importlib.util.spec_from_file_location("_simplex_fast", so_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module.project


@njit(cache=True, fastmath=True)
def _project_sort(y):
    n = y.size
    u = np.sort(y).copy()
    cssv = 0.0
    theta = 0.0
    for j in range(n):
        value = u[n - 1 - j]
        cssv += value
        candidate = (cssv - 1.0) / (j + 1)
        if value > candidate:
            theta = candidate
    out = np.empty(n, dtype=np.float64)
    for i in range(n):
        value = y[i] - theta
        out[i] = value if value > 0.0 else 0.0
    return out


@njit(cache=True, fastmath=True)
def _median3(a, b, c):
    if a < b:
        if b < c:
            return b
        if a < c:
            return c
        return a
    if a < c:
        return a
    if b < c:
        return c
    return b


@njit(cache=True, fastmath=True)
def _project_select(y):
    n = y.size
    work = y.copy()
    m = n
    total = 0.0
    count = 0

    while m > 0:
        pivot = _median3(work[0], work[m >> 1], work[m - 1])
        left = 0
        i = 0
        right = m - 1
        sum_greater = 0.0

        while i <= right:
            value = work[i]
            if value > pivot:
                work[i] = work[left]
                work[left] = value
                sum_greater += value
                left += 1
                i += 1
            elif value < pivot:
                work[i] = work[right]
                work[right] = value
                right -= 1
            else:
                i += 1

        equal_count = right - left + 1
        ge_count = left + equal_count
        ge_sum = sum_greater + pivot * equal_count

        if (total + ge_sum) - (count + ge_count) * pivot < 1.0:
            total += ge_sum
            count += ge_count
            lower_count = m - ge_count
            lower_start = ge_count
            for j in range(lower_count):
                work[j] = work[lower_start + j]
            m = lower_count
        else:
            m = left

    theta = (total - 1.0) / count
    out = np.empty(n, dtype=np.float64)
    for i in range(n):
        value = y[i] - theta
        out[i] = value if value > 0.0 else 0.0
    return out


@njit(cache=True, fastmath=True)
def _project_topk(y, k):
    n = y.size
    if k > n:
        k = n
    top = np.empty(k, dtype=np.float64)
    size = 0

    for i in range(n):
        value = y[i]
        if size < k:
            j = size - 1
            while j >= 0 and top[j] < value:
                top[j + 1] = top[j]
                j -= 1
            top[j + 1] = value
            size += 1
        elif value > top[k - 1]:
            j = k - 2
            while j >= 0 and top[j] < value:
                top[j + 1] = top[j]
                j -= 1
            top[j + 1] = value

    cssv = 0.0
    theta = 0.0
    active = 0
    for j in range(k):
        value = top[j]
        cssv += value
        candidate = (cssv - 1.0) / (j + 1)
        if value > candidate:
            theta = candidate
            active = j + 1

    success = False
    if active < k:
        if top[active] <= theta + 1e-12:
            success = True
    elif k == n:
        success = True

    out = np.empty(n, dtype=np.float64)
    if success:
        for i in range(n):
            value = y[i] - theta
            out[i] = value if value > 0.0 else 0.0
    return out, success


@njit(cache=True, fastmath=True)
def _project_small(y):
    n = y.size
    if n == 1:
        out = np.empty(1, dtype=np.float64)
        out[0] = 1.0
        return out

    values = np.empty(n, dtype=np.float64)
    for i in range(n):
        value = y[i]
        j = i - 1
        while j >= 0 and values[j] < value:
            values[j + 1] = values[j]
            j -= 1
        values[j + 1] = value

    cssv = 0.0
    theta = 0.0
    for j in range(n):
        value = values[j]
        cssv += value
        candidate = (cssv - 1.0) / (j + 1)
        if value > candidate:
            theta = candidate

    out = np.empty(n, dtype=np.float64)
    for i in range(n):
        value = y[i] - theta
        out[i] = value if value > 0.0 else 0.0
    return out


class Solver:
    def __init__(self):
        try:
            self._c_project = _load_c_projector()
        except Exception:
            self._c_project = None
        sample = np.array([0.2, -0.7, 1.5, 0.1], dtype=np.float64)
        _project_small(sample[:2])
        _project_sort(sample)
        _project_select(sample)
        _project_topk(sample, 3)

    def solve(self, problem, **kwargs) -> Any:
        raw = problem["y"]
        if self._c_project is not None and isinstance(raw, (list, tuple)):
            n = len(raw)
            out = np.empty(n, dtype=np.float64)
            if n <= 64:
                k = 64
            elif n <= 4096:
                k = 8
            else:
                k = 16
            self._c_project(raw, out, k)
            return {"solution": out}

        y = np.asarray(raw, dtype=np.float64).ravel()
        n = y.size
        if n <= 24:
            x = _project_small(y)
        elif n <= 64:
            x = _project_sort(y)
        else:
            x, success = _project_topk(y, 64)
            if not success:
                x = _project_select(y)
        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.6075s
Total Solver Time:   0.2525s
Raw Speedup:         26.1639 x
Final Reward (Score): 26.1639
---------------------------
..

==================================== 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 238.59s (0:03:58) =========================
