from __future__ import division, print_function, absolute_import
from collections import defaultdict
import numpy as np
from treetime import config as ttconf
from .seq_utils import alphabets, profile_maps, alphabet_synonyms

class GTR(object):
    """
    Defines General-Time-Reversible model of character evolution.
    """

    def __init__(self, alphabet='nuc', prof_map=None, logger=None):
        """
        Initialize empty evolutionary model.

        Parameters
        ----------

         alphabet : str, numpy.array
            Alphabet of the sequence. If a string is passed, it is understood as
            an alphabet name. In this case, the alphabet and its profile map are pulled
            from :py:obj:`treetime.seq_utils`.

            If a numpy array of characters is passed, a new alphabet is constructed,
            and the default profile map is atached to it.

         prof_map : dict
            Dictionary linking characters in the sequence to the likelihood
            of observing characters in the alphabet. This is used to
            implement ambiguous characters like 'N'=[1,1,1,1] which are
            equally likely to be any of the 4 nucleotides. Standard profile_maps
            are defined in file seq_utils.py. If None is provided, no ambigous
            characters are supported.

         logger : callable
            Custom logging function that should take arguments (msg, level, warn=False),
            where msg is a string and level an integer to be compared against verbose.

        """
        self.debug=False
        if isinstance(alphabet, str):
            if alphabet not in alphabet_synonyms:
                raise AttributeError("Unknown alphabet type specified")
            else:
                tmp_alphabet = alphabet_synonyms[alphabet]
                self.alphabet = alphabets[tmp_alphabet]
                self.profile_map = profile_maps[tmp_alphabet]
        else:
            # not a predefine alphabet
            self.alphabet = alphabet
            if prof_map is None: # generate trivial unambiguous profile map is none is given
                self.profile_map = {s:x for s,x in zip(self.alphabet, np.eye(len(self.alphabet)))}
            else:
                self.profile_map = prof_map

        if logger is None:
            def logger_default(*args,**kwargs):
                """standard logging function if none provided"""
                if self.debug:
                    print(*args)
            self.logger = logger_default
        else:
            self.logger = logger

        self.ambiguous = None
        self.gap_index = None
        self.n_states = len(self.alphabet)
        self.assign_gap_and_ambiguous()

        # NEEDED TO BREAK RATE MATRIX DEGENERACY AND FORCE NP TO RETURN REAL ORTHONORMAL EIGENVECTORS
        # ugly hack, but works and shouldn't affect results
        tmp_rng_state = np.random.get_state()
        np.random.seed(12345)
        self.break_degen = np.random.random(size=(self.n_states, self.n_states))*1e-6
        np.random.set_state(tmp_rng_state)

        # init all matrices with dummy values
        self.logger("GTR: init with dummy values!", 3)
        self.v = None # right eigenvectors
        self.v_inv = None # left eigenvectors
        self.eigenvals =None # eigenvalues
        self.assign_rates()

    def assign_gap_and_ambiguous(self):
        n_states = len(self.alphabet)
        self.logger("GTR: with alphabet: "+str(self.alphabet),1)
        # determine if a character exists that corresponds to no info, i.e. all one profile
        if any([x.sum()==n_states for x in self.profile_map.values()]):
            amb_states = [c for c,x in self.profile_map.items() if x.sum()==n_states]
            self.ambiguous = 'N' if 'N' in amb_states else amb_states[0]
            self.logger("GTR: ambiguous character: "+self.ambiguous,2)
        else:
            self.ambiguous=None

        # check for a gap symbol
        try:
            self.gap_index = list(self.alphabet).index('-')
        except:
            self.logger("GTR: no gap symbol!", 4, warn=True)
            self.gap_index=None

    @property
    def Q(self):
        """function that return the product of the transition matrix
           and the equilibrium frequencies to obtain the rate matrix
           of the GTR model
        """
        return (self.W*self.Pi).T

