#!/usr/bin/env python
"""Unit tests for fast unifrac."""
from __future__ import division

from numpy import array, logical_not, argsort
from cogent.util.unit_test import TestCase, main
from cogent.parse.tree import DndParser
from cogent.maths.unifrac.fast_tree import (count_envs, index_tree, index_envs,
    get_branch_lengths)
from cogent.maths.unifrac.fast_unifrac import (reshape_by_name,
    meta_unifrac, shuffle_tipnames, weight_equally, weight_by_num_tips, 
    weight_by_branch_length, weight_by_num_seqs, get_all_env_names,
    consolidate_skipping_missing_matrices, consolidate_missing_zero,
    consolidate_missing_one, consolidate_skipping_missing_values,
    UniFracTreeNode, mcarlo_sig, num_comps, fast_unifrac, 
    fast_unifrac_whole_tree, PD_whole_tree, PD_generic_whole_tree,
    TEST_ON_TREE, TEST_ON_ENVS, TEST_ON_PAIRWISE, shared_branch_length,
    shared_branch_length_to_root, fast_unifrac_one_sample)
from numpy.random import permutation 

__author__ = "Rob Knight and Micah Hamady"
__copyright__ = "Copyright 2007-2012, The Cogent Project"
__credits__ = ["Rob Knight", "Micah Hamady", "Daniel McDonald", 
"Justin Kuczynski"]
__license__ = "GPL"
__version__ = "1.5.3"
__maintainer__ = "Rob Knight, Micah Hamady"
__email__ = "rob@spot.colorado.edu, hamady@colorado.edu"
__status__ = "Prototype"

