import pytest

from apb.sequence.mutation import (
    BLOSUM62,
    BLOSUM90,
    Mutation,
    apply_mutations,
    find_mutations,
    get_multi_mutation_string,
)


def test_blosum_lookups():
    # https://en.wikipedia.org/wiki/BLOSUM.
    assert BLOSUM62.get_score("H", "S") == -1
    assert BLOSUM62.get_score("K", "C") == -3
    assert BLOSUM62.get_score("N", "P") == -2

    # https://github.com/biopython/biopython/blob/master/Bio/Align/substitution_matrices/data/BLOSUM90.
    assert BLOSUM90.get_score("H", "S") == -2
    assert BLOSUM90.get_score("K", "C") == -4
    assert BLOSUM90.get_score("N", "P") == -3


def test_mutation_init():
    mut = Mutation("A", 10, "V")
    assert mut.before == "A"
    assert mut.pos == 10
    assert mut.after == "V"


def test_mutation_repr():
    mut = Mutation("A", 10, "V")
    assert repr(mut) == "A10V"


def test_mutation_as_str():
    mut = Mutation("A", 10, "V")
    assert mut.as_str() == "A10V"


def test_mutation_as_pikaa():
    mut = Mutation("A", 10, "V")
    assert mut.as_pikaa() == "10 A PIKAA V"
    assert mut.as_pikaa(chain="B") == "10 B PIKAA V"


def test_mutation_reverse():
    mut = Mutation("A", 10, "V")
    rev = mut.reverse()
    assert rev.before == "V"
    assert rev.pos == 10
    assert rev.after == "A"


def test_mutation_blosum90():
    mut = Mutation("H", 5, "S")
    assert mut.blosum90() == -2

    mut = Mutation("K", 10, "C")
    assert mut.blosum90() == -4

    mut = Mutation("N", 15, "P")
    assert mut.blosum90() == -3


def test_mutation_blosum62():
    mut = Mutation("H", 5, "S")
    assert mut.blosum62() == -1

    mut = Mutation("K", 10, "C")
    assert mut.blosum62() == -3

    mut = Mutation("N", 15, "P")
    assert mut.blosum62() == -2


def test_mutation_from_str():
    mut = Mutation.from_str("A10V")
    assert mut.before == "A"
    assert mut.pos == 10
    assert mut.after == "V"

    mut = Mutation.from_str("K123C")
    assert mut.before == "K"
    assert mut.pos == 123
    assert mut.after == "C"


def test_mutation_from_multistr():
    muts = Mutation.from_multistr("A10V:K20C:N30P")
    assert len(muts) == 3
    assert muts[0].as_str() == "A10V"
    assert muts[1].as_str() == "K20C"
    assert muts[2].as_str() == "N30P"

    # Works with custom delimiter.
    muts = Mutation.from_multistr("A10V|K20C", delim="|")
    assert len(muts) == 2
    assert muts[0].as_str() == "A10V"
    assert muts[1].as_str() == "K20C"


def test_apply_mutations():
    seq = "ACDEFGHIKLMNPQRSTVWY"

    muts = [Mutation("A", 1, "V")]
    result = apply_mutations(seq, muts)
    assert result == "VCDEFGHIKLMNPQRSTVWY"
    #                 |
    #               mut 1

    muts = [Mutation("A", 1, "V"), Mutation("L", 10, "R"), Mutation("Y", 20, "F")]
    result = apply_mutations(seq, muts)
    assert result == "VCDEFGHIKRMNPQRSTVWF"
    #                 |        |         |
    #               mut 1    mut 2     mut 3

    # Cannot mutate outside sequence's length
    with pytest.raises(IndexError, match="outside sequence range"):
        apply_mutations(seq, [Mutation("A", 25, "V")])

    with pytest.raises(IndexError, match="outside sequence range"):
        apply_mutations(seq, [Mutation("A", 0, "V")])

    # Cannot mutate if starting amino acid doesn't match
    with pytest.raises(ValueError, match="sequence doesn't match mutation"):
        apply_mutations(seq, [Mutation("C", 1, "V")])


def test_find_mutations():
    seq1 = "ACDEFGHIKLMNPQRSTVWY"
    #       |        |         |
    seq2 = "VCDEFGHIKRMNPQRSTVWF"

    muts = find_mutations(seq1, seq2)
    assert len(muts) == 3
    assert muts[0].as_str() == "A1V"
    assert muts[1].as_str() == "L10R"
    assert muts[2].as_str() == "Y20F"

    # No mutations
    muts = find_mutations(seq1, seq1)
    assert len(muts) == 0

    # Test strict length checking
    with pytest.raises(ValueError):
        find_mutations("ACD", "ACDE")


def test_get_multi_mutation_string():
    muts = [Mutation("A", 1, "V"), Mutation("K", 10, "R"), Mutation("Y", 20, "F")]

    result = get_multi_mutation_string(muts)
    assert result == "A1V:K10R:Y20F"

    # Custom delimiter works.
    result = get_multi_mutation_string(muts, delim="|")
    assert result == "A1V|K10R|Y20F"

    # Empty list is ok.
    result = get_multi_mutation_string([])
    assert result == ""


def test_apply_find_roundtrip():
    """Test that applying mutations and then finding them gives back the same mutations."""
    original_seq = "ACDEFGHIKLMNPQRSTVWY"

    mutations_to_apply = [
        Mutation("A", 1, "V"),
        Mutation("F", 5, "Q"),
        Mutation("K", 9, "R"),
        Mutation("L", 10, "R"),
        Mutation("Y", 20, "F"),
    ]

    mutated_seq = apply_mutations(original_seq, mutations_to_apply)
    found_mutations = find_mutations(original_seq, mutated_seq)

    for applied, found in zip(mutations_to_apply, found_mutations, strict=False):
        assert applied == found


def test_mutation_ordering_multistring():
    """Test that mutations maintain order when converted to/from multistring format."""
    # Create mutations in a specific order
    mutations = [
        Mutation("Y", 20, "F"),
        Mutation("F", 5, "Q"),
        Mutation("A", 1, "V"),
        Mutation("K", 9, "R"),
        Mutation("L", 10, "R"),
    ]

    multistring = get_multi_mutation_string(mutations)
    assert multistring == "Y20F:F5Q:A1V:K9R:L10R"

    parsed_mutations = Mutation.from_multistr(multistring)
    for original, parsed in zip(mutations, parsed_mutations, strict=False):
        assert original == parsed