######################################################################
## constructor methods
######################################################################

    def __str__(self):
        '''
        String representation of the GTR model for pretty printing
        '''
        multi_site = len(self.Pi.shape)==2

        if multi_site:
            eq_freq_str = "Average substitution rate (mu): "+str(np.round(self.average_rate,6))+'\n'
        else:
            eq_freq_str = "Substitution rate (mu): "+str(np.round(self.mu,6))+'\n'

        if not multi_site:
            eq_freq_str += "\nEquilibrium frequencies (pi_i):\n"
            for a,p in zip(self.alphabet, self.Pi):
                eq_freq_str+='  '+str(a)+': '+str(np.round(p,4))+'\n'

        W_str = "\nSymmetrized rates from j->i (W_ij):\n"
        W_str+='\t'+'\t'.join(map(str, self.alphabet))+'\n'
        for a,Wi in zip(self.alphabet, self.W):
            W_str+= '  '+str(a)+'\t'+'\t'.join([str(np.round(max(0,p),4)) for p in Wi])+'\n'

        if not multi_site:
            Q_str = "\nActual rates from j->i (Q_ij):\n"
            Q_str+='\t'+'\t'.join(map(str, self.alphabet))+'\n'
            for a,Qi in zip(self.alphabet, self.Q):
                Q_str+= '  '+str(a)+'\t'+'\t'.join([str(np.round(max(0,p),4)) for p in Qi])+'\n'

        return eq_freq_str + W_str + Q_str

    def assign_rates(self, mu=1.0, pi=None, W=None):
        """
        Overwrite the GTR model given the provided data

        Parameters
        ----------

         mu : float
            Substitution rate

         W : nxn matrix
            Substitution matrix

         pi : n vector
            Equilibrium frequencies

        """
        n = len(self.alphabet)
        self.mu = mu

        if pi is not None and len(pi)==n:
            Pi = np.array(pi)
        else:
            if pi is not None and len(pi)!=n:
                self.logger("length of equilibrium frequency vector does not match alphabet length", 4, warn=True)
                self.logger("Ignoring input equilibrium frequencies", 4, warn=True)
            Pi = np.ones(shape=(n,))

        self.Pi = Pi/np.sum(Pi)

        if W is None or W.shape!=(n,n):
            if (W is not None) and W.shape!=(n,n):
                self.logger("Substitution matrix size does not match alphabet size", 4, warn=True)
                self.logger("Ignoring input substitution matrix", 4, warn=True)
            # flow matrix
            W = np.ones((n,n))
            np.fill_diagonal(W, 0.0)
            np.fill_diagonal(W, - W.sum(axis=0))
        else:
            W=np.array(W)

        self.W = 0.5*(W+W.T)
        self._check_fix_Q(fixed_mu=True)
        self._eig()

    @classmethod
    def custom(cls, mu=1.0, pi=None, W=None, **kwargs):
        """
        Create a GTR model by specifying the matrix explicitly

        Parameters
        ----------

         mu : float
            Substitution rate

         W : nxn matrix
            Substitution matrix

         pi : n vector
            Equilibrium frequencies

         **kwargs:
            Key word arguments to be passed

        Keyword Args
        ------------

         alphabet : str
            Specify alphabet when applicable. If the alphabet specification is
            required, but no alphabet is specified, the nucleotide alphabet will be used as
            default.

        """
        gtr = cls(**kwargs)
        gtr.assign_rates(mu=mu, pi=pi, W=W)
        return gtr

    @staticmethod
    def standard(model, **kwargs):
        """
        Create standard model of molecular evolution.

        Parameters
        ----------

         model : str
            Model to create. See list of available models below
         **kwargs:
            Key word arguments to be passed to the model

        **Available models**

        - JC69:

          Jukes-Cantor 1969 model. This model assumes equal frequencies
          of the nucleotides and equal transition rates between nucleotide states.
          For more info, see: Jukes and Cantor (1969).
          Evolution of Protein Molecules. New York: Academic Press. pp. 21-132.
          To create this model, use:

          :code:`mygtr = GTR.standard(model='jc69', mu=<my_mu>, alphabet=<my_alph>)`

          :code:`my_mu` - substitution rate (float)

          :code:`my_alph` - alphabet (str: :code:`'nuc'` or  :code:`'nuc_nogap'`)

        - K80:

          Kimura 1980 model. Assumes equal concentrations across nucleotides, but
          allows different rates between transitions and transversions. The ratio
          of the transversion/transition rates is given by kappa parameter.
          For more info, see
          Kimura (1980),  J. Mol. Evol. 16 (2): 111-120. doi:10.1007/BF01731581.
          Current implementation of the model does not account for the gaps.

          :code:`mygtr = GTR.standard(model='k80', mu=<my_mu>, kappa=<my_kappa>)`

          :code:`mu` - overall substitution rate (float)

          :code:`kappa` - ratio of transversion/transition rates (float)

        - F81:

          Felsenstein 1981 model. Assumes non-equal concentrations across nucleotides,
          but the transition rate between all states is assumed to be equal. See
          Felsenstein (1981), J. Mol. Evol. 17  (6): 368-376. doi:10.1007/BF01734359
          for details.

          :code:`mygtr = GTR.standard(model='F81', mu=<mu>, pi=<pi>, alphabet=<alph>)`

          :code:`mu` -  substitution rate  (float)

          :code:`pi`  - : nucleotide concentrations (numpy.array)

          :code:`alphabet' -  alphabet to use. (:code:`'nuc'` or  :code:`'nuc_nogap'`)

        - HKY85:

          Hasegawa, Kishino and Yano 1985 model. Allows different concentrations of the
          nucleotides (as in F81) + distinguishes between transition/transversion substitutions
          (similar to K80). Link:
          Hasegawa, Kishino, Yano (1985), J. Mol. Evol. 22 (2): 160-174. doi:10.1007/BF02101694

          Current implementation of the model does not account for the gaps

          :code:`mygtr = GTR.standard(model='HKY85', mu=<mu>, pi=<pi>, kappa=<kappa>)`

          :code:`mu` -  substitution rate  (float)

          :code:`pi`  - : nucleotide concentrations (numpy.array)

          :code:`kappa` - ratio of transversion/transition rates (float)

        - T92:

          Tamura 1992 model. Extending Kimura  (1980) model for the case where a
          G+C-content bias exists. Link:
          Tamura K (1992),  Mol.  Biol. Evol. 9 (4): 678-687.  DOI: 10.1093/oxfordjournals.molbev.a040752

          Current implementation of the model does not account for the gaps

          :code:`mygtr = GTR.standard(model='T92', mu=<mu>, pi_GC=<pi_gc>, kappa=<kappa>)`

          :code:`mu` -  substitution rate  (float)

          :code:`pi_GC`  - : relative GC content

          :code:`kappa` - ratio of transversion/transition rates (float)

        - TN93:

          Tamura and Nei 1993. The model distinguishes between the two different types of
          transition: (A <-> G) is allowed to have a different rate to (C<->T).
          Transversions have the same rate. The frequencies of the nucleotides are allowed
          to be different. Link: Tamura, Nei (1993), MolBiol Evol. 10 (3): 512-526.
          DOI:10.1093/oxfordjournals.molbev.a040023

          :code:`mygtr = GTR.standard(model='TN93', mu=<mu>, kappa1=<k1>, kappa2=<k2>)`

          :code:`mu` -  substitution rate  (float)

          :code:`kappa1`  - relative A<-->C, A<-->T, T<-->G and G<-->C rates (float)

          :code:`kappa` - relative C<-->T rate (float)

          .. Note::
                Rate of A<-->G substitution is set to one. All other rates
                (kappa1, kappa2) are specified relative to this rate

        """
        from .nuc_models import JC69, K80, F81, HKY85, T92, TN93
        from .aa_models  import JTT92

        if model.lower() in ['jc', 'jc69', 'jukes-cantor', 'jukes-cantor69', 'jukescantor', 'jukescantor69']:
            return JC69(**kwargs)
        elif model.lower() in ['k80', 'kimura80', 'kimura1980']:
            return K80(**kwargs)
        elif model.lower() in ['f81', 'felsenstein81', 'felsenstein1981']:
            return F81(**kwargs)
        elif model.lower() in ['hky', 'hky85', 'hky1985']:
            return HKY85(**kwargs)
        elif model.lower() in ['t92', 'tamura92', 'tamura1992']:
            return T92(**kwargs)
        elif model.lower() in ['tn93', 'tamura_nei_93', 'tamuranei93']:
            return TN93(**kwargs)
        elif model.lower() in ['jtt', 'jtt92']:
            return JTT92(**kwargs)
        else:
            raise KeyError("The GTR model '{}' is not in the list of available models."
                "".format(model))

    @classmethod
    def random(cls, mu=1.0, alphabet='nuc'):
        """
        Creates a random GTR model

        Parameters
        ----------

         mu : float
            Substitution rate

         alphabet : str
            Alphabet name (should be standard: 'nuc', 'nuc_gap', 'aa', 'aa_gap')

        """

        alphabet=alphabets[alphabet]
        gtr = cls(alphabet)
        n = gtr.alphabet.shape[0]
        pi = 1.0*np.random.randint(0,100,size=(n))
        W = 1.0*np.random.randint(0,100,size=(n,n)) # with gaps

        gtr.assign_rates(mu=mu, pi=pi, W=W)
        return gtr

    @classmethod
    def infer(cls, nij, Ti, root_state, fixed_pi=None, pc=5.0, gap_limit=0.01, **kwargs):
        """
        Infer a GTR model by specifying the number of transitions and time spent in each
        character. The basic equation that is being solved is

        :math:`n_{ij} = pi_i W_{ij} T_j`

        where :math:`n_{ij}` are the transitions, :math:`pi_i` are the equilibrium
        state frequencies, :math:`W_{ij}` is the "substitution attempt matrix",
        while :math:`T_i` is the time on the tree spent in character state
        :math:`i`. To regularize the process, we add pseudocounts and also need
        to account for the fact that the root of the tree is in a particular
        state. the modified equation is

        :math:`n_{ij} + pc = pi_i W_{ij} (T_j+pc+root\_state)`

        Parameters
        ----------

         nij : nxn matrix
            The number of times a change in character state is observed
            between state j and i

         Ti :n vector
            The time spent in each character state

         root_state : n vector
            The number of characters in state i in the sequence
            of the root node.

         pc : float
            Pseudocounts, this determines the lower cutoff on the rate when
            no substitutions are observed

         **kwargs:
            Key word arguments to be passed

        Keyword Args
        ------------

         alphabet : str
            Specify alphabet when applicable. If the alphabet specification
            is required, but no alphabet is specified, the nucleotide alphabet will be used as default.

        """
        from scipy import linalg as LA
        gtr = cls(**kwargs)
        gtr.logger("GTR: model inference ",1)
        dp = 1e-5
        Nit = 40
        pc_mat = pc*np.ones_like(nij)
        np.fill_diagonal(pc_mat, 0.0)
        count = 0
        pi_old = np.zeros_like(Ti)
        if fixed_pi is None:
            pi = np.ones_like(Ti)
        else:
            pi = np.copy(fixed_pi)
        pi/=pi.sum()
        W_ij = np.ones_like(nij)
        mu = nij.sum()/Ti.sum()
        # if pi is fixed, this will immediately converge
        while LA.norm(pi_old-pi) > dp and count < Nit:
            gtr.logger(' '.join(map(str, ['GTR inference iteration',count,'change:',LA.norm(pi_old-pi)])), 3)
            count += 1
            pi_old = np.copy(pi)
            W_ij = (nij+nij.T+2*pc_mat)/mu/(np.outer(pi,Ti) + np.outer(Ti,pi)
                                                    + ttconf.TINY_NUMBER + 2*pc_mat)

            np.fill_diagonal(W_ij, 0)
            scale_factor = np.einsum('i,ij,j',pi,W_ij,pi)

            W_ij = W_ij/scale_factor
            if fixed_pi is None:
                pi = (np.sum(nij+pc_mat,axis=1)+root_state)/(ttconf.TINY_NUMBER + mu*np.dot(W_ij,Ti)+root_state.sum()+np.sum(pc_mat, axis=1))
                pi /= pi.sum()
            mu = nij.sum()/(ttconf.TINY_NUMBER + np.sum(pi * (W_ij.dot(Ti))))
        if count >= Nit:
            gtr.logger('WARNING: maximum number of iterations has been reached in GTR inference',3, warn=True)
            if LA.norm(pi_old-pi) > dp:
                gtr.logger('the iterative scheme has not converged',3,warn=True)
            elif np.abs(1-np.max(pi.sum(axis=0))) > dp:
                gtr.logger('the iterative scheme has converged, but proper normalization was not reached',3,warn=True)
        if gtr.gap_index is not None:
            if pi[gtr.gap_index]<gap_limit:
                gtr.logger('The model allows for gaps which are estimated to occur at a low fraction of %1.3e'%pi[gtr.gap_index]+
                       '\n\t\tthis can potentially result in artificats.'+
                       '\n\t\tgap fraction will be set to %1.4f'%gap_limit,2,warn=True)
            pi[gtr.gap_index] = gap_limit
            pi /= pi.sum()

        gtr.assign_rates(mu=mu, W=W_ij, pi=pi)
        return gtr

