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

"""Cython acceleration for pareto_core.pareto_rating.

Public entry-point
------------------
rate_population_arrays_cy(X, mode_i, within_mode_i) ->
    front_id, within_score, order, offsets, items

- All objectives are maximized.
- Deterministic: all tie-breaks by index.

within_mode_i
-------------
0: crowding distance (NSGA-II)

The function returns a packed representation of fronts:
- offsets: int64 array of length (n_fronts+1)
- items: int64 array of length N; items[offsets[f]:offsets[f+1]] are indices for front f

This representation avoids per-front Python allocations.
"""

from libc.math cimport INFINITY
from libc.stdlib cimport qsort, malloc, free
from libc.string cimport memset

import numpy as np
cimport numpy as cnp

# NOTE: C-API types (F64/I64/I32/U8) are declared in the companion .pxd file.


cdef struct Pair:
    double val
    int idx


cdef int _pair_cmp(const void* a, const void* b) noexcept nogil:
    cdef Pair* pa = <Pair*>a
    cdef Pair* pb = <Pair*>b
    if pa.val < pb.val:
        return -1
    if pa.val > pb.val:
        return 1
    if pa.idx < pb.idx:
        return -1
    if pa.idx > pb.idx:
        return 1
    return 0


cdef struct OrderKey:
    long front
    double neg_score
    long idx


cdef int _order_cmp(const void* a, const void* b) noexcept nogil:
    cdef OrderKey* ka = <OrderKey*>a
    cdef OrderKey* kb = <OrderKey*>b
    if ka.front < kb.front:
        return -1
    if ka.front > kb.front:
        return 1
    if ka.neg_score < kb.neg_score:
        return -1
    if ka.neg_score > kb.neg_score:
        return 1
    if ka.idx < kb.idx:
        return -1
    if ka.idx > kb.idx:
        return 1
    return 0


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


cdef void _transform_to_ranks_fill(
    double[:, ::1] X,
    double[:, ::1] R,
    Pair* pairs,
) noexcept nogil:
    """Fill R with stable ranks of X in [0,1], per-column.

    Sorting is by (value asc, index asc).
    rank = pos/(N-1) (or 0 when N==1).
    """
    cdef int N = X.shape[0]
    cdef int M = X.shape[1]
    cdef double denom = (N - 1) if N > 1 else 1.0
    cdef int i, m, pos

    for m in range(M):
        for i in range(N):
            pairs[i].val = X[i, m]
            pairs[i].idx = i
        qsort(pairs, N, sizeof(Pair), _pair_cmp)
        for pos in range(N):
            R[pairs[pos].idx, m] = (<double>pos) / denom


cdef void _build_dom_and_counts(
    double[:, ::1] X,
    U8[:, ::1] dom,
    I32[::1] n_dom,
) noexcept nogil:
    """Compute dom matrix and domination counts.

    dom[i,j]=1 iff i dominates j (maximize): all(X[i]>=X[j]) and any(X[i]>X[j]).
    n_dom[j] = number of i dominating j.
    """
    cdef int N = X.shape[0]
    cdef int M = X.shape[1]
    cdef int i, j, m
    cdef bint all_ge
    cdef bint any_gt
    cdef double ai, bj

    for i in range(N):
        n_dom[i] = 0

    for i in range(N):
        for j in range(N):
            if i == j:
                dom[i, j] = 0
                continue
            all_ge = True
            any_gt = False
            for m in range(M):
                ai = X[i, m]
                bj = X[j, m]
                if ai < bj:
                    all_ge = False
                    break
                if ai > bj:
                    any_gt = True
            if all_ge and any_gt:
                dom[i, j] = 1
                n_dom[j] += 1
            else:
                dom[i, j] = 0


