<commit_before># ----------------------------------------------------------------------------
# Copyright (c) 2015--, ghost-tree development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the LICENSE file, distributed with this software.
# ----------------------------------------------------------------------------
import unittest
import os
from io import StringIO

import skbio
from skbio import Sequence

from ghosttree.scaffold.hybridtree import _make_nr_foundation_alignment
from ghosttree.scaffold.hybridtree import _create_taxonomy_dict
from ghosttree.scaffold.hybridtree import _make_mini_otu_files
from ghosttree.scaffold.hybridtree import _extension_genus_accession_dict


class TestScaffoldExtensionsIntofoundation(unittest.TestCase):
    def setUp(self):
        self.otu_clusters = StringIO(otu_clusters)
        self.extension_taxonomy = StringIO(extension_taxonomy)
        self.extension_taxonomy_unids = StringIO(extension_taxonomy_unids)
        self.extension_taxonomy_none = StringIO(extension_taxonomy_none)
        self.extension_seqs = StringIO(extension_seqs)
        self.foundation_alignment = StringIO(foundation_alignment)
        self.ghost_tree_fp = StringIO(ghost_tree_fp)
        self.extension_genus_dic_few = extension_genus_dic_few
        self.extension_genus_dic_none = extension_genus_dic_none
        self.key_node = key_node
        self.graft_letter_g = 'g'
        self.graft_level_6 = 6
        self.graft_level_5 = 5
        self.foundation_taxonomy = StringIO(foundation_taxonomy)
        self.foundation_newick = StringIO(foundation_newick)

    def test_make_nr_foundation_alignment_few(self):
        result = _make_nr_foundation_alignment(self.foundation_alignment,
                                               self.extension_genus_dic_few,
                                               self.graft_letter_g)

        self.assertEqual(list(result), [
            Sequence("AAA---", metadata={"id": "PBB1",
                                         "description": "Phoma"}),
            Sequence("AAG---", metadata={"id": "CBB3",
                                         "description": "Candida"}),
        ])

    def test_make_nr_foundation_alignment_none(self):
        result = _make_nr_foundation_alignment(self.foundation_alignment,
                                               self.extension_genus_dic_none,
                                               self.graft_letter_g)
        self.assertEqual(list(result), [])

    # def test_newick_file_few_extensions(self):
    #     result = _make_nr_foundation_newick(self.foundation_newick,
    #                                         self.extension_genus_dic_few)
    #
    #     print('result', result)

    def test_create_taxonomy_dic_many_genus(self):

        test = {'P1': 'Phoma', 'P2': 'Phoma', 'C1': 'Candida', 'C2': 'Candida',
                'C3': 'Candida', 'C4': 'Candida', 'C5': 'Candida',
                'M1': 'Mucor', 'M2': 'Mucor', 'M3': 'Mucor', 'M4': 'Mucor'}

        result = _create_taxonomy_dict(self.extension_taxonomy,
                                       self.graft_level_6)
        print('\n\n\nlalala', result)
        self.assertDictEqual(result, test)

    def test_create_taxonomy_dic_none(self):
        with self.assertRaises(ValueError):
            list(_create_taxonomy_dict(self.extension_taxonomy_none,
                                       9))

    def test_make_mini_otu_files(self):
        os.system("mkdir tmp")
        extension_seqs = list(skbio.io.read(self.extension_seqs,
                                            format='fasta'))

        _make_mini_otu_files(self.key_node,
                             self.extension_genus_dic_few,
                             extension_seqs)

        result = str(list(skbio.io.read('tmp/mini_seq_gt.fasta',
                                        format='fasta'))[0])

        os.system("rm -r tmp")
        self.assertEqual(result, 'TTAAAAAA')

    def test_extension_genus_accession_dic(self):
        test = {'Candida': ['C1', 'C2', 'C3', 'M1', 'C4', 'C5', 'M4'],
                'Mucor': ['M2', 'M3'],
                'Phoma': ['P1', 'P2']}
        result = _extension_genus_accession_dict(self.otu_clusters,
                                                 self.extension_taxonomy,
                                                 self.graft_level_6)
        self.assertDictEqual(result, test)

    def test_extension_genus_accession_dic_unidentifieds_level_6_genus(self):
        test = {'Candida': ['C1', 'C2', 'C3', 'M1'],
                'Mucor': ['M2', 'M3', 'C4', 'C5', 'M4'],
                'Phoma': ['P1', 'P2']}
        result = _extension_genus_accession_dict(self.otu_clusters,
                                                 self.extension_taxonomy_unids,
                                                 self.graft_level_6)
        self.assertDictEqual(result, test)

    def test_extension_genus_accession_dic_unidentifieds_level_5_family(self):
        test = {'Saccharomycetaceae': ['C1', 'C2', 'C3', 'M1', 'C4', 'C5',
                'M4'], 'Didymellaceae': ['P1', 'P2'], 'Mucoraceae':
                ['M2', 'M3']}
        result = _extension_genus_accession_dict(self.otu_clusters,
                                                 self.extension_taxonomy_unids,
                                                 self.graft_level_5)

        self.assertDictEqual(result, test)