########################################################################
### prepare model
########################################################################
    def _check_fix_Q(self, fixed_mu=False):
        """
        Check the main diagonal of Q and fix it in case it does not corresond
        the definition of the rate matrix. Should be run every time when creating
        custom GTR model.
        """
        # fix Q
        self.Pi /= self.Pi.sum() # correct the Pi manually
        # NEEDED TO BREAK RATE MATRIX DEGENERACY AND FORCE NP TO RETURN REAL ORTHONORMAL EIGENVECTORS
        self.W += self.break_degen + self.break_degen.T
        # fix W
        np.fill_diagonal(self.W, 0)
        Wdiag = -(self.Q).sum(axis=0)/self.Pi
        np.fill_diagonal(self.W, Wdiag)
        scale_factor = -np.sum(np.diagonal(self.Q)*self.Pi)
        self.W /= scale_factor
        if not fixed_mu:
            self.mu *= scale_factor
        if (self.Q.sum(axis=0) < 1e-10).sum() <  self.alphabet.shape[0]: # fix failed
            print ("Cannot fix the diagonal of the GTR rate matrix. Should be all zero", self.Q.sum(axis=0))
            import ipdb; ipdb.set_trace()
            raise ArithmeticError("Cannot fix the diagonal of the GTR rate matrix.")

    def _eig(self):
        """
        Perform eigendecompositon of the rate matrix and stores the left- and right-
        matrices to convert the sequence profiles to the GTR matrix eigenspace
        and hence to speed-up the computations.
        """
        # eigendecomposition of the rate matrix
        eigvals, eigvecs = np.linalg.eig(self.Q)
        self.v = np.real(eigvecs)
        self.v_inv = np.linalg.inv(self.v)
        self.eigenvals = np.real(eigvals)

    def _eig_sym(self):
        """
        Perform eigendecompositon of the rate matrix and stores the left- and right-
        matrices to convert the sequence profiles to the GTR matrix eigenspace
        and hence to speed-up the computations.
        """
        # eigendecomposition of the rate matrix
        tmpp = np.sqrt(self.Pi)
        symQ = self.W*np.outer(tmpp, tmpp)
        eigvals, eigvecs = np.linalg.eigh(symQ)
        tmp_v = eigvecs.T*tmpp
        one_norm = np.sum(np.abs(tmp_v), axis=1)
        self.v = tmp_v.T/one_norm
        self.v_inv = (eigvecs*one_norm).T/tmpp
        self.eigenvals = eigvals

    def compress_sequence_pair(self, seq_p, seq_ch, pattern_multiplicity=None,
                               ignore_gaps=False):
        '''
        Make a compressed representation of a pair of sequences, only counting
        the number of times a particular pair of states (e.g. (A,T)) is observed
        in the aligned sequences of parent and child.

        Parameters
        ----------

         seq_p: numpy array
            Parent sequence as numpy array of chars

         seq_ch: numpy array
            Child sequence as numpy array of chars

         pattern_multiplicity : numpy array
            If sequences are reduced by combining identical alignment patterns,
            these multplicities need to be accounted for when counting the number
            of mutations across a branch. If None, all pattern are assumed to
            occur exactly once.

         ignore_gap: bool
            Whether or not to include gapped positions of the alignment
            in the multiplicity count

        Returns
        -------
          seq_pair : list
            :code:`[(0,1), (2,2), (3,4)]` list of parent_child state pairs
            as indices in the alphabet

          multiplicity : numpy array
            Number of times a particular pair is observed

        '''
        if pattern_multiplicity is None:
            pattern_multiplicity = np.ones_like(seq_p, dtype=float)

        from collections import Counter
        if seq_ch.shape != seq_p.shape:
            raise ValueError("GTR.compress_sequence_pair: Sequence lengths do not match!")

        if len(self.alphabet)<10: # for small alphabet, repeatedly check array for all state pairs
            pair_count = []
            bool_seqs_p = []
            bool_seqs_ch = []
            for seq, bs in [(seq_p,bool_seqs_p), (seq_ch, bool_seqs_ch)]:
                for ni,nuc in enumerate(self.alphabet):
                    bs.append(seq==nuc)

            for n1,nuc1 in enumerate(self.alphabet):
                if (self.gap_index is None) or (not ignore_gaps) or (n1!=self.gap_index):
                    for n2,nuc2 in enumerate(self.alphabet):
                        if (self.gap_index is None) or (not ignore_gaps) or (n2!=self.gap_index):
                            count = ((bool_seqs_p[n1]&bool_seqs_ch[n2])*pattern_multiplicity).sum()
                            if count: pair_count.append(((n1,n2), count))
        else: # enumerate state pairs of the sequence for large alphabets
            num_seqs = []
            for seq in [seq_p, seq_ch]: # for each sequence (parent and child) construct a numerical sequence [0,5,3,1,2,3...]
                tmp = np.ones_like(seq, dtype=int)
                for ni,nuc in enumerate(self.alphabet):
                    tmp[seq==nuc] = ni  # set each position corresponding to a state to the corresponding index
                num_seqs.append(tmp)
            pair_count = defaultdict(int)
            if ignore_gaps:  # if gaps are ignored skip positions where one or the other sequence is gapped
                for i in range(len(seq_p)):
                    if self.gap_index!=num_seqs[0][i] and self.gap_index!=num_seqs[1][i]:
                        pair_count[(num_seqs[0][i],num_seqs[1][i])]+=pattern_multiplicity[i]
            else: # otherwise, just count
                for i in range(len(seq_p)):
                    pair_count[(num_seqs[0][i],num_seqs[1][i])]+=pattern_multiplicity[i]
            pair_count = pair_count.items()

        return (np.array([x[0] for x in pair_count], dtype=int),    # [(child_nuc, parent_nuc),()...]
                np.array([x[1] for x in pair_count], dtype=int))    # multiplicity of each parent/child nuc pair

