"""
Additional MSA test cases beyond the C++ ground-truth fixture.

These tests focus on edge-cases and small MSAs where we can reason about
expected behavior (especially around all-gap column trimming and unique-gap
interval counting).
"""

from __future__ import annotations

from typing import Sequence

import pytest

import msastats


def _stats_by_name(msa: Sequence[str]) -> dict[str, float]:
    """Return `msastats.calculate_msa_stats` results as a name->value mapping."""
    stats = msastats.calculate_msa_stats(msa)
    names = msastats.stats_names()
    return dict(zip(names, stats))

def _make_seq(length: int, gap_intervals: Sequence[tuple[int, int]]) -> str:
    if length <= 0:
        raise ValueError("length must be positive")

    seq = ["A"] * length
    for start, end in gap_intervals:
        if start < 0 or end < 0:
            raise ValueError("gap interval indices must be non-negative")
        if end < start:
            raise ValueError("gap interval end must be >= start")
        if end >= length:
            raise ValueError("gap interval end must be within sequence length")
        for i in range(start, end + 1):
            seq[i] = "-"

    return "".join(seq)


def test_no_gaps_all_sequences_identical() -> None:
    msa = ["ACGT", "ACGT", "ACGT"]
    s = _stats_by_name(msa)

    assert s["MSA_LEN"] == 4
    assert s["TOT_NUM_GAPS"] == 0
    assert s["TOT_NUM_UNIQUE_GAPS"] == 0
    assert s["AVG_GAP_SIZE"] == 0.0
    assert s["AVG_UNIQUE_GAP_SIZE"] == 0.0
    assert s["LONGEST_UNALIGNED_SEQ"] == 4
    assert s["SHORTEST_UNALIGNED_SEQ"] == 4
    assert s["MSA_POSITION_WITH_0_GAPS"] == 4
    assert s["MSA_POSITION_WITH_1_GAPS"] == 0
    assert s["MSA_POSITION_WITH_2_GAPS"] == 0
    assert s["MSA_POSITION_WITH_N_MINUS_1_GAPS"] == 0


def test_single_gap_event_in_one_sequence() -> None:
    msa = ["ACGT", "A-GT", "ACGT"]
    s = _stats_by_name(msa)

    assert s["TOT_NUM_GAPS"] == 1
    assert s["TOT_NUM_UNIQUE_GAPS"] == 1
    assert s["NUM_GAPS_LEN_ONE"] == 1
    assert s["NUM_GAPS_LEN_ONE_IN_ONE_SEQ"] == 1
    assert s["AVG_GAP_SIZE"] == 1.0
    assert s["AVG_UNIQUE_GAP_SIZE"] == 1.0
    assert s["MSA_POSITION_WITH_1_GAPS"] == 1
    assert s["MSA_POSITION_WITH_0_GAPS"] == 3


def test_all_gap_column_is_trimmed_and_not_counted_as_indel() -> None:
    # Middle column is all-gap (every sequence has '-'), so it should be trimmed
    # before unique gap interval detection. That means zero indels are reported.
    msa = ["A-A", "A-A", "A-A"]
    s = _stats_by_name(msa)

    assert s["MSA_LEN"] == 3
    assert s["TOT_NUM_GAPS"] == 0
    assert s["TOT_NUM_UNIQUE_GAPS"] == 0
    assert s["AVG_GAP_SIZE"] == 0.0
    assert s["AVG_UNIQUE_GAP_SIZE"] == 0.0
    assert s["MSA_POSITION_WITH_0_GAPS"] == 2
    assert s["MSA_POSITION_WITH_1_GAPS"] == 0
    assert s["MSA_POSITION_WITH_2_GAPS"] == 0


def test_all_gap_only_alignment_has_no_indels_after_trim() -> None:
    msa = ["----", "----", "----"]
    s = _stats_by_name(msa)

    assert s["MSA_LEN"] == 4
    assert s["TOT_NUM_GAPS"] == 0
    assert s["TOT_NUM_UNIQUE_GAPS"] == 0
    assert s["LONGEST_UNALIGNED_SEQ"] == 0
    assert s["SHORTEST_UNALIGNED_SEQ"] == 0
    assert s["MSA_POSITION_WITH_0_GAPS"] == 0
    assert s["MSA_POSITION_WITH_1_GAPS"] == 0
    assert s["MSA_POSITION_WITH_2_GAPS"] == 0
    assert s["MSA_POSITION_WITH_N_MINUS_1_GAPS"] == 0


