#! /usr/bin/env python
# Matthew Wakefield Feb 2004

from __future__ import with_statement
import unittest
import os
import warnings

from cogent import LoadSeqs, LoadTree
import cogent.evolve.parameter_controller, cogent.evolve.substitution_model
from cogent.maths import optimisers

__author__ = "Peter Maxwell"
__copyright__ = "Copyright 2007-2012, The Cogent Project"
__credits__ = ["Peter Maxwell", "Gavin Huttley", "Matthew Wakefield"]
__license__ = "GPL"
__version__ = "1.5.3"
__maintainer__ = "Gavin Huttley"
__email__ = "gavin.huttley@anu.edu.au"
__status__ = "Production"

base_path = os.getcwd()
data_path = os.path.join(base_path, 'data')

good_rule_sets = [
    [
    {'par_name' : 'length','is_independent':True},
    ],
    [
    {'par_name' : 'length','is_independent':True},
    ],
    [
    {'par_name' : 'length','is_clade' :True, 'is_independent':True, 'edges' : ['a','b']},
    ],
    [
    {'par_name' : 'length','is_independent':True, 'edges' : ['a','c','e']},
    ],
    [
    {'par_name' : 'length','is_independent':True, 'edge' : 'a'},
    ],
]
bad_rule_sets = [
    [
    {'par_name' : 'length','is_clade' :True, 'edges' : ['b','f'],},
    ],
]
class test_parameter_controller(unittest.TestCase):
    """Tesing Parameter Controller"""
    def setUp(self):
        #length all edges 1 except c=2.  b&d transitions all other transverions
        self.al = LoadSeqs(
            data={'a':'tata', 'b':'tgtc', 'c':'gcga', 'd':'gaac', 'e':'gagc',})
        self.tree = LoadTree(treestring='((a,b),(c,d),e);')
        self.model = cogent.evolve.substitution_model.Nucleotide(
            do_scaling=True, equal_motif_probs=True, model_gaps=True)
        
    def test_scoped_local(self):
        model = cogent.evolve.substitution_model.Nucleotide(
                do_scaling=True, equal_motif_probs=True, model_gaps=True,
                predicates = {'kappa':'transition'})
        lf = model.makeLikelihoodFunction(self.tree)
        lf.setConstantLengths()
        lf.setAlignment(self.al)
        null = lf.getNumFreeParams()
        lf.setParamRule(par_name='kappa',
                            is_independent=True,
                            edges=['b','d'])
        self.assertEqual(null+2, lf.getNumFreeParams())
    
    def test_setMotifProbs(self):
        """Mprobs supplied to the parameter controller"""
        model = cogent.evolve.substitution_model.Nucleotide(
            model_gaps=True, motif_probs=None)
        lf = model.makeLikelihoodFunction(self.tree, 
                motif_probs_from_align=False)
                
        mprobs = {'A':0.1,'C':0.2,'G':0.2,'T':0.5,'-':0.0}
        lf.setMotifProbs(mprobs)
        self.assertEqual(lf.getMotifProbs(), mprobs)
        
        lf.setMotifProbsFromData(self.al[:1], is_constant=True)
        self.assertEqual(lf.getMotifProbs()['G'], 0.6)
        
        lf.setMotifProbsFromData(self.al[:1], pseudocount=1)
        self.assertNotEqual(lf.getMotifProbs()['G'], 0.6)
        
        # test with consideration of ambiguous states
        al = LoadSeqs(data = {'seq1': 'ACGTAAGNA', 'seq2': 'ACGTANGTC',
                               'seq3': 'ACGTACGTG'})
        lf.setMotifProbsFromData(al, include_ambiguity=True, is_constant=True)
        motif_probs = dict(lf.getMotifProbs())
        correct_probs = {'A': 8.5/27, 'C': 5.5/27, '-': 0.0, 'T': 5.5/27,
                         'G': 7.5/27}
        self.assertEqual(motif_probs, correct_probs)
        self.assertEqual(sum(motif_probs.values()), 1.0)

    def test_setMultiLocus(self):
        """2 loci each with own mprobs"""
        model = cogent.evolve.substitution_model.Nucleotide(motif_probs=None)
        lf = model.makeLikelihoodFunction(self.tree, 
                motif_probs_from_align=False, loci=["a", "b"])
                
        mprobs_a = dict(A=.2, T=.2, C=.3, G=.3)
        mprobs_b = dict(A=.1, T=.2, C=.3, G=.4)
        
        for is_constant in [False, True]:
            lf.setMotifProbs(mprobs_a, is_constant=is_constant)
            s = str(lf)
            lf.setMotifProbs(mprobs_b, locus="b")
            self.assertEqual(lf.getMotifProbs(locus="a"), mprobs_a)
            self.assertEqual(lf.getMotifProbs(locus="b"), mprobs_b)
            s = str(lf)
            #lf.setParamRule('mprobs', is_independent=False)

    def test_setParamRules(self):
        lf = self.model.makeLikelihoodFunction(self.tree)
        def do_rules(rule_set):
            for rule in rule_set:
                lf.setParamRule(**rule)
        for rule_set in good_rule_sets:
            lf.setDefaultParamRules()
            do_rules(rule_set)
        for rule_set in bad_rule_sets:
            lf.setDefaultParamRules()
            self.assertRaises((KeyError, TypeError,
                    AssertionError, ValueError), do_rules, rule_set)
    
    def test_setLocalClock(self):
        pass

    def test_setConstantLengths(self):
        t = LoadTree(treestring='((a:1,b:2):3,(c:4,d:5):6,e:7);')
        lf = self.model.makeLikelihoodFunction(t)#self.tree)
        lf.setParamRule('length', is_constant=True)
        # lf.setConstantLengths(t)
        lf.setAlignment(self.al)
        self.assertEqual(lf.getParamValue('length', 'b'), 2)
        self.assertEqual(lf.getParamValue('length', 'd'), 5)

    def test_pairwise_clock(self):
        al = LoadSeqs(data={'a':'agct','b':'ggct'})
        tree = LoadTree(treestring='(a,b);')
        model = cogent.evolve.substitution_model.Dinucleotide(
                do_scaling=True, equal_motif_probs=True, model_gaps=True,
                mprob_model='tuple')
        lf = model.makeLikelihoodFunction(tree)
        lf.setLocalClock('a','b')
        lf.setAlignment(al)
        lf.optimise(local=True)
        rd = lf.getParamValueDict(['edge'], params=['length'])
        self.assertAlmostEquals(lf.getLogLikelihood(),-10.1774488956)
        self.assertEqual(rd['length']['a'],rd['length']['b'])

    def test_local_clock(self):
        lf = self.model.makeLikelihoodFunction(self.tree)
        lf.setLocalClock('c','d')
        lf.setAlignment(self.al)
        lf.optimise(local=True, 
                tolerance=1e-8, max_restarts=2)
        rd = lf.getParamValueDict(['edge'], params=['length'])
        self.assertAlmostEquals(lf.getLogLikelihood(),-27.84254174)
        self.assertEqual(rd['length']['c'],rd['length']['d'])
        self.assertNotEqual(rd['length']['a'],rd['length']['e'])
        
    def test_complex_parameter_rules(self):
            # This test has many local minima and so does not cope
            # with changes to optimiser details.
        model = cogent.evolve.substitution_model.Nucleotide(
                do_scaling=True, equal_motif_probs=True, model_gaps=True,
                predicates = {'kappa':'transition'})
        lf = model.makeLikelihoodFunction(self.tree)
        lf.setParamRule(par_name='kappa',
                                is_independent=True)
        lf.setParamRule(par_name='kappa',
                                is_independent=False,
                                edges=['b','d'])
        lf.setConstantLengths(LoadTree(
                            treestring='((a:1,b:1):1,(c:2,d:1):1,e:1);'))
        #print self.pc
        lf.setAlignment(self.al)
        lf.optimise(local=True)
        rd = lf.getParamValueDict(['edge'], params=['kappa'])
        self.assertAlmostEquals(lf.getLogLikelihood(),-27.3252, 3)
        self.assertEqual(rd['kappa']['b'],rd['kappa']['d'])
        self.assertNotEqual(rd['kappa']['a'],rd['kappa']['b'])
        
    def test_bounds(self):
        """Test setting upper and lower bounds for parameters"""
        lf = self.model.makeLikelihoodFunction(self.tree)
        lf.setParamRule('length', value=3, lower=0, upper=5)
        
        # Out of bounds value should warn and keep bounded
        with warnings.catch_warnings(record=True) as w:
            lf.setParamRule('length', lower=0, upper=2)
            self.assertTrue(len(w), 'No warning issued')
        self.assertEqual(lf.getParamValue('length', edge='a'), 2)
        
        # upper < lower bounds should fail
        self.assertRaises(ValueError, lf.setParamRule, 
            'length', lower=2, upper=0)
        
        
if __name__ == '__main__':
    unittest.main()

    
    
        
