"""Tests for TreeHandler"""
import os
import unittest
from nexus.reader import NexusReader
from nexus.exceptions import NexusFormatException
from nexus.handlers.tree import TreeHandler

EXAMPLE_DIR = os.path.join(os.path.dirname(__file__), '../examples')

class Test_TreeHandler_SimpleTreefile(unittest.TestCase):
    def setUp(self):
        self.nex = NexusReader(os.path.join(EXAMPLE_DIR, 'example.trees'))

    def test_block_find(self):
        # did we get a tree block?
        assert 'trees' in self.nex.blocks

    def test_treecount(self):
        # did we find 3 trees?
        assert len(self.nex.blocks['trees'].trees) == 3
        assert self.nex.blocks['trees'].ntrees == 3
        assert len(self.nex.trees.trees) == 3
        assert self.nex.trees.ntrees == 3

    def test_taxa(self):
        expected = [
            'Chris', 'Bruce', 'Tom', 'Henry', 'Timothy', 'Mark', 'Simon',
            'Fred', 'Kevin', 'Roger', 'Michael', 'Andrew', 'David'
        ]
        assert len(self.nex.trees.taxa) == len(expected)
        for taxon in expected:
            assert taxon in self.nex.trees.taxa

    def test_iterable(self):
        assert list(self.nex.blocks['trees'])
        assert list(self.nex.trees)

    def test_write(self):
        written = self.nex.trees.write()
        with open(os.path.join(EXAMPLE_DIR, 'example.trees'), 'r') as handle:
            expected = handle.read()
        # remove leading header which isn't generated by .write()
        expected = expected.lstrip("#NEXUS\n\n")
        assert expected == written
    
    def test_repr(self):
        assert repr(self.nex.trees) == '<NexusTreeBlock: 3 trees>'

    def test_write_produces_end(self):
        assert "end;" in self.nex.trees.write()
    