########################################################################
### evolution functions
########################################################################
    def prob_t_compressed(self, seq_pair, multiplicity, t, return_log=False):
        '''
        Calculate the probability of observing a sequence pair at a distance t,
        for compressed sequences

        Parameters
        ----------

          seq_pair : numpy array
            :code:`np.array([(0,1), (2,2), ()..])` as indicies of
            pairs of aligned positions. (e.g. 'A'==0, 'C'==1 etc).
            This only lists all occuring parent-child state pairs, order is irrelevant

          multiplicity : numpy array
            The number of times a parent-child state pair is observed.
            This allows compression of the sequence representation

          t : float
            Length of the branch separating parent and child

          return_log : bool
            Whether or not to exponentiate the result

        '''
        if t<0:
            logP = -ttconf.BIG_NUMBER
        else:
            tmp_eQT = self.expQt(t)
            bad_indices=(tmp_eQT==0)
            logQt = np.log(tmp_eQT + ttconf.TINY_NUMBER*(bad_indices))
            logQt[np.isnan(logQt) | np.isinf(logQt) | bad_indices] = -ttconf.BIG_NUMBER
            logP = np.sum(logQt[seq_pair[:,1], seq_pair[:,0]]*multiplicity)

        return logP if return_log else np.exp(logP)

    def prob_t(self, seq_p, seq_ch, t, pattern_multiplicity = None,
               return_log=False, ignore_gaps=True):
        """
        Compute the probability to observe seq_ch (child sequence) after time t starting from seq_p
        (parent sequence).

        Parameters
        ----------

         seq_p : character array
            Parent sequence

         seq_c : character array
            Child sequence

         t : double
            Time (branch len) separating the profiles.

         pattern_multiplicity : numpy array
            If sequences are reduced by combining identical alignment patterns,
            these multplicities need to be accounted for when counting the number
            of mutations across a branch. If None, all pattern are assumed to
            occur exactly once.

         return_log : bool
            It True, return log-probability.

        Returns
        -------
         prob : np.array
            Resulting probability

        """
        seq_pair, multiplicity = self.compress_sequence_pair(seq_p, seq_ch,
                                        pattern_multiplicity=pattern_multiplicity, ignore_gaps=ignore_gaps)
        return self.prob_t_compressed(seq_pair, multiplicity, t, return_log=return_log)

    def optimal_t(self, seq_p, seq_ch, pattern_multiplicity=None, ignore_gaps=False):
        '''
        Find the optimal distance between the two sequences

        Parameters
        ----------

         seq_p : character array
            Parent sequence

         seq_c : character array
            Child sequence

         pattern_multiplicity : numpy array
            If sequences are reduced by combining identical alignment patterns,
            these multplicities need to be accounted for when counting the number
            of mutations across a branch. If None, all pattern are assumed to
            occur exactly once.

         ignore_gaps : bool
            If True, ignore gaps in distance calculations

        '''
        seq_pair, multiplicity = self.compress_sequence_pair(seq_p, seq_ch,
                                                            pattern_multiplicity = pattern_multiplicity,
                                                            ignore_gaps=ignore_gaps)
        return self.optimal_t_compressed(seq_pair, multiplicity)

    def optimal_t_compressed(self, seq_pair, multiplicity, profiles=False, tol=1e-10):
        """
        Find the optimal distance between the two sequences, for compressed sequences

        Parameters
        ----------

         seq_pair : compressed_sequence_pair
            Compressed representation of sequences along a branch, either
            as tuple of state pairs or as tuple of profiles.

         multiplicity : array
            Number of times each state pair in seq_pair appears (if profile==False)

            Number of times an alignment pattern is observed (if profiles==True)

         profiles : bool, default False
            The standard branch length optimization assumes fixed sequences at
            either end of the branch. With profiles==True, optimization is performed
            while summing over all possible states of the nodes at either end of the
            branch. Note that the meaning/format of seq_pair and multiplicity
            depend on the value of profiles.

        """

        def _neg_prob(t, seq_pair, multiplicity):
            """
            Probability to observe a child given the the parent state, transition
            matrix, and the time of evolution (branch length).

            Parameters
            ----------

             t : double
                Branch length (time between sequences)

             parent :  numpy.array
                Parent sequence

             child : numpy.array
                Child sequence

             tm :  GTR
                Model of evolution

            Returns
            -------

             prob : double
                Negative probability of the two given sequences
                to be separated by the time t.
            """
            if profiles:
                res = -1.0*self.prob_t_profiles(seq_pair, multiplicity,t**2, return_log=True)
                return res
            else:
                return -1.0*self.prob_t_compressed(seq_pair, multiplicity,t**2, return_log=True)

        try:
            from scipy.optimize import minimize_scalar
            opt = minimize_scalar(_neg_prob,
                    bounds=[-np.sqrt(ttconf.MAX_BRANCH_LENGTH),np.sqrt(ttconf.MAX_BRANCH_LENGTH)],
                    args=(seq_pair, multiplicity), tol=tol)
            new_len = opt["x"]**2
            if 'success' not in opt:
                opt['success'] = True
                self.logger("WARNING: the optimization result does not contain a 'success' flag:"+str(opt),4, warn=True)
        except:
            import scipy
            print('legacy scipy', scipy.__version__)
            from scipy.optimize import fminbound
            new_len = fminbound(_neg_prob,
                      -np.sqrt(ttconf.MAX_BRANCH_LENGTH),np.sqrt(ttconf.MAX_BRANCH_LENGTH),
                       args=(seq_pair, multiplicity))
            new_len = new_len**2
            opt={'success':True}

        if new_len > .9 * ttconf.MAX_BRANCH_LENGTH:
            self.logger("WARNING: GTR.optimal_t_compressed -- The branch length seems to be very long!", 4, warn=True)

        if opt["success"] != True:
            # return hamming distance: number of state pairs where state differs/all pairs
            new_len =  np.sum(multiplicity[seq_pair[:,1]!=seq_pair[:,0]])/np.sum(multiplicity)

        return new_len

    def prob_t_profiles(self, profile_pair, multiplicity, t,
                        return_log=False, ignore_gaps=True):
        '''
        Calculate the probability of observing a node pair at a distance t

        Parameters
        ----------

          profile_pair: numpy arrays
            Probability distributions of the nucleotides at either
            end of the branch. pp[0] = parent, pp[1] = child

          multiplicity : numpy array
            The number of times an alignment pattern is observed

          t : float
            Length of the branch separating parent and child

          ignore_gaps: bool
            If True, ignore mutations to and from gaps in distance calculations

          return_log : bool
            Whether or not to exponentiate the result

        '''
        if t<0:
            logP = -ttconf.BIG_NUMBER
        else:
            Qt = self.expQt(t)
            if len(Qt.shape)==3:
                res = np.einsum('ai,ija,aj->a', profile_pair[1], Qt, profile_pair[0])
            else:
                res = np.einsum('ai,ij,aj->a', profile_pair[1], Qt, profile_pair[0])
            if ignore_gaps and (self.gap_index is not None): # calculate the probability that neither outgroup/node has a gap
                non_gap_frac = (1-profile_pair[0][:,self.gap_index])*(1-profile_pair[1][:,self.gap_index])
                # weigh log LH by the non-gap probability
                logP = np.sum(multiplicity*np.log(res)*non_gap_frac)
            else:
                logP = np.sum(multiplicity*np.log(res))

        return logP if return_log else np.exp(logP)

    def propagate_profile(self, profile, t, return_log=False):
        """
        Compute the probability of the sequence state of the parent
        at time (t+t0, backwards), given the sequence state of the
        child (profile) at time t0.

        Parameters
        ----------

         profile : numpy.array
            Sequence profile. Shape = (L, a),
            where L - sequence length, a - alphabet size.

         t : double
            Time to propagate

         return_log: bool
            If True, return log-probability

        Returns
        -------

         res : np.array
            Profile of the sequence after time t in the past.
            Shape = (L, a), where L - sequence length, a - alphabet size.

        """
        Qt = self.expQt(t)
        res = profile.dot(Qt)

        return np.log(res) if return_log else res

    def evolve(self, profile, t, return_log=False):
        """
        Compute the probability of the sequence state of the child
        at time t later, given the parent profile.

        Parameters
        ----------

         profile : numpy.array
            Sequence profile. Shape = (L, a),
            where L - sequence length, a - alphabet size.

         t : double
            Time to propagate

         return_log: bool
            If True, return log-probability

        Returns
        -------

         res : np.array
            Profile of the sequence after time t in the future.
            Shape = (L, a), where L - sequence length, a - alphabet size.

        """
        Qt = self.expQt(t).T
        res = profile.dot(Qt)
        return np.log(res) if return_log else res

    def _exp_lt(self, t):
        """
        Parameters
        ----------

         t : float
            time to propagate

        Returns
        --------

         exp_lt : numpy.array
            Array of values exp(lambda(i) * t),
            where (i) - alphabet index (the eigenvalue number).
        """
        return np.exp(self.mu * t * self.eigenvals)

    def expQt(self, t):
        '''
        Parameters
        ----------

         t : float
            Time to propagate

        Returns
        --------

         expQt : numpy.array
            Matrix exponential of exo(Qt)
        '''
        eLambdaT = np.diag(self._exp_lt(t)) # vector length = a
        Qs = self.v.dot(eLambdaT.dot(self.v_inv))   # This is P(nuc1 | given nuc_2)
        return np.maximum(0,Qs)

    def expQs(self, s):
        eLambdaT = np.diag(self._exp_lt(s**2)) # vector length = a
        Qs = self.v.dot(eLambdaT.dot(self.v_inv))   # This is P(nuc1 | given nuc_2)
        return np.maximum(0,Qs)

    def expQsds(self, s):
        '''
        Returns
        -------
        Qtds :  Returns 2 V_{ij} \lambda_j s e^{\lambda_j s**2 } V^{-1}_{jk}
                This is the derivative of the branch probability with respect to s=\sqrt(t)
        '''
        lambda_eLambdaT = np.diag(2.0*self._exp_lt(s**2)*self.eigenvals*s) # vector length = a
        Qsds = self.v.dot(lambda_eLambdaT.dot(self.v_inv))
        return Qsds

    def expQsdsds(self, s):
        '''
        Returns
        -------
        Qtdtdt :  Returns V_{ij} \lambda_j^2 e^{\lambda_j s**2} V^{-1}_{jk}
                This is the second derivative of the branch probability wrt time
        '''
        t=s**2
        elt = self._exp_lt(t)
        lambda_eLambdaT = np.diag(elt*(4.0*t*self.eigenvals**2 + 2.0*self.eigenvals))
        Qsdsds = self.v.dot(lambda_eLambdaT.dot(self.v_inv))
        return Qsdsds

    def sequence_logLH(self,seq, pattern_multiplicity=None):
        """
        Returns the log-likelihood of sampling a sequence from equilibrium frequency.
        Expects a sequence as numpy array

        Parameters
        ----------

         seq : numpy array
            Compressed sequence as an array of chars

         pattern_multiplicity : numpy_array
            The number of times each position in sequence is observed in the
            initial alignment. If None, sequence is assumed to be not compressed

        """
        if pattern_multiplicity is None:
            pattern_multiplicity = np.ones_like(seq, dtype=float)
        return np.sum([np.sum((seq==state)*pattern_multiplicity*np.log(self.Pi[si]))
                      for si,state in enumerate(self.alphabet)])

    def average_rate(self):
        return self.mu*np.einsum('i,ij,j',self.Pi, self.W, self.Pi)

    def save_to_npz(self, outfile):
        full_gtr = self.mu * np.dot(self.Pi, self.W)
        desc=np.array(["GTR matrix description\n", "Substitution rate: " + str(self.mu)])
        np.savez(outfile,   description=desc,
                            full_gtr=full_gtr,
                            char_dist=self.Pi,
                            flow_matrix=self.W)

    def save_to_json(self, zip):
        d = {
        "full_gtr": self.mu * np.dot(self.Pi, self.W),
        "Substitution rate" : self.mu,
        "Equilibrium character composition": self.Pi,
        "Flow rate matrix": self.W
        }