key_node = 'Phoma'

extension_genus_dic_few = {'Candida': ['C1', 'C2', 'C3', 'P2'],
                           'Phoma': ['P1']}
extension_genus_dic_none = {'Asperggg': ['C2', 'C3', 'P2'],
                            'Nope': ['P1', 'M1']}

otu_clusters = """C1\tC1\tC2\tC3\tM1
M2\tM2\tM3
P1\tP1\tP2<commit_msg>fix unit test; handles otu cluster tie breakers<commit_after># ----------------------------------------------------------------------------
# Copyright (c) 2015--, ghost-tree development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the LICENSE file, distributed with this software.
# ----------------------------------------------------------------------------
import unittest
import os
from io import StringIO

import skbio
from skbio import Sequence

from ghosttree.scaffold.hybridtree import _make_nr_foundation_alignment
from ghosttree.scaffold.hybridtree import _create_taxonomy_dict
from ghosttree.scaffold.hybridtree import _make_mini_otu_files
from ghosttree.scaffold.hybridtree import _extension_genus_accession_dict


class TestScaffoldExtensionsIntofoundation(unittest.TestCase):
    def setUp(self):
        self.otu_clusters = StringIO(otu_clusters)
        self.extension_taxonomy = StringIO(extension_taxonomy)
        self.extension_taxonomy_unids = StringIO(extension_taxonomy_unids)
        self.extension_taxonomy_none = StringIO(extension_taxonomy_none)
        self.extension_seqs = StringIO(extension_seqs)
        self.foundation_alignment = StringIO(foundation_alignment)
        self.ghost_tree_fp = StringIO(ghost_tree_fp)
        self.extension_genus_dic_few = extension_genus_dic_few
        self.extension_genus_dic_none = extension_genus_dic_none
        self.key_node = key_node
        self.graft_letter_g = 'g'
        self.graft_level_6 = 6
        self.graft_level_5 = 5
        self.foundation_taxonomy = StringIO(foundation_taxonomy)
        self.foundation_newick = StringIO(foundation_newick)

    def test_make_nr_foundation_alignment_few(self):
        result = _make_nr_foundation_alignment(self.foundation_alignment,
                                               self.extension_genus_dic_few,
                                               self.graft_letter_g)

        self.assertEqual(list(result), [
            Sequence("AAA---", metadata={"id": "PBB1",
                                         "description": "Phoma"}),
            Sequence("AAG---", metadata={"id": "CBB3",
                                         "description": "Candida"}),
        ])

    def test_make_nr_foundation_alignment_none(self):
        result = _make_nr_foundation_alignment(self.foundation_alignment,
                                               self.extension_genus_dic_none,
                                               self.graft_letter_g)
        self.assertEqual(list(result), [])

    # def test_newick_file_few_extensions(self):
    #     result = _make_nr_foundation_newick(self.foundation_newick,
    #                                         self.extension_genus_dic_few)
    #
    #     print('result', result)

    def test_create_taxonomy_dic_many_genus(self):

        test = {'P1': 'Phoma', 'P2': 'Phoma', 'C1': 'Candida', 'C2': 'Candida',
                'C3': 'Candida', 'C4': 'Candida', 'C5': 'Candida',
                'M1': 'Mucor', 'M2': 'Mucor', 'M3': 'Mucor', 'M4': 'Mucor'}

        result = _create_taxonomy_dict(self.extension_taxonomy,
                                       self.graft_level_6)
        print('\n\n\nlalala', result)
        self.assertDictEqual(result, test)

    def test_create_taxonomy_dic_none(self):
        with self.assertRaises(ValueError):
            list(_create_taxonomy_dict(self.extension_taxonomy_none,
                                       9))

    def test_make_mini_otu_files(self):
        os.system("mkdir tmp")
        extension_seqs = list(skbio.io.read(self.extension_seqs,
                                            format='fasta'))

        _make_mini_otu_files(self.key_node,
                             self.extension_genus_dic_few,
                             extension_seqs)

        result = str(list(skbio.io.read('tmp/mini_seq_gt.fasta',
                                        format='fasta'))[0])

        os.system("rm -r tmp")
        self.assertEqual(result, 'TTAAAAAA')

    def test_extension_genus_accession_dic(self):
        test = {'Candida': ['C1', 'C2', 'C3', 'M1', 'C4', 'C5', 'M4'],
                'Mucor': ['M2', 'M3'],
                'Phoma': ['P1', 'P2']}
        result = _extension_genus_accession_dict(self.otu_clusters,
                                                 self.extension_taxonomy,
                                                 self.graft_level_6)
        self.assertDictEqual(result, test)

    def test_extension_genus_accession_dic_unidentifieds_level_6_genus(self):
        test = {'Candida': ['C1', 'C2', 'C3', 'M1'],
                'Mucor': ['M2', 'M3', 'C4', 'C5', 'M4'],
                'Phoma': ['P1', 'P2']}
        result = _extension_genus_accession_dict(self.otu_clusters,
                                                 self.extension_taxonomy_unids,
                                                 self.graft_level_6)
        self.assertDictEqual(result, test)

    def test_extension_genus_accession_dic_unidentifieds_level_5_family(self):
        # There is a tie breaker here so data can switch to either
        # Saccharomycetaceae OR Mucoraceae randomly (no priority).
        test_case_1 = {'Saccharomycetaceae': ['C1', 'C2', 'C3', 'M1', 'C4',
                       'C5', 'M4'], 'Didymellaceae': ['P1', 'P2'],
                       'Mucoraceae': ['M2', 'M3']}

        test_case_2 = {'Mucoraceae': ['M2', 'M3', 'C4', 'C5', 'M4'],
                       'Saccharomycetaceae': ['C1', 'C2', 'C3', 'M1'],
                       'Didymellaceae': ['P1', 'P2']}

        result = _extension_genus_accession_dict(self.otu_clusters,
                                                 self.extension_taxonomy_unids,
                                                 self.graft_level_5)

        try:
            self.assertDictEqual(result, test_case_1)
        except AssertionError:
            self.assertDictEqual(result, test_case_2)


key_node = 'Phoma'

extension_genus_dic_few = {'Candida': ['C1', 'C2', 'C3', 'P2'],
                           'Phoma': ['P1']}
extension_genus_dic_none = {'Asperggg': ['C2', 'C3', 'P2'],
                            'Nope': ['P1', 'M1']}

otu_clusters = """C1\tC1\tC2\tC3\tM1
M2\tM2\tM3
P1\tP1\tP2
{'fraction_non_alphanumeric': 0.13213513063737728, 'fraction_numerical': 0.015060742220003328, 'mean_word_length': 1.873296677026058, 'pattern_counts': {'":': 8, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 62, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3859586', 'n_tokens_mistral': 3745, 'n_tokens_neox': 3575, 'n_words': 584}