cdef int _non_dominated_sort_core(
    double[:, ::1] X,
    U8[:, ::1] dom,
    I32[::1] n_dom,
    I32[::1] cur,
    I32[::1] nxt,
    I64[::1] front_id,
    I64[::1] items,
    I64[::1] offsets,
) noexcept nogil:
    """Compute non-dominated sorting.

    Returns n_fronts. Packed fronts are written into `items` and `offsets`.
    Each front's items are emitted in ascending index order (deterministic).
    """
    cdef int N = X.shape[0]
    cdef int i, j, q, a, b, tmp
    cdef int cur_k = 0
    cdef int nxt_k
    cdef int f = 0
    cdef int items_pos = 0

    for i in range(N):
        front_id[i] = -1

    _build_dom_and_counts(X, dom, n_dom)

    # current front (n_dom==0), ascending by i
    for i in range(N):
        if n_dom[i] == 0:
            cur[cur_k] = i
            cur_k += 1

    offsets[0] = 0

    while cur_k > 0:
        # emit current
        for j in range(cur_k):
            i = cur[j]
            front_id[i] = f
            items[items_pos] = i
            items_pos += 1

        offsets[f + 1] = items_pos

        # compute next candidates
        nxt_k = 0
        for j in range(cur_k):
            i = cur[j]
            for q in range(N):
                if dom[i, q] != 0:
                    n_dom[q] -= 1
                    if n_dom[q] == 0:
                        nxt[nxt_k] = q
                        nxt_k += 1

        if nxt_k == 0:
            f += 1
            break

        # sort nxt ascending (insertion sort; N<=500)
        for a in range(1, nxt_k):
            tmp = nxt[a]
            b = a - 1
            while b >= 0 and nxt[b] > tmp:
                nxt[b + 1] = nxt[b]
                b -= 1
            nxt[b + 1] = tmp

        # copy nxt -> cur
        for j in range(nxt_k):
            cur[j] = nxt[j]
        cur_k = nxt_k

        f += 1

    return f


cdef void _crowding_distance_front(
    double[:, ::1] X,
    I64[::1] front_items,
    int start,
    int end,
    double[::1] within_score,
    Pair* pairs,
) noexcept nogil:
    """Compute NSGA-II crowding distance for a packed front slice [start,end)."""
    cdef int K = end - start
    cdef int j, m
    cdef int idx
    cdef double minv, maxv, denom

    if K <= 0:
        return
    if K == 1:
        within_score[front_items[start]] = INFINITY
        return

    # reset
    for j in range(start, end):
        within_score[front_items[j]] = 0.0

    cdef int M = X.shape[1]

    for m in range(M):
        # build sortable pairs (value, idx), stable by idx
        for j in range(K):
            idx = <int>front_items[start + j]
            pairs[j].val = X[idx, m]
            pairs[j].idx = idx

        qsort(pairs, K, sizeof(Pair), _pair_cmp)

        minv = pairs[0].val
        maxv = pairs[K - 1].val
        if maxv == minv:
            # degenerate objective contributes nothing and creates no artificial extremes
            continue

        within_score[pairs[0].idx] = INFINITY
        within_score[pairs[K - 1].idx] = INFINITY

        denom = maxv - minv
        for j in range(1, K - 1):
            idx = pairs[j].idx
            if within_score[idx] == INFINITY:
                continue
            within_score[idx] += (pairs[j + 1].val - pairs[j - 1].val) / denom