if __name__ == "__main__":
    pass


from __future__ import print_function, division, absolute_import
import numpy as np
from scipy import optimize as sciopt
from Bio import Phylo
from treetime import config as ttconf
from .utils import tree_layout
from .clock_tree import ClockTree

rerooting_mechanisms = ["min_dev", "best", "least-squares"]
deprecated_rerooting_mechanisms = {"residual":"least-squares", "res":"least-squares",
                                   "min_dev_ML": "min_dev", "ML":"least-squares"}

class TreeTime(ClockTree):
    """
    TreeTime is a wrapper class to ClockTree that adds additional functionality
    such as reroot, detection and exclusion of outliers, resolution of polytomies
    using temporal information, and relaxed molecular clock models
    """

    def __init__(self, *args,**kwargs):
        """
        TreeTime constructor

        Parameters
        -----------
         *args
            Arguments to construct ClockTree

         **kwargs
            Keyword arguments to construct the GTR model

        """
        super(TreeTime, self).__init__(*args, **kwargs)

    def run(self, root=None, infer_gtr=True, relaxed_clock=None, n_iqd = None,
            resolve_polytomies=True, max_iter=0, Tc=None, fixed_clock_rate=None,
            time_marginal=False, sequence_marginal=False, branch_length_mode='auto',
            vary_rate=False, use_covariation=False, **kwargs):

        """
        Run TreeTime reconstruction. Based on the input parameters, it divides
        the analysis into semi-independent jobs and conquers them one-by-one,
        gradually optimizing the tree given the temporal constarints and leaf
        node sequences.

        Parameters
        ----------
        root : str
           Try to find better root position on a given tree. If string is passed,
           the root will be searched according to the specified method. If none,
           use tree as-is.

           See :py:meth:`treetime.TreeTime.reroot` for available rooting methods.

        infer_gtr : bool
           If True, infer GTR model

        relaxed_clock : dic
           If not None, use autocorrelated molecular clock model. Specify the
           clock parameters as :code:`{slack:<slack>, coupling:<coupling>}` dictionary.

        n_iqd : int
           If not None, filter tree nodes which do not obey the molecular clock
           for the particular tree. The nodes, which deviate more than
           :code:`n_iqd` interquantile intervals from the molecular clock
           regression will be marked as 'BAD' and not used in the TreeTime
           analysis

        resolve_polytomies : bool
           If True, attempt to resolve multiple mergers

        max_iter : int
           Maximum number of iterations to optimize the tree

        Tc : float, str
           If not None, use coalescent model to correct the branch lengths by
           introducing merger costs.

           If Tc is float, it is interpreted as the coalescence time scale

           If Tc is str, it should be one of (:code:`opt`, :code:`const`, :code:`skyline`)

        fixed_clock_rate : float
           Fixed clock rate to be used. If None, infer clock rate from the molecular clock.

        time_marginal : bool
           If True, perform a final round of marginal reconstruction of the node's positions.

        sequence_marginal : bool, optional
            use marginal reconstruction for ancestral sequences

        branch_length_mode : str
           Should be one of: :code:`joint`, :code:`marginal`, :code:`input`.

           If 'input', rely on the branch lengths in the input tree and skip directly
           to the maximum-likelihood ancestral sequence reconstruction.
           Otherwise, perform preliminary sequence reconstruction using parsimony
           algorithm and do branch length optimization

        vary_rate : bool or float, optional
            redo the time tree estimation for rates +/- one standard deviation.
            if a float is passed, it is interpreted as standard deviation,
            otherwise this standard deviation is estimated from the root-to-tip regression

        use_covariation : bool, optional
            default False, if False, rate estimates will be performed using simple
            regression ignoring phylogenetic covaration between nodes.

        **kwargs
           Keyword arguments needed by the downstream functions

        Returns
        -------
        TreeTime error/succces code : str
            return value depending on success or error

        """

        # register the specified covaration mode
        self.use_covariation = use_covariation

        if (self.tree is None) or (self.aln is None and self.seq_len is None):
            self.logger("TreeTime.run: ERROR, alignment or tree are missing", 0)
            return ttconf.ERROR
        if (self.aln is None):
            branch_length_mode='input'

        self._set_branch_length_mode(branch_length_mode)

        # determine how to reconstruct and sample sequences
        seq_kwargs = {"marginal_sequences":sequence_marginal or (self.branch_length_mode=='marginal'),
                      "sample_from_profile":"root"}

        tt_kwargs = {'clock_rate':fixed_clock_rate, 'time_marginal':False}
        tt_kwargs.update(kwargs)

        seq_LH = 0
        if "fixed_pi" in kwargs:
            seq_kwargs["fixed_pi"] = kwargs["fixed_pi"]
        if "do_marginal" in kwargs:
            time_marginal=kwargs["do_marginal"]

        # initially, infer ancestral sequences and infer gtr model if desired
        if self.branch_length_mode=='input':
            if self.aln:
                self.infer_ancestral_sequences(infer_gtr=infer_gtr, **seq_kwargs)
                self.prune_short_branches()
        else:
            self.optimize_sequences_and_branch_length(infer_gtr=infer_gtr,
                                                      max_iter=1, prune_short=True, **seq_kwargs)
        avg_root_to_tip = np.mean([x.dist2root for x in self.tree.get_terminals()])

        # optionally reroot the tree either by oldest, best regression or with a specific leaf
        if n_iqd or root=='clock_filter':
            if "plot_rtt" in kwargs and kwargs["plot_rtt"]:
                plot_rtt=True
            else:
                plot_rtt=False
            reroot_mechanism = 'least-squares' if root=='clock_filter' else root
            if self.clock_filter(reroot=reroot_mechanism, n_iqd=n_iqd, plot=plot_rtt)==ttconf.ERROR:
                return ttconf.ERROR
        elif root is not None:
            if self.reroot(root=root)==ttconf.ERROR:
                return ttconf.ERROR

        if self.branch_length_mode=='input':
            if self.aln:
                self.infer_ancestral_sequences(**seq_kwargs)
        else:
            self.optimize_sequences_and_branch_length(max_iter=1, prune_short=False,
                                                      **seq_kwargs)

        # infer time tree and optionally resolve polytomies
        self.logger("###TreeTime.run: INITIAL ROUND",0)
        self.make_time_tree(**tt_kwargs)

        if self.aln:
            seq_LH = self.tree.sequence_marginal_LH if seq_kwargs['marginal_sequences'] else self.tree.sequence_joint_LH
        self.LH =[[seq_LH, self.tree.positional_joint_LH, 0]]

        if root is not None and max_iter:
            if self.reroot(root='least-squares' if root=='clock_filter' else root)==ttconf.ERROR:
                return ttconf.ERROR

        # iteratively reconstruct ancestral sequences and re-infer
        # time tree to ensure convergence.
        niter = 0
        ndiff = 0
        need_new_time_tree=False
        while niter < max_iter:
            self.logger("###TreeTime.run: ITERATION %d out of %d iterations"%(niter+1,max_iter),0)
            # add coalescent prior
            if Tc:
                if Tc=='skyline' and niter<max_iter-1:
                    tmpTc='const'
                else:
                    tmpTc=Tc
                self.add_coalescent_model(tmpTc, **kwargs)
                need_new_time_tree = True

            # estimate a relaxed molecular clock
            if relaxed_clock:
                print("relaxed_clock", relaxed_clock)
                self.relaxed_clock(**relaxed_clock)
                need_new_time_tree = True

            n_resolved=0
            if resolve_polytomies:
                # if polytomies are found, rerun the entire procedure
                n_resolved = self.resolve_polytomies()
                if n_resolved:
                    self.prepare_tree()
                    if self.branch_length_mode!='input': # otherwise reoptimize branch length while preserving branches without mutations
                        self.optimize_sequences_and_branch_length(prune_short=False,
                                                          max_iter=0, **seq_kwargs)

                    need_new_time_tree = True

            if need_new_time_tree:
                self.make_time_tree(**tt_kwargs)
                if self.aln:
                    ndiff = self.infer_ancestral_sequences('ml',**seq_kwargs)
            else: # no refinements, just iterate
                if self.aln:
                    ndiff = self.infer_ancestral_sequences('ml',**seq_kwargs)
                self.make_time_tree(**tt_kwargs)

            self.tree.coalescent_joint_LH = self.merger_model.total_LH() if Tc else 0.0

            if self.aln:
                seq_LH = self.tree.sequence_marginal_LH if seq_kwargs['marginal_sequences'] else self.tree.sequence_joint_LH
            self.LH.append([seq_LH, self.tree.positional_joint_LH, self.tree.coalescent_joint_LH])
            niter+=1

            if ndiff==0 and n_resolved==0 and Tc!='skyline':
                self.logger("###TreeTime.run: CONVERGED",0)
                break

        # if the rate is too be varied and the rate estimate has a valid confidence interval
        # rerun the estimation for variations of the rate
        if vary_rate:
            if type(vary_rate)==float:
                res = self.calc_rate_susceptibility(rate_std=vary_rate, params=tt_kwargs)
            elif self.clock_model['valid_confidence']:
                res = self.calc_rate_susceptibility(params=tt_kwargs)
            else:
                res = ttconf.ERROR

            if res==ttconf.ERROR:
                self.logger("TreeTime.run: rate variation failed and can't be used for confidence estimation", 1, warn=True)

        # if marginal reconstruction requested, make one more round with marginal=True
        # this will set marginal_pos_LH, which to be used as error bar estimations
        if time_marginal:
            self.logger("###TreeTime.run: FINAL ROUND - confidence estimation via marginal reconstruction", 0)
            tt_kwargs['time_marginal']=time_marginal
            self.make_time_tree(**tt_kwargs)

        # explicitly print out which branches are bad and whose dates don't correspond to the input dates
        bad_branches =[n for n in self.tree.get_terminals()
                       if n.bad_branch and n.raw_date_constraint]
        if bad_branches:
            self.logger("TreeTime: The following tips don't fit the clock model, "
                        "please remove them from the tree. Their dates have been reset:",0,warn=True)
            for n in bad_branches:
                self.logger("%s, input date: %s, apparent date: %1.2f"%(n.name, str(n.raw_date_constraint), n.numdate),0,warn=True)

        return ttconf.SUCCESS

    def _set_branch_length_mode(self, branch_length_mode):
        '''
        if branch_length mode is not explicitly set, set according to
        empirical branch length distribution in input tree

        Parameters
        ----------

         branch_length_mode : str, 'input', 'joint', 'marginal'
            if the maximal branch length in the tree is longer than 0.05, this will
            default to 'input'. Otherwise set to 'joint'
        '''
        if branch_length_mode in ['joint', 'marginal', 'input']:
            self.branch_length_mode = branch_length_mode
        elif self.aln:
            bl_dis = [n.branch_length for n in self.tree.find_clades() if n.up]
            max_bl = np.max(bl_dis)
            if max_bl>0.1:
                bl_mode = 'input'
            else:
                bl_mode = 'joint'
            self.logger("TreeTime._set_branch_length_mode: maximum branch length is %1.3e, using branch length mode %s"%(max_bl, bl_mode),1)
            self.branch_length_mode = bl_mode
        else:
            self.branch_length_mode = 'input'

    def clock_filter(self, reroot='least-squares', n_iqd=None, plot=False):
        '''
        Labels outlier branches that don't seem to follow a molecular clock
        and excludes them from subsequent molecular clock estimation and
        the timetree propagation.

        Parameters
        ----------
         reroot : str
            Method to find the best root in the tree (see :py:meth:`treetime.TreeTime.reroot` for options)

         n_iqd : int
            Number of iqd intervals. The outlier nodes are those which do not fall
            into :math:`IQD\cdot n_iqd` interval (:math:`IQD` is the interval between
            75\ :sup:`th` and 25\ :sup:`th` percentiles)

            If None, the default (3) assumed

         plot : bool
            If True, plot the results

        '''
        if n_iqd is None:
            n_iqd = ttconf.NIQD
        if type(reroot) is list and len(reroot)==1:
            reroot=str(reroot[0])

        terminals = self.tree.get_terminals()
        if reroot:
            if self.reroot(root='least-squares' if reroot=='best' else reroot, covariation=False)==ttconf.ERROR:
                return ttconf.ERROR
        else:
            self.get_clock_model(covariation=False)

        clock_rate = self.clock_model['slope']
        icpt = self.clock_model['intercept']
        res = {}
        for node in terminals:
            if hasattr(node, 'raw_date_constraint') and  (node.raw_date_constraint is not None):
                res[node] = node.dist2root - clock_rate*np.mean(node.raw_date_constraint) - icpt

        residuals = np.array(list(res.values()))
        iqd = np.percentile(residuals,75) - np.percentile(residuals,25)
        for node,r in res.items():
            if abs(r)>n_iqd*iqd and node.up.up is not None:
                self.logger('TreeTime.ClockFilter: marking %s as outlier, residual %f interquartile distances'%(node.name,r/iqd), 3, warn=True)
                node.bad_branch=True
            else:
                node.bad_branch=False

        # redo root estimation after outlier removal
        if reroot and self.reroot(root=reroot)==ttconf.ERROR:
                return ttconf.ERROR

        if plot:
            self.plot_root_to_tip()

        return ttconf.SUCCESS

    def plot_root_to_tip(self, add_internal=False, label=True, ax=None):
        """
        Plot root-to-tip regression

        Parameters
        ----------
        add_internal : bool
           If true, plot inte`rnal node positions

        label : bool
           If true, label the plots

        ax : matplotlib axes
           If not None, use the provided matplotlib axes to plot the results
        """
        Treg = self.setup_TreeRegression()
        if self.clock_model and 'cov' in self.clock_model:
            cf = self.clock_model['valid_confidence']
        else:
            cf = False
        Treg.clock_plot(ax=ax, add_internal=add_internal, confidence=cf, n_sigma=2,
                        regression=self.clock_model)

    def reroot(self, root='least-squares', force_positive=True, covariation=None):
        """
        Find best root and re-root the tree to the new root

        Parameters
        ----------

         root : str
            Which method should be used to find the best root. Available methods are:

            :code:`best`, `least-squares` - minimize squared residual or likelihood of root-to-tip regression

            :code:`min_dev` - minimize variation of root-to-tip distance

            :code:`oldest` - reroot on the oldest node

            :code:`<node_name>` - reroot to the node with name :code:`<node_name>`

            :code:`[<node_name1>, <node_name2>, ...]` - reroot to the MRCA of these nodes

          force_positive : bool
            only consider positive rates when searching for the optimal root

          covariation : bool
             account for covariation in root-to-tip regression
        """
        if type(root) is list and len(root)==1:
            root=str(root[0])

        if root=='best':
            root='least-squares'

        use_cov = self.use_covariation if covariation is None else covariation

        self.logger("TreeTime.reroot: with method or node: %s"%root,0)
        for n in self.tree.find_clades():
            n.branch_length=n.mutation_length

        if (type(root) is str) and \
           (root in rerooting_mechanisms or root in deprecated_rerooting_mechanisms):
            if root in deprecated_rerooting_mechanisms:
                if "ML" in root:
                    use_cov=True
                self.logger('TreeTime.reroot: rerooting mechanisms %s has been renamed to %s'
                             %(root, deprecated_rerooting_mechanisms[root]), 1, warn=True)
                root = deprecated_rerooting_mechanisms[root]

            self.logger("TreeTime.reroot: rerooting will %s covariance and shared ancestry."%("account for" if use_cov else "ignore"),0)
            new_root = self._find_best_root(covariation=use_cov,
                                            slope = 0.0 if root.startswith('min_dev') else None,
                                            force_positive=force_positive and (not root.startswith('min_dev')))
        else:
            if isinstance(root,Phylo.BaseTree.Clade):
                new_root = root
            elif isinstance(root, list):
                new_root = self.tree.common_ancestor(*root)
            elif root in self._leaves_lookup:
                new_root = self._leaves_lookup[root]
            elif root=='oldest':
                new_root = sorted([n for n in self.tree.get_terminals()
                                   if n.raw_date_constraint is not None],
                                   key=lambda x:np.mean(x.raw_date_constraint))[0]
            else:
                self.logger('TreeTime.reroot -- ERROR: unsupported rooting mechanisms or root not found',0,warn=True)
                return ttconf.ERROR

            #this forces a bifurcating root, as we want. Branch lengths will be reoptimized anyway.
            #(Without outgroup_branch_length, gives a trifurcating root, but this will mean
            #mutations may have to occur multiple times.)
            self.tree.root_with_outgroup(new_root, outgroup_branch_length=new_root.branch_length/2)
            self.get_clock_model(covariation=use_cov)

        if new_root == ttconf.ERROR:
            return ttconf.ERROR

        self.logger("TreeTime.reroot: Tree was re-rooted to node "
                    +('new_node' if new_root.name is None else new_root.name), 2)

        self.tree.root.branch_length = self.one_mutation
        self.tree.root.clock_length = self.one_mutation
        self.tree.root.raw_date_constraint = None
        if hasattr(new_root, 'time_before_present'):
            self.tree.root.time_before_present = new_root.time_before_present
        if hasattr(new_root, 'numdate'):
            self.tree.root.numdate = new_root.numdate
        # set root.gamma bc root doesn't have a branch_length_interpolator but gamma is needed
        if not hasattr(self.tree.root, 'gamma'):
            self.tree.root.gamma = 1.0
        for n in self.tree.find_clades():
            n.mutation_length = n.branch_length
            if not hasattr(n, 'clock_length'):
                n.clock_length = n.branch_length
        self.prepare_tree()

        self.get_clock_model(covariation=self.use_covariation)

        return ttconf.SUCCESS

    def resolve_polytomies(self, merge_compressed=False):
        """
        Resolve the polytomies on the tree.

        The function scans the tree, resolves polytomies if present,
        and re-optimizes the tree with new topology. Note that polytomies are only
        resolved if that would result in higher likelihood. Sometimes, stretching
        two or more branches that carry several mutations is less costly than
        an additional branch with zero mutations (long branches are not stiff,
        short branches are).

        Parameters
        ----------
         merge_compressed : bool
            If True, keep compressed branches as polytomies. If False,
            return a strictly binary tree.

        Returns
        --------
         poly_found : int
            The number of polytomies found

        """
        self.logger("TreeTime.resolve_polytomies: resolving multiple mergers...",1)
        poly_found=0

        for n in self.tree.find_clades():
            if len(n.clades) > 2:
                prior_n_clades = len(n.clades)
                self._poly(n, merge_compressed)
                poly_found+=prior_n_clades - len(n.clades)

        obsolete_nodes = [n for n in self.tree.find_clades() if len(n.clades)==1 and n.up is not None]
        for node in obsolete_nodes:
            self.logger('TreeTime.resolve_polytomies: remove obsolete node '+node.name,4)
            if node.up is not None:
                self.tree.collapse(node)

        if poly_found:
            self.logger('TreeTime.resolve_polytomies: introduces %d new nodes'%poly_found,3)
        else:
            self.logger('TreeTime.resolve_polytomies: No more polytomies to resolve',3)
        return poly_found

    def _poly(self, clade, merge_compressed):

        """
        Function to resolve polytomies for a given parent node. If the
        number of the direct decendants is less than three (not a polytomy), does
        nothing. Otherwise, for each pair of nodes, assess the possible LH increase
        which could be gained by merging the two nodes. The increase in the LH is
        basically the tradeoff between the gain of the LH due to the changing the
        branch lenghts towards the optimal values and the decrease due to the
        introduction of the new branch with zero optimal length.
        """

        from .branch_len_interpolator import BranchLenInterpolator

        zero_branch_slope = self.gtr.mu*self.seq_len

        def _c_gain(t, n1, n2, parent):
            """
            cost gain if nodes n1, n2 are joined and their parent is placed at time t
            cost gain = (LH loss now) - (LH loss when placed at time t)
            """
            cg2 = n2.branch_length_interpolator(parent.time_before_present - n2.time_before_present) - n2.branch_length_interpolator(t - n2.time_before_present)
            cg1 = n1.branch_length_interpolator(parent.time_before_present - n1.time_before_present) - n1.branch_length_interpolator(t - n1.time_before_present)
            cg_new = - zero_branch_slope * (parent.time_before_present - t) # loss in LH due to the new branch
            return -(cg2+cg1+cg_new)

        def cost_gain(n1, n2, parent):
            """
            cost gained if the two nodes would have been connected.
            """
            try:
                cg = sciopt.minimize_scalar(_c_gain,
                    bounds=[max(n1.time_before_present,n2.time_before_present), parent.time_before_present],
                    method='Bounded',args=(n1,n2, parent))
                return cg['x'], - cg['fun']
            except:
                self.logger("TreeTime._poly.cost_gain: optimization of gain failed", 3, warn=True)
                return parent.time_before_present, 0.0

        def merge_nodes(source_arr, isall=False):
            mergers = np.array([[cost_gain(n1,n2, clade) if i1<i2 else (0.0,-1.0)
                                    for i1,n1 in enumerate(source_arr)]
                                for i2, n2 in enumerate(source_arr)])
            LH = 0
            while len(source_arr) > 1 + int(isall):
                # max possible gains of the cost when connecting the nodes:
                # this is only a rough approximation because it assumes the new node positions
                # to be optimal
                new_positions = mergers[:,:,0]
                cost_gains = mergers[:,:,1]
                # set zero to large negative value and find optimal pair
                np.fill_diagonal(cost_gains, -1e11)
                idxs = np.unravel_index(cost_gains.argmax(),cost_gains.shape)
                if (idxs[0] == idxs[1]) or cost_gains.max()<0:
                    self.logger("TreeTime._poly.merge_nodes: node is not fully resolved "+clade.name,4)
                    return LH

                n1, n2 = source_arr[idxs[0]], source_arr[idxs[1]]
                LH += cost_gains[idxs]

                new_node = Phylo.BaseTree.Clade()

                # fix positions and branch lengths
                new_node.time_before_present = new_positions[idxs]
                new_node.branch_length = clade.time_before_present - new_node.time_before_present
                new_node.clades = [n1,n2]
                n1.branch_length = new_node.time_before_present - n1.time_before_present
                n2.branch_length = new_node.time_before_present - n2.time_before_present

                # set parameters for the new node
                new_node.up = clade
                n1.up = new_node
                n2.up = new_node
                if hasattr(clade, "cseq"):
                    new_node.cseq = clade.cseq
                    self._store_compressed_sequence_to_node(new_node)

                new_node.mutations = []
                new_node.mutation_length = 0.0
                new_node.branch_length_interpolator = BranchLenInterpolator(new_node, self.gtr, one_mutation=self.one_mutation,
                                                                            branch_length_mode = self.branch_length_mode)
                clade.clades.remove(n1)
                clade.clades.remove(n2)
                clade.clades.append(new_node)
                self.logger('TreeTime._poly.merge_nodes: creating new node as child of '+clade.name,3)
                self.logger("TreeTime._poly.merge_nodes: Delta-LH = " + str(cost_gains[idxs].round(3)), 3)

                # and modify source_arr array for the next loop
                if len(source_arr)>2: # if more than 3 nodes in polytomy, replace row/column
                    for ii in np.sort(idxs)[::-1]:
                        tmp_ind = np.arange(mergers.shape[0])!=ii
                        mergers = mergers[tmp_ind].swapaxes(0,1)
                        mergers = mergers[tmp_ind].swapaxes(0,1)

                    source_arr.remove(n1)
                    source_arr.remove(n2)
                    new_gains = np.array([[cost_gain(n1,new_node, clade) for n1 in source_arr]])
                    mergers = np.vstack((mergers, new_gains)).swapaxes(0,1)

                    source_arr.append(new_node)
                    new_gains = np.array([[cost_gain(n1,new_node, clade) for n1 in source_arr]])
                    mergers = np.vstack((mergers, new_gains)).swapaxes(0,1)
                else: # otherwise just recalculate matrix
                    source_arr.remove(n1)
                    source_arr.remove(n2)
                    source_arr.append(new_node)
                    mergers = np.array([[cost_gain(n1,n2, clade) for n1 in source_arr]
                                       for n2 in source_arr])

            return LH

        stretched = [c for c  in clade.clades if c.mutation_length < c.clock_length]
        compressed = [c for c in clade.clades if c not in stretched]

        if len(stretched)==1 and merge_compressed is False:
            return 0.0

        LH = merge_nodes(stretched, isall=len(stretched)==len(clade.clades))
        if merge_compressed and len(compressed)>1:
            LH += merge_nodes(compressed, isall=len(compressed)==len(clade.clades))

        return LH

    def print_lh(self, joint=True):
        """
        Print the total likelihood of the tree given the constrained leaves

        Parameters
        ----------

         joint : bool
            If true, print joint LH, else print marginal LH

        """
        try:
            u_lh = self.tree.unconstrained_sequence_LH
            if joint:
                s_lh = self.tree.sequence_joint_LH
                t_lh = self.tree.positional_joint_LH
                c_lh = self.tree.coalescent_joint_LH
            else:
                s_lh = self.tree.sequence_marginal_LH
                t_lh = self.tree.positional_marginal_LH
                c_lh = 0

            print ("###  Tree Log-Likelihood  ###\n"
                " Sequence log-LH without constraints: \t%1.3f\n"
                " Sequence log-LH with constraints:    \t%1.3f\n"
                " TreeTime sequence log-LH:            \t%1.3f\n"
                " Coalescent log-LH:                   \t%1.3f\n"
               "#########################"%(u_lh, s_lh,t_lh, c_lh))
        except:
            print("ERROR. Did you run the corresponding inference (joint/marginal)?")

    def add_coalescent_model(self, Tc, **kwargs):
        """Add a coalescent model to the tree and optionally optimze

        Parameters
        ----------
        Tc : float,str
            If this is a float, it will be interpreted as the inverse merger
            rate in molecular clock units, if its is a
        """
        from .merger_models import Coalescent
        self.logger('TreeTime.run: adding coalescent prior with Tc='+str(Tc),1)
        self.merger_model = Coalescent(self.tree,
                                       date2dist=self.date2dist, logger=self.logger)

        if Tc=='skyline': # restrict skyline model optimization to last iteration
            self.merger_model.optimize_skyline(**kwargs)
            self.logger("optimized a skyline ", 2)
        else:
            if Tc in ['opt', 'const']:
                self.merger_model.optimize_Tc()
                self.logger("optimized Tc to %f"%self.merger_model.Tc.y[0], 2)
            else:
                try:
                    self.merger_model.set_Tc(Tc)
                except:
                    self.logger("setting of coalescent time scale failed", 1, warn=True)

        self.merger_model.attach_to_tree()

    def relaxed_clock(self, slack=None, coupling=None, **kwargs):
        """
        Allow the mutation rate to vary on the tree (relaxed molecular clock).
        Changes of the mutation rates from one branch to another are penalized.
        In addition, deviation of the mutation rate from the mean rate is
        penalized.

        Parameters
        ----------
         slack : float
            Maximum change in substitution rate between parent and child nodes

         coupling : float
            Maximum difference in substitution rates in sibling nodes

        """
        if slack is None: slack=ttconf.MU_ALPHA
        if coupling is None: coupling=ttconf.MU_BETA
        self.logger("TreeTime.relaxed_clock: slack=%f, coupling=%f"%(slack, coupling),2)

        c=1.0/self.one_mutation
        for node in self.tree.find_clades(order='postorder'):
            opt_len = node.mutation_length
            act_len = node.clock_length if hasattr(node, 'clock_length') else node.branch_length

            # opt_len \approx 1.0*len(node.mutations)/node.profile.shape[0] but calculated via gtr model
            # stiffness is the expectation of the inverse variance of branch length (one_mutation/opt_len)
            # contact term: stiffness*(g*bl - bl_opt)^2 + slack(g-1)^2 =
            #               (slack+bl^2) g^2 - 2 (bl*bl_opt+1) g + C= k2 g^2 + k1 g + C
            node._k2 = slack + c*act_len**2/(opt_len+self.one_mutation)
            node._k1 = -2*(c*act_len*opt_len/(opt_len+self.one_mutation) + slack)
            # coupling term: \sum_c coupling*(g-g_c)^2 + Cost_c(g_c|g)
            # given g, g_c needs to be optimal-> 2*coupling*(g-g_c) = 2*child.k2 g_c  + child.k1
            # hence g_c = (coupling*g - 0.5*child.k1)/(coupling+child.k2)
            # substituting yields
            for child in node.clades:
                denom = coupling+child._k2
                node._k2 += coupling*(1.0-coupling/denom)**2 + child._k2*coupling**2/denom**2
                node._k1 += (coupling*(1.0-coupling/denom)*child._k1/denom \
                            - coupling*child._k1*child._k2/denom**2 \
                            + coupling*child._k1/denom)

        for node in self.tree.find_clades(order='preorder'):
            if node.up is None:
                node.gamma = max(0.1, -0.5*node._k1/node._k2)
            else:
                if node.up.up is None:
                    g_up = node.up.gamma
                else:
                    g_up = node.up.branch_length_interpolator.gamma
                node.branch_length_interpolator.gamma = max(0.1,(coupling*g_up - 0.5*node._k1)/(coupling+node._k2))

