# cython: language_level=3
# cython: boundscheck=False
# cython: wraparound=False
# cython: cdivision=True
# cython: nonecheck=False

"""Monte-Carlo Pareto ranking core (Cython).

Goal
----
Given per-point objective distributions (Normal with mean `mu` and variance
`sigma2`), repeatedly sample populations, compute Pareto ranking via
`pareto_core` (front id + within-front crowding), and accumulate statistics.

This module is designed as a hot-path kernel for async PBT selection.

Key properties
--------------
- Deterministic given the same seed.
- No allocations inside the sampling loop.
- Calls Pareto ranking on C-level via `rate_population_inplace`.

Public API
----------
mc_rank_counts(mu, sigma2, n_samples, percentile=0.2, seed=0,
              mode_i=1, within_mode_i=0)
    -> (win_count, lose_count, rank_sum)

mc_rank_select(...)
    -> (winners, losers, win_count, lose_count, rank_sum)

mc_rank_select_detailed(...)
    -> (winners, losers, win_count, lose_count, rank_sum,
        place_stats, front_stats, within_stats)

    Each *_stats is a dict with keys:
        mean, median, std, q25, q75

    The detailed path uses constant-memory streaming quantile estimators
    (P² algorithm) to avoid storing all (n_samples x N) values.
"""

from libc.math cimport sqrt, log, cos, sin
from libc.math cimport fabs
from libc.stdint cimport uint64_t
from libc.stdlib cimport malloc, free
from libc.math cimport fabs

import numpy as np
cimport numpy as cnp

from zen_fronts.pareto_core._pareto_rating_cy cimport (
    rate_population_inplace,
    Pair,
    OrderKey,
    I64,
    I32,
    U8,
)


ctypedef cnp.float64_t F64
ctypedef cnp.int32_t I32T


# ------------------------------ RNG (xoshiro256**) ------------------------------

cdef inline uint64_t _rotl(const uint64_t x, int k) noexcept nogil:
    return (x << k) | (x >> (64 - k))


cdef inline uint64_t _splitmix64(uint64_t* state) noexcept nogil:
    cdef uint64_t z
    state[0] = state[0] + <uint64_t>0x9E3779B97F4A7C15
    z = state[0]
    z = (z ^ (z >> 30)) * <uint64_t>0xBF58476D1CE4E5B9
    z = (z ^ (z >> 27)) * <uint64_t>0x94D049BB133111EB
    return z ^ (z >> 31)


cdef struct Xoshiro256ss:
    uint64_t s0
    uint64_t s1
    uint64_t s2
    uint64_t s3
    bint has_spare
    double spare


cdef inline void _xoshiro_seed(Xoshiro256ss* st, uint64_t seed) noexcept nogil:
    cdef uint64_t sm = seed
    st.s0 = _splitmix64(&sm)
    st.s1 = _splitmix64(&sm)
    st.s2 = _splitmix64(&sm)
    st.s3 = _splitmix64(&sm)
    st.has_spare = False
    st.spare = 0.0


cdef inline uint64_t _xoshiro_next_u64(Xoshiro256ss* st) noexcept nogil:
    cdef uint64_t result
    cdef uint64_t t
    result = _rotl(st.s1 * <uint64_t>5, 7) * <uint64_t>9
    t = st.s1 << 17
    st.s2 ^= st.s0
    st.s3 ^= st.s1
    st.s1 ^= st.s2
    st.s0 ^= st.s3
    st.s2 ^= t
    st.s3 = _rotl(st.s3, 45)
    return result


cdef inline double _xoshiro_u01(Xoshiro256ss* st) noexcept nogil:
    # Convert to double in [0,1) using top 53 bits.
    cdef uint64_t x = _xoshiro_next_u64(st)
    return ((x >> 11) * 1.1102230246251565e-16)  # 1/2**53


cdef inline double _xoshiro_randn(Xoshiro256ss* st) noexcept nogil:
    # Box-Muller transform.
    cdef double u1, u2, r, theta
    if st.has_spare:
        st.has_spare = False
        return st.spare
    u1 = _xoshiro_u01(st)
    if u1 < 1e-300:
        u1 = 1e-300
    u2 = _xoshiro_u01(st)
    r = sqrt(-2.0 * log(u1))
    theta = 6.2831853071795864769 * u2  # 2*pi
    st.spare = r * sin(theta)
    st.has_spare = True
    return r * cos(theta)


# ------------------------------ helpers ------------------------------

cdef inline cnp.ndarray[F64, ndim=2] _as_contig_f64_2d(object A_obj):
    cdef cnp.ndarray arr = np.asarray(A_obj, dtype=np.float64)
    if arr.ndim != 2:
        raise ValueError("expected a 2D array")
    if not arr.flags["C_CONTIGUOUS"]:
        arr = np.ascontiguousarray(arr)
    return <cnp.ndarray[F64, ndim=2]>arr


cdef inline void _sort5(double* a) noexcept nogil:
    """In-place insertion sort for 5 doubles."""
    cdef int i, j
    cdef double key
    for i in range(1, 5):
        key = a[i]
        j = i - 1
        while j >= 0 and a[j] > key:
            a[j + 1] = a[j]
            j -= 1
        a[j + 1] = key


