#! /usr/bin/env python

##############################################################################
##  DendroPy Phylogenetic Computing Library.
##
##  Copyright 2010-2015 Jeet Sukumaran and Mark T. Holder.
##  All rights reserved.
##
##  See "LICENSE.rst" for terms and conditions of usage.
##
##  If you use this work or any portion thereof in published work,
##  please cite it as:
##
##     Sukumaran, J. and M. T. Holder. 2010. DendroPy: a Python library
##     for phylogenetic computing. Bioinformatics 26: 1569-1571.
##
##############################################################################

"""
Testing of calculation of and operations with split bitmask hashes.
"""

import warnings
import unittest
import re
import sys
try:
    from StringIO import StringIO # Python 2 legacy support: StringIO in this module is the one needed (not io)
except ImportError:
    from io import StringIO # Python 3

from dendropy.test.support import pathmap
from dendropy.test.support import paupsplitsreference
from dendropy.test.support.dendropytest import ExtendedTestCase
from dendropy.utility import messaging
from dendropy.utility import bitprocessing
from dendropy.interop import paup
from dendropy.calculate import treecompare
import dendropy

_LOG = messaging.get_logger(__name__)

class SplitDistributionTestCases(ExtendedTestCase):

    def check_splits_distribution(self,
            tree_filename,
            splits_filename,
            use_tree_weights,
            is_rooted,
            expected_num_trees,
            ):
        if is_rooted is None:
            key_column_index = 2 # default to unrooted: normalized split bitmask
        elif is_rooted:
            key_column_index = 1 # leafset_bitmask / unnormalized split bitmask
        else:
            key_column_index = 2 # normalized split bitmask
        splits_ref = paupsplitsreference.get_splits_reference(
                splits_filename=splits_filename,
                key_column_index=key_column_index,
                )
        # print("* {} ({})".format(tree_filename, splits_filename))
        tree_filepath = pathmap.tree_source_path(tree_filename)
        trees = dendropy.TreeList.get_from_path(
                tree_filepath,
                "nexus",
                store_tree_weights=use_tree_weights)
        sd = dendropy.SplitDistribution(
                taxon_namespace=trees.taxon_namespace,
                use_tree_weights=use_tree_weights)
        for tree in trees:
            sd.count_splits_on_tree(tree)

        # trees counted ...
        self.assertEqual(sd.total_trees_counted, len(trees))
        # frequencies have not yet been calculated
        self.assertEqual(sd._trees_counted_for_freqs, 0)
        self.assertFalse(sd.is_mixed_rootings_counted())
        if is_rooted:
            self.assertTrue(sd.is_all_counted_trees_rooted())
        else:
            self.assertFalse(sd.is_all_counted_trees_rooted())
            self.assertTrue(sd.is_all_counted_trees_treated_as_unrooted() or sd.is_all_counted_trees_strictly_unrooted())

        # splits_distribution also counts trivial splits, so this will not work
        # self.assertEqual(len(splits_ref), len(sd))

        expected_nontrivial_splits = list(splits_ref.keys())
        observed_splits = set(sd.split_counts.keys())
        visited_splits = []
        # for k in sorted(observed_splits):
        #     print("{}: {}, {}".format(k, sd.split_counts[k], sd[k]))
        all_taxa_bitmask = sd.taxon_namespace.all_taxa_bitmask()
        for split in expected_nontrivial_splits:
            self.assertAlmostEqual(sd.split_counts[split], splits_ref[split]["count"], 2,
                    "{} (using '{}'): {}".format(tree_filename, splits_filename, split))
            self.assertAlmostEqual(sd[split], splits_ref[split]["frequency"], 2,
                    "{} (using '{}'): {}".format(tree_filename, splits_filename, split))
            self.assertAlmostEqual(sd.split_frequencies[split], splits_ref[split]["frequency"], 2,
                    "{} (using '{}'): {}".format(tree_filename, splits_filename, split))
            if split in observed_splits:
                observed_splits.remove(split)
            visited_splits.append(split)
        self.assertEqual(len(visited_splits), len(expected_nontrivial_splits))

        # ensure remaining splits (not given in PAUP splits file) are trivial ones (which are not tracked by PAUP)
        for split in observed_splits:
            self.assertTrue(dendropy.Bipartition.is_trivial_bitmask(split, all_taxa_bitmask))

    def test_group1(self):
        sources = [
                ("cetaceans.mb.no-clock.mcmc.trees"    , 251, False, False), # Trees explicitly unrooted
                ("cetaceans.mb.no-clock.mcmc.weighted-01.trees" , 251, False , True), # Weighted
                ("cetaceans.mb.no-clock.mcmc.weighted-02.trees" , 251, False , True), # Weighted
                ("cetaceans.mb.no-clock.mcmc.weighted-03.trees" , 251, False , True), # Weighted
                ("cetaceans.mb.strict-clock.mcmc.trees", 251, True , False), # Trees explicitly rooted
                ("cetaceans.mb.strict-clock.mcmc.weighted-01.trees" , 251, True , True), # Weighted
                ("cetaceans.mb.strict-clock.mcmc.weighted-02.trees" , 251, True , True), # Weighted
                ("cetaceans.mb.strict-clock.mcmc.weighted-03.trees" , 251, True , True), # Weighted
                ("issue_mth_2009-02-03.rooted.nexus"   , 100, True , False), # 100 trees (frequency column not reported by PAUP)
                ("issue_mth_2009-02-03.unrooted.nexus" , 100, False , False), # 100 trees (frequency column not reported by PAUP)
                ("cetaceans.raxml.bootstraps.trees"    , 250, None , False), # No tree rooting statement; PAUP defaults to rooted, DendroPy defaults to unrooted
                ("cetaceans.raxml.bootstraps.weighted-01.trees"    , 250, None , False), # No tree rooting statement; PAUP defaults to rooted, DendroPy defaults to unrooted
                ("cetaceans.raxml.bootstraps.weighted-02.trees"    , 250, None , False), # No tree rooting statement; PAUP defaults to rooted, DendroPy defaults to unrooted
                ("cetaceans.raxml.bootstraps.weighted-03.trees"    , 250, None , False), # No tree rooting statement; PAUP defaults to rooted, DendroPy defaults to unrooted
        ]
        splits_filename_template = "{stemname}.is-rooted-{is_rooted}.use-tree-weights-{use_weights}.burnin-{burnin}.splits.txt"
        for tree_filename, num_trees, treefile_is_rooted, treefile_is_weighted in sources:
            stemname = tree_filename
            for use_weights in (False, True, None):
                expected_is_rooted = treefile_is_rooted
                splits_filename = splits_filename_template.format(
                        stemname=stemname,
                        is_rooted=expected_is_rooted,
                        use_weights=use_weights,
                        burnin=0)
                self.check_splits_distribution(
                        tree_filename=tree_filename,
                        splits_filename=splits_filename,
                        is_rooted=treefile_is_rooted,
                        use_tree_weights=use_weights,
                        expected_num_trees=num_trees)

