# ----------------------------------------------------------------------------
# Copyright (c) 2013--, scikit-bio development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
# ----------------------------------------------------------------------------

from __future__ import absolute_import, division, print_function

from unittest import TestCase, main

import numpy as np

from skbio.io._fileobject import StringIO
from skbio import TreeNode
from skbio.tree import majority_rule
from skbio.tree._majority_rule import (_walk_clades, _filter_clades,
                                       _build_trees)


class MajorityRuleTests(TestCase):
    def test_majority_rule(self):
        trees = [
            TreeNode.read(StringIO("(A,(B,(H,(D,(J,(((G,E),(F,I)),C))))));")),
            TreeNode.read(StringIO("(A,(B,(D,((J,H),(((G,E),(F,I)),C)))));")),
            TreeNode.read(StringIO("(A,(B,(D,(H,(J,(((G,E),(F,I)),C))))));")),
            TreeNode.read(StringIO("(A,(B,(E,(G,((F,I),((J,(H,D)),C))))));")),
            TreeNode.read(StringIO("(A,(B,(E,(G,((F,I),(((J,H),D),C))))));")),
            TreeNode.read(StringIO("(A,(B,(E,((F,I),(G,((J,(H,D)),C))))));")),
            TreeNode.read(StringIO("(A,(B,(E,((F,I),(G,(((J,H),D),C))))));")),
            TreeNode.read(StringIO("(A,(B,(E,((G,(F,I)),((J,(H,D)),C)))));")),
            TreeNode.read(StringIO("(A,(B,(E,((G,(F,I)),(((J,H),D),C)))));"))]

        exp = TreeNode.read(StringIO("(((E,(G,(F,I),(C,(D,J,H)))),B),A);"))
        obs = majority_rule(trees)
        self.assertEqual(exp.compare_subsets(obs[0]), 0.0)
        self.assertEqual(len(obs), 1)

        tree = obs[0]
        exp_supports = sorted([9.0, 9.0, 9.0, 6.0, 6.0, 6.0])
        obs_supports = sorted([n.support for n in tree.non_tips()])
        self.assertEqual(obs_supports, exp_supports)

        obs = majority_rule(trees, weights=np.ones(len(trees)) * 2)
        self.assertEqual(exp.compare_subsets(obs[0]), 0.0)
        self.assertEqual(len(obs), 1)

        tree = obs[0]
        exp_supports = sorted([18.0, 18.0, 12.0, 18.0, 12.0, 12.0])
        obs_supports = sorted([n.support for n in tree.non_tips()])

        with self.assertRaises(ValueError):
            majority_rule(trees, weights=[1, 2])

    def test_majority_rule_multiple_trees(self):
        trees = [
            TreeNode.read(StringIO("((a,b),(c,d),(e,f));")),
            TreeNode.read(StringIO("(a,(c,d),b,(e,f));")),
            TreeNode.read(StringIO("((c,d),(e,f),b);")),
            TreeNode.read(StringIO("(a,(c,d),(e,f));"))]

        trees = majority_rule(trees)
        self.assertEqual(len(trees), 4)

        exp = set([
                  frozenset(['a']),
                  frozenset(['b']),
                  frozenset([None, 'c', 'd']),
                  frozenset([None, 'e', 'f'])])

        obs = set([frozenset([n.name for n in t.traverse()]) for t in trees])
        self.assertEqual(obs, exp)

    def test_walk_clades(self):
        trees = [TreeNode.read(StringIO("((A,B),(D,E));")),
                 TreeNode.read(StringIO("((A,B),(D,(E,X)));"))]
        exp_clades = [
            (frozenset(['A']), 2.0),
            (frozenset(['B']), 2.0),
            (frozenset(['A', 'B']), 2.0),
            (frozenset(['D', 'E']), 1.0),
            (frozenset(['D', 'E', 'A', 'B']), 1.0),
            (frozenset(['D']), 2.0),
            (frozenset(['E']), 2.0),
            (frozenset(['X']), 1.0),
            (frozenset(['E', 'X']), 1.0),
            (frozenset(['D', 'E', 'X']), 1.0),
            (frozenset(['A', 'B', 'D', 'E', 'X']), 1.0)]

        exp_lengths_nolength = {
            frozenset(['A']): None,
            frozenset(['B']): None,
            frozenset(['A', 'B']): None,
            frozenset(['D', 'E']): None,
            frozenset(['D', 'E', 'A', 'B']): None,
            frozenset(['D']): None,
            frozenset(['E']): None,
            frozenset(['X']): None,
            frozenset(['E', 'X']): None,
            frozenset(['D', 'E', 'X']): None,
            frozenset(['A', 'B', 'D', 'E', 'X']): None}

        exp_lengths = {
            frozenset(['A']): 2.0,
            frozenset(['B']): 2.0,
            frozenset(['A', 'B']): 2.0,
            frozenset(['D', 'E']): 1.0,
            frozenset(['D', 'E', 'A', 'B']): 1.0,
            frozenset(['D']): 2.0,
            frozenset(['E']): 2.0,
            frozenset(['X']): 1.0,
            frozenset(['E', 'X']): 1.0,
            frozenset(['D', 'E', 'X']): 1.0,
            frozenset(['A', 'B', 'D', 'E', 'X']): 1.0}

        obs_clades, obs_lengths = _walk_clades(trees, np.ones(len(trees)))
        self.assertEqual(set(obs_clades), set(exp_clades))
        self.assertEqual(obs_lengths, exp_lengths_nolength)

        for t in trees:
            for n in t.traverse(include_self=True):
                n.length = 2.0

        obs_clades, obs_lengths = _walk_clades(trees, np.ones(len(trees)))

        self.assertEqual(set(obs_clades), set(exp_clades))
        self.assertEqual(obs_lengths, exp_lengths)

    def test_filter_clades(self):
        clade_counts = [(frozenset(['A', 'B']), 8),
                        (frozenset(['A', 'C']), 7),
                        (frozenset(['A']), 6),
                        (frozenset(['B']), 5)]
        obs = _filter_clades(clade_counts, 2)
        exp = {frozenset(['A', 'B']): 8,
               frozenset(['A']): 6,
               frozenset(['B']): 5}
        self.assertEqual(obs, exp)

        clade_counts = [(frozenset(['A']), 8),
                        (frozenset(['B']), 7),
                        (frozenset(['C']), 7),
                        (frozenset(['A', 'B']), 6),
                        (frozenset(['A', 'B', 'C']), 5),
                        (frozenset(['D']), 2)]
        obs = _filter_clades(clade_counts, 4)
        exp = {frozenset(['A']): 8,
               frozenset(['B']): 7,
               frozenset(['C']): 7,
               frozenset(['A', 'B']): 6,
               frozenset(['A', 'B', 'C']): 5}
        self.assertEqual(obs, exp)

    def test_build_trees(self):
        clade_counts = {frozenset(['A', 'B']): 6,
                        frozenset(['A']): 7,
                        frozenset(['B']): 8}
        edge_lengths = {frozenset(['A', 'B']): 1,
                        frozenset(['A']): 2,
                        frozenset(['B']): 3}
        tree = _build_trees(clade_counts, edge_lengths, 'foo')[0]
        self.assertEqual(tree.foo, 6)
        tree_foos = set([c.foo for c in tree.children])
        tree_lens = set([c.length for c in tree.children])
        self.assertEqual(tree_foos, set([7, 8]))
        self.assertEqual(tree_lens, set([2, 3]))


if __name__ == '__main__':
    main()
