import pytest
from utils.lca_determination.determine_lca import create_lca
from utils.lca_determination import weighted_lca_class, lca_utils, determine_lca
from utils.lca_determination.lca_utils import ProcessedTaskData, CheckLineageConsistence
from utils.lca_determination.weighted_lca_class import WeightedLCA

testdata_pgid_to_accs_dict = {0: ['A_1', 'A_2', 'A_3'], 1: ['B_1', 'B_2', 'B_3'], 2: ['C_1', 'C_2', 'C_3'],
                              3:  ['D_1', 'D_2', 'D_3'],  4:  ['E_1', 'E_2'], 5: ['F1', 'F2', 'F3']}
testdata_accs_to_spectra_dict_proteins = {'A_1': {'A_1'}, 'A_2': {'A_2'}, 'A_3': {'A_3'}, 'B_1': {'B_1'}, 'B_2': {'B_2'},
                                          'B_3': {'B_3'}, 'C_1': {'C_1'}, 'C_2': {'C_2'}, 'C_3': {'C_3'},'D_1': {'D_1'},
                                          'D_2': {'D_2'}, 'D_3': {'D_3'}, 'E_1': {'E_1'}, 'E_2': {'E_2'},
                                          'F1': {'F11'}, 'F2': {'F2'}, 'F3': {'F3'}}
testdata_accs_to_spectra_dict_spectra = {'A_1': {'s_1', 's_2', 's3'}, 'A_2': {'s_2', 's3'}, 'A_3': {'s_1'}}
testdata_acc_to_lineages_dict_group_all_same = {
    'A_1': [['Bacteria',	'Firmicutes', 'Bacilli', 'Lactobacillales', 'Streptococcaceae', 'Lactococcus',
            'Lactococcus lactis'], ['','','','','','','']],
    'A_2': [['Bacteria',	'Firmicutes', 'Bacilli', 'Lactobacillales', 'Streptococcaceae', 'Lactococcus',
            'Lactococcus lactis'], ['','','','','','','']],
    'A_3': [['Bacteria',	'Firmicutes', 'Bacilli', 'Lactobacillales', 'Streptococcaceae', 'Lactococcus',
            'Lactococcus lactis'], ['','','','','','','']]
}
testdata_acc_to_lineages_dict_group_all_unclassified = {
    'B_1': [['unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified',
            'unclassified'], ['','','','','','','']],
    'B_2': [['unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified',
            'unclassified'], ['','','','','','','']],
    'B_3': [['unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified',
            'unclassified'], ['','','','','','','']],
}
testdata_acc_to_lineages_dict_group_different = {
    'C_1': [['Bacteria', 'Actinobacteria', 'Actinobacteria', 'Bifidobacteriales', 'Bifidobacteriaceae', 'Bifidobacterium', 'Bifidobacterium longum'],
            ['Bacteria', 'Actinobacteria', 'Actinobacteria', 'Bifidobacteriales', 'Bifidobacteriaceae', 'Bifidobacterium', 	'Bifidobacterium longum CAG:69']],
    'C_2': [['Bacteria', 'Proteobacteria', 'Gammaproteobacteria', 'Pseudomonadales', 'Moraxellaceae', 'Acinetobacter', 'Acinetobacter baumannii'],
            ['Bacteria', 'Proteobacteria', 'Gammaproteobacteria', 'Pseudomonadales', 'Pseudomonadaceae', 'Pseudomonas', 'Pseudomonas aeruginosa'],
            ['Bacteria', 'Proteobacteria', 'Gammaproteobacteria', 'Pseudomonadales', 'Pseudomonadaceae', 'Pseudomonas', 'Pseudomonas sp. HMSC058A10'],
            ['Bacteria', 'Proteobacteria', 'Gammaproteobacteria', 'Pseudomonadales', 'Pseudomonadaceae', 'Pseudomonas',	'Pseudomonas sp. HMSC058C05'],
            ['Bacteria', 'Proteobacteria', 'Gammaproteobacteria', 'Pseudomonadales', 'Pseudomonadaceae', 'Pseudomonas',	'Pseudomonas sp. HMSC059F05'],
            ['Bacteria', 'Proteobacteria', 'Gammaproteobacteria', 'Pseudomonadales', 'Pseudomonadaceae', 'Pseudomonas',	'Pseudomonas sp. HMSC05H02'],
            ['Bacteria', 'Proteobacteria', 'Gammaproteobacteria', 'Pseudomonadales', 'Pseudomonadaceae', 'Pseudomonas', 'sp. HMSC060F12']],
    'C_3': [['Eukaryota', 'Ascomycota', 'Sordariomycetes',	'Sordariales', 'Sordariaceae', 'Neurospora', 'Neurospora crassa']]
}
testdata_acc_to_lineages_dict_group_one_classified = {
    'D_1': [['Bacteria', 'Actinobacteria', 'Actinobacteria', 'Bifidobacteriales', 'Bifidobacteriaceae', 'Bifidobacterium', 'Bifidobacterium longum'],
            ['','','','','','','']],
    'D_2': [['unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified',
             'unclassified'], ['','','','','','','']],
    'D_3': [['unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified',
             'unclassified'], ['','','','','','','']]
}