if not paup.DENDROPY_PAUP_INTEROPERABILITY:
    _LOG.warn("PAUP interoperability not available: skipping split counting tests")
else:

    class SplitCountTest(ExtendedTestCase):

        @classmethod
        def setUpClass(cls):
            if sys.version_info.major < 3:
                cls.assertRaisesRegex = cls.assertRaisesRegexp

        def check_split_counting(self,
                tree_filename,
                test_as_rooted,
                parser_rooting_interpretation,
                test_ignore_tree_weights=False,
                dp_ignore_tree_weights=False,
                ):
            tree_filepath = pathmap.tree_source_path(tree_filename)
            ps = paup.PaupService()
            paup_sd = ps.get_split_distribution_from_files(
                    tree_filepaths=[tree_filepath],
                    is_rooted=test_as_rooted,
                    use_tree_weights=not test_ignore_tree_weights,
                    burnin=0,
                    taxa_definition_filepath=tree_filepath
                    )
            taxon_namespace = paup_sd.taxon_namespace
            dp_sd = dendropy.SplitDistribution(taxon_namespace=taxon_namespace)
            dp_sd.ignore_edge_lengths = True
            dp_sd.ignore_node_ages = True
            dp_sd.ignore_tree_weights = dp_ignore_tree_weights
            taxa_mask = taxon_namespace.all_taxa_bitmask()
            taxon_namespace.is_mutable = False
            trees = dendropy.TreeList.get_from_path(tree_filepath,
                    "nexus",
                    rooting=parser_rooting_interpretation,
                    taxon_namespace=taxon_namespace)
            for tree in trees:
                self.assertIs(tree.taxon_namespace, taxon_namespace)
                self.assertIs(tree.taxon_namespace, dp_sd.taxon_namespace)
                dp_sd.count_splits_on_tree(
                        tree,
                        is_bipartitions_updated=False)
            self.assertEqual(dp_sd.total_trees_counted, paup_sd.total_trees_counted)
            taxa_mask = taxon_namespace.all_taxa_bitmask()
            for split in dp_sd.split_counts:
                if not dendropy.Bipartition.is_trivial_bitmask(split, taxa_mask):
                    # if split not in paup_sd.split_counts:
                    #     print("{}: {}".format(split, split in paup_sd.split_counts))
                    #     s2 = taxon_namespace.normalize_bitmask(split)
                    #     print("{}: {}".format(s2, s2 in paup_sd.split_counts))
                    #     s3 = ~split & taxon_namespace.all_taxa_bitmask()
                    #     print("{}: {}".format(s3, s3 in paup_sd.split_counts))
                    self.assertIn(split, paup_sd.split_counts, "split not found")
                    self.assertEqual(dp_sd.split_counts[split], paup_sd.split_counts[split], "incorrect split frequency")
                    del paup_sd.split_counts[split]
            remaining_splits = list(paup_sd.split_counts.keys())
            for split in remaining_splits:
                if dendropy.Bipartition.is_trivial_bitmask(split, taxa_mask):
                    del paup_sd.split_counts[split]
            self.assertEqual(len(paup_sd.split_counts), 0)

        def test_basic_split_count_with_incorrect_rootings_raises_error(self):
            assertion_error_regexp1 = re.compile("(incorrect split frequency|split not found)")
            test_cases = (
                ('pythonidae.reference-trees.nexus', True, "force-unrooted", assertion_error_regexp1),
                ('feb032009.trees.nexus', False, "force-rooted", assertion_error_regexp1),
                )
            for test_case, test_as_rooted, parser_rooting_interpretation, assertion_error_regexp in test_cases:
                with self.assertRaisesRegex(AssertionError, assertion_error_regexp):
                    self.check_split_counting(
                            test_case,
                            test_as_rooted=test_as_rooted,
                            parser_rooting_interpretation=parser_rooting_interpretation)

        def test_basic_split_count_with_incorrect_weight_treatment_raises_error(self):
            assertion_error_regexp1 = re.compile("incorrect split frequency")
            test_cases = (
                    ("cetaceans.mb.no-clock.mcmc.weighted-01.trees", False),
                    ("cetaceans.mb.strict-clock.mcmc.weighted-01.trees", True),
                )
            for test_case, test_as_rooted in test_cases:
                with self.assertRaisesRegex(AssertionError, assertion_error_regexp1):
                    self.check_split_counting(
                            test_case,
                            test_as_rooted=test_as_rooted,
                            parser_rooting_interpretation="default-rooted",
                            test_ignore_tree_weights=False,
                            dp_ignore_tree_weights=False,
                            )

        def test_basic_split_counting_under_different_rootings(self):
            test_cases = (
                'pythonidae.reference-trees.nexus',
                'feb032009.trees.nexus',
                'maj-rule-bug1.trees.nexus',
                'maj-rule-bug2.trees.nexus',
                )
            for is_rooted in (True, False):
                if is_rooted:
                    rooting = "force-rooted"
                else:
                    rooting = "force-unrooted"
                for test_case in test_cases:
                    self.check_split_counting(
                            test_case,
                            test_as_rooted=is_rooted,
                            parser_rooting_interpretation=rooting)