###############################################################################
### rerooting
###############################################################################

    def _find_best_root(self, covariation=True, force_positive=True, slope=0, **kwarks):
        '''
        Determine the node that, when the tree is rooted on this node, results
        in the best regression of temporal constraints and root to tip distances.

        Parameters
        ----------

         infer_gtr : bool
            If True, infer new GTR model after re-root

         covariation : bool
            account for covariation structure when rerooting the tree

         force_positive : bool
            only accept positive evolutionary rate estimates when rerooting the tree

        '''
        for n in self.tree.find_clades():
            n.branch_length=n.mutation_length
        self.logger("TreeTime._find_best_root: searching for the best root position...",2)
        Treg = self.setup_TreeRegression(covariation=covariation)
        return Treg.optimal_reroot(force_positive=force_positive, slope=slope)['node']

def plot_vs_years(tt, step = None, ax=None, confidence=None, ticks=True, **kwargs):
    '''
    Converts branch length to years and plots the time tree on a time axis.

    Parameters
    ----------
     tt : TreeTime object
        A TreeTime instance after a time tree is inferred

     step : int
        Width of shaded boxes indicating blocks of years. Will be inferred if not specified.
        To switch off drawing of boxes, set to 0

     ax : matplotlib axes
        Axes to be used to plot, will create new axis if None

     confidence : tuple, float
        Draw confidence intervals. This assumes that marginal time tree inference was run.
        Confidence intervals are either specified as an interval of the posterior distribution
        like (0.05, 0.95) or as the weight of the maximal posterior region , e.g. 0.9

     **kwargs : dict
        Key word arguments that are passed down to Phylo.draw

    '''
    import matplotlib.pyplot as plt
    tt.branch_length_to_years()
    nleafs = tt.tree.count_terminals()

    if ax is None:
        fig = plt.figure(figsize=(12,10))
        ax = plt.subplot(111)
    else:
        fig = None
    # draw tree
    if "label_func" not in kwargs:
        kwargs["label_func"] = lambda x:x.name if (x.is_terminal() and nleafs<30) else ""
    Phylo.draw(tt.tree, axes=ax, **kwargs)

    offset = tt.tree.root.numdate - tt.tree.root.branch_length
    date_range = np.max([n.numdate for n in tt.tree.get_terminals()])-offset

    # estimate year intervals if not explicitly specified
    if step is None or (step>0 and date_range/step>100):
        step = 10**np.floor(np.log10(date_range))
        if date_range/step<2:
            step/=5
        elif date_range/step<5:
            step/=2
        step = max(1.0/12,step)

    # set axis labels
    if step:
        dtick = step
        min_tick = step*(offset//step)
        extra = dtick if dtick<date_range else dtick
        tick_vals = np.arange(min_tick, min_tick+date_range+extra, dtick)
        xticks = tick_vals - offset
    else:
        xticks = ax.get_xticks()
        dtick = xticks[1]-xticks[0]
        shift = offset - dtick*(offset//dtick)
        xticks -= shift
        tick_vals = [x+offset-shift for x in xticks]

    ax.set_xticks(xticks)
    ax.set_xticklabels(map(str, tick_vals))
    ax.set_xlabel('year')
    ax.set_ylabel('')
    ax.set_xlim((0,date_range))

    # put shaded boxes to delineate years
    if step:
        ylim = ax.get_ylim()
        xlim = ax.get_xlim()
        from matplotlib.patches import Rectangle
        for yi,year in enumerate(np.arange(np.floor(tick_vals[0]), tick_vals[-1]+.01, step)):
            pos = year - offset
            r = Rectangle((pos, ylim[1]-5),
                          step, ylim[0]-ylim[1]+10,
                          facecolor=[0.7+0.1*(1+yi%2)] * 3,
                          edgecolor=[1,1,1])
            ax.add_patch(r)
            if year in tick_vals and pos>=xlim[0] and pos<=xlim[1] and ticks:
                label_str = str(step*(year//step)) if step<1 else  str(int(year))
                ax.text(pos,ylim[0]-0.04*(ylim[1]-ylim[0]), label_str,
                        horizontalalignment='center')
        ax.set_axis_off()

    # add confidence intervals to the tree graph -- grey bars
    if confidence:
        tree_layout(tt.tree)
        if not hasattr(tt.tree.root, "marginal_inverse_cdf"):
            print("marginal time tree reconstruction required for confidence intervals")
            return ttconf.ERROR
        elif type(confidence) is float:
            cfunc = tt.get_max_posterior_region
        elif len(confidence)==2:
            cfunc = tt.get_confidence_interval
        else:
            print("confidence needs to be either a float (for max posterior region) or a two numbers specifying lower and upper bounds")
            return ttconf.ERROR

        for n in tt.tree.find_clades():
            pos = cfunc(n, confidence)
            ax.plot(pos-offset, np.ones(len(pos))*n.ypos, lw=3, c=(0.5,0.5,0.5))
    return fig, ax

def treetime_to_newick(tt, outf):
    Phylo.write(tt.tree, outf, 'newick')

if __name__=="__main__":
    pass