cdef inline void _p2_init(
    double* q,          # len 5
    I32T* npos,         # len 5
    double* ndes,       # len 5 (desired positions)
    double p,           # target quantile in (0,1)
    double s0, double s1, double s2, double s3, double s4,
) noexcept nogil:
    """Initialize one P² quantile estimator from 5 sorted samples."""
    q[0] = s0
    q[1] = s1
    q[2] = s2
    q[3] = s3
    q[4] = s4

    npos[0] = 1
    npos[1] = 2
    npos[2] = 3
    npos[3] = 4
    npos[4] = 5

    # Desired positions for n=5
    # See: Jain & Chlamtac (1985), P² algorithm.
    ndes[0] = 1.0
    ndes[1] = 1.0 + 2.0 * p
    ndes[2] = 1.0 + 4.0 * p
    ndes[3] = 3.0 + 2.0 * p
    ndes[4] = 5.0


cdef inline double _p2_update(
    double* q,          # len 5
    I32T* npos,         # len 5
    double* ndes,       # len 5
    double p,
    double x,
) noexcept nogil:
    """Update one P² estimator with a new sample. Returns current estimate."""
    cdef int k, i
    cdef double d
    cdef int di
    cdef double qhat

    # Find cell k
    if x < q[0]:
        q[0] = x
        k = 0
    elif x >= q[4]:
        q[4] = x
        k = 3
    else:
        k = 0
        for i in range(4):
            if q[i] <= x < q[i + 1]:
                k = i
                break

    # Increment positions of markers above k
    for i in range(k + 1, 5):
        npos[i] += 1

    # Increment desired positions
    # dn = [0, p/2, p, (1+p)/2, 1]
    ndes[1] += 0.5 * p
    ndes[2] += p
    ndes[3] += 0.5 * (1.0 + p)
    ndes[4] += 1.0

    # Adjust marker heights
    for i in range(1, 4):
        d = ndes[i] - <double>npos[i]
        if (d >= 1.0 and (npos[i + 1] - npos[i]) > 1) or (d <= -1.0 and (npos[i - 1] - npos[i]) < -1):
            di = 1 if d >= 1.0 else -1

            # Parabolic prediction
            qhat = q[i] + di / (npos[i + 1] - npos[i - 1]) * (
                (npos[i] - npos[i - 1] + di) * (q[i + 1] - q[i]) / (npos[i + 1] - npos[i]) +
                (npos[i + 1] - npos[i] - di) * (q[i] - q[i - 1]) / (npos[i] - npos[i - 1])
            )

            if q[i - 1] < qhat < q[i + 1]:
                q[i] = qhat
            else:
                # Linear
                q[i] = q[i] + di * (q[i + di] - q[i]) / (npos[i + di] - npos[i])

            npos[i] += di

    return q[2]