class CladeMaskTest(unittest.TestCase):

    def runTest(self):
        # rooted tree: so clade bitmasks
        tree_list = dendropy.TreeList.get_from_stream(
            StringIO("""[&R]((t5:0.161175,t6:0.161175):0.392293,((t4:0.104381,(t2:0.075411,t1:0.075411):1):0.065840,t3:0.170221):0.383247);"""),
            "newick")
        for i in tree_list:
            _LOG.debug(i._get_indented_form())
            i.encode_bipartitions()
            _LOG.debug(i._get_indented_form(splits=True))
            i._debug_check_tree(splits=True, logger_obj=_LOG)
        root1 = tree_list[0].seed_node
        root1e = root1.edge
        self.assertEqual(bitprocessing.indexes_of_set_bits(root1e.split_bitmask), list(range(6)))
        self.assertEqual(bitprocessing.indexes_of_set_bits(root1e.split_bitmask, one_based=True), list(range(1,7)))
        self.assertEqual(bitprocessing.indexes_of_set_bits(root1e.split_bitmask, fill_bitmask=21, one_based=True), [1, 3, 5])
        self.assertEqual(bitprocessing.indexes_of_set_bits(root1e.split_bitmask, fill_bitmask=21), [0, 2, 4])
        self.assertEqual(bitprocessing.num_set_bits(root1e.split_bitmask), 6)

        fc1 = root1.child_nodes()[0]
        fc1e = fc1.edge
        self.assertEqual(bitprocessing.indexes_of_set_bits(fc1e.split_bitmask), [0, 1])
        self.assertEqual(bitprocessing.indexes_of_set_bits(fc1e.split_bitmask, one_based=True), [1, 2])
        self.assertEqual(bitprocessing.indexes_of_set_bits(fc1e.split_bitmask, fill_bitmask=0x15, one_based=True), [1])
        self.assertEqual(bitprocessing.indexes_of_set_bits(fc1e.split_bitmask, fill_bitmask=0x15), [0])
        self.assertEqual(bitprocessing.num_set_bits(fc1e.split_bitmask), 2)

