import numpy as np
import pytest

from pynw import needleman_wunsch


def test_returns_score_and_index_arrays() -> None:
    score_matrix = np.array(
        [[2.0, -1.0, -1.0], [-1.0, 2.0, -1.0], [-1.0, -1.0, 2.0]],
        dtype=np.float64,
    )

    score, s1, s2 = needleman_wunsch(score_matrix)

    assert score == 6.0
    assert s1.dtype == np.intp
    assert s2.dtype == np.intp
    np.testing.assert_array_equal(s1, [0, 1, 2])
    np.testing.assert_array_equal(s2, [0, 1, 2])


def test_rejects_non_2d_input() -> None:
    with pytest.raises(ValueError, match=r"2-dimensional"):
        needleman_wunsch(np.array([1.0, 2.0]))

    with pytest.raises(ValueError, match=r"2-dimensional"):
        needleman_wunsch(np.zeros((1, 1, 1)))


def test_accepts_non_float_input_by_casting() -> None:
    score_matrix = np.array([[2, -1], [-1, 2]], dtype=np.int32)

    score, s1, s2 = needleman_wunsch(score_matrix)

    assert score == 4.0
    np.testing.assert_array_equal(s1, [0, 1])
    np.testing.assert_array_equal(s2, [0, 1])


def test_empty_matrix_returns_empty_alignment() -> None:
    score, s1, s2 = needleman_wunsch(np.empty((0, 0), dtype=np.float64))

    assert score == 0.0
    assert len(s1) == 0
    assert len(s2) == 0


def test_one_empty_dimension_returns_all_gaps() -> None:
    score, s1, s2 = needleman_wunsch(
        np.empty((0, 3), dtype=np.float64), gap_penalty_seq1=-2
    )
    assert score == -6.0
    np.testing.assert_array_equal(s1, [-1, -1, -1])
    np.testing.assert_array_equal(s2, [0, 1, 2])

    score2, s1b, s2b = needleman_wunsch(
        np.empty((2, 0), dtype=np.float64), gap_penalty_seq2=-3
    )
    assert score2 == -6.0
    np.testing.assert_array_equal(s1b, [0, 1])
    np.testing.assert_array_equal(s2b, [-1, -1])


def test_tie_break_prefers_diagonal_when_all_moves_tie() -> None:
    score, s1, s2 = needleman_wunsch(np.array([[0.0]], dtype=np.float64), 0.0, 0.0)

    assert score == 0.0
    np.testing.assert_array_equal(s1, [0])
    np.testing.assert_array_equal(s2, [0])


def test_tie_break_prefers_up_over_left_when_diagonal_worse() -> None:
    score, s1, s2 = needleman_wunsch(np.array([[-2.0]], dtype=np.float64), 0.0, 0.0)

    assert score == 0.0
    np.testing.assert_array_equal(s1, [-1, 0])
    np.testing.assert_array_equal(s2, [0, -1])


def test_asymmetric_gaps_change_traceback_shape() -> None:
    score_matrix = np.array([[-1.0, 2.0], [2.0, -1.0]], dtype=np.float64)

    score, s1, s2 = needleman_wunsch(
        score_matrix, gap_penalty_seq1=-0.5, gap_penalty_seq2=-3.0
    )

    assert score == -1.5
    np.testing.assert_array_equal(s1, [-1, 0, 1])
    np.testing.assert_array_equal(s2, [0, 1, -1])