cpdef tuple mc_rank_counts(
    object mu_obj,
    object sigma2_obj,
    int n_samples,
    double percentile=0.2,
    unsigned long long seed=0,
    int mode_i=1,
    int within_mode_i=0,
):
    """Accumulate MC win/lose counts and rank sums.

    Parameters
    ----------
    mu_obj:
        (N,M) float64 means.
    sigma2_obj:
        (N,M) float64 variances (>=0 recommended).
    n_samples:
        Monte-Carlo samples.
    percentile:
        Fraction for top/bottom selection per sample; default 0.2.
    seed:
        Seed for deterministic RNG.
    mode_i:
        0=values, 1=ranks (recommended for robustness).
    within_mode_i:
        0=crowding.

    Returns
    -------
    (win_count, lose_count, rank_sum)
        win_count/lose_count: int32 arrays of shape (N,)
        rank_sum: int64 array of shape (N,) (sum of positions in global order)
    """

    cdef cnp.ndarray[F64, ndim=2] mu = _as_contig_f64_2d(mu_obj)
    cdef cnp.ndarray[F64, ndim=2] sigma2 = _as_contig_f64_2d(sigma2_obj)

    if mu.shape[0] != sigma2.shape[0] or mu.shape[1] != sigma2.shape[1]:
        raise ValueError("mu and sigma2 must have the same shape")

    cdef int N = mu.shape[0]
    cdef int M = mu.shape[1]

    if N == 0:
        return (
            np.zeros(0, dtype=np.int32),
            np.zeros(0, dtype=np.int32),
            np.zeros(0, dtype=np.int64),
        )

    if n_samples <= 0:
        raise ValueError("n_samples must be > 0")

    if percentile <= 0.0 or percentile > 1.0:
        raise ValueError("percentile must be in (0,1]")

    cdef int k = <int>np.ceil(percentile * N)
    if k < 1:
        k = 1
    if k > N:
        k = N

    # Precompute sigma = sqrt(max(sigma2, 0)).
    cdef cnp.ndarray[F64, ndim=2] sigma = np.empty((N, M), dtype=np.float64)
    cdef double[:, ::1] s2_v = sigma2
    cdef double[:, ::1] sig_v = sigma
    cdef int i, m
    with nogil:
        for i in range(N):
            for m in range(M):
                if s2_v[i, m] <= 0.0:
                    sig_v[i, m] = 0.0
                else:
                    sig_v[i, m] = sqrt(s2_v[i, m])

    # Allocate once: sample matrix + pareto buffers.
    cdef cnp.ndarray[F64, ndim=2] X = np.empty((N, M), dtype=np.float64)
    cdef cnp.ndarray[F64, ndim=2] R = np.empty((N, M), dtype=np.float64)

    cdef cnp.ndarray[I64, ndim=1] front_id = np.empty(N, dtype=np.int64)
    cdef cnp.ndarray[F64, ndim=1] within = np.empty(N, dtype=np.float64)
    cdef cnp.ndarray[I64, ndim=1] order = np.empty(N, dtype=np.int64)
    cdef cnp.ndarray[I64, ndim=1] items = np.empty(N, dtype=np.int64)
    cdef cnp.ndarray[I64, ndim=1] offsets = np.empty(N + 1, dtype=np.int64)

    cdef cnp.ndarray[U8, ndim=2] dom = np.empty((N, N), dtype=np.uint8)
    cdef cnp.ndarray[I32, ndim=1] n_dom = np.empty(N, dtype=np.int32)
    cdef cnp.ndarray[I32, ndim=1] cur = np.empty(N, dtype=np.int32)
    cdef cnp.ndarray[I32, ndim=1] nxt = np.empty(N, dtype=np.int32)

    cdef cnp.ndarray[I32, ndim=1] win_count = np.zeros(N, dtype=np.int32)
    cdef cnp.ndarray[I32, ndim=1] lose_count = np.zeros(N, dtype=np.int32)
    cdef cnp.ndarray[I64, ndim=1] rank_sum = np.zeros(N, dtype=np.int64)

    cdef Pair* pairs = <Pair*>malloc(max(N, 1) * sizeof(Pair))
    cdef OrderKey* keys = <OrderKey*>malloc(max(N, 1) * sizeof(OrderKey))
    if pairs == NULL or keys == NULL:
        if pairs != NULL:
            free(pairs)
        if keys != NULL:
            free(keys)
        raise MemoryError()

    # Bind typed views once (avoid view creation under nogil).
    cdef double[:, ::1] mu_v = mu
    cdef double[:, ::1] sigm_v = sigma
    cdef double[:, ::1] X_v = X
    cdef double[:, ::1] R_v = R

    cdef I64[::1] front_v = front_id
    cdef double[::1] within_v = within
    cdef I64[::1] order_v = order
    cdef I64[::1] items_v = items
    cdef I64[::1] offsets_v = offsets
    cdef U8[:, ::1] dom_v = dom
    cdef I32[::1] n_dom_v = n_dom
    cdef I32[::1] cur_v = cur
    cdef I32[::1] nxt_v = nxt

    cdef I32[::1] win_v = win_count
    cdef I32[::1] lose_v = lose_count
    cdef I64[::1] rsum_v = rank_sum

    cdef Xoshiro256ss rng
    cdef int s, pos, idx
    cdef int n_fronts

    _xoshiro_seed(&rng, <uint64_t>seed)

    try:
        with nogil:
            for s in range(n_samples):
                # 1) sample X ~ Normal(mu, sigma)
                for i in range(N):
                    for m in range(M):
                        if sigm_v[i, m] == 0.0:
                            X_v[i, m] = mu_v[i, m]
                        else:
                            X_v[i, m] = mu_v[i, m] + sigm_v[i, m] * _xoshiro_randn(&rng)

                # 2) pareto ranking (C-level)
                n_fronts = rate_population_inplace(
                    X_v, mode_i, within_mode_i,
                    front_v, within_v, order_v,
                    offsets_v, items_v,
                    dom_v, n_dom_v, cur_v, nxt_v,
                    pairs, keys,
                    R_v,
                )
                if n_fronts < 0:
                    # error code: break; Python wrapper will raise.
                    break

                # 3) update counts for top-k / bottom-k and rank sums
                for pos in range(k):
                    idx = <int>order_v[pos]
                    win_v[idx] += 1

                for pos in range(N - k, N):
                    idx = <int>order_v[pos]
                    lose_v[idx] += 1

                for pos in range(N):
                    idx = <int>order_v[pos]
                    rsum_v[idx] += pos

        if n_fronts < 0:
            raise ValueError(f"pareto_core returned error code: {n_fronts}")

        return win_count, lose_count, rank_sum

    finally:
        free(pairs)
        free(keys)


cpdef tuple mc_rank_select(
    object mu_obj,
    object sigma2_obj,
    int n_samples,
    double percentile=0.2,
    unsigned long long seed=0,
    int mode_i=1,
    int within_mode_i=0,
):
    """Compute winners/losers according to MC top/bottom percentiles.

    Returns
    -------
    (winners, losers, win_count, lose_count, rank_sum)
        winners/losers are int64 arrays of shape (k,).
    """
    win_count, lose_count, rank_sum = mc_rank_counts(
        mu_obj, sigma2_obj, n_samples,
        percentile=percentile,
        seed=seed,
        mode_i=mode_i,
        within_mode_i=within_mode_i,
    )

    cdef int N = win_count.shape[0]
    if N == 0:
        empty = np.zeros(0, dtype=np.int64)
        return empty, empty, win_count, lose_count, rank_sum

    cdef int k = <int>np.ceil(percentile * N)
    if k < 1:
        k = 1
    if k > N:
        k = N

    idx = np.arange(N, dtype=np.int64)

    # winners: max win_count, then min rank_sum, then min index
    winners_order = np.lexsort((idx, rank_sum, -win_count))
    winners = winners_order[:k].astype(np.int64, copy=False)

    # losers: max lose_count, then max rank_sum, then min index
    losers_order = np.lexsort((idx, -rank_sum, -lose_count))
    losers = losers_order[:k].astype(np.int64, copy=False)

    return winners, losers, win_count, lose_count, rank_sum


