import pytest

from tral.hmm.hmm import HMM
from tral.hmm.hmm_viterbi import hmm_path_to_non_aligned_tandem_repeat_units, \
    hmm_path_to_maximal_complete_tandem_repeat_units
from tral.repeat import repeat
from tral.repeat import repeat_align


@pytest.mark.no_external_software_required
def test_hmm_path_to_non_aligned_tandem_repeat_units():

    TEST = {"Double": ["AAAA", ["M1", "M2", "M1", "M2"], 2, ["AA", "AA"]],
            "Double_converted": ["ABCD", ["M2", "M1", "M2", "M1"], 2, ["AB", "CD"]],
            "Double_converted_complex": ["NNAAAACC",
                                         ["N", "N", "M2", "M1", "M2", "M1", "C", "C"],
                                         2,
                                         ["AA", "AA"]],
            "Double_converted_complex_insertions": ["NNAIAAACC",
                                                    ["N", "N", "M2", "I2", "M1", "M2", "M1", "C", "C"],
                                                    2,
                                                    ["AIA", "AA"]],
            "Single": ["AAAA", ["M1", "M1", "M1", "M1"], 1, ["A", "A", "A", "A"]],
            "Single_Complex": ["NNAAIIIAA",
                               ["N", "N", "M1", "M1", "I1", "I1", "I1", "M1", "M1"],
                               1,
                               ["A", "AIII", "A", "A"]],
            "Long": ["GYRADKLADKLADKL",
                     ["N", "N", "N", "M1", "M2", "M3", "M4", "M1", "M2", "M3", "M4", "M1", "M2", "M3", "M4"],
                     4,
                     ["ADKL", "ADKL", "ADKL"]],
            }
    for test, p in TEST.items():
        test_repeat_msa = hmm_path_to_non_aligned_tandem_repeat_units(sequence=p[0], path=p[1], l_effective=p[2])
        assert test_repeat_msa == p[3]


def test_viterbi():

    # {Test_name: [Original_TR_MSA, Sequence, Viterbi_path, Refined_TR_MSA], ... }
    TEST = {"Single": [["A", "A", "A"], "AAAAAA", ["M1", "M1", "M1", "M1", "M1", "M1"], ["A", "A", "A", "A", "A", "A"]],
            "Double": [["AA", "AA"], "AAAAAA", ["M1", "M2", "M1", "M2", "M1", "M2"], ["AA", "AA", "AA"]],
            "Long": [["ADKL", "ADKL"],
                     "GYRADKLADKLADKL",
                     ["N", "N", "N", "M1", "M2", "M3", "M4", "M1", "M2", "M3", "M4", "M1", "M2", "M3", "M4"],
                     ["ADKL", "ADKL", "ADKL"]]
            }

    for test, p in TEST.items():
        test_repeat = repeat.Repeat(msa=p[0])
        test_hmm = HMM.create(input_format="repeat", repeat=test_repeat)

        for iHMM in [test_hmm]:
            # Detect TRs on self.seq with hmm using the Viterbi algorithm.
            most_likely_path = iHMM.viterbi(p[1])
            assert type(most_likely_path) == list
            assert most_likely_path == p[2]

            unaligned_msa = hmm_path_to_non_aligned_tandem_repeat_units(p[1], most_likely_path, iHMM.l_effective)
            assert unaligned_msa == p[3]

            aligned_msa = repeat_align.realign_repeat(unaligned_msa)
            assert aligned_msa == p[3]

# The following test functions were copy-pasted from hmm_viterbi.py. They need to be either
# strongly adapted, or - if not useful - erased.


@pytest.mark.notfixed
@pytest.mark.no_external_software_required
def test_conversion_multiple():
    """ Test <hmm_path_to_maximal_complete_tandem_repeat_units>.
        However, no positive set is defined right now.
    """
    assert 0, "Test not fixed"
    l_effective = 3
    lPaths = [['N', 'N', 'N', 'M0', 'M1', 'M2', 'M0', 'I1', 'M1', 'M2', 'M0', 'M1', 'M2', 'N', 'N', 'N'],
              ['N', 'N', 'N', 'M1', 'M2', 'M0', 'M1', 'M2', 'M0', 'M1', 'M2', 'M0', 'N', 'N', 'N'],
              ['N', 'N', 'N', 'M2', 'M0', 'M1', 'M2', 'M0', 'M1', 'M2', 'M0', 'M1', 'N', 'N', 'N']]
    lSequence = ['XXXABCADBCABCXXX', 'XXXBCABCABCAXXX', 'XXXCABCABCABXXX']
    lMSA = hmm_path_to_maximal_complete_tandem_repeat_units(lSequence, lPaths, l_effective)
    assert lMSA is not None


@pytest.mark.notfixed
@pytest.mark.no_external_software_required
def test_create_viterbi():
    """ This test needs to be fixed, the HMM cannot be initialised in this way.
    """
    assert 0, "Test not fixed"
    my_hmm = HMM()
    my_hmm.states = ['H', 'F']

    # Initialisation
    my_hmm.p_0 = {"H": 0.6, "F": 0.4}

    # Feed Values to p_t
    #my_hmm.p_t["START"] = {"H": 0.6, "F": 0.4}
    my_hmm.p_t["H"] = {"H": 0.7, "F": 0.3}
    my_hmm.p_t["F"] = {"H": 0.4, "F": 0.6}

    # emissions
    my_hmm.emissions = ["D", "C", "N"]

    # emission probabilities
    my_hmm.p_e['H'] = {"D": 0.1, "C": 0.4, "N": 0.5}
    my_hmm.p_e['F'] = {"D": 0.6, "C": 0.3, "N": 0.1}

    my_viterbi = Viterbi(my_hmm, "NCD")
    if my_viterbi.viterbi() == "HHF":
        print("Test SUCCESSFUL")
    else:
        print("Test FAILED")


@pytest.mark.notfixed
def test_init_with_repeat():
    """ This test needs to be fixed.
    """
    assert 0, "Test not fixed"
    my_TR = repeat.Repeat(begin=0, msa=['A-G', 'ACG', 'ACG'], sequence_type='DNA')
    my_hmm = HMM.create(repeat=my_TR)
    from . import sequence
    my_sequence = sequence.Sequence()
    my_viterbi = Viterbi(my_hmm, my_sequence.sequence)
    v = my_viterbi.viterbi()
    "".join(v)