class CountBitsTest(unittest.TestCase):

    def runTest(self):
        self.assertEqual(bitprocessing.num_set_bits(21), 3)

class LowestBitTest(unittest.TestCase):

    def runTest(self):
        for n, expected in enumerate([0, 1, 2, 1, 4, 1, 2, 1, 8, 1, 2, 1, 4, 1, 2, 1, 16]):
            self.assertEqual(bitprocessing.least_significant_set_bit(n), expected)

class IsTrivialTest(unittest.TestCase):

    def runTest(self):
        y = True
        n = False
        for i, r in enumerate([y, y, y, n, y, n, n, y, y, n, n, y, n, y, y, y, y, y, y, n, y, n, n, y, y, n, n, y, n, y, y, y, ]):
            self.assertEqual(r, dendropy.Bipartition.is_trivial_bitmask(i, 0xF))
        for i, r in enumerate([y, y, y, n, y, n, n, n, y, n, n, n, n, n, n, y, y, n, n, n, n, n, n, y, n, n, n, y, n, y, y, y, ]):
            self.assertEqual(r, dendropy.Bipartition.is_trivial_bitmask(i, 0x1F))
                              #0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5  6  7  8  9  0  1
        for i, r in enumerate([y, y, y, n, y, n, n, y, y, y, y, n, y, n, n, y, y, n, n, y, n, y, y, y, y, n, n, y, n, y, y, y, ]):
            self.assertEqual(r, dendropy.Bipartition.is_trivial_bitmask(i, 0x17))