def test_all_gap_column_inside_gap_run_reduces_reported_length() -> None:
    # Column 2 is all-gap (trimmed), while columns 1 and 3 are gaps only in seq1.
    # After trimming column 2, seq1's "---" becomes "--", so the reported indel
    # length should be 2 (not 3).
    msa = [
        "A---A",
        "AB-CA",
        "AD-EA",
    ]
    s = _stats_by_name(msa)

    assert s["MSA_LEN"] == 5
    assert s["TOT_NUM_GAPS"] == 1
    assert s["TOT_NUM_UNIQUE_GAPS"] == 1
    assert s["NUM_GAPS_LEN_TWO"] == 1
    assert s["NUM_GAPS_LEN_THREE"] == 0
    assert s["NUM_GAPS_LEN_TWO_IN_ONE_SEQ"] == 1
    assert s["AVG_GAP_SIZE"] == 2.0
    assert s["AVG_UNIQUE_GAP_SIZE"] == 2.0
    assert s["MSA_POSITION_WITH_0_GAPS"] == 2
    assert s["MSA_POSITION_WITH_1_GAPS"] == 2


def test_gap_run_at_start_and_end_are_counted() -> None:
    msa = [
        "--AA",
        "AAAA",
        "AA--",
    ]
    s = _stats_by_name(msa)

    assert s["MSA_LEN"] == 4
    assert s["TOT_NUM_GAPS"] == 2
    assert s["TOT_NUM_UNIQUE_GAPS"] == 2
    assert s["NUM_GAPS_LEN_TWO"] == 2
    assert s["NUM_GAPS_LEN_TWO_IN_ONE_SEQ"] == 2
    assert s["AVG_GAP_SIZE"] == 2.0
    assert s["AVG_UNIQUE_GAP_SIZE"] == 2.0
    assert s["LONGEST_UNALIGNED_SEQ"] == 4
    assert s["SHORTEST_UNALIGNED_SEQ"] == 2
    assert s["MSA_POSITION_WITH_1_GAPS"] == 4


def test_n_minus_1_gaps_column_and_unique_interval_counting() -> None:
    msa = [
        "AAAA",
        "A-AA",
        "A-AA",
        "A-AA",
        "A-AA",
    ]
    s = _stats_by_name(msa)

    assert s["TOT_NUM_GAPS"] == 4
    assert s["TOT_NUM_UNIQUE_GAPS"] == 1
    assert s["NUM_GAPS_LEN_ONE"] == 4
    assert s["NUM_GAPS_LEN_ONE_IN_ALL_EXCEPT_ONE"] == 1
    assert s["AVG_GAP_SIZE"] == 1.0
    assert s["AVG_UNIQUE_GAP_SIZE"] == 1.0
    assert s["MSA_POSITION_WITH_N_MINUS_1_GAPS"] == 1


def test_recompute_stats_is_idempotent() -> None:
    calc = msastats.MsaStatsCalculator(["ACGT", "A-GT", "ACGT"])
    calc.recompute_stats()
    first = calc.get_stat_vec()
    calc.recompute_stats()
    second = calc.get_stat_vec()

    assert first == second


def test_two_sequences_n_minus_1_equals_one_sequence() -> None:
    # With 2 sequences, "n-1" is 1. The same unique gap interval should be counted
    # as both "in one seq" and "in all except one".
    msa = ["A-A", "AAA"]
    s = _stats_by_name(msa)

    assert s["TOT_NUM_GAPS"] == 1
    assert s["TOT_NUM_UNIQUE_GAPS"] == 1
    assert s["NUM_GAPS_LEN_ONE"] == 1
    assert s["NUM_GAPS_LEN_ONE_IN_ONE_SEQ"] == 1
    assert s["NUM_GAPS_LEN_ONE_IN_ALL_EXCEPT_ONE"] == 1
    assert s["MSA_POSITION_WITH_1_GAPS"] == 1
    # Column-bucket counters are mutually exclusive: the implementation classifies
    # this as a "1 gap" column (not also as an "n-1 gaps" column when n=2).
    assert s["MSA_POSITION_WITH_N_MINUS_1_GAPS"] == 0


def test_multiple_gap_runs_in_one_sequence() -> None:
    msa = [
        "A-A-A",
        "AAAAA",
        "AAAAA",
    ]
    s = _stats_by_name(msa)

    assert s["TOT_NUM_GAPS"] == 2
    assert s["TOT_NUM_UNIQUE_GAPS"] == 2
    assert s["NUM_GAPS_LEN_ONE"] == 2
    assert s["NUM_GAPS_LEN_ONE_IN_ONE_SEQ"] == 2
    assert s["MSA_POSITION_WITH_1_GAPS"] == 2
    assert s["MSA_POSITION_WITH_0_GAPS"] == 3