class Test_TreeHandler_TranslatedTreefile(unittest.TestCase):
    def setUp(self):
        self.nex = NexusReader(
            os.path.join(EXAMPLE_DIR, 'example-translated.trees')
        )

    def test_block_find(self):
        # did we get a tree block?
        assert 'trees' in self.nex.blocks

    def test_treecount(self):
        # did we find 3 trees?
        assert len(self.nex.blocks['trees'].trees) == 3
        assert self.nex.blocks['trees'].ntrees == 3
        assert len(self.nex.trees.trees) == 3
        assert self.nex.trees.ntrees == 3

    def test_iterable(self):
        assert list(self.nex.blocks['trees'])
        assert list(self.nex.trees)

    def test_taxa(self):
        expected = [
            'Chris', 'Bruce', 'Tom', 'Henry', 'Timothy', 'Mark', 'Simon',
            'Fred', 'Kevin', 'Roger', 'Michael', 'Andrew', 'David'
        ]
        assert len(self.nex.trees.taxa) == len(expected)
        for taxon in expected:
            assert taxon in self.nex.trees.taxa

    def test_was_translated_flag_set(self):
        assert self.nex.trees.was_translated

    def test_parsing_sets_translators(self):
        assert len(self.nex.trees.translators) == 13

    def test_been_detranslated_flag_set(self):
        assert not self.nex.trees._been_detranslated
        self.nex.trees.detranslate()
        assert self.nex.trees._been_detranslated

    def test_write(self):
        assert not self.nex.trees._been_detranslated
        written = self.nex.trees.write()
        with open(os.path.join(EXAMPLE_DIR, 'example-translated.trees'), 'r') as handle:
            expected = handle.read()
        # remove leading header which isn't generated by .write()
        expected = expected.lstrip("#NEXUS\n\n")
        # remove tabs since we reformat things a bit
        expected = expected.replace("\t", "").strip()
        written = written.replace("\t", "").strip()
        # handle the workaround for the beast bug
        expected = expected.replace("12 David;", "12 David\n;")
        assert expected == written, "%s\n----\n%s" % (expected, written)

    def test_no_error_on_multiple_translate(self):
        assert not self.nex.trees._been_detranslated
        self.nex.trees.detranslate()
        assert self.nex.trees._been_detranslated
        self.nex.trees.detranslate()  # should not cause an error

    def test_detranslate(self):
        assert not self.nex.trees._been_detranslated
        self.nex.trees.detranslate()
        # should NOW be the same as tree 0 in example.trees
        other_tree_file = NexusReader(
            os.path.join(EXAMPLE_DIR, 'example.trees')
        )
        assert other_tree_file.trees[0] == self.nex.trees[0]

    def test_treelabel(self):
        nex = NexusReader()
        nex.read_string("""
        #NEXUS
    
        begin trees;
            translate
                0 Tom,
                1 Simon,
                2 Fred;
                tree TREEONE = (0,1,2);
        end;
        """)
        assert len(nex.trees.trees) == 1
        assert nex.trees.trees == ['tree TREEONE = (0,1,2);']
        
    def test_no_treelabel(self):
        nex = NexusReader()
        nex.read_string("""
        #NEXUS
    
        begin trees;
            translate
                0 Tom,
                1 Simon,
                2 Fred;
                tree = (0,1,2);
        end;
        """)
        assert len(nex.trees.trees) == 1
        assert nex.trees.trees == ['tree = (0,1,2);']
    
    def test_rooted(self):
        nex = NexusReader()
        nex.read_string("""
        #NEXUS
    
        begin trees;
            translate
                0 Tom,
                1 Simon,
                2 Fred;
                tree [&] = (0,1,2);
        end;
        """)
        assert len(nex.trees.trees) == 1
        assert nex.trees.trees == ['tree [&] = (0,1,2);']
    
    def test_unrooted(self):
        nex = NexusReader()
        nex.read_string("""
        #NEXUS
    
        begin trees;
            translate
                0 Tom,
                1 Simon,
                2 Fred;
                tree [&U] = (0,1,2);
        end;
        """)
        assert len(nex.trees.trees) == 1
        assert nex.trees.trees == ['tree [&U] = (0,1,2);']
        
    
    def test_labelled_unrooted(self):
        nex = NexusReader()
        nex.read_string("""
        #NEXUS
    
        begin trees;
            translate
                0 Tom,
                1 Simon,
                2 Fred;
                tree unrooted [U] = (0,1,2);
        end;
        """)
        assert len(nex.trees.trees) == 1
        assert nex.trees.trees == ['tree unrooted [U] = (0,1,2);']
        
    
    def test_ok_starting_with_zero(self):
        nex = NexusReader()
        nex.read_string("""
        #NEXUS
    
        begin trees;
            translate
                0 Tom,
                1 Simon,
                2 Fred;
                tree tree = (0,1,2)
        end;
        """)
        assert len(nex.trees.translators) == 3
        assert '0' in nex.trees.translators
        assert '1' in nex.trees.translators
        assert '2' in nex.trees.translators

    def test_ok_starting_with_one(self):
        nex = NexusReader()
        nex.read_string("""
        #NEXUS
    
        begin trees;
            translate
                1 Tom,
                2 Simon,
                3 Fred;
                tree tree = (1,2,3)
        end;
        """)
        assert len(nex.trees.translators) == 3
        assert '1' in nex.trees.translators
        assert '2' in nex.trees.translators
        assert '3' in nex.trees.translators
        
        
    def test_error_on_duplicate_taxa_id(self):
        with self.assertRaises(NexusFormatException):
            NexusReader().read_string("""
            #NEXUS
        
            begin trees;
                translate
                    0 Tom,
                    1 Simon,
                    1 Tom;
                    tree tree = (0,1,2)
            end;
            """)
    
    def test_error_on_duplicate_taxa(self):
        with self.assertRaises(NexusFormatException):
            NexusReader().read_string("""
            #NEXUS
        
            begin trees;
                translate
                    0 Tom,
                    1 Simon,
                    2 Tom;
                    tree tree = (0,1,2)
            end;
            """)