cdef int rate_population_inplace(
    double[:, ::1] X,            # (N,M) objectives, maximize
    int mode_i,                  # 0=values, 1=ranks
    int within_mode_i,           # 0=crowding
    I64[::1] front_id,           # (N,)
    double[::1] within_score,    # (N,)
    I64[::1] order,              # (N,)
    I64[::1] offsets,            # (N+1,)
    I64[::1] items,              # (N,)
    U8[:, ::1] dom,              # (N,N)
    I32[::1] n_dom,              # (N,)
    I32[::1] cur,                # (N,)
    I32[::1] nxt,                # (N,)
    Pair* pairs,                 # len>=N
    OrderKey* keys,              # len>=N
    double[:, ::1] R,            # (N,M) scratch for ranks
) noexcept nogil:
    """In-place Pareto rating kernel for cross-extension (C-API) use.

    This function does **no allocations** and is safe under `nogil`.
    All buffers must be pre-allocated by the caller.

    Returns
    -------
    n_fronts (>=1) or a negative error code.
    """
    cdef int N = X.shape[0]
    cdef int n_fronts
    cdef int f
    cdef int i

    cdef double[:, ::1] Xv

    if within_mode_i != 0:
        return -2

    if mode_i == 0:
        Xv = X
    elif mode_i == 1:
        _transform_to_ranks_fill(X, R, pairs)
        Xv = R
    else:
        return -1

    # reset within scores for determinism (caller may pass uninitialized memory)
    for i in range(N):
        within_score[i] = 0.0

    n_fronts = _non_dominated_sort_core(
        Xv, dom, n_dom, cur, nxt, front_id, items, offsets
    )

    for f in range(n_fronts):
        _crowding_distance_front(
            Xv,
            items,
            <int>offsets[f],
            <int>offsets[f + 1],
            within_score,
            pairs,
        )

    # global order by (front asc, within desc, idx asc)
    for i in range(N):
        keys[i].front = <long>front_id[i]
        keys[i].neg_score = -(<double>within_score[i])
        keys[i].idx = <long>i

    qsort(keys, N, sizeof(OrderKey), _order_cmp)

    for i in range(N):
        order[i] = <I64>keys[i].idx

    return n_fronts