def test_length_one_bucket_counters_for_n4() -> None:
    # Build an MSA with three distinct length-1 gap intervals:
    # - One interval in exactly 1 sequence
    # - One interval in exactly 2 sequences
    # - One interval in exactly n-1=3 sequences
    length = 6
    msa = [
        _make_seq(length, [(1, 1), (3, 3)]),          # gaps at cols 1 and 3
        _make_seq(length, [(3, 3), (5, 5)]),          # gaps at cols 3 and 5
        _make_seq(length, [(5, 5)]),                  # gap at col 5
        _make_seq(length, [(5, 5)]),                  # gap at col 5
    ]
    s = _stats_by_name(msa)

    assert s["MSA_LEN"] == length
    assert s["TOT_NUM_GAPS"] == 6
    assert s["TOT_NUM_UNIQUE_GAPS"] == 3
    assert s["NUM_GAPS_LEN_ONE"] == 6
    assert s["NUM_GAPS_LEN_ONE_IN_ONE_SEQ"] == 1
    assert s["NUM_GAPS_LEN_ONE_IN_TWO_SEQS"] == 1
    assert s["NUM_GAPS_LEN_ONE_IN_ALL_EXCEPT_ONE"] == 1
    assert s["AVG_GAP_SIZE"] == 1.0
    assert s["AVG_UNIQUE_GAP_SIZE"] == 1.0
    assert s["MSA_POSITION_WITH_0_GAPS"] == 3
    assert s["MSA_POSITION_WITH_1_GAPS"] == 1
    assert s["MSA_POSITION_WITH_2_GAPS"] == 1
    assert s["MSA_POSITION_WITH_N_MINUS_1_GAPS"] == 1


def test_bucket_counters_across_lengths_and_counts() -> None:
    # Exercise length=1,2,3,>=4 with different "shared-by-N-seqs" bucket counters.
    #
    # n=5 sequences, length=15:
    # - len 1 interval shared by 1 seq
    # - len 2 interval shared by 2 seqs
    # - len 3 interval shared by n-1=4 seqs
    # - len 4 interval shared by 2 seqs
    length = 15

    interval_len1 = (1, 1)
    interval_len2 = (3, 4)
    interval_len3 = (6, 8)
    interval_len4 = (10, 13)

    msa = [
        _make_seq(length, [interval_len1, interval_len2, interval_len3, interval_len4]),
        _make_seq(length, [interval_len2, interval_len3, interval_len4]),
        _make_seq(length, [interval_len3]),
        _make_seq(length, [interval_len3]),
        _make_seq(length, []),
    ]
    s = _stats_by_name(msa)

    assert s["MSA_LEN"] == length
    assert s["TOT_NUM_GAPS"] == 9
    assert s["TOT_NUM_UNIQUE_GAPS"] == 4
    assert s["NUM_GAPS_LEN_ONE"] == 1
    assert s["NUM_GAPS_LEN_TWO"] == 2
    assert s["NUM_GAPS_LEN_THREE"] == 4
    assert s["NUM_GAPS_LEN_AT_LEAST_FOUR"] == 2
    assert s["NUM_GAPS_LEN_ONE_IN_ONE_SEQ"] == 1
    assert s["NUM_GAPS_LEN_TWO_IN_TWO_SEQS"] == 1
    assert s["NUM_GAPS_LEN_THREE_IN_ALL_EXCEPT_ONE"] == 1
    assert s["NUM_GAPS_LEN_AT_LEAST_FOUR_IN_TWO_SEQS"] == 1
    assert s["AVG_GAP_SIZE"] == pytest.approx(25.0 / 9.0)
    assert s["AVG_UNIQUE_GAP_SIZE"] == pytest.approx(10.0 / 4.0)
    assert s["LONGEST_UNALIGNED_SEQ"] == 15
    assert s["SHORTEST_UNALIGNED_SEQ"] == 5
    assert s["MSA_POSITION_WITH_0_GAPS"] == 5
    assert s["MSA_POSITION_WITH_1_GAPS"] == 1
    assert s["MSA_POSITION_WITH_2_GAPS"] == 6
    assert s["MSA_POSITION_WITH_N_MINUS_1_GAPS"] == 3


def test_calculate_msa_stats_accepts_tuple_input() -> None:
    stats = msastats.calculate_msa_stats(("ACGT", "A-GT", "ACGT"))
    assert isinstance(stats, list)
    assert len(stats) == 27