class Test_TreeHandler_BEAST_Format(unittest.TestCase):

    def setUp(self):
        self.nex = NexusReader(
            os.path.join(EXAMPLE_DIR, 'example-beast.trees')
        )

    def test_read_BEAST_format(self):
        assert self.nex.trees[0].startswith('tree STATE_201000')

    def test_block_find(self):
        # did we get a tree block?
        assert 'trees' in self.nex.blocks

    def test_taxa(self):
        expected = [
            "R1", "B2", "S3", "T4", "A5", "E6", "U7", "T8", "T9", "F10", "U11",
            "T12", "N13", "F14", "K15", "N16", "I17", "L18", "S19", "T20",
            "V21", "R22", "M23", "H24", "M25", "M26", "M27", "R28", "T29",
            "M30", "P31", "T32", "R33", "P34", "R35", "W36", "F37", "F38"
        ]
        assert len(self.nex.trees.taxa) == len(expected)
        for taxon in expected:
            assert taxon in self.nex.trees.taxa

    def test_treecount(self):
        assert len(self.nex.blocks['trees'].trees) == 1
        assert self.nex.blocks['trees'].ntrees == 1
        assert len(self.nex.trees.trees) == 1
        assert self.nex.trees.ntrees == 1

    def test_flag_set(self):
        assert self.nex.trees.was_translated

    def test_parsing_sets_translators(self):
        assert len(self.nex.trees.translators) == 38

    def test_detranslate_BEAST_format_extended(self):
        self.nex.trees.detranslate()
        for index, taxon in self.nex.trees.translators.items():
            # check if the taxon name is present in the tree...
            assert taxon in self.nex.trees[0], \
                "Expecting taxon %s in tree description" % taxon
        assert self.nex.trees._been_detranslated


class Test_TreeHandler_translate_regex(unittest.TestCase):

    def setUp(self):
        th = TreeHandler()
        self.findall = th._findall_chunks

    def test_tree(self):
        expected = {
            0: {
                'start': '(',
                'taxon': 'Chris',
                'comment': None,
                'branch': None,
                'end': ','
            },
            1: {
                'start': ',',
                'taxon': 'Bruce',
                'comment': None,
                'branch': None,
                'end': ')'
            },
            2: {
                'start': ',',
                'taxon': 'Tom',
                'comment': None,
                'branch': None,
                'end': ')'
            },
        }
        found = self.findall("tree a = ((Chris,Bruce),Tom);")
        assert len(found) == 3
        for match in expected:
            for key in expected[match]:
                if expected[match][key] != found[match][key]:
                    raise AssertionError(
                        "Expected %s for %s, got %s" %
                        (expected[match][key], key, found[match][key])
                    )

    def test_tree_digits(self):
        expected = {
            0: {
                'start': '(',
                'taxon': '1',
                'comment': None,
                'branch': None,
                'end': ','
            },
            1: {
                'start': ',',
                'taxon': '2',
                'comment': None,
                'branch': None,
                'end': ')'
            },
            2: {
                'start': ',',
                'taxon': '3',
                'comment': None,
                'branch': None,
                'end': ')'
            },
        }
        found = self.findall("tree a = ((1,2),3);")
        assert len(found) == 3
        for match in expected:
            for key in expected[match]:
                if expected[match][key] != found[match][key]:
                    raise AssertionError(
                        "Expected %s for %s, got %s" %
                        (expected[match][key], key, found[match][key])
                    )

    def test_tree_with_branchlengths(self):
        expected = {
            0: {
                'start': '(',
                'taxon': '1',
                'comment': None,
                'branch': '0.1',
                'end': ','
            },
            1: {
                'start': ',',
                'taxon': '2',
                'comment': None,
                'branch': '0.2',
                'end': ')'
            },
            2: {
                'start': ',',
                'taxon': '3',
                'comment': None,
                'branch': '0.3',
                'end': ')'
            },
        }
        found = self.findall("tree a = ((1:0.1,2:0.2):0.9,3:0.3):0.9;")
        assert len(found) == 3
        for match in expected:
            for key in expected[match]:
                if expected[match][key] != found[match][key]:
                    raise AssertionError(
                        "Expected %s for %s, got %s" %
                        (expected[match][key], key, found[match][key])
                    )

    def test_tree_complex(self):
        expected = {
            0: {
                'start': '(',
                'taxon': '1',
                'comment': '[&var=1]',
                'branch': '0.1',
                'end': ','
            },
            1: {
                'start': ',',
                'taxon': '2',
                'comment': '[&var=2]',
                'branch': '0.2',
                'end': ')'
            },
            2: {
                'start': ',',
                'taxon': '3',
                'comment': '[&var=4]',
                'branch': '0.3',
                'end': ')'
            },
        }
        found = self.findall(
            "tree a = ((1:[&var=1]0.1,2:[&var=2]0.2):[&var=3]0.9,3:[&var=4]0.3):[&var=5]0.9;"
        )
        assert len(found) == 3
        for match in expected:
            for key in expected[match]:
                if expected[match][key] != found[match][key]:
                    raise AssertionError(
                        "Expected %s for %s, got %s" %
                        (expected[match][key], key, found[match][key])
                    )


