# This source code is part of the Biotite package and is distributed
# under the 3-Clause BSD License. Please see 'LICENSE.rst' for further
# information.

import numpy as np
import pytest
import biotite.sequence as seq
import biotite.sequence.align as align


@pytest.fixture
def kmer_alphabet():
    return align.KmerAlphabet(seq.ProteinSequence.alphabet, 3)


np.random.seed(0)
N = 10
@pytest.mark.parametrize("ref_kmer, threshold", zip(
        np.random.randint(10000,  size=N),
        np.random.randint(-5, 15, size=N)
))
def test_score_threshold_rule(kmer_alphabet, ref_kmer, threshold):
    """
    Test if the similar k-mers given by :class:`ScoreThresholdRule`
    are equal to k-mers generated by a brute-force approach.
    """
    matrix = align.SubstitutionMatrix.std_protein_matrix()
    
    ref_kmer_sequence = seq.ProteinSequence()
    ref_kmer_sequence.code = kmer_alphabet.split(ref_kmer)
    
    ref_sim_kmer_set = set()
    # Iterate through all possible k-mers 
    for kmer in range(len(kmer_alphabet)):
        kmer_sequence = seq.ProteinSequence()
        kmer_sequence.code = kmer_alphabet.split(kmer)
        score = align.align_ungapped(
            ref_kmer_sequence, kmer_sequence, matrix, score_only=True
        )
        # Add k-mer to list if the threshold score is reached
        if score >= threshold:
            ref_sim_kmer_set.add(kmer)
    
    test_rule = align.ScoreThresholdRule(matrix, threshold)
    test_sim_kmer_set = set(test_rule.similar_kmers(kmer_alphabet, ref_kmer))

    assert test_sim_kmer_set == ref_sim_kmer_set


def test_invalid_matrix(kmer_alphabet):
    """
    Expect an error when the k-mer alphabet is incompatible with the
    substitution matrix.
    """
    test_rule = align.ScoreThresholdRule(
        align.SubstitutionMatrix.std_nucleotide_matrix(), 0
    )
    with pytest.raises(ValueError):
        test_rule.similar_kmers(kmer_alphabet, 0)


@pytest.mark.parametrize("invalid_kmer", [-1, 1000000])
def test_invalid_kmer(kmer_alphabet, invalid_kmer):
    """
    Expect an error when an invalid k-mer code is given as input.
    """
    test_rule = align.ScoreThresholdRule(
        align.SubstitutionMatrix.std_protein_matrix(), 0
    )
    with pytest.raises(seq.AlphabetError):
        test_rule.similar_kmers(kmer_alphabet, invalid_kmer)