"""sem_cython12.wrapper - thin Python API over the compiled core.

Numpy-facing entry points for OpenMP-parallel numerical routines.
Each function accepts numpy arrays, marshals them to contiguous
float64 form, and returns numpy arrays.

Thread count is shared globally via ``set_num_threads`` / the
``SEM_NUM_THREADS`` environment variable; default is roughly 50% of
the host's logical cores.
"""

from __future__ import annotations

import numpy as np

# Fail gracefully when the compiled extension is not present.
try:
    from sem_cython12.sem_core12 import (
        cy12_get_num_threads,
        cy12_set_num_threads,
        cy12_batch_max_similarity,
        cy12_concept_support_matrix,
        cy12_pareto_core_mask,
        cy12_one_sided_mask,
        cy12_non_redundant_witnesses,
        cy12_pairwise_distances,
        cy12_nn_distances,
        cy12_extend_frontier_kernel,
    )
    _HAS_CYTHON = True
except ImportError:
    _HAS_CYTHON = False


def available() -> bool:
    """Return True iff the compiled extension imported successfully."""
    return _HAS_CYTHON


def backend() -> str:
    """Return ``'cython12'`` when the compiled extension is loaded,
    otherwise ``'python-fallback'``.  In the fallback case the
    numpy-facing helpers below raise RuntimeError; callers should
    branch on ``available()`` and provide their own pure-Python path
    when False."""
    return "cython12" if _HAS_CYTHON else "python-fallback"


def get_num_threads() -> int:
    """Return the active OpenMP worker count."""
    if not _HAS_CYTHON:
        return 1
    return int(cy12_get_num_threads())


def set_num_threads(n: int) -> None:
    """Set the OpenMP worker count.  No-op when the extension is
    unavailable.  ``n`` must be >= 1."""
    if not _HAS_CYTHON:
        return
    cy12_set_num_threads(int(n))


# ---- numpy-facing helpers ---------------------------------------


def batch_max_similarity(X_query, X_members, lam: float) -> np.ndarray:
    """For each row of ``X_query`` (shape ``(Q, D)``), return a
    similarity score in ``[0, 1]`` summarising how close the row is
    to the most similar row of ``X_members`` (shape ``(M, D)``).
    ``lam`` is the kernel scale (> 0) - smaller values produce a
    narrower kernel.  Output shape: ``(Q,)``.  Empty ``X_members``
    yields a zeros vector."""
    if not _HAS_CYTHON:
        raise RuntimeError("sem_cython12 not available; import guarded caller")
    Xq = np.ascontiguousarray(X_query, dtype=np.float64)
    Xm = np.ascontiguousarray(X_members, dtype=np.float64)
    if Xm.size == 0:
        return np.zeros(Xq.shape[0], dtype=np.float64)
    return cy12_batch_max_similarity(Xq, Xm, float(lam))


def concept_support_matrix(X_query, member_mats, lam: float) -> np.ndarray:
    """``batch_max_similarity`` evaluated against ``K`` independent
    reference sets in one call.  ``X_query`` has shape ``(Q, D)`` and
    ``member_mats`` is a list of ``(M_k, D)`` arrays (``M_k`` may
    differ per ``k``).  Output is a ``(Q, K)`` matrix; column ``k``
    is the similarity vector against ``member_mats[k]``.  Empty
    member matrices contribute a column of zeros."""
    if not _HAS_CYTHON:
        raise RuntimeError("sem_cython12 not available; import guarded caller")
    Xq = np.ascontiguousarray(X_query, dtype=np.float64)
    return cy12_concept_support_matrix(Xq, list(member_mats), float(lam))


def pareto_core_mask(S) -> np.ndarray:
    """Given a real-valued matrix ``S`` of shape ``(N, k)``, return a
    byte mask of shape ``(N,)`` flagging the rows that survive a
    multi-objective best-tradeoff filter (the "Pareto core") in the
    maximisation order."""
    if not _HAS_CYTHON:
        raise RuntimeError("sem_cython12 not available; import guarded caller")
    return cy12_pareto_core_mask(np.ascontiguousarray(S, dtype=np.float64))


def one_sided_mask(S) -> np.ndarray:
    """Given ``S`` of shape ``(N, k)``, return a byte mask of shape
    ``(N, k)`` flagging row/column pairs ``(w, i)`` where row ``w``
    contributes meaningfully on column ``i`` only.  Used together
    with ``pareto_core_mask`` to build ``non_redundant_witnesses``."""
    if not _HAS_CYTHON:
        raise RuntimeError("sem_cython12 not available; import guarded caller")
    return cy12_one_sided_mask(np.ascontiguousarray(S, dtype=np.float64))


def non_redundant_witnesses(S) -> np.ndarray:
    """Given ``S`` of shape ``(N, k)``, return a 1-D int32 array of
    row indices that survive both the best-tradeoff filter and the
    one-sided filter - i.e., rows that contribute meaningfully across
    multiple columns rather than peaking on one alone."""
    if not _HAS_CYTHON:
        raise RuntimeError("sem_cython12 not available; import guarded caller")
    return cy12_non_redundant_witnesses(np.ascontiguousarray(S, dtype=np.float64))


def pairwise_distances(X) -> np.ndarray:
    """Symmetric ``(N, N)`` distance matrix between rows of ``X``
    (shape ``(N, D)``).  Computed with parallel upper-triangle
    iteration; the lower triangle is mirrored."""
    if not _HAS_CYTHON:
        raise RuntimeError("sem_cython12 not available; import guarded caller")
    return cy12_pairwise_distances(np.ascontiguousarray(X, dtype=np.float64))


def nn_distances(X) -> np.ndarray:
    """For each row of ``X`` (shape ``(N, D)``), return the minimum
    POSITIVE distance to any other row.  Rows for which no
    positive-distance neighbour exists receive ``inf`` so the caller
    can filter them with ``np.isfinite``."""
    if not _HAS_CYTHON:
        raise RuntimeError("sem_cython12 not available; import guarded caller")
    return cy12_nn_distances(np.ascontiguousarray(X, dtype=np.float64))


def extend_frontier_kernel(cur_centers, cur_radii, new_emb, cur_arity: int):
    """Fused bulk update for an incremental aggregation step.  Given
    ``F`` reference items - each summarised by a centre vector
    (``cur_centers``, shape ``(F, D)``) and a radius
    (``cur_radii``, shape ``(F,)``) representing the dispersion of
    ``cur_arity`` underlying points - and ``A`` candidate new
    contributions (``new_emb``, shape ``(A, D)``), produce all
    ``F * A`` updated (centre, radius) pairs that result from
    appending one candidate to one reference item.

    Returns ``(flat_new_centers, flat_radii)`` with shapes
    ``((F*A, D), (F*A,))`` packed in row-major (f, a) order.
    OpenMP-parallel over the outer ``f`` index."""
    if not _HAS_CYTHON:
        raise RuntimeError("sem_cython12 not available; import guarded caller")
    return cy12_extend_frontier_kernel(
        np.ascontiguousarray(cur_centers, dtype=np.float64),
        np.ascontiguousarray(cur_radii, dtype=np.float64),
        np.ascontiguousarray(new_emb, dtype=np.float64),
        int(cur_arity),
    )