class Test_TreeHandler__detranslate_tree(unittest.TestCase):
    def test_no_change(self):
        translatetable = {'0': 'Chris', '1': 'Bruce', '2': 'Tom'}
        oldtree = "tree a = ((Chris,Bruce),Tom);"
        newtree = "tree a = ((Chris,Bruce),Tom);"
        trans = TreeHandler()._detranslate_tree(oldtree, translatetable)
        assert trans == newtree, \
            "Unable to correctly NOT translate a simple tree"

    def test_no_change_branchlengths(self):
        translatetable = {'0': 'Chris', '1': 'Bruce', '2': 'Tom'}
        oldtree = "tree a = ((Chris:0.1,Bruce:0.2):0.3,Tom:0.4);"
        newtree = "tree a = ((Chris:0.1,Bruce:0.2):0.3,Tom:0.4);"
        trans = TreeHandler()._detranslate_tree(oldtree, translatetable)
        assert trans == newtree, \
            "Unable to correctly NOT translate a tree with branchlengths"

    def test_change(self):
        translatetable = {'0': 'Chris', '1': 'Bruce', '2': 'Tom'}
        oldtree = "tree a = ((0,1),2);"
        newtree = "tree a = ((Chris,Bruce),Tom);"
        trans = TreeHandler()._detranslate_tree(oldtree, translatetable)
        assert trans == newtree, \
            "Unable to correctly detranslate a simple tree"

    def test_change_branchlengths(self):
        translatetable = {'0': 'Chris', '1': 'Bruce', '2': 'Tom'}
        oldtree = "tree a = ((0:0.1,1:0.2):0.3,2:0.4);"
        newtree = "tree a = ((Chris:0.1,Bruce:0.2):0.3,Tom:0.4);"
        trans = TreeHandler()._detranslate_tree(oldtree, translatetable)
        assert trans == newtree, \
            "Unable to correctly detranslate a tree with branchlengths"
    
    def test_change_comment(self):
        translatetable = {'0': 'Chris', '1': 'Bruce', '2': 'Tom'}
        oldtree = "tree a = ((0[x],1[y]),2[z]);"
        newtree = "tree a = ((Chris[x],Bruce[y]),Tom[z]);"
        trans = TreeHandler()._detranslate_tree(oldtree, translatetable)
        assert trans == newtree, \
            "Unable to correctly detranslate a tree with branchlengths"
    
    def test_BEAST_format(self):
        translatetable = {'1': 'Chris', '2': 'Bruce', '3': 'Tom'}
        oldtree = "tree STATE_0 [&lnP=-584.441] = [&R] ((1:[&rate=1.0]48.056,3:[&rate=1.0]48.056):[&rate=1.0]161.121,2:[&rate=1.0]209.177);"
        newtree = "tree STATE_0 [&lnP=-584.441] = [&R] ((Chris:[&rate=1.0]48.056,Tom:[&rate=1.0]48.056):[&rate=1.0]161.121,Bruce:[&rate=1.0]209.177);"
        trans = TreeHandler()._detranslate_tree(oldtree, translatetable)
        assert trans == newtree, \
            "Unable to correctly detranslate a BEAST tree"
