"""Global sequence alignment via Needleman-Wunsch dynamic programming.

``needleman_wunsch(score_matrix, ...)`` accepts a precomputed ``(n, m)``
similarity matrix and returns index arrays (similar to
``scipy.optimize.linear_sum_assignment``). The DP fill and traceback run
entirely in Rust.
"""

from __future__ import annotations

import numpy as np

from pynw._pynw_core import nw_traceback_indices

# Tie-breaking priority: DIAG > UP > LEFT.  This is enforced during the DP
# fill by choosing DIAG as the default and only overwriting on strict >.
# The effect is compact alignments that prefer substitutions over gaps when
# costs are equal.  The optimal *score* is unique regardless of tie-breaking,
# but the alignment path is not -- different tools may return different
# co-optimal alignments.


def needleman_wunsch(
    score_matrix: np.ndarray,
    gap_penalty_seq1: float = -1.0,
    gap_penalty_seq2: float = -1.0,
) -> tuple[float, np.ndarray, np.ndarray]:
    """Solve the global sequence alignment problem.

    Parameters
    ----------
    score_matrix : array_like, shape (n, m)
        ``score_matrix[i, j]`` is the similarity score for aligning
        element *i* of sequence 1 with element *j* of sequence 2.
    gap_penalty_seq1 : float, default -1.0
        Penalty added when a gap is inserted in sequence 1
        (sequence 2 advances).
    gap_penalty_seq2 : float, default -1.0
        Penalty added when a gap is inserted in sequence 2
        (sequence 1 advances).

    Returns
    -------
    score : float
        The optimal alignment score.
    seq1_idx : ndarray of intp
        Index into sequence 1 at each alignment position, or ``-1``
        for a gap.
    seq2_idx : ndarray of intp
        Index into sequence 2 at each alignment position, or ``-1``
        for a gap.

    Examples
    --------
    >>> import numpy as np
    >>> sm = np.array([[1, -1, -1],
    ...                [-1, -1, -1],
    ...                [-1, -1,  1]], dtype=float)
    >>> score, s1, s2 = needleman_wunsch(sm, gap_penalty_seq1=-2, gap_penalty_seq2=-2)
    >>> s1
    array([0, 1, 2])
    >>> s2
    array([0, 1, 2])

    The match-score component is composable, like
    ``scipy.optimize.linear_sum_assignment``::

        matched = (s1 >= 0) & (s2 >= 0)
        score_matrix[s1[matched], s2[matched]].sum()   # match contribution

    Notes
    -----
    Tie-breaking is deterministic: ``DIAG > UP > LEFT``.  The optimal
    *score* is unique, but multiple alignments may achieve it.  This
    rule prefers substitutions over gaps, producing compact alignments.
    Other tools may return different co-optimal alignments.

    ``NaN`` and ``Inf`` values in the score matrix are not rejected but
    will propagate through the dynamic programming computation, likely
    producing meaningless results.

    """
    sm = np.asarray(score_matrix, dtype=np.float64)
    if sm.ndim != 2:
        msg = f"score_matrix must be 2-dimensional, got {sm.ndim}-dimensional array"
        raise ValueError(msg)

    return nw_traceback_indices(sm, gap_penalty_seq1, gap_penalty_seq2)