testdata_acc_to_lineages_dict_group_one_nonesense = {
    'E_1': [['Bacteria', 'Actinobacteria', 'Actinobacteria', 'Bifidobacteriales', 'Bifidobacteriaceae', 'Bifidobacterium', 'Bifidobacterium longum'],
            ['','','','','','','']],
    'E_2': [['unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'vector_x']]
}

testdata_acc_to_lineages_dict_group_error_prone = {
    'F1': ["Viruses", "unclassified", "unclassified", "unclassified", "Iridoviridae", "Ranavirus", "Frog virus 3"],
    'F2': ["Archaea", "unclassified", "unclassified", "unclassified", "unclassified", "unclassified", "uncultured ammonia-oxidizing archaeon"],
    'F3': ["Archaea", "unclassified", "unclassified", "unclassified", "unclassified", "unclassified", "uncultured ammonia-oxidizing archaeon"]
}

testdata_acc_to_lineages_dict_all = {**testdata_acc_to_lineages_dict_group_all_same, **testdata_acc_to_lineages_dict_group_all_unclassified,
                                 **testdata_acc_to_lineages_dict_group_different, **testdata_acc_to_lineages_dict_group_one_classified,
                                 **testdata_acc_to_lineages_dict_group_one_nonesense, **testdata_acc_to_lineages_dict_group_error_prone}
testdata_acc_to_lineages_dict = [testdata_acc_to_lineages_dict_group_all_same, testdata_acc_to_lineages_dict_group_all_unclassified,
                                 testdata_acc_to_lineages_dict_group_different, testdata_acc_to_lineages_dict_group_one_classified,
                                 testdata_acc_to_lineages_dict_group_one_nonesense, testdata_acc_to_lineages_dict_group_error_prone]

test_result_group_lca_501 = {'lca': {0: ['Bacteria', 'Firmicutes', 'Bacilli', 'Lactobacillales', 'Streptococcaceae', 'Lactococcus', 'Lactococcus lactis'],
                             1: ['unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified'],
                             2: ['Bacteria', 'various', 'various', 'various', 'various', 'various', 'various'],
                             3: ['various', 'various', 'various', 'various', 'various', 'various', 'various'],
                             4: ['various', 'various', 'various', 'various', 'various', 'various', 'various'],
                                     5: ['Archaea', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'uncultured ammonia-oxidizing archaeon']}
                   }
test_result_group_lca_support_501 = {'lca': {0: ['3/3', '3/3', '3/3', '3/3', '3/3', '3/3', '3/3'],
                                      1: ['0/0', '0/0', '0/0', '0/0', '0/0', '0/0', '0/0'],
                                      2: ['2/3', '(1)/3', '(1)/3', '(1)/3', '(1)/3', '(1)/3', '(1)/3'],
                                      3: ['(1)/3', '(1)/3', '(1)/3', '(1)/3', '(1)/3', '(1)/3', '(1)/3'],
                                      4: ['(1)/2', '(1)/2', '(1)/2', '(1)/2', '(1)/2', '(1)/2', '(1)/2'],
                                             5: ['2/3', '2/3', '2/3', '2/3', '2/3', '2/3', '2/3']}}

test_result_lca_1 = {'lca': {0: ['Bacteria', 'Firmicutes', 'Bacilli', 'Lactobacillales', 'Streptococcaceae', 'Lactococcus', 'Lactococcus lactis'],
                                     1: ['unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified'],
                                     2: ['various', 'various', 'various', 'various', 'various', 'various', 'various'],
                                     3: ['various', 'various', 'various', 'various', 'various', 'various', 'various'],
                                     4: ['various', 'various', 'various', 'various', 'various', 'various', 'various'],
                                    5: ['various', 'unclassified', 'unclassified', 'unclassified', 'various', 'various', 'various']},
                             'democratic_lca': {0: ['Bacteria', 'Firmicutes', 'Bacilli', 'Lactobacillales', 'Streptococcaceae', 'Lactococcus', 'Lactococcus lactis'],
                                                1: ['unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified'],
                                                2: ['Bacteria', 'Actinobacteria', 'Actinobacteria', 'Bifidobacteriales', 'Bifidobacteriaceae', 'Bifidobacterium', 'Bifidobacterium longum'],
                                                3: ['Bacteria', 'Actinobacteria', 'Actinobacteria', 'Bifidobacteriales', 'Bifidobacteriaceae', 'Bifidobacterium', 'Bifidobacterium longum'],
                                                4: ['Bacteria', 'Actinobacteria', 'Actinobacteria', 'Bifidobacteriales', 'Bifidobacteriaceae', 'Bifidobacterium', 'Bifidobacterium longum'],
                                                5: ['Archaea', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'unclassified', 'uncultured ammonia-oxidizing archaeon']}
                             }
test_result_lca_support_1 = {'lca': {0: ['3/3', '3/3', '3/3', '3/3', '3/3', '3/3', '3/3'],
                                            1: ['0/0', '0/0', '0/0', '0/0', '0/0', '0/0', '0/0'],
                                            2: ['(2)/3', '(1)/3', '(1)/3', '(1)/3', '(1)/3', '(1)/3', '(1)/3'],
                                            3: ['(1)/3', '(1)/3', '(1)/3', '(1)/3', '(1)/3', '(1)/3', '(1)/3'],
                                            4: ['(1)/2', '(1)/2', '(1)/2', '(1)/2', '(1)/2', '(1)/2', '(1)/2'],
                                            5: ['(2)/3', '(2)/3', '(2)/3', '(2)/3', '(2)/3', '(2)/3', '(2)/3']},
                                      'democratic_lca': {0: ['3/7', '3/3', '3/3', '3/3', '3/3', '3/3', '3/3'],
                                                         1: ['0/0', '0/0', '0/0', '0/0', '0/0', '0/0', '0/0'],
                                                         2: ['2/7', '1/3', '1/3', '1/3', '1/3', '1/3', '1/3'],
                                                         3: ['1/7', '1/3', '1/3', '1/3', '1/3', '1/3', '1/3'],
                                                         4: ['1/7', '1/3', '1/3', '1/3', '1/3', '1/3', '1/3'],
                                                         5: ['2/2', '2/2', '2/2', '2/2', '2/2', '2/2', '2/2']}}
test_result_min_nb_annotations_2 = {4: ['Bacteria', 'Actinobacteria', 'Actinobacteria', 'Bifidobacteriales', 'Bifidobacteriaceae', 'Bifidobacterium', 'Bifidobacterium longum']}
test_result_min_nb_annotations_sp = {4: ['(1)/2', '(1)/2', '(1)/2', '(1)/2', '(1)/2', '(1)/2', '(1)/2']}


@pytest.mark.parametrize("method", ['lca', 'democratic_lca'])
@pytest.mark.parametrize("threshold", [0.501, 1])
@pytest.mark.parametrize("pgid", [0, 1, 2, 3, 4, 5])
def test_lca(pgid, threshold, method):
    with_spectra = False
    is_tax_task = True
    acc_to_lin = testdata_acc_to_lineages_dict[pgid]
    democratic_annotation_weight_df = None
    data = ProcessedTaskData(testdata_pgid_to_accs_dict, testdata_accs_to_spectra_dict_proteins, acc_to_lin,
                             ignore_unclassified=False, minimum_number_of_annotations=0)
    if method == 'democratic_lca':
        lca = WeightedLCA('', testdata_acc_to_lineages_dict_all, testdata_pgid_to_accs_dict,
                          testdata_accs_to_spectra_dict_proteins, method, ' ', is_tax_task,
                          threshold=threshold, ignore_unclassified=False, minimum_number_of_annotations=0)
        democratic_annotation_weight_df = lca.democratic_annotation_weight_df

    lca_per_group = weighted_lca_class.GroupLCA(pgid, 7, data.get_group_data(pgid), method, False,
                                                democratic_annotation_weight_df, threshold)
    result = lca_per_group.get_lca_and_support()
    if threshold == 1:
        assert test_result_lca_1[method][pgid] == result[0]
        assert test_result_lca_support_1[method][pgid] == result[1]
    if threshold == 0.501 and method == 'lca':
        assert test_result_group_lca_501[method][pgid] == result[0]
        assert test_result_group_lca_support_501[method][pgid] == result[1]
