"""Pareto rating core.

This module provides a fast, deterministic implementation of:

- Non-dominated sorting (Pareto fronts), assuming *all* objectives are maximized.
- Within-front ranking via crowding distance (NSGA-II).
- Optional preprocessing that converts objective values to population ranks
  (stable by index).

The intended use is PBT selection: rank individuals by
    (front_id asc, within_score desc, index asc).

Determinism guarantees
----------------------
- All internal sorts use an explicit tie-break by index.
- In `mode="ranks"`, ties in objective values are broken by index, producing
  unique ranks.

Performance
-----------
For N≈200–500 and M≤4 this code is already fast. A Cython acceleration layer is
supported via `zen_fronts.pareto_core._pareto_rating_cy`.

If the Cython extension is available (compiled), `rate_population` will use it
by default. The pure-Python implementation remains as a reference and fallback.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Literal, Tuple, Union

import numpy as np


Mode = Literal["values", "ranks"]

# Within-front ranking modes. Integers are preferred for speed and future ports.
WITHIN_CROWDING = 0
WITHIN_KNN = 1

WithinMode = Union[int, Literal["crowding", "knn"]]


@dataclass(frozen=True)
class ParetoRating:
    """Result of Pareto rating."""

    front_id: np.ndarray  # (N,) int64
    within_score: np.ndarray  # (N,) float64
    order: np.ndarray  # (N,) int64
    fronts: List[np.ndarray]  # list of 1D int64 arrays, each sorted by index


# --- Optional Cython acceleration ------------------------------------------------

_HAVE_CYTHON = False
try:
    from ._pareto_rating_cy import (  # type: ignore
        rate_population_arrays_cy as _rate_population_arrays_cy,
    )

    _HAVE_CYTHON = True
except Exception:
    _HAVE_CYTHON = False


def transform_to_ranks(X: np.ndarray) -> np.ndarray:
    """Convert each objective column to stable normalized ranks in [0,1].

    For each objective m, the population is sorted by X[:, m] ascending with a
    tie-break by index.

    The resulting position pos (0..N-1) is normalized as pos/(N-1)
    (or 0 when N==1).
    """

    X = np.asarray(X, dtype=np.float64)
    if X.ndim != 2:
        raise ValueError(f"X must be 2D array, got shape={X.shape}")

    N, M = X.shape
    if N == 0:
        return X.copy()

    denom = float(N - 1) if N > 1 else 1.0
    idx = np.arange(N, dtype=np.int64)

    R = np.empty_like(X, dtype=np.float64)
    for m in range(M):
        col = X[:, m]
        order = np.lexsort((idx, col))  # primary: col asc, tie: idx asc
        pos = np.empty(N, dtype=np.float64)
        pos[order] = np.arange(N, dtype=np.float64)
        R[:, m] = pos / denom

    return R


def non_dominated_sort(X: np.ndarray) -> Tuple[np.ndarray, List[np.ndarray]]:
    """Perform non-dominated sorting (all objectives maximized).

    Returns
    -------
    front_id:
        int64 array of shape (N,) with front numbers (0 is best).
    fronts:
        list of int64 arrays; each array contains indices belonging to that
        front, sorted ascending for determinism.

    Notes
    -----
    This is a pure-Python/Numpy reference implementation. When the Cython
    extension is available, `rate_population()` uses it instead.
    """

    X = np.asarray(X, dtype=np.float64)
    if X.ndim != 2:
        raise ValueError(f"X must be 2D array, got shape={X.shape}")

    N, M = X.shape
    front_id = np.full(N, -1, dtype=np.int64)
    if N == 0:
        return front_id, []

    # dom[i, j] == True if i dominates j.
    ge = X[:, None, :] >= X[None, :, :]
    gt = X[:, None, :] > X[None, :, :]
    dom = np.all(ge, axis=2) & np.any(gt, axis=2)
    np.fill_diagonal(dom, False)

    n_dom = dom.sum(axis=0).astype(np.int64)
    S = [np.nonzero(dom[i])[0].astype(np.int64) for i in range(N)]

    fronts: List[np.ndarray] = []

    current = np.nonzero(n_dom == 0)[0].astype(np.int64)
    current.sort()
    f = 0
    while current.size > 0:
        fronts.append(current)
        front_id[current] = f

        next_candidates: List[int] = []
        for p in current.tolist():
            for q in S[p].tolist():
                n_dom[q] -= 1
                if n_dom[q] == 0:
                    next_candidates.append(q)

        if not next_candidates:
            break

        current = np.array(sorted(set(next_candidates)), dtype=np.int64)
        f += 1

    return front_id, fronts


def crowding_distance(X: np.ndarray, front: np.ndarray, out: np.ndarray) -> None:
    """Compute NSGA-II crowding distance for one front.

    Parameters
    ----------
    X:
        (N, M) float64 objectives, maximized.
    front:
        1D int64 indices of individuals in the front.
    out:
        (N,) float64 array to be updated in place; only entries in `front` are written.

    Notes
    -----
    - For each objective, boundary points receive +inf if the objective is non-degenerate.
    - If an objective is degenerate within the front (max==min), it contributes 0 and
      produces no boundary infinities.
    - Sorting per objective uses tie-break by index for determinism.
    """

    X = np.asarray(X, dtype=np.float64)
    front = np.asarray(front, dtype=np.int64)
    if out.dtype != np.float64:
        raise ValueError("out must be float64")

    K = int(front.size)
    if K == 0:
        return

    out[front] = 0.0

    if K == 1:
        out[front[0]] = np.inf
        return

    _, M = X.shape
    idx = front.copy()

    for m in range(M):
        vals = X[idx, m]
        ord_local = np.lexsort((idx, vals))
        sorted_idx = idx[ord_local]
        sorted_vals = vals[ord_local]

        minv = float(sorted_vals[0])
        maxv = float(sorted_vals[-1])
        if maxv == minv:
            continue

        out[sorted_idx[0]] = np.inf
        out[sorted_idx[-1]] = np.inf

        denom = maxv - minv
        for j in range(1, K - 1):
            i_mid = sorted_idx[j]
            if np.isinf(out[i_mid]):
                continue
            prevv = float(sorted_vals[j - 1])
            nextv = float(sorted_vals[j + 1])
            out[i_mid] += (nextv - prevv) / denom


def _rate_population_python(
    X: np.ndarray,
    *,
    mode: Mode,
    within_mode: WithinMode,
) -> ParetoRating:
    X = np.asarray(X, dtype=np.float64)

    if mode == "ranks":
        Xw = transform_to_ranks(X)
    elif mode == "values":
        Xw = X
    else:
        raise ValueError(f"Unknown mode: {mode!r}")

    front_id, fronts = non_dominated_sort(Xw)

    N = Xw.shape[0]
    within_score = np.zeros(N, dtype=np.float64)

    if within_mode == "crowding":
        within_mode_i = WITHIN_CROWDING
    elif within_mode == "knn":
        within_mode_i = WITHIN_KNN
    elif isinstance(within_mode, int):
        within_mode_i = within_mode
    else:
        raise ValueError(f"Unknown within_mode: {within_mode!r}")

    if within_mode_i == WITHIN_CROWDING:
        for fr in fronts:
            crowding_distance(Xw, fr, within_score)
    elif within_mode_i == WITHIN_KNN:
        raise NotImplementedError(
            "within_mode=knn is reserved but not implemented yet. "
            "If you need it, specify k, distance normalization, and tie-breaks."
        )
    else:
        raise ValueError(f"Unknown within_mode int code: {within_mode_i}")

    idx = np.arange(N, dtype=np.int64)
    order = np.lexsort((idx, -within_score, front_id))

    return ParetoRating(front_id=front_id, within_score=within_score, order=order, fronts=fronts)


def rate_population(
    X: np.ndarray,
    *,
    mode: Mode = "values",
    within_mode: WithinMode = WITHIN_CROWDING,
    use_cython: bool = True,
) -> ParetoRating:
    """Rank individuals by Pareto fronts and within-front crowding distance.

    Ordering is deterministic:
        (front_id asc, within_score desc, index asc)

    Parameters
    ----------
    X:
        (N,M) float64 objectives (all maximized).
    mode:
        "values" for direct objectives, "ranks" to convert each objective to stable ranks.
    within_mode:
        Within-front tie-break mode.
    use_cython:
        If True (default) and Cython extension is available, use it.

    Returns
    -------
    ParetoRating
        front_id, within_score, order, and per-front index lists.
    """

    if use_cython and _HAVE_CYTHON:
        X = np.asarray(X, dtype=np.float64)
        if mode == "values":
            mode_i = 0
        elif mode == "ranks":
            mode_i = 1
        else:
            raise ValueError(f"Unknown mode: {mode!r}")

        if within_mode == "crowding":
            within_mode_i = WITHIN_CROWDING
        elif within_mode == "knn":
            within_mode_i = WITHIN_KNN
        elif isinstance(within_mode, int):
            within_mode_i = within_mode
        else:
            raise ValueError(f"Unknown within_mode: {within_mode!r}")

        front_id, within_score, order, offsets, items = _rate_population_arrays_cy(
            X, mode_i, within_mode_i
        )

        # Build list-of-arrays fronts from offsets/items.
        offsets = np.asarray(offsets, dtype=np.int64)
        items = np.asarray(items, dtype=np.int64)
        fronts: List[np.ndarray] = []
        for k in range(int(offsets.size) - 1):
            a = int(offsets[k])
            b = int(offsets[k + 1])
            fronts.append(items[a:b].copy())

        return ParetoRating(
            front_id=np.asarray(front_id, dtype=np.int64),
            within_score=np.asarray(within_score, dtype=np.float64),
            order=np.asarray(order, dtype=np.int64),
            fronts=fronts,
        )

    return _rate_population_python(X, mode=mode, within_mode=within_mode)