class unifrac_tests(TestCase):
    """Tests of top-level functions."""
    def setUp(self):
        """Define some standard trees."""
        self.t_str = '((a:1,b:2):4,(c:3,(d:1,e:1):2):3)'
        self.t = DndParser(self.t_str, UniFracTreeNode)
        self.env_str = """
a   A   1
a   C   2
b   A   1
b   B   1
c   B   1
d   B   3
e   C   1"""
        self.env_counts = count_envs(self.env_str.splitlines())
        self.missing_env_str = """
a   A   1
a   C   2
e   C   1"""
        self.missing_env_counts = count_envs(self.missing_env_str.splitlines())
        self.extra_tip_str = """
q   A   1
w   C   2
e   A   1
r   B   1
t   B   1
y   B   3
u   C   1"""
        self.extra_tip_counts = count_envs(self.extra_tip_str.splitlines())
        self.wrong_tip_str = """
q   A   1
w   C   2
r   B   1
t   B   1
y   B   3
u   C   1"""
        self.wrong_tip_counts = count_envs(self.wrong_tip_str.splitlines())

        self.t2_str = '(((a:1,b:1):1,c:5):2,d:4)'
        self.t2 = DndParser(self.t2_str, UniFracTreeNode)
        self.env2_str = """
a   B   1
b   A   1
c   A   2
c   C   2
d   B   1
d   C   1"""
        self.env2_counts = count_envs(self.env2_str.splitlines())
        self.trees = [self.t, self.t2]
        self.envs = [self.env_counts, self.env2_counts]

        self.mc_1 = array([.5, .4, .3, .2, .1, .6, .7, .8, .9, 1.0])
       
        # from old EnvsNode tests
        self.old_t_str = '((org1:0.11,org2:0.22,(org3:0.12,org4:0.23)g:0.33)b:0.2,(org5:0.44,org6:0.55)c:0.3,org7:0.4)'

        self.old_t = DndParser(self.old_t_str, UniFracTreeNode)
        self.old_env_str = """
org1    env1    1
org1    env2    1
org2    env2    1
org3    env2    1
org4    env3    1
org5    env1    1
org6    env1    1
org7    env3    1
"""
        self.old_env_counts = count_envs(self.old_env_str.splitlines())
        self.old_node_index, self.old_nodes = index_tree(self.old_t)
        self.old_count_array, self.old_unique_envs, self.old_env_to_index, \
            self.old_node_to_index = index_envs(self.old_env_counts, self.old_node_index)
        self.old_branch_lengths = get_branch_lengths(self.old_node_index)

    def test_shared_branch_length(self):
        """Should return the correct shared branch length by env"""
        t_str = "(((a:1,b:2):3,c:4),(d:5,e:6,f:7):8);"
        envs = """
a A 1
b A 1
c A 1
d A 1
e A 1
f B 1
"""
        env_counts = count_envs(envs.splitlines())
        t = DndParser(t_str, UniFracTreeNode)
        exp = {('A',):21.0,('B',):7.0}
        obs = shared_branch_length(t, env_counts, 1)
        self.assertEqual(obs, exp)

        exp = {('A','B'):8.0}
        obs = shared_branch_length(t, env_counts, 2)
        self.assertEqual(obs, exp)

        self.assertRaises(ValueError, shared_branch_length, t, env_counts, 3)

    def test_shared_branch_length_to_root(self):
        """Should return the correct shared branch length by env to root"""
        t_str = "(((a:1,b:2):3,c:4),(d:5,e:6,f:7):8);"
        envs = """
a A 1
b A 1
c A 1
d A 1
e A 1
f B 1 
"""
        env_counts = count_envs(envs.splitlines())
        t = DndParser(t_str, UniFracTreeNode)
        exp = {'A':29.0,'B':15.0}
        obs = shared_branch_length_to_root(t, env_counts)
        self.assertEqual(obs, exp)


    def test_fast_unifrac(self):
        """Should calc unifrac values for whole tree."""
        #Note: results not tested for correctness here as detailed tests
        #in fast_tree module.
        res = fast_unifrac(self.t, self.env_counts)
        res = fast_unifrac(self.t, self.missing_env_counts)
        res = fast_unifrac(self.t, self.extra_tip_counts)
        self.assertRaises(ValueError,  fast_unifrac, self.t, \
            self.wrong_tip_counts)
            
    def test_fast_unifrac_one_sample(self):
        """ fu one sample should match whole unifrac result, for env 'B'"""
        # first get full unifrac matrix
        res = fast_unifrac(self.t, self.env_counts)
        dmtx, env_order =  res['distance_matrix']
        dmtx_vec = dmtx[env_order.index('B')]
        dmtx_vec = dmtx_vec[argsort(env_order)]
        
        # then get one sample unifrac vector
        one_sam_dvec, one_sam_env_order = \
            fast_unifrac_one_sample('B', self.t, self.env_counts)
        one_sam_dvec = one_sam_dvec[argsort(one_sam_env_order)]
        self.assertFloatEqual(one_sam_dvec, dmtx_vec)

    def test_fast_unifrac_one_sample2(self):
        """fu one sam should match whole weighted unifrac result, for env 'B'"""
        # first get full unifrac matrix
        res = fast_unifrac(self.t, self.env_counts, weighted=True)
        dmtx, env_order =  res['distance_matrix']
        dmtx_vec = dmtx[env_order.index('B')]
        dmtx_vec = dmtx_vec[argsort(env_order)]
        
        # then get one sample unifrac vector
        one_sam_dvec, one_sam_env_order = \
            fast_unifrac_one_sample('B', self.t, self.env_counts,weighted=True)
        one_sam_dvec = one_sam_dvec[argsort(one_sam_env_order)]
        self.assertFloatEqual(one_sam_dvec, dmtx_vec)
        
    def test_fast_unifrac_one_sample3(self):
        """fu one sam should match missing env unifrac result, for env 'B'"""
        # first get full unifrac matrix
        res = fast_unifrac(self.t, self.missing_env_counts, weighted=False)
        dmtx, env_order =  res['distance_matrix']
        dmtx_vec = dmtx[env_order.index('C')]
        dmtx_vec = dmtx_vec[argsort(env_order)]
        
        # then get one sample unifrac vector
        one_sam_dvec, one_sam_env_order = \
            fast_unifrac_one_sample('C', self.t, 
            self.missing_env_counts,weighted=False)
        one_sam_dvec = one_sam_dvec[argsort(one_sam_env_order)]
        self.assertFloatEqual(one_sam_dvec, dmtx_vec)
        
        # and should raise valueerror when 'B'
        self.assertRaises(ValueError, fast_unifrac_one_sample, 'B', self.t, 
            self.missing_env_counts,weighted=False)
            


    def test_fast_unifrac_whole_tree(self):
        """ should correctly compute one p-val for whole tree """
        # "should test with fake permutation but 
        # using same as old envs nodefor now"
        result = []
        num_to_do = 10
        for i in range(num_to_do):
            real_ufracs, sim_ufracs = fast_unifrac_whole_tree(self.old_t, \
                self.old_env_counts, 1000, permutation_f=permutation)
            rawp, corp = mcarlo_sig(sum(real_ufracs), [sum(x) for x in \
                sim_ufracs], 1, tail='high')
            result.append(rawp)
        self.assertSimilarMeans(result, 0.047)

    def test_unifrac_explicit(self):
        """unifrac should correctly compute correct values.
        
        environment M contains only tips not in tree, tip j is in no envs
        values were calculated by hand
        """
        t1 = DndParser('((a:1,b:2):4,((c:3, j:17),(d:1,e:1):2):3)', \
            UniFracTreeNode) # note c,j is len 0 node
        #           /-------- /-a
        # ---------|          \-b
        #          |          /-------- /-c
        #           \--------|          \-j
        #                     \-------- /-d
        #                               \-e

        env_str = """
        a   A   1
        a   C   2
        b   A   1
        b   B   1
        c   B   1
        d   B   3
        e   C   1
        m   M   88"""
        env_counts = count_envs(env_str.splitlines())
        self.assertFloatEqual(fast_unifrac(t1,env_counts)['distance_matrix'], \
            (array(
            [[0,10/16, 8/13],
            [10/16,0,8/17],
            [8/13,8/17,0]]),['A','B','C']))
        # changing tree topology relative to c,j tips shouldn't change 
        # anything
        t2 = DndParser('((a:1,b:2):4,((c:2, j:16):1,(d:1,e:1):2):3)', \
            UniFracTreeNode)
        self.assertFloatEqual(fast_unifrac(t2,env_counts)['distance_matrix'], \
            (array(
            [[0,10/16, 8/13],
            [10/16,0,8/17],
            [8/13,8/17,0]]),['A','B','C']))

    def test_unifrac_make_subtree(self):
        """unifrac result should not depend on make_subtree
        
        environment M contains only tips not in tree, tip j, k is in no envs
        one clade is missing entirely
        values were calculated by hand
        we also test that we still have a valid tree at the end
        """
        t1 = DndParser('((a:1,b:2):4,((c:3, (j:1,k:2)mt:17),(d:1,e:1):2):3)',\
            UniFracTreeNode) # note c,j is len 0 node
        #           /-------- /-a
        # ---------|          \-b
        #          |          /-------- /-c
        #           \--------|          \mt------ /-j
        #                    |                    \-k
        #                     \-------- /-d
        #                               \-e
        # 

        env_str = """
        a   A   1
        a   C   2
        b   A   1
        b   B   1
        c   B   1
        d   B   3
        e   C   1
        m   M   88"""
        env_counts = count_envs(env_str.splitlines())
        self.assertFloatEqual(fast_unifrac(t1,env_counts,make_subtree=False)['distance_matrix'], \
            (array(
            [[0,10/16, 8/13],
            [10/16,0,8/17],
            [8/13,8/17,0]]),['A','B','C']))
        self.assertFloatEqual(fast_unifrac(t1,env_counts,make_subtree=True)['distance_matrix'], \
            (array(
            [[0,10/16, 8/13],
            [10/16,0,8/17],
            [8/13,8/17,0]]),['A','B','C']))
        # changing tree topology relative to c,j tips shouldn't change anything
        t2 = DndParser('((a:1,b:2):4,((c:2, (j:1,k:2)mt:17):1,(d:1,e:1):2):3)', \
            UniFracTreeNode)
        self.assertFloatEqual(fast_unifrac(t2,env_counts,make_subtree=False)['distance_matrix'], \
            (array(
            [[0,10/16, 8/13],
            [10/16,0,8/17],
            [8/13,8/17,0]]),['A','B','C']))
        self.assertFloatEqual(fast_unifrac(t2,env_counts,make_subtree=True)['distance_matrix'], \
            (array(
            [[0,10/16, 8/13],
            [10/16,0,8/17],
            [8/13,8/17,0]]),['A','B','C']))

        # ensure we haven't meaningfully changed the tree 
        # by passing it to unifrac
        t3 = DndParser('((a:1,b:2):4,((c:3, (j:1,k:2)mt:17),(d:1,e:1):2):3)',\
            UniFracTreeNode) # note c,j is len 0 node
        t1_tips = [tip.Name for tip in t1.tips()]
        t1_tips.sort()
        t3_tips = [tip.Name for tip in t3.tips()]
        t3_tips.sort()
        
        self.assertEqual(t1_tips, t3_tips)
        tipj3 = t3.getNodeMatchingName('j')
        tipb3 = t3.getNodeMatchingName('b')
        tipj1 = t1.getNodeMatchingName('j')
        tipb1 = t1.getNodeMatchingName('b')
        self.assertFloatEqual(tipj1.distance(tipb1), tipj3.distance(tipb3))


    def test_PD_whole_tree(self):
        """PD_whole_tree should correctly compute PD for test tree.
        
        environment M contains only tips not in tree, tip j is in no envs
        """
        t1 = DndParser('((a:1,b:2):4,((c:3, j:17),(d:1,e:1):2):3)', \
            UniFracTreeNode)
        env_str = """
        a   A   1
        a   C   2
        b   A   1
        b   B   1
        c   B   1
        d   B   3
        e   C   1
        m   M   88"""
        env_counts = count_envs(env_str.splitlines())
        self.assertEqual(PD_whole_tree(t1,env_counts), \
            (['A','B','C'], array([7.,15.,11.])))

    def test_PD_generic_whole_tree(self):
        """PD_generic_whole_tree should correctly compute PD for test tree."""
        self.t1 = DndParser('((a:1,b:2):4,(c:3,(d:1,e:1):2):3)', \
            UniFracTreeNode)
        self.env_str = """
        a   A   1
        a   C   2
        b   A   1
        b   B   1
        c   B   1
        d   B   3
        e   C   1"""
        env_counts = count_envs(self.env_str.splitlines())
        self.assertEqual(PD_generic_whole_tree(self.t1,self.env_counts), \
            (['A','B','C'], array([7.,15.,11.])))


    def test_mcarlo_sig(self):
        """test_mcarlo_sig should calculate monte carlo sig high/low"""
        self.assertEqual(mcarlo_sig(.5, self.mc_1, 1, 'high'), (5.0/10, 5.0/10))
        self.assertEqual(mcarlo_sig(.5, self.mc_1, 1, 'low'), (4.0/10, 4.0/10))
        self.assertEqual(mcarlo_sig(.5, self.mc_1, 5, 'high'), (5.0/10, 1.0))
        self.assertEqual(mcarlo_sig(.5, self.mc_1, 5, 'low'), (4.0/10, 1.0))
        self.assertEqual(mcarlo_sig(0, self.mc_1, 1, 'low'), (0.0, "<=%.1e" % (1.0/10)))
        self.assertEqual(mcarlo_sig(100, self.mc_1, 10, 'high'), (0.0, "<=%.1e" % (1.0/10)))

    
    def test_num_comps(self):
        """ test num comps """
        self.assertEqual(num_comps(5), sum([i for i in range(1, 5)]))
        self.assertEqual(num_comps(15), sum([i for i in range(1, 15)]))
        self.assertEqual(num_comps(10000), sum([i for i in range(1, 10000)]))
        self.assertEqual(num_comps(1833), sum([i for i in range(1, 1833)]))

    def test_shuffle_tipnames(self):
        """shuffle_tipnames should return copy of tree w/ labels permuted"""
        #Note: this should never fail but is technically still stochastic
        #5! is 120 so repeating 5 times should fail about 1 in 10^10.
        for i in range(5):
            try:
                t = DndParser(self.t_str)
                result = shuffle_tipnames(t)
                orig_names = [n.Name for n in t.tips()]
                new_names = [n.Name for n in result.tips()]
                self.assertIsPermutation(orig_names, new_names)
                return
            except AssertionError:
                continue
        raise AssertionError, "Produced same permutation in 5 tries: broken?"

    def test_weight_equally(self):
        """weight_equally should return unit weight per tree"""
        self.assertEqual(weight_equally(self.trees, self.envs),
            array([1,1]))

    def test_weight_by_num_tips(self):
        """weight_by_num_tips should return tips per tree"""
        self.assertEqual(weight_by_num_tips(self.trees, self.envs),
            array([5, 4]))

    def test_weight_by_branch_length(self):
        """weight_by_branch_length should return branch length per tree"""
        self.assertEqual(weight_by_branch_length(self.trees, self.envs),
            array([17, 14]))

    def test_weight_by_num_seqs(self):
        """weight_by_num_seqs should return num seqs per tree"""
        self.assertEqual(weight_by_num_seqs(self.trees, self.envs),
            array([10, 8]))

    def test_get_all_env_names(self):
        """get_all_env_names should get all names from counts"""
        self.assertEqual(get_all_env_names(self.env_counts), 
            set('ABC'))

    def test_consolidate_skipping_missing_matrices(self):
        """consolidate_skipping_missing_matrices should skip those missing data"""
        m1 = array([[1,2],[3,4]])
        m2 = array([[1,2,3],[4,5,6],[7,8,9]])
        m3 = array([[2,2,2],[3,3,3],[4,4,4]])
        matrices = [m1,m2, m3]
        env_names = map(list, ['AB', 'ABC', 'ABC'])
        weights = [1, 2, 3]
        all_names =list('ABC')
        result = consolidate_skipping_missing_matrices(matrices, env_names, weights,
            all_names)
        self.assertFloatEqual(result, .4*m2 + .6*m3)

    def test_consolidate_missing_zero(self):
        """consolidate_missing_zero should fill missing values to zero"""
        m1 = array([[1,2],[3,4]])
        m2 = array([[1,2,3],[4,5,6],[7,8,9]])
        m3 = array([[2,2,2],[3,3,3],[4,4,4]])
        matrices = [m1,m2, m3]
        env_names = map(list, ['AB', 'ABC', 'ABC'])
        weights = [1, 2, 3]
        weights = array(weights, float)
        weights/=weights.sum()
        all_names =list('ABC')
        transformed_m1 = array([[1,2,0],[3,4,0],[0,0,0]])
        result = consolidate_missing_zero(matrices, env_names, weights,
            all_names)
        self.assertFloatEqual(result, (1/6.)*transformed_m1 + (2/6.)*m2 + (3/6.)*m3)

    def test_consolidate_missing_one(self):
        """consolidate_missing_one should fill missing off-diags to one"""
        m1 = array([[1,2],[3,4]])
        m2 = array([[1,2,3],[4,5,6],[7,8,9]])
        m3 = array([[2,2,2],[3,3,3],[4,4,4]])
        matrices = [m1,m2, m3]
        env_names = map(list, ['AB', 'ABC', 'ABC'])
        weights = [1, 2, 3]
        weights = array(weights, float)
        weights/=weights.sum()
        all_names =list('ABC')
        transformed_m1 = array([[1,2,1],[3,4,1],[1,1,0]])
        result = consolidate_missing_one(matrices, env_names, weights,
            all_names)
        self.assertFloatEqual(result, (1/6.)*transformed_m1 + (2/6.)*m2 + (3/6.)*m3)

    def test_consolidate_skipping_missing_values(self):
        """consolidate_skipping_missing_values should average over filled values"""
        m1 = array([[1,2],[3,4]])
        m2 = array([[1,2,3],[4,5,6],[7,8,9]])
        m3 = array([[2,2,2],[3,3,3],[4,4,4]])
        matrices = [m1,m2, m3]
        env_names = map(list, ['AB', 'ABC', 'ABC'])
        weights = [1., 2, 3]
        weights = array(weights)
        weights/=weights.sum()
        all_names =list('ABC')
        expected = array([[ 1/6.*1 + 2/6.*1 + 3/6.*2,
                            1/6.*2 + 2/6.*2 + 3/6.*2,
                            2/5.*3 + 3/5.*2],
                          [ 1/6.*3 + 2/6.*4 + 3/6.*3,
                            1/6.*4 + 2/6.*5 + 3/6.*3,
                            2/5.*6 + 3/5.*3],
                          [ 2/5.*7 + 3/5.*4,
                            2/5.*8 + 3/5.*4,
                            2/5.*9 + 3/5.*4]])
        result = consolidate_skipping_missing_values(matrices, env_names, weights,
            all_names)
        self.assertFloatEqual(result, expected)

    def test_reshape_by_name(self):
        """reshape_by_name should reshape matrix from old to new names"""
        old = array([[0,1,2],[3,4,5],[6,7,8]])
        old_names = 'ABC'
        new_names = 'xCyBA'
        exp = array([[0,0,0,0,0],[0,8,0,7,6],[0,0,0,0,0],\
            [0,5,0,4,3],[0,2,0,1,0]])
        self.assertEqual(reshape_by_name(old, old_names, new_names), exp)
        result = reshape_by_name(old, old_names, new_names, masked=True)
        result.fill_value=0
        self.assertEqual(result._data * logical_not(result._mask), exp)

    def test_meta_unifrac(self):
        """meta_unifrac should give correct result on sample trees"""
        tree_list = [self.t, self.t2]
        envs_list = [self.env_counts, self.env2_counts]
        result = meta_unifrac(tree_list, envs_list, weight_equally,
            modes=["distance_matrix"])

        u1_distances = array([[0, 10/16.,8/13.],[10/16.,0,8/17.],\
                            [8/13.,8/17.,0]])
        u2_distances = array([[0,11/14.,6/13.],[11/14.,0,7/13.],[6/13.,7/13., 0]])
        exp = (u1_distances + u2_distances)/2
        self.assertFloatEqual(result['distance_matrix'], (exp, list('ABC')))


if __name__ == '__main__':
    main()