cpdef tuple mc_rank_select_detailed(
    object mu_obj,
    object sigma2_obj,
    int n_samples,
    double percentile=0.2,
    unsigned long long seed=0,
    int mode_i=1,
    int within_mode_i=0,
    int quantiles_mode_i=2,
    long long quantiles_budget=2500000,
    double quality_w_place=1.0,
    double quality_w_front=1.0,
    double quality_w_within=1.0,
    double quality_sign=-1.0,
):
    """Monte-Carlo winner/loser selection + per-individual distribution stats.

    Returns
    -------
    winners_i, losers_i, win_count, lose_count, rank_sum, place_stats, front_stats, within_stats, quality_score

    Each *_stats is a dict with keys: mean, median, std, q25, q75.

    Notes
    -----
    This kernel keeps **constant memory** per individual for quantile estimates
    using the P² (Jain-Chlamtac) streaming estimator. Mean/std are exact.
    """

    cdef cnp.ndarray[F64, ndim=2] mu = _as_contig_f64_2d(mu_obj)
    cdef cnp.ndarray[F64, ndim=2] sigma2 = _as_contig_f64_2d(sigma2_obj)

    if mu.shape[0] != sigma2.shape[0] or mu.shape[1] != sigma2.shape[1]:
        raise ValueError("mu and sigma2 must have the same shape")

    cdef int N = mu.shape[0]
    cdef int M = mu.shape[1]

    if N == 0:
        empty_i = np.zeros(0, dtype=np.int64)
        z32 = np.zeros(0, dtype=np.int32)
        z64 = np.zeros(0, dtype=np.int64)
        zf = np.zeros(0, dtype=np.float64)
        empty_stats = {"mean": zf, "median": zf, "std": zf, "q25": zf, "q75": zf}
        return empty_i, empty_i, z32, z32, z64, empty_stats, empty_stats, empty_stats, zf

    if n_samples <= 0:
        raise ValueError("n_samples must be > 0")
    if percentile <= 0.0 or percentile > 1.0:
        raise ValueError("percentile must be in (0,1]")

    cdef int k = <int>np.ceil(percentile * N)
    if k < 1:
        k = 1
    if k > N:
        k = N

    # Precompute sigma = sqrt(max(sigma2, 0)).
    cdef cnp.ndarray[F64, ndim=2] sigma = np.empty((N, M), dtype=np.float64)
    cdef double[:, ::1] s2_v = sigma2
    cdef double[:, ::1] sig_v = sigma
    cdef int i, m
    with nogil:
        for i in range(N):
            for m in range(M):
                if s2_v[i, m] <= 0.0:
                    sig_v[i, m] = 0.0
                else:
                    sig_v[i, m] = sqrt(s2_v[i, m])

    # Allocate once: sample matrix + pareto buffers.
    cdef cnp.ndarray[F64, ndim=2] X = np.empty((N, M), dtype=np.float64)
    cdef cnp.ndarray[F64, ndim=2] R = np.empty((N, M), dtype=np.float64)

    cdef cnp.ndarray[I64, ndim=1] front_id = np.empty(N, dtype=np.int64)
    cdef cnp.ndarray[F64, ndim=1] within_score = np.empty(N, dtype=np.float64)
    cdef cnp.ndarray[I64, ndim=1] order = np.empty(N, dtype=np.int64)
    cdef cnp.ndarray[I64, ndim=1] items = np.empty(N, dtype=np.int64)
    cdef cnp.ndarray[I64, ndim=1] offsets = np.empty(N + 1, dtype=np.int64)

    cdef cnp.ndarray[U8, ndim=2] dom = np.empty((N, N), dtype=np.uint8)
    cdef cnp.ndarray[I32, ndim=1] n_dom = np.empty(N, dtype=np.int32)
    cdef cnp.ndarray[I32, ndim=1] cur = np.empty(N, dtype=np.int32)
    cdef cnp.ndarray[I32, ndim=1] nxt = np.empty(N, dtype=np.int32)

    cdef cnp.ndarray[I32, ndim=1] win_count = np.zeros(N, dtype=np.int32)
    cdef cnp.ndarray[I32, ndim=1] lose_count = np.zeros(N, dtype=np.int32)
    cdef cnp.ndarray[I64, ndim=1] rank_sum = np.zeros(N, dtype=np.int64)

    # Online mean/std (exact) for place/front/within.
    cdef cnp.ndarray[F64, ndim=1] place_mean = np.zeros(N, dtype=np.float64)
    cdef cnp.ndarray[F64, ndim=1] place_M2 = np.zeros(N, dtype=np.float64)
    cdef cnp.ndarray[F64, ndim=1] front_mean = np.zeros(N, dtype=np.float64)
    cdef cnp.ndarray[F64, ndim=1] front_M2 = np.zeros(N, dtype=np.float64)
    cdef cnp.ndarray[F64, ndim=1] within_mean = np.zeros(N, dtype=np.float64)
    cdef cnp.ndarray[F64, ndim=1] within_M2 = np.zeros(N, dtype=np.float64)

    # Quantile strategy:
    #   0 = exact (store all samples)
    #   1 = streaming P² (constant memory; requires n_samples>=5)
    #   2 = auto (exact for small N*n_samples, else streaming)
    cdef long long budget = quantiles_budget  # threshold on (N*n_samples) for exact quantiles in auto mode
    cdef bint use_p2
    cdef int S_init

    if quantiles_mode_i == 0:
        use_p2 = False
        S_init = n_samples
    elif quantiles_mode_i == 1:
        use_p2 = n_samples >= 5
        S_init = 5 if use_p2 else n_samples
    else:
        if n_samples < 5:
            use_p2 = False
            S_init = n_samples
        else:
            if budget <= 0:
                use_p2 = True
                S_init = 5
            elif (<long long>N) * (<long long>n_samples) <= budget:
                use_p2 = False
                S_init = n_samples
            else:
                use_p2 = True
                S_init = 5

    cdef cnp.ndarray[F64, ndim=2] place_buf = np.empty((S_init, N), dtype=np.float64)
    cdef cnp.ndarray[F64, ndim=2] front_buf = np.empty((S_init, N), dtype=np.float64)
    cdef cnp.ndarray[F64, ndim=2] within_buf = np.empty((S_init, N), dtype=np.float64)

    # P² states: (N,5) for each quantile and each metric.
    cdef cnp.ndarray[F64, ndim=2] q_place25
    cdef cnp.ndarray[F64, ndim=2] q_place50
    cdef cnp.ndarray[F64, ndim=2] q_place75
    cdef cnp.ndarray[I32T, ndim=2] npos_place25
    cdef cnp.ndarray[I32T, ndim=2] npos_place50
    cdef cnp.ndarray[I32T, ndim=2] npos_place75
    cdef cnp.ndarray[F64, ndim=2] des_place25
    cdef cnp.ndarray[F64, ndim=2] des_place50
    cdef cnp.ndarray[F64, ndim=2] des_place75

    cdef cnp.ndarray[F64, ndim=2] q_front25
    cdef cnp.ndarray[F64, ndim=2] q_front50
    cdef cnp.ndarray[F64, ndim=2] q_front75
    cdef cnp.ndarray[I32T, ndim=2] npos_front25
    cdef cnp.ndarray[I32T, ndim=2] npos_front50
    cdef cnp.ndarray[I32T, ndim=2] npos_front75
    cdef cnp.ndarray[F64, ndim=2] des_front25
    cdef cnp.ndarray[F64, ndim=2] des_front50
    cdef cnp.ndarray[F64, ndim=2] des_front75

    cdef cnp.ndarray[F64, ndim=2] q_within25
    cdef cnp.ndarray[F64, ndim=2] q_within50
    cdef cnp.ndarray[F64, ndim=2] q_within75
    cdef cnp.ndarray[I32T, ndim=2] npos_within25
    cdef cnp.ndarray[I32T, ndim=2] npos_within50
    cdef cnp.ndarray[I32T, ndim=2] npos_within75
    cdef cnp.ndarray[F64, ndim=2] des_within25
    cdef cnp.ndarray[F64, ndim=2] des_within50
    cdef cnp.ndarray[F64, ndim=2] des_within75

    if use_p2:
        q_place25 = np.empty((N, 5), dtype=np.float64)
        q_place50 = np.empty((N, 5), dtype=np.float64)
        q_place75 = np.empty((N, 5), dtype=np.float64)
        npos_place25 = np.empty((N, 5), dtype=np.int32)
        npos_place50 = np.empty((N, 5), dtype=np.int32)
        npos_place75 = np.empty((N, 5), dtype=np.int32)
        des_place25 = np.empty((N, 5), dtype=np.float64)
        des_place50 = np.empty((N, 5), dtype=np.float64)
        des_place75 = np.empty((N, 5), dtype=np.float64)

        q_front25 = np.empty((N, 5), dtype=np.float64)
        q_front50 = np.empty((N, 5), dtype=np.float64)
        q_front75 = np.empty((N, 5), dtype=np.float64)
        npos_front25 = np.empty((N, 5), dtype=np.int32)
        npos_front50 = np.empty((N, 5), dtype=np.int32)
        npos_front75 = np.empty((N, 5), dtype=np.int32)
        des_front25 = np.empty((N, 5), dtype=np.float64)
        des_front50 = np.empty((N, 5), dtype=np.float64)
        des_front75 = np.empty((N, 5), dtype=np.float64)

        q_within25 = np.empty((N, 5), dtype=np.float64)
        q_within50 = np.empty((N, 5), dtype=np.float64)
        q_within75 = np.empty((N, 5), dtype=np.float64)
        npos_within25 = np.empty((N, 5), dtype=np.int32)
        npos_within50 = np.empty((N, 5), dtype=np.int32)
        npos_within75 = np.empty((N, 5), dtype=np.int32)
        des_within25 = np.empty((N, 5), dtype=np.float64)
        des_within50 = np.empty((N, 5), dtype=np.float64)
        des_within75 = np.empty((N, 5), dtype=np.float64)

    cdef Pair* pairs = <Pair*>malloc(max(N, 1) * sizeof(Pair))
    cdef OrderKey* keys = <OrderKey*>malloc(max(N, 1) * sizeof(OrderKey))
    if pairs == NULL or keys == NULL:
        if pairs != NULL:
            free(pairs)
        if keys != NULL:
            free(keys)
        raise MemoryError()

    # Per-sample front counters for within-front ranks.
    cdef cnp.ndarray[I32T, ndim=1] front_counts = np.empty(N, dtype=np.int32)

    # Bind typed views once.
    cdef double[:, ::1] mu_v = mu
    cdef double[:, ::1] sigm_v = sigma
    cdef double[:, ::1] X_v = X
    cdef double[:, ::1] R_v = R

    cdef I64[::1] front_v = front_id
    cdef double[::1] within_score_v = within_score
    cdef I64[::1] order_v = order
    cdef I64[::1] items_v = items
    cdef I64[::1] offsets_v = offsets
    cdef U8[:, ::1] dom_v = dom
    cdef I32[::1] n_dom_v = n_dom
    cdef I32[::1] cur_v = cur
    cdef I32[::1] nxt_v = nxt

    cdef I32[::1] win_v = win_count
    cdef I32[::1] lose_v = lose_count
    cdef I64[::1] rsum_v = rank_sum

    cdef double[::1] pl_mean_v = place_mean
    cdef double[::1] pl_M2_v = place_M2
    cdef double[::1] fr_mean_v = front_mean
    cdef double[::1] fr_M2_v = front_M2
    cdef double[::1] wi_mean_v = within_mean
    cdef double[::1] wi_M2_v = within_M2

    cdef double[:, ::1] pl_buf_v = place_buf
    cdef double[:, ::1] fr_buf_v = front_buf
    cdef double[:, ::1] wi_buf_v = within_buf

    cdef double[:, ::1] qpl25_v
    cdef double[:, ::1] qpl50_v
    cdef double[:, ::1] qpl75_v
    cdef I32T[:, ::1] npl25_v
    cdef I32T[:, ::1] npl50_v
    cdef I32T[:, ::1] npl75_v
    cdef double[:, ::1] dpl25_v
    cdef double[:, ::1] dpl50_v
    cdef double[:, ::1] dpl75_v

    cdef double[:, ::1] qfr25_v
    cdef double[:, ::1] qfr50_v
    cdef double[:, ::1] qfr75_v
    cdef I32T[:, ::1] nfr25_v
    cdef I32T[:, ::1] nfr50_v
    cdef I32T[:, ::1] nfr75_v
    cdef double[:, ::1] dfr25_v
    cdef double[:, ::1] dfr50_v
    cdef double[:, ::1] dfr75_v

    cdef double[:, ::1] qwi25_v
    cdef double[:, ::1] qwi50_v
    cdef double[:, ::1] qwi75_v
    cdef I32T[:, ::1] nwi25_v
    cdef I32T[:, ::1] nwi50_v
    cdef I32T[:, ::1] nwi75_v
    cdef double[:, ::1] dwi25_v
    cdef double[:, ::1] dwi50_v
    cdef double[:, ::1] dwi75_v

    # Output views for quantiles (assigned later).
    cdef double[::1] pl_med_v
    cdef double[::1] pl_q25_v
    cdef double[::1] pl_q75_v
    cdef double[::1] fr_med_v
    cdef double[::1] fr_q25_v
    cdef double[::1] fr_q75_v
    cdef double[::1] wi_med_v
    cdef double[::1] wi_q25_v
    cdef double[::1] wi_q75_v

    if use_p2:
        qpl25_v = q_place25
        qpl50_v = q_place50
        qpl75_v = q_place75
        npl25_v = npos_place25
        npl50_v = npos_place50
        npl75_v = npos_place75
        dpl25_v = des_place25
        dpl50_v = des_place50
        dpl75_v = des_place75

        qfr25_v = q_front25
        qfr50_v = q_front50
        qfr75_v = q_front75
        nfr25_v = npos_front25
        nfr50_v = npos_front50
        nfr75_v = npos_front75
        dfr25_v = des_front25
        dfr50_v = des_front50
        dfr75_v = des_front75

        qwi25_v = q_within25
        qwi50_v = q_within50
        qwi75_v = q_within75
        nwi25_v = npos_within25
        nwi50_v = npos_within50
        nwi75_v = npos_within75
        dwi25_v = des_within25
        dwi50_v = des_within50
        dwi75_v = des_within75

    cdef I32T[::1] fc_v = front_counts

    cdef Xoshiro256ss rng
    cdef int s, pos, idx
    cdef int n_fronts
    cdef double n_d
    cdef double x_place, x_front, x_within
    cdef double delta
    cdef I32T fid
    cdef I32T within_rank
    cdef double tmp[5]
    cdef double s0, s1, s2, s3, s4

    _xoshiro_seed(&rng, <uint64_t>seed)

    try:
        with nogil:
            for s in range(n_samples):
                # 1) sample X ~ Normal(mu, sigma)
                for i in range(N):
                    for m in range(M):
                        if sigm_v[i, m] == 0.0:
                            X_v[i, m] = mu_v[i, m]
                        else:
                            X_v[i, m] = mu_v[i, m] + sigm_v[i, m] * _xoshiro_randn(&rng)

                # 2) pareto ranking
                n_fronts = rate_population_inplace(
                    X_v, mode_i, within_mode_i,
                    front_v, within_score_v, order_v,
                    offsets_v, items_v,
                    dom_v, n_dom_v, cur_v, nxt_v,
                    pairs, keys,
                    R_v,
                )
                if n_fronts < 0:
                    break

                # 3) within-front counters
                for i in range(n_fronts):
                    fc_v[i] = 0

                # sample count (1-indexed) for Welford
                n_d = <double>(s + 1)

                # 4) update selection counters + distributions
                for pos in range(N):
                    idx = <int>order_v[pos]

                    # top/bottom counters
                    if pos < k:
                        win_v[idx] += 1
                    if pos >= (N - k):
                        lose_v[idx] += 1
                    rsum_v[idx] += pos

                    fid = <I32T>front_v[idx]
                    within_rank = fc_v[fid]
                    fc_v[fid] = within_rank + 1

                    x_place = <double>pos
                    x_front = <double>fid
                    x_within = <double>within_rank

                    # Exact mean/std (Welford)
                    delta = x_place - pl_mean_v[idx]
                    pl_mean_v[idx] += delta / n_d
                    pl_M2_v[idx] += delta * (x_place - pl_mean_v[idx])

                    delta = x_front - fr_mean_v[idx]
                    fr_mean_v[idx] += delta / n_d
                    fr_M2_v[idx] += delta * (x_front - fr_mean_v[idx])

                    delta = x_within - wi_mean_v[idx]
                    wi_mean_v[idx] += delta / n_d
                    wi_M2_v[idx] += delta * (x_within - wi_mean_v[idx])

                    # Quantile summaries
                    if s < S_init:
                        pl_buf_v[s, idx] = x_place
                        fr_buf_v[s, idx] = x_front
                        wi_buf_v[s, idx] = x_within
                    elif use_p2:
                        # Update P² for each quantile (place)
                        _p2_update(&qpl25_v[idx, 0], &npl25_v[idx, 0], &dpl25_v[idx, 0], 0.25, x_place)
                        _p2_update(&qpl50_v[idx, 0], &npl50_v[idx, 0], &dpl50_v[idx, 0], 0.50, x_place)
                        _p2_update(&qpl75_v[idx, 0], &npl75_v[idx, 0], &dpl75_v[idx, 0], 0.75, x_place)

                        _p2_update(&qfr25_v[idx, 0], &nfr25_v[idx, 0], &dfr25_v[idx, 0], 0.25, x_front)
                        _p2_update(&qfr50_v[idx, 0], &nfr50_v[idx, 0], &dfr50_v[idx, 0], 0.50, x_front)
                        _p2_update(&qfr75_v[idx, 0], &nfr75_v[idx, 0], &dfr75_v[idx, 0], 0.75, x_front)

                        _p2_update(&qwi25_v[idx, 0], &nwi25_v[idx, 0], &dwi25_v[idx, 0], 0.25, x_within)
                        _p2_update(&qwi50_v[idx, 0], &nwi50_v[idx, 0], &dwi50_v[idx, 0], 0.50, x_within)
                        _p2_update(&qwi75_v[idx, 0], &nwi75_v[idx, 0], &dwi75_v[idx, 0], 0.75, x_within)

                # Initialize P² states once we collected first 5 samples.
                if use_p2 and s == 4:
                    for i in range(N):
                        # place
                        tmp[0] = pl_buf_v[0, i]
                        tmp[1] = pl_buf_v[1, i]
                        tmp[2] = pl_buf_v[2, i]
                        tmp[3] = pl_buf_v[3, i]
                        tmp[4] = pl_buf_v[4, i]
                        _sort5(&tmp[0])
                        s0 = tmp[0]; s1 = tmp[1]; s2 = tmp[2]; s3 = tmp[3]; s4 = tmp[4]
                        _p2_init(&qpl25_v[i, 0], &npl25_v[i, 0], &dpl25_v[i, 0], 0.25, s0, s1, s2, s3, s4)
                        _p2_init(&qpl50_v[i, 0], &npl50_v[i, 0], &dpl50_v[i, 0], 0.50, s0, s1, s2, s3, s4)
                        _p2_init(&qpl75_v[i, 0], &npl75_v[i, 0], &dpl75_v[i, 0], 0.75, s0, s1, s2, s3, s4)

                        # front
                        tmp[0] = fr_buf_v[0, i]
                        tmp[1] = fr_buf_v[1, i]
                        tmp[2] = fr_buf_v[2, i]
                        tmp[3] = fr_buf_v[3, i]
                        tmp[4] = fr_buf_v[4, i]
                        _sort5(&tmp[0])
                        s0 = tmp[0]; s1 = tmp[1]; s2 = tmp[2]; s3 = tmp[3]; s4 = tmp[4]
                        _p2_init(&qfr25_v[i, 0], &nfr25_v[i, 0], &dfr25_v[i, 0], 0.25, s0, s1, s2, s3, s4)
                        _p2_init(&qfr50_v[i, 0], &nfr50_v[i, 0], &dfr50_v[i, 0], 0.50, s0, s1, s2, s3, s4)
                        _p2_init(&qfr75_v[i, 0], &nfr75_v[i, 0], &dfr75_v[i, 0], 0.75, s0, s1, s2, s3, s4)

                        # within
                        tmp[0] = wi_buf_v[0, i]
                        tmp[1] = wi_buf_v[1, i]
                        tmp[2] = wi_buf_v[2, i]
                        tmp[3] = wi_buf_v[3, i]
                        tmp[4] = wi_buf_v[4, i]
                        _sort5(&tmp[0])
                        s0 = tmp[0]; s1 = tmp[1]; s2 = tmp[2]; s3 = tmp[3]; s4 = tmp[4]
                        _p2_init(&qwi25_v[i, 0], &nwi25_v[i, 0], &dwi25_v[i, 0], 0.25, s0, s1, s2, s3, s4)
                        _p2_init(&qwi50_v[i, 0], &nwi50_v[i, 0], &dwi50_v[i, 0], 0.50, s0, s1, s2, s3, s4)
                        _p2_init(&qwi75_v[i, 0], &nwi75_v[i, 0], &dwi75_v[i, 0], 0.75, s0, s1, s2, s3, s4)

        if n_fronts < 0:
            raise ValueError(f"pareto_core returned error code: {n_fronts}")

    finally:
        free(pairs)
        free(keys)

    # Winners/losers from aggregated counters (Python-side lexsort for stable tie-breaks).
    idx_arr = np.arange(N, dtype=np.int64)
    winners_order = np.lexsort((idx_arr, rank_sum, -win_count))
    losers_order = np.lexsort((idx_arr, -rank_sum, -lose_count))
    winners = winners_order[:k].astype(np.int64, copy=False)
    losers = losers_order[:k].astype(np.int64, copy=False)

    # Finalize std (ddof=0)
    cdef cnp.ndarray[F64, ndim=1] place_std = np.empty(N, dtype=np.float64)
    cdef cnp.ndarray[F64, ndim=1] front_std = np.empty(N, dtype=np.float64)
    cdef cnp.ndarray[F64, ndim=1] within_std = np.empty(N, dtype=np.float64)

    cdef double[::1] pl_std_v = place_std
    cdef double[::1] fr_std_v = front_std
    cdef double[::1] wi_std_v = within_std
    with nogil:
        for i in range(N):
            pl_std_v[i] = sqrt(pl_M2_v[i] / <double>n_samples)
            fr_std_v[i] = sqrt(fr_M2_v[i] / <double>n_samples)
            wi_std_v[i] = sqrt(wi_M2_v[i] / <double>n_samples)

    # Quantiles/median
    if not use_p2:
        place_median = np.median(place_buf, axis=0)
        place_q25 = np.quantile(place_buf, 0.25, axis=0)
        place_q75 = np.quantile(place_buf, 0.75, axis=0)

        front_median = np.median(front_buf, axis=0)
        front_q25 = np.quantile(front_buf, 0.25, axis=0)
        front_q75 = np.quantile(front_buf, 0.75, axis=0)

        within_median = np.median(within_buf, axis=0)
        within_q25 = np.quantile(within_buf, 0.25, axis=0)
        within_q75 = np.quantile(within_buf, 0.75, axis=0)
    else:
        place_median = np.empty(N, dtype=np.float64)
        place_q25 = np.empty(N, dtype=np.float64)
        place_q75 = np.empty(N, dtype=np.float64)
        front_median = np.empty(N, dtype=np.float64)
        front_q25 = np.empty(N, dtype=np.float64)
        front_q75 = np.empty(N, dtype=np.float64)
        within_median = np.empty(N, dtype=np.float64)
        within_q25 = np.empty(N, dtype=np.float64)
        within_q75 = np.empty(N, dtype=np.float64)

        pl_med_v = place_median
        pl_q25_v = place_q25
        pl_q75_v = place_q75
        fr_med_v = front_median
        fr_q25_v = front_q25
        fr_q75_v = front_q75
        wi_med_v = within_median
        wi_q25_v = within_q25
        wi_q75_v = within_q75

        with nogil:
            for i in range(N):
                pl_q25_v[i] = qpl25_v[i, 2]
                pl_med_v[i] = qpl50_v[i, 2]
                pl_q75_v[i] = qpl75_v[i, 2]

                fr_q25_v[i] = qfr25_v[i, 2]
                fr_med_v[i] = qfr50_v[i, 2]
                fr_q75_v[i] = qfr75_v[i, 2]

                wi_q25_v[i] = qwi25_v[i, 2]
                wi_med_v[i] = qwi50_v[i, 2]
                wi_q75_v[i] = qwi75_v[i, 2]


    # Composite quality score (larger is better by default with quality_sign=-1).
    cdef cnp.ndarray[F64, ndim=1] quality_score = np.empty(N, dtype=np.float64)
    cdef double[::1] qsc_v = quality_score
    with nogil:
        for i in range(N):
            qsc_v[i] = quality_sign * (
                quality_w_place * pl_mean_v[i] +
                quality_w_front * fr_mean_v[i] +
                quality_w_within * wi_mean_v[i]
            )

    q_mode = 1 if use_p2 else 0

    place_stats = {
        "mean": place_mean,
        "quantiles_mode_used": 1 if use_p2 else 0,
        "median": np.asarray(place_median, dtype=np.float64),
        "std": place_std,
        "q25": np.asarray(place_q25, dtype=np.float64),
        "q75": np.asarray(place_q75, dtype=np.float64),
    }
    front_stats = {
        "mean": front_mean,
        "quantiles_mode_used": 1 if use_p2 else 0,
        "median": np.asarray(front_median, dtype=np.float64),
        "std": front_std,
        "q25": np.asarray(front_q25, dtype=np.float64),
        "q75": np.asarray(front_q75, dtype=np.float64),
    }
    within_stats = {
        "mean": within_mean,
        "quantiles_mode_used": 1 if use_p2 else 0,
        "median": np.asarray(within_median, dtype=np.float64),
        "std": within_std,
        "q25": np.asarray(within_q25, dtype=np.float64),
        "q75": np.asarray(within_q75, dtype=np.float64),
    }

    return winners, losers, win_count, lose_count, rank_sum, place_stats, front_stats, within_stats, quality_score