cpdef tuple _rate_population_inplace_debug(object X_obj, int mode_i, int within_mode_i):
    """Debug helper: call `rate_population_inplace` once and return arrays.

    This is not used by the high-level API but is valuable for unit tests and
    for validating cross-extension calls.
    """
    cdef cnp.ndarray[F64, ndim=2] X_arr = _as_contig_f64_2d(X_obj)
    cdef int N = X_arr.shape[0]

    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_arr = np.empty((N, N), dtype=np.uint8)
    cdef cnp.ndarray[I32, ndim=1] n_dom_arr = np.empty(N, dtype=np.int32)
    cdef cnp.ndarray[I32, ndim=1] cur_arr = np.empty(N, dtype=np.int32)
    cdef cnp.ndarray[I32, ndim=1] nxt_arr = np.empty(N, dtype=np.int32)

    cdef cnp.ndarray[F64, ndim=2] R_arr = np.empty_like(X_arr, 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()

    # IMPORTANT: bind typed memoryviews *before* releasing the GIL.
    # Passing ndarrays directly to a memoryview-typed signature under `nogil`
    # can trigger unsafe buffer acquisition and crash.
    cdef double[:, ::1] X_v = X_arr
    cdef I64[::1] front_v = front_id
    cdef double[::1] within_v = within_score
    cdef I64[::1] order_v = order
    cdef I64[::1] offsets_v = offsets
    cdef I64[::1] items_v = items
    cdef U8[:, ::1] dom_v = dom_arr
    cdef I32[::1] n_dom_v = n_dom_arr
    cdef I32[::1] cur_v = cur_arr
    cdef I32[::1] nxt_v = nxt_arr
    cdef double[:, ::1] R_v = R_arr

    cdef int n_fronts
    try:
        with nogil:
            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:
            raise ValueError(f"rate_population_inplace error code: {n_fronts}")
        return front_id, within_score, order, offsets[: n_fronts + 1].copy(), items
    finally:
        free(pairs)
        free(keys)


cpdef tuple rate_population_arrays_cy(object X_obj, int mode_i, int within_mode_i):
    """Compute Pareto ranking and return raw arrays.

    Parameters
    ----------
    X_obj:
        Array-like (N,M) objectives; all maximized.
    mode_i:
        0=values, 1=ranks.
    within_mode_i:
        0=crowding.

    Returns
    -------
    (front_id, within_score, order, offsets, items)
        offsets is length (n_fronts+1); items is length N and contains packed fronts.
    """
    cdef cnp.ndarray[F64, ndim=2] X_arr = _as_contig_f64_2d(X_obj)
    cdef int N = X_arr.shape[0]

    # outputs
    cdef cnp.ndarray[I64, ndim=1] front_id = np.empty(N, dtype=np.int64)
    cdef cnp.ndarray[F64, ndim=1] within_score = np.zeros(N, dtype=np.float64)
    cdef cnp.ndarray[I64, ndim=1] items = np.empty(N, dtype=np.int64)
    cdef cnp.ndarray[I64, ndim=1] offsets = np.zeros(N + 1, dtype=np.int64)  # worst-case fronts=N

    # work buffers for sorting/front building
    cdef cnp.ndarray[U8, ndim=2] dom_arr = np.zeros((N, N), dtype=np.uint8)
    cdef cnp.ndarray[I32, ndim=1] n_dom_arr = np.empty(N, dtype=np.int32)
    cdef cnp.ndarray[I32, ndim=1] cur_arr = np.empty(N, dtype=np.int32)
    cdef cnp.ndarray[I32, ndim=1] nxt_arr = np.empty(N, dtype=np.int32)

    # shared work buffer for sorting pairs (rank-mode + crowding)
    cdef Pair* pairs = <Pair*>malloc(max(N, 1) * sizeof(Pair))
    if pairs == NULL:
        raise MemoryError()

    # typed views (declared once)
    cdef double[:, ::1] Xv
    cdef double[:, ::1] Xv0
    cdef double[:, ::1] Rv0
    cdef U8[:, ::1] dom
    cdef I32[::1] n_dom
    cdef I32[::1] cur
    cdef I32[::1] nxt
    cdef I64[::1] front_id_v
    cdef I64[::1] items_v
    cdef I64[::1] offsets_v
    cdef double[::1] within_v

    cdef cnp.ndarray[F64, ndim=2] Xw_arr
    cdef cnp.ndarray[F64, ndim=2] R_arr

    cdef int n_fronts
    cdef int f
    cdef int i

    cdef cnp.ndarray[I64, ndim=1] order
    cdef OrderKey* keys

    try:
        if within_mode_i != 0:
            raise ValueError("unknown within_mode_i")

        # prepare Xw (rank-mode optional)
        if mode_i == 0:
            Xw_arr = X_arr
        elif mode_i == 1:
            R_arr = np.empty_like(X_arr, dtype=np.float64)
            Xw_arr = R_arr
            Xv0 = X_arr
            Rv0 = R_arr
            with nogil:
                _transform_to_ranks_fill(Xv0, Rv0, pairs)
        else:
            raise ValueError("unknown mode_i")

        # bind views
        Xv = Xw_arr
        dom = dom_arr
        n_dom = n_dom_arr
        cur = cur_arr
        nxt = nxt_arr
        front_id_v = front_id
        items_v = items
        offsets_v = offsets
        within_v = within_score

        # compute fronts and within scores
        with nogil:
            n_fronts = _non_dominated_sort_core(
                Xv, dom, n_dom, cur, nxt, front_id_v, items_v, offsets_v
            )
            for f in range(n_fronts):
                _crowding_distance_front(
                    Xv,
                    items_v,
                    <int>offsets_v[f],
                    <int>offsets_v[f + 1],
                    within_v,
                    pairs,
                )

        # trim offsets to (n_fronts+1)
        offsets_trim = offsets[: n_fronts + 1].copy()

        # global order by (front asc, within desc, idx asc)
        order = np.empty(N, dtype=np.int64)
        keys = <OrderKey*>malloc(max(N, 1) * sizeof(OrderKey))
        if keys == NULL:
            raise MemoryError()
        try:
            for i in range(N):
                keys[i].front = <long>front_id[i]
                keys[i].neg_score = -(<double>within_score[i])
                keys[i].idx = <long>i
            with nogil:
                qsort(keys, N, sizeof(OrderKey), _order_cmp)
            for i in range(N):
                order[i] = <I64>keys[i].idx
        finally:
            free(keys)

        return front_id, within_score, order, offsets_trim, items

    finally:
        free(pairs)