class IncompleteLeafSetSplitTest(unittest.TestCase):

    def check(self, title, src_prefix):
        tns = dendropy.TaxonNamespace()
        input_ds = dendropy.DataSet.get_from_path(
                src=pathmap.tree_source_path(src_prefix + ".dendropy-pruned.nex"),
                schema='nexus',
                attached_taxon_namespace=tns)
        input_taxa = input_ds.taxon_namespaces[0]
        output_ds = dendropy.DataSet.get_from_path(
                src=pathmap.tree_source_path(src_prefix + ".paup-pruned.nex"),
                schema='nexus',
                taxon_namespace=input_taxa)
        for set_idx, src_trees in enumerate(input_ds.tree_lists):
            src_trees = input_ds.tree_lists[set_idx]
            ref_trees = output_ds.tree_lists[set_idx]
            for tree_idx, src_tree in enumerate(src_trees):
                _LOG.debug("%s Set %d/%d, Tree %d/%d" % (title, set_idx+1, len(input_ds.tree_lists), tree_idx+1, len(src_trees)))
                ref_tree = ref_trees[tree_idx]
                # tree_dist = paup.symmetric_difference(src_tree, ref_tree)
                # d = src_tree.symmetric_difference(ref_tree)
                # if d > 0:
                #     print d
                self.assertEqual(treecompare.symmetric_difference(src_tree, ref_tree), 0)

    def testUnrooted(self):
        self.check("Unrooted", "incomplete_leaves_unrooted")

    def testRooted(self):
        self.check("Rooted", "incomplete_leaves_rooted")

    def testPrunedThenEncoding(self):
        inp = StringIO('''(a,b,c,(d,e));
        (b,d,(c,e));''')
        first, second = dendropy.TreeList.get_from_stream(inp, schema='newick')
        # prune tree 1 to have the same leaf set as tree 2.
        #   this removes the first taxon in the taxon list "A"
        retain_list = set([node.taxon for node in second.leaf_nodes()])
        exclude_list = [node for node in first.leaf_nodes() if node.taxon not in retain_list]
        for nd in exclude_list:
            first.prune_subtree(nd)
        # the trees are now (b,c,(d,e)) and (b,d,(c,e)) so the symmetric diff is 2
        self.assertEqual(2, treecompare.symmetric_difference(first, second))

class TestTreeSplitSupportCredibilityScoring(unittest.TestCase):

    def setUp(self):
        self.trees = dendropy.TreeList.get_from_path(
                pathmap.tree_source_path("issue_mth_2009-02-03.rooted.nexus"),
                "nexus")
        self.split_distribution = dendropy.SplitDistribution(taxon_namespace=self.trees.taxon_namespace)
        for tree in self.trees:
            self.split_distribution.count_splits_on_tree(
                    tree,
                    is_bipartitions_updated=False)

    def test_product_of_split_support_on_tree(self):
        t1 = self.trees[70]
        self.assertAlmostEqual(
                self.split_distribution.log_product_of_split_support_on_tree(t1),
                -33.888380488585284)

    def test_sum_of_split_support_on_tree(self):
        t1 = self.trees[73]
        self.assertAlmostEqual(
                self.split_distribution.sum_of_split_support_on_tree(t1),
                30.89000000000001)

    def test_sum_of_split_support_on_tree2(self):
        t1 = self.trees[73]
        self.assertAlmostEqual(
                self.split_distribution.sum_of_split_support_on_tree(t1, include_external_splits=True),
                30.89000000000001 + len(self.trees.taxon_namespace))

if __name__ == "__main__":
    unittest.main()
    # if paup.DENDROPY_PAUP_INTEROPERABILITY:
    #     unittest.main()
    # else:
    #     _LOG.warn("PAUP interoperability not available: skipping split counting tests")
