import numpy as np
from Bio import Phylo

tavgii, davgii, tsqii, dtavgii, dsqii, sii = 0,1,2,3,4,5

def base_regression(Q, slope=None):
    """
    this function calculates the regression coefficients for a
    given vector containing the averages of tip and branch
    quantities.

    Parameters
    ----------
    Q : numpy.array
        vector with
    slope : None, optional
        Description

    Returns
    -------
    TYPE
        Description
    """
    if slope is None:
        slope = (Q[dtavgii] - Q[tavgii]*Q[davgii]/Q[sii]) \
                /(Q[tsqii] - Q[tavgii]**2/Q[sii])
        only_intercept=False
    else:
        only_intercept=True

    intercept = (Q[davgii] - Q[tavgii]*slope)/Q[sii]

    if only_intercept:
        return {'slope':slope, 'intercept':intercept,
                'chisq': 0.5*(Q[dsqii]/Q[sii] - Q[davgii]**2/Q[sii]**2)}

    chisq = 0.5*(Q[dsqii] - Q[davgii]**2/Q[sii]
                - (Q[dtavgii] - Q[davgii]*Q[tavgii]/Q[sii])**2/(Q[tsqii]
                - Q[tavgii]**2/Q[sii]))

    estimator_hessian = np.array([[Q[tsqii], Q[tavgii]], [Q[tavgii], Q[sii]]])

    return {'slope':slope, 'intercept':intercept,
            'chisq':chisq, 'hessian':estimator_hessian,
            'cov':np.linalg.inv(estimator_hessian)}

class TreeRegression(object):
    """TreeRegression
    This class implements an efficient regression method
    for quantity associated with tips and one that changes
    in an additive manner along the branches of the tree,
    e.g. the distance to the root. This implemented
    algorithm take into account the correlation structure
    of the data under the assumptions that variance increase
    linearly along branches as well.
    """
    def __init__(self, tree_in, tip_value = None,
                 branch_value = None, branch_variance = None):
        """
        Parameters
        ----------
         T : (Bio.Phylo.tree)
            Tree for which the covariances and regression
            are to be calculated.

         tip_value : (callable)
            function that for each tip returns the value to
            be used in the regression.

         branch_value : (callable)
            function that for each node of the tree returns
            the contribution of this branch to the value of
            the subtending tips.

         variance_function : (callable)
            function that for each node of the tree returns
            the accumulated variance
        """
        super(TreeRegression, self).__init__()
        self.tree = tree_in
        # prep tree
        for li, l in enumerate(self.tree.get_terminals()):
            l._ii = np.array([li])
        total_bl = 0
        for n in self.tree.get_nonterminals(order='postorder'):
            n._ii = np.concatenate([c._ii for c in n])
            n._ii.sort()
            for c in n:
                c.up=n
                total_bl+=c.branch_length
        self.tree.root.up=None
        self.N = self.tree.root._ii.shape[0]

        if tip_value is None:
            self.tip_value = lambda x:np.mean(x.numdate) if x.is_terminal() else None
        else:
            self.tip_value = tip_value

        if branch_value is None:
            self.branch_value = lambda x:x.branch_length
        else:
            self.branch_value = branch_value

        if branch_variance is None:
            # provide a default equal to the branch_length (Poisson) and add
            # a tenth of the average branch length to avoid numerical instabilities and division by 0.
            self.branch_variance = lambda x:x.branch_length + 0.05*total_bl/self.N
        else:
            self.branch_variance = branch_variance

    def Cov(self):
        """
        calculate the covariance matrix of the tips assuming variance
        has accumulated along branches of the tree accoriding to the
        the provided
        Returns
        -------

         M : (np.array)
            covariance matrix with tips arranged standard transersal order.
        """
        # accumulate the covariance matrix by adding 'squares'
        M = np.zeros((self.N, self.N))
        for n in self.tree.find_clades():
            if n == self.tree.root:
                continue
            M[np.meshgrid(n._ii, n._ii)] += self.branch_variance(n)
        return M

    def CovInv(self):
        """
        Inverse of the covariance matrix

        Returns
        -------

         H : (np.array)
            inverse of the covariance matrix.
        """
        self.recurse(full_matrix=True)
        return self.tree.root.cinv

    def recurse(self, full_matrix=False):
        """
        recursion to calculate inverse covariance matrix

        Parameters
        ----------
        full_matrix : bool, optional
            if True, the entire inverse matrix is calculated. otherwise, only the weighing vector.
        """
        for n in self.tree.get_nonterminals(order='postorder'):
            n_leaves = len(n._ii)
            if full_matrix: M = np.zeros((n_leaves, n_leaves), dtype=float)
            r = np.zeros(n_leaves, dtype=float)
            c_count = 0
            for c in n:
                ssq = self.branch_variance(c)
                nc = len(c._ii)
                if c.is_terminal():
                    if full_matrix:
                        M[c_count, c_count] = 1.0/ssq
                    r[c_count] = 1.0/ssq
                else:
                    if full_matrix:
                        M[c_count:c_count+nc, c_count:c_count+nc] = c.cinv - ssq*np.outer(c.r,c.r)/(1+ssq*c.s)
                    r[c_count:c_count+nc] = c.r/(1+ssq*c.s)
                c_count += nc

            if full_matrix: n.cinv = M
            n.r = r #M.sum(axis=1)
            n.s = n.r.sum()

    def _calculate_averages(self):
        """
        calculate the weighted sums of the tip and branch values and
        their second moments.
        """
        for n in self.tree.get_nonterminals(order='postorder'):
            Q = np.zeros(6, dtype=float)
            for c in n:
                tv = self.tip_value(c)
                bv = self.branch_value(c)
                var = self.branch_variance(c)
                Q += self.propagate_averages(c, tv, bv, var)
            n.Q=Q

        for n in self.tree.find_clades(order='preorder'):
            O = np.zeros(6, dtype=float)
            if n==self.tree.root:
                n.Qtot = n.Q
                continue

            for c in n.up:
                if c==n:
                    continue

                tv = self.tip_value(c)
                bv = self.branch_value(c)
                var = self.branch_variance(c)
                O += self.propagate_averages(c, tv, bv, var)

            if n.up!=self.tree.root:
                c = n.up
                tv = self.tip_value(c)
                bv = self.branch_value(c)
                var = self.branch_variance(c)
                O += self.propagate_averages(c, tv, bv, var, outgroup=True)
            n.O = O

            if not n.is_terminal():
                tv = self.tip_value(n)
                bv = self.branch_value(n)
                var = self.branch_variance(n)
                n.Qtot = n.Q + self.propagate_averages(n, tv, bv, var, outgroup=True)

    def propagate_averages(self, n, tv, bv, var, outgroup=False):
        """
        This function implements the propagation of the means,
        variance, and covariances along a branch. It operates
        both towards the root and tips.

        Parameters
        ----------
         n : (node)
            the branch connecting this node to its parent is used
            for propagation
         tv : (float)
            tip value. Only required if not is terminal
         bl : (float)
            branch value. The increment of the tree associated quantity'

         var : (float)
            the variance increment along the branch

        Returns
        -------
         Q : (np.array)
            a vector of length 6 containing the updated quantities
        """
        if n.is_terminal() and outgroup==False:
            if tv is None or np.isinf(tv) or np.isnan(tv):
                res = np.array([0, 0, 0, 0, 0, 0])
            elif var==0:
                res = np.array([np.inf, np.inf, np.inf, np.inf, np.inf, np.inf])
            else:
                res = np.array([
                    tv/var,
                    bv/var,
                    tv**2/var,
                    bv*tv/var,
                    bv**2/var,
                    1.0/var], dtype=float)
        else:
            tmpQ = n.O if outgroup else n.Q
            denom = 1.0/(1+var*tmpQ[sii])
            res = np.array([
                tmpQ[tavgii]*denom,
                (tmpQ[davgii] + bv*tmpQ[sii])*denom,
                tmpQ[tsqii] - var*tmpQ[tavgii]**2*denom,
                tmpQ[dtavgii] + tmpQ[tavgii]*bv - var*tmpQ[tavgii]*(tmpQ[davgii] + bv*tmpQ[sii])*denom,
                tmpQ[dsqii] + 2*bv*tmpQ[davgii] + bv**2*tmpQ[sii] - var*(tmpQ[davgii]**2 + 2*bv*tmpQ[davgii]*tmpQ[sii] + bv**2*tmpQ[sii]**2)*denom,
                tmpQ[sii]*denom]
            )

        return res

    def explained_variance(self):
        """calculate standard explained variance

        Returns
        -------
        float
            r-value of the root-to-tip distance and time.
            independent of regression model, but dependent on root choice
        """
        self.tree.root._v=0
        for n in self.tree.get_nonterminals(order='preorder'):
            for c in n:
                c._v = n._v + self.branch_value(c)
        raw = np.array([(self.tip_value(n), n._v) for n in self.tree.get_terminals()
                         if self.tip_value(n) is not None])
        return np.corrcoef(raw.T)[0,1]

    def regression(self, slope=None):
        """regress tip values against branch values

        Parameters
        ----------
        slope : None, optional
            if given, the slope isn't optimized

        Returns
        -------
        dict
            regression parameters
        """
        self._calculate_averages()

        clock_model = base_regression(self.tree.root.Q, slope)
        clock_model['r_val'] = self.explained_variance()

        return clock_model

    def find_best_root(self, force_positive=True, slope=None):
        """
        determine the position on the tree that minimizes the bilinear
        product of the inverse covariance and the data vectors.

        Returns
        -------
         best_root : (dict)
            dictionary with the node, the fraction `x` at which the branch
            is to be split, and the regression parameters
        """
        self._calculate_averages()
        best_root = {"chisq": np.inf}
        for n in self.tree.find_clades():
            if n==self.tree.root:
                continue

            tv = self.tip_value(n)
            bv = self.branch_value(n)
            var = self.branch_variance(n)
            x, chisq = self._optimal_root_along_branch(n, tv, bv, var, slope=slope)

            if (chisq<best_root["chisq"]):
                tmpQ = self.propagate_averages(n, tv, bv*x, var*x) \
                     + self.propagate_averages(n, tv, bv*(1-x), var*(1-x), outgroup=True)
                reg = base_regression(tmpQ, slope=slope)
                if reg["slope"]>=0 or (force_positive==False):
                    best_root = {"node":n, "split":x}
                    best_root.update(reg)

        if 'node' not in best_root:
            print("TreeRegression.find_best_root: No valid root found!", force_positive)
            return None

        if 'hessian' in best_root:
            # calculate differentials with respect to x
            deriv = []
            n = best_root["node"]
            tv = self.tip_value(n)
            bv = self.branch_value(n)
            var = self.branch_variance(n)
            for dx in [-0.001, 0.001]:
                y = min(1.0, max(0.0, best_root["split"]+dx))
                tmpQ = self.propagate_averages(n, tv, bv*y, var*y) \
                     + self.propagate_averages(n, tv, bv*(1-y), var*(1-y), outgroup=True)
                reg = base_regression(tmpQ, slope=slope)
                deriv.append([y,reg['chisq'], tmpQ[tavgii], tmpQ[davgii]])

            estimator_hessian = np.zeros((3,3))
            estimator_hessian[:2,:2] = best_root['hessian']
            estimator_hessian[2,2] = (deriv[0][1] + deriv[1][1] - 2.0*best_root['chisq'])/(deriv[0][0] - deriv[1][0])**2
            # estimator_hessian[2,0] = (deriv[0][2] - deriv[1][2])/(deriv[0][0] - deriv[1][0])
            # estimator_hessian[2,1] = (deriv[0][3] - deriv[1][3])/(deriv[0][0] - deriv[1][0])
            estimator_hessian[0,2] = estimator_hessian[2,0]
            estimator_hessian[1,2] = estimator_hessian[2,1]
            best_root['hessian'] = estimator_hessian
            best_root['cov'] = np.linalg.inv(estimator_hessian)

        return best_root

    def _optimal_root_along_branch(self, n, tv, bv, var, slope=None):
        from scipy.optimize import minimize_scalar
        def chisq(x):
            tmpQ = self.propagate_averages(n, tv, bv*x, var*x) \
                 + self.propagate_averages(n, tv, bv*(1-x), var*(1-x), outgroup=True)
            return base_regression(tmpQ, slope=slope)['chisq']

        chisq_prox = np.inf if n.is_terminal() else base_regression(n.Qtot)['chisq']
        chisq_dist = np.inf if n==self.tree.root else base_regression(n.up.Qtot)['chisq']

        grid = np.linspace(0.001,0.999,6)
        chisq_grid = np.array([chisq(x) for x in grid])
        min_chisq = chisq_grid.min()
        if chisq_prox<=min_chisq:
            return 0.0, chisq_prox
        elif chisq_dist<=min_chisq:
            return 1.0, chisq_dist
        else:
            ii = np.argmin(chisq_grid)
            bounds = (0 if ii==0 else grid[ii-1], 1.0 if ii==len(grid)-1 else grid[ii+1])
            sol = minimize_scalar(chisq, bounds=bounds, method="bounded")
            if sol["success"]:
                return sol['x'], sol['fun']
            else:
                return np.nan, np.inf

    def optimal_reroot(self, force_positive=True, slope=None):
        """
        determine the best root and reroot the tree to this value.
        Note that this can change the parent child relations of the tree
        and values associated with branches rather than nodes
        (e.g. confidence) might need to be re-evaluated afterwards

        Parameters
        ----------
        force_positive : bool, optional
            if True, the search for a root will only consider positive rate estimates

        slope : float, optional
            if given, it will find the optimal root given a fixed rate. If slope==0, this
            corresponds to minimal root-to-tip variance rooting (min_dev)

        Returns
        -------
        dict
            regression parameters
        """
        best_root = self.find_best_root(force_positive=force_positive, slope=slope)
        best_node = best_root["node"]

        x = best_root["split"]
        if x<1e-5:
            new_node = best_node
        elif x>1.0-1e-5:
            new_node = best_node.up
        else:
            # create new node in the branch and root the tree to it
            new_node = Phylo.BaseTree.Clade()

            # insert the new node in the middle of the branch
            # by simple re-wiring the links on the both sides of the branch
            # and fix the branch lengths
            new_node.branch_length = best_node.branch_length*(1-x)
            new_node.up = best_node.up
            new_node.clades = [best_node]
            new_node.up.clades = [k if k!=best_node else new_node
                                  for k in best_node.up.clades]

            best_node.branch_length *= x
            best_node.up = new_node

        new_node.rtt_regression = best_root
        self.tree.root_with_outgroup(new_node)

        self.tree.ladderize()
        for n in self.tree.get_nonterminals(order='postorder'):
            for c in n:
                c.up=n
        return best_root

    def clock_plot(self, add_internal=False, ax=None, regression=None,
                   confidence=True, n_sigma = 2, fs=14):
        """Plot root-to-tip distance vs time as a basic time-tree diagnostic

        Parameters
        ----------
        add_internal : bool, optional
            add internal nodes. this will only work if the tree has been dated already
        ax : None, optional
            an matplotlib axis to plot into. if non provided, a new figure is opened
        regression : None, optional
            a dict containing parameters of a root-to-tip vs time regression as
            returned by the function base_regression
        confidence : bool, optional
            add confidence area to the regression line
        n_sigma : int, optional
            number of standard deviations for the confidence area.
        fs : int, optional
            fontsize

        """
        import matplotlib.pyplot as plt
        if ax is None:
            plt.figure()
            ax=plt.subplot(111)

        self.tree.root._v=0
        for n in self.tree.get_nonterminals(order='preorder'):
            for c in n:
                c._v = n._v + self.branch_value(c)

        tips = self.tree.get_terminals()
        internal = self.tree.get_nonterminals()

        # get values of terminals
        xi = np.array([self.tip_value(n) for n in tips])
        yi = np.array([n._v for n in tips])
        ind = np.array([n.bad_branch  if hasattr(n, 'bad_branch') else False for n in tips])
        if add_internal:
            xi_int = np.array([n.numdate for n in internal])
            yi_int = np.array([n._v for n in internal])
            ind_int = np.array([n.bad_branch  if hasattr(n, 'bad_branch') else False  for n in internal])

        if regression:
            # plot regression line
            t_mrca = -regression['intercept']/regression['slope']
            if add_internal:
                time_span = np.max(xi_int[~ind_int]) - np.min(xi_int[~ind_int])
                x_vals = np.array([max(np.min(xi_int[~ind_int]), t_mrca) - 0.1*time_span, np.max(xi[~ind])+0.05*time_span])
            else:
                time_span = np.max(xi[~ind]) - np.min(xi[~ind])
                x_vals = np.array([max(np.min(xi[~ind]), t_mrca) - 0.1*time_span, np.max(xi[~ind]+0.05*time_span)])

            # plot confidence interval
            if confidence and 'cov' in regression:
                x_vals = np.linspace(x_vals[0], x_vals[1], 100)
                y_vals = regression['slope']*x_vals + regression['intercept']
                dev = n_sigma*np.array([np.sqrt(regression['cov'][:2,:2].dot(np.array([x, 1])).dot(np.array([x,1]))) for x in x_vals])
                dev_slope = n_sigma*np.sqrt(regression['cov'][0,0])
                ax.fill_between(x_vals, y_vals-dev, y_vals+dev, alpha=0.2)
                dp = np.array([regression['intercept']/regression['slope']**2,-1./regression['slope']])
                dev_rtt = n_sigma*np.sqrt(regression['cov'][:2,:2].dot(dp).dot(dp))

            else:
                dev_rtt = None
                dev_slope = None

            ax.plot(x_vals, regression['slope']*x_vals + regression['intercept'],
                    label = r"$y=\alpha + \beta t$"+"\n"+
                            r"$\beta=$%1.2e"%(regression["slope"])
                            + ("+/- %1.e"%dev_slope if dev_slope else "") +
                            "\nroot date: %1.1f"%(-regression['intercept']/regression['slope']) +
                            ("+/- %1.2f"%dev_rtt if dev_rtt else ""))

        ax.scatter(xi[~ind], yi[~ind], label=("tips" if add_internal else None))
        if ind.sum():
            try:
                # note: this is treetime specific
                tmp_x = np.array([np.mean(n.raw_date_constraint) if n.raw_date_constraint else None
                                  for n in self.tree.get_terminals()])
                ax.scatter(tmp_x[ind], yi[ind], label="ignored tips", c='r')
            except:
                pass
        if add_internal:
            ax.scatter(xi_int[~ind_int], yi_int[~ind_int], label="internal nodes")

        ax.set_ylabel('root-to-tip distance', fontsize=fs)
        ax.set_xlabel('date', fontsize=fs)
        ax.ticklabel_format(useOffset=False)
        ax.tick_params(labelsize=fs*0.8)
        ax.set_ylim([0, 1.1*np.max(yi)])
        plt.tight_layout()
        plt.legend(fontsize=fs*0.8)

if __name__ == '__main__':
    import matplotlib.pyplot as plt
    import time
    plt.ion()
    # tree_file = '../data/H3N2_NA_allyears_NA.20.nwk'
    # date_file = '../data/H3N2_NA_allyears_NA.20.metadata.csv'
    tree_file = '../data/ebola.nwk'
    date_file = '../data/ebola.metadata.csv'

    T = Phylo.read(tree_file, 'newick')
    #T.root_with_outgroup('A/Canterbury/58/2000|CY009150|09/05/2000|New_Zealand||H3N2/8-1416')

    dates = {}
    with open(date_file) as ifile:
        ifile.readline()
        for line in ifile:
            if line[0]!='#':
                fields = line.strip().split(',')
                dates[fields[0]] = float(fields[1])

    for l in T.get_terminals():
        l.numdate = dates[l.name]
    branch_variance = lambda x:(x.branch_length+(0.0005 if x.is_terminal() else 0.0))/19000.0
    #branch_variance = lambda x:(x.branch_length+(0.005 if x.is_terminal() else 0.0))/1700.0
    #branch_variance = lambda x:1.0 if x.is_terminal() else 0.0
    tstart = time.time()
    mtc = TreeRegression(T, branch_variance = branch_variance)
    print(time.time()-tstart)
    reg = mtc.optimal_reroot()
    print(time.time()-tstart)
    print(reg)

    plt.figure()
    ti = []
    rtt = []
    T.root.rtt=0
    for n in T.get_nonterminals(order='preorder'):
        for c in n:
            c.rtt = n.rtt + c.branch_length
    for l in T.get_terminals():
        ti.append(l.numdate)
        rtt.append(l.rtt)

    ti = np.array(ti)
    rtt = np.array(rtt)
    plt.plot(ti, rtt)
    plt.plot(ti, reg["slope"]*ti + reg["intercept"])
    plt.show()

    Phylo.draw(T)


#!/usr/local/bin/python
# -*- coding: utf-8 -*-
from __future__ import division, print_function, absolute_import
import numpy as np
from .seq_utils import alphabets, profile_maps
from .gtr import GTR

def JC69 (mu=1.0, alphabet="nuc", **kwargs):
    """
    Jukes-Cantor 1969 model. This model assumes equal concentrations
    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

    Parameters
    -----------

     mu : float
        substitution rate

     alphabet : str
        specify alphabet to use.
        Available alphabets are:

         'nuc' - nucleotides only, gaps ignored

         'nuc_gap' - nucleotide alphabet with gaps, gaps can be ignored optionally

    """
    num_chars = len(alphabets[alphabet])
    W, pi = np.ones((num_chars,num_chars)), np.ones(num_chars)
    gtr = GTR(alphabet=alphabet)
    gtr.assign_rates(mu=mu, pi=pi, W=W)
    return gtr

def K80(mu=1., kappa=0.1, **kwargs):
    """
    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.

    Parameters
    -----------

     mu : float
        Overall substitution rate

     kappa : float
        Ratio of transversion/transition rates
    """

    num_chars = len(alphabets['nuc_nogap'])
    pi = np.ones(len(alphabets['nuc_nogap']), dtype=float)/len(alphabets['nuc_nogap'])
    W = _create_transversion_transition_W(kappa)
    gtr = GTR(alphabet=alphabets['nuc_nogap'])
    gtr.assign_rates(mu=mu, pi=pi, W=W)
    return gtr

def F81(mu=1.0, pi=None, alphabet="nuc", **kwargs):
    """
    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.

    Current implementation of the model does not account for the gaps (treatment of
    gaps as characters is possible if specify alphabet='nuc_gap').

    Parameters
    -----------

     mu : float
        Substitution rate

     pi : numpy.array
        Nucleotide concentrations

     alphabet : str
        Alphabet to use. POsiible values are: ['nuc', 'nuc_gap'] Default 'nuc', which discounts al gaps.
        'nuc_gap' alphabet enables treatmen of gaps as characters.
    """

    if pi is None:
        pi=0.25*np.ones(4, dtype=float)
    num_chars = len(alphabets[alphabet])

    pi = np.array(pi, dtype=float)

    if num_chars != len(pi) :
        pi = np.ones((num_chars, ), dtype=float)
        print ("GTR: Warning!The number of the characters in the alphabet does not match the "
            "shape of the vector of equilibrium frequencies Pi -- assuming equal frequencies for all states.")

    W = np.ones((num_chars,num_chars))
    pi /= (1.0 * np.sum(pi))
    gtr = GTR(alphabet=alphabets[alphabet])
    gtr.assign_rates(mu=mu, pi=pi, W=W)
    return gtr

def HKY85(mu=1.0, pi=None, kappa=0.1, **kwargs):
    """
    Hasegawa, Kishino and Yano 1985 model. Allows different concentrations of the
    nucleotides (as in F81) + distinguishes between transition/transversionsubstitutions
    (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

    Parameters
    -----------

     mu : float
        Substitution rate

     pi : numpy.array
        Nucleotide concentrations

     kappa : float
        Ratio of transversion/transition substitution rates

    """

    if pi is None:
        pi=0.25*np.ones(4, dtype=float)
    num_chars = len(alphabets['nuc_nogap'])
    if num_chars != pi.shape[0] :
        pi = np.ones((num_chars, ), dtype=float)
        print ("GTR: Warning!The number of the characters in the alphabet does not match the "
            "shape of the vector of equilibrium frequencies Pi -- assuming equal frequencies for all states.")

    W = _create_transversion_transition_W(kappa)
    pi /= pi.sum()
    gtr = GTR(alphabet=alphabets['nuc_nogap'])
    gtr.assign_rates(mu=mu, pi=pi, W=W)
    return gtr

def T92(mu=1.0, pi_GC=0.5, kappa=0.1, **kwargs):
    """
    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

    Parameters
    -----------

     mu : float
        substitution rate

     pi_GC : float
        relative GC content

     kappa : float
        relative transversion/transition rate

    """

    W = _create_transversion_transition_W(kappa)
    # A C G T
    if pi_CG >=1.:
        raise ValueError("The relative CG content specified is larger than 1.0!")
    pi = np.array([(1.-pi_CG)*0.5, pi_CG*0.5, pi_CG*0.5, (1-pi_CG)*0.5])
    gtr = GTR(alphabet=alphabets['nuc_nogap'])
    gtr.assign_rates(mu=mu, pi=pi, W=W)
    return gtr

def TN93(mu=1.0, kappa1=1., kappa2=1., pi=None, **kwargs):
    """
    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

    Parameters
    -----------

     mu : float
        Substitution rate

     kappa1 : float
        relative A<-->C, A<-->T, T<-->G and G<-->C rates

     kappa2 : float
        relative C<-->T rate

    Note
    ----

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

    """

    if pi is None:
        pi=0.25*np.ones(4, dtype=float)
    W = np.ones((4,4))
    W = np.array([
        [1,      kappa1, 1,      kappa1],
        [kappa1, 1,      kappa1, kappa2],
        [1,      kappa1, 1,      kappa1],
        [kappa1, kappa2, kappa1,  1]], dtype=float)

    pi /=pi.sum()
    num_chars = len(alphabets['nuc_nogap'])
    if num_chars != pi.shape[0] :
        pi = np.ones((num_chars, ), dtype=float)
        print ("GTR: Warning!The number of the characters in the alphabet does not match the "
            "shape of the vector of equilibrium frequencies Pi -- assuming equal frequencies for all states.")

    gtr = GTR(alphabet=alphabets['nuc'])
    gtr.assign_rates(mu=mu, pi=pi, W=W)
    return gtr

def _create_transversion_transition_W(kappa):
    """
    Alphabet = [A, C, G, T]
    """
    W = np.ones((4,4))
    W[0, 2]=W[1, 3]=W[2, 0]=W[3,1]=kappa
    return W

if __name__ == '__main__':
    pass


"""
methods to calculate merger models for a time tree
"""
from __future__ import print_function, division, absolute_import
import numpy as np
import scipy.special as sf
from scipy.interpolate import interp1d
from Bio import AlignIO, Phylo
from scipy.interpolate import interp1d
from collections import Iterable
from treetime import config as ttconf

class Coalescent(object):
    """docstring for Coalescent"""
    def __init__(self, tree, Tc=0.001, logger=None, date2dist=None):
        super(Coalescent, self).__init__()
        self.tree = tree
        self.calc_branch_count()
        self.set_Tc(Tc)
        self.date2dist = date2dist
        if logger is None:
            def f(*args):
                print(*args)
            self.logger = f
        else:
            self.logger = logger

    def set_Tc(self, Tc, T=None):
        '''
        initialize the merger model with a coalescent time

        Args:
            - Tc:   a float or an iterable, if iterable another argument T of same shape is required
            - T:    an array like of same shape as Tc that specifies the time pivots corresponding to Tc
        Returns:
            - None
        '''
        if isinstance(Tc, Iterable):
            if len(Tc)==len(T):
                x = np.concatenate(([-ttconf.BIG_NUMBER], T, [ttconf.BIG_NUMBER]))
                y = np.concatenate(([Tc[0]], Tc, [Tc[-1]]))
                self.Tc = interp1d(x,y)
            else:
                self.logger("need Tc values and Timepoints of equal length",2,warn=True)
                self.Tc = interp1d([-ttconf.BIG_NUMBER, ttconf.BIG_NUMBER], [1e-5, 1e-5])
        else:
            self.Tc = interp1d([-ttconf.BIG_NUMBER, ttconf.BIG_NUMBER],
                               [Tc+ttconf.TINY_NUMBER, Tc+ttconf.TINY_NUMBER])
        self.calc_integral_merger_rate()

    def calc_branch_count(self):
        '''
        calculates an interpolation object that maps time to the number of
        concurrent branches in the tree. The result is stored in self.nbranches
        '''

        # make a list of (time, merger or loss event) by root first iteration
        self.tree_events = np.array(sorted([(n.time_before_present, len(n.clades)-1)
                                for n in self.tree.find_clades() if not n.bad_branch],
                                key=lambda x:-x[0]))

        # collapse multiple events at one time point into sum of changes
        from collections import defaultdict
        dn_branch = defaultdict(int)
        for (t, dn) in self.tree_events:
            dn_branch[t]+=dn
        unique_mergers = np.array(sorted(dn_branch.items(), key = lambda x:-x[0]))

        # calculate the branch count at each point summing the delta branch counts
        nbranches = [[ttconf.BIG_NUMBER, 1], [unique_mergers[0,0]+ttconf.TINY_NUMBER, 1]]
        for ti, (t, dn) in enumerate(unique_mergers[:-1]):
            new_n = nbranches[-1][1]+dn
            next_t = unique_mergers[ti+1,0]+ttconf.TINY_NUMBER
            nbranches.append([t, new_n])
            nbranches.append([next_t, new_n])

        new_n += unique_mergers[-1,1]
        nbranches.append([next_t, new_n])
        nbranches.append([-ttconf.BIG_NUMBER, new_n])
        nbranches=np.array(nbranches)

        self.nbranches = interp1d(nbranches[:,0], nbranches[:,1], kind='linear')

    def calc_integral_merger_rate(self):
        '''
        calculates the integral int_0^t (k(t')-1)/2Tc(t') dt' and stores it as
        self.integral_merger_rate. This differences of this quantity evaluated at
        different times points are the cost of a branch.
        '''
        # integrate the piecewise constant branch count function.
        tvals = np.unique(self.nbranches.x[1:-1])
        rate = self.branch_merger_rate(tvals)
        avg_rate = 0.5*(rate[1:] + rate[:-1])
        cost = np.concatenate(([0],np.cumsum(np.diff(tvals)*avg_rate)))
        # make interpolation objects for the branch count and its integral
        # the latter is scaled by 0.5/Tc
        # need to add extra point at very large time before present to
        # prevent 'out of interpolation range' errors
        self.integral_merger_rate = interp1d(np.concatenate(([-ttconf.BIG_NUMBER], tvals,[ttconf.BIG_NUMBER])),
                                  np.concatenate(([cost[0]], cost,[cost[-1]])), kind='linear')

    def branch_merger_rate(self, t):
        # returns the rate at which one particular branch merges with any other branch
        # note that we always have a positive merger rate by capping the
        # number of branches at 0.5 from below. in these regions, the
        # function should only be called if the tree changes.
        return 0.5*np.maximum(0.5,self.nbranches(t)-1.0)/self.Tc(t)

    def total_merger_rate(self, t):
        # returns the rate at which any branch merges with any other branch
        # not that we always have a positive merger rate by capping the
        # number of branches at 0.5 from below. in these regions, the
        # function should only be called if the tree changes.
        nlineages = np.maximum(0.5,self.nbranches(t)-1.0)
        return 0.5*nlineages*(nlineages+1)/self.Tc(t)

    def cost(self, t_node, branch_length, multiplicity=2.0):
        '''
        returns the cost associated with a branch starting at t_node
        t_node is time before present, the branch goes back in time

        Args:
            - t_node:           time of the node
            - branch_length:    branch length, determines when this branch merges with sister
            - multiplicity:     2 if merger is binary, higher if this is a polytomy
        '''
        merger_time = t_node+branch_length
        return self.integral_merger_rate(merger_time) - self.integral_merger_rate(t_node)\
                 - np.log(self.total_merger_rate(merger_time))*(multiplicity-1.0)/multiplicity

    def attach_to_tree(self):
        '''
        attaches the the merger cost to each branch length interpolator in the tree.
        '''
        for clade in self.tree.find_clades():
            if clade.up is not None:
                clade.branch_length_interpolator.merger_cost = self.cost

    def total_LH(self):
        LH = 0.0 #np.log(self.total_merger_rate([node.time_before_present for node in self.tree.get_nonterminals()])).sum()
        for node in self.tree.find_clades():
            if node.up:
                LH -= self.cost(node.time_before_present, node.branch_length)
        return LH

    def optimize_Tc(self):
        '''
        determines the coalescent time scale that optimizes the coalescent likelihood of the tree
        '''
        from scipy.optimize import minimize_scalar
        initial_Tc = self.Tc
        def cost(Tc):
            self.set_Tc(Tc)
            return -self.total_LH()

        sol = minimize_scalar(cost, bounds=[ttconf.TINY_NUMBER,10.0])
        if "success" in sol and sol["success"]:
            self.set_Tc(sol['x'])
        else:
            self.logger("merger_models:optimze_Tc: optimization of coalescent time scale failed: " + str(sol), 0, warn=True)
            self.set_Tc(initial_Tc.y, T=initial_Tc.x)

    def optimize_skyline(self, n_points=20, stiffness=2.0, method = 'SLSQP',
                         tol=0.03, regularization=10.0, **kwarks):
        '''
        optimize the trajectory of the merger rate 1./T_c to maximize the
        coalescent likelihood.
        parameters:
            n_points    --  number of pivots of the Tc interpolation object
            stiffness   --  penalty for rapid changes in log(Tc)
            methods     --  method used to optimize
            tol         --  optimization tolerance
            regularization --  cost of moving logTc outsize of the range [-100,0]
                               merger rate is measured in branch length units, no
                               plausible rates should never be outside this window
        '''
        self.logger("Coalescent:optimize_skyline:... current LH: %f"%self.total_LH(),2)
        from scipy.optimize import minimize
        initial_Tc = self.Tc
        tvals = np.linspace(self.tree_events[0,0], self.tree_events[-1,0], n_points)
        def cost(logTc):
            # cap log Tc to avoid under or overflow and nan in logs
            self.set_Tc(np.exp(np.maximum(-200,np.minimum(100,logTc))), tvals)
            neglogLH = -self.total_LH() + stiffness*np.sum(np.diff(logTc)**2) \
                       + np.sum((logTc>0)*logTc*regularization)\
                       - np.sum((logTc<-100)*logTc*regularization)
            return neglogLH

        sol = minimize(cost, np.ones_like(tvals)*np.log(self.Tc.y.mean()), method=method, tol=tol)
        if "success" in sol and sol["success"]:
            dlogTc = 0.1
            opt_logTc = sol['x']
            dcost = []
            for ii in range(len(opt_logTc)):
                tmp = opt_logTc.copy()
                tmp[ii]+=dlogTc
                cost_plus = cost(tmp)
                tmp[ii]-=2*dlogTc
                cost_minus = cost(tmp)
                dcost.append([cost_minus, cost_plus])

            dcost = np.array(dcost)
            optimal_cost = cost(opt_logTc)
            self.confidence = -dlogTc/(2*optimal_cost - dcost[:,0] - dcost[:,1])
            self.logger("Coalescent:optimize_skyline:...done. new LH: %f"%self.total_LH(),2)
        else:
            self.set_Tc(initial_Tc.y, T=initial_Tc.x)
            self.logger("Coalescent:optimize_skyline:...failed:"+str(sol),0, warn=True)

    def skyline_empirical(self, gen=1.0, n_points = 20):
        '''
        returns the skyline, i.e., an estimate of the inverse rate of coalesence.
        Here, the skyline is estimated from a sliding window average of the observed
        mergers, i.e., without reference to the coalescence likelihood.
        parameters:
            gen -- number of generations per year.
        '''

        mergers = self.tree_events[:,1]>0
        merger_tvals = self.tree_events[mergers,0]
        nlineages = self.nbranches(merger_tvals-ttconf.TINY_NUMBER)
        expected_merger_density = nlineages*(nlineages-1)*0.5

        nmergers = len(mergers)
        et = merger_tvals
        ev = 1.0/expected_merger_density
        # reduce the window size if there are few events in the tree
        if 2*n_points>len(expected_merger_density):
            n_points = len(ev)//4

        # smoothes with a sliding window over data points
        avg = np.sum(ev)/np.abs(et[0]-et[-1])
        dt = et[0]-et[-1]
        mid_points = np.concatenate(([et[0]-0.5*(et[1]-et[0])],
                                      0.5*(et[1:] + et[:-1]),
                                     [et[-1]+0.5*(et[-1]-et[-2])]))

        # this smoothes the ratio of expected and observed merger rate
        self.Tc_inv = interp1d(mid_points[n_points:-n_points],
                        [np.sum(ev[(et>=l)&(et<u)])/(u-l+dt/nmergers)
                        for u,l in zip(mid_points[:-2*n_points],mid_points[2*n_points:])])

        return interp1d(self.date2dist.to_numdate(self.Tc_inv.x), gen/self.date2dist.clock_rate/self.Tc_inv.y)

    def skyline_inferred(self, gen=1.0, confidence=False):
        '''
        return the skyline, i.e., an estimate of the inverse rate of coalesence.
        This function merely returns the merger rate self.Tc that was set or
        estimated by other means. If it was determined using self.optimize_skyline,
        the returned skyline will maximize the coalescent likelihood.
        parameters:
            gen -- number of generations per year. Unit of time is branch length,
                   hence this needs to be the inverse substitution rate per generation
            confidence -- False, or number of standard deviations of confidence intervals
        '''
        if len(self.Tc.x)<=2:
            print("no skyline has been inferred, returning constant population size")
            return gen/self.date2dist.clock_rate*self.Tc.y[-1]

        skyline = interp1d(self.date2dist.to_numdate(self.Tc.x[1:-1]), gen/self.date2dist.clock_rate*self.Tc.y[1:-1])
        if confidence and hasattr(self, 'confidence'):
            conf = [skyline.y*np.exp(-confidence*self.confidence), skyline.y*np.exp(confidence*self.confidence)]
            return skyline, conf
        else:
            return skyline


import numpy as np

alphabet_synonyms = {'nuc':'nuc', 'nucleotide':'nuc', 'aa':'aa', 'aminoacid':'aa',
                     'nuc_nogap':'nuc_nogap', 'nucleotide_nogap':'nuc_nogap',
                     'aa_nogap':'aa_nogap', 'aminoacid_nogap':'aa_nogap',
                     'DNA':'nuc', 'DNA_nogap':'nuc_nogap'}

alphabets = {
            "nuc":           np.array(['A', 'C', 'G', 'T', '-']),

            "nuc_nogap":np.array(['A', 'C', 'G', 'T']),

            "aa":            np.array(['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K',
                                       'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V',
                                       'W', 'Y', '*', '-']),

            "aa_nogap": np.array(['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K',
                                       'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V',
                                       'W', 'Y'])
            }

profile_maps = {
'nuc':{
    'A': np.array([1, 0, 0, 0, 0], dtype='float'),
    'C': np.array([0, 1, 0, 0, 0], dtype='float'),
    'G': np.array([0, 0, 1, 0, 0], dtype='float'),
    'T': np.array([0, 0, 0, 1, 0], dtype='float'),
    '-': np.array([0, 0, 0, 0, 1], dtype='float'),
    'N': np.array([1, 1, 1, 1, 1], dtype='float'),
    'X': np.array([1, 1, 1, 1, 1], dtype='float'),
    'R': np.array([1, 0, 1, 0, 0], dtype='float'),
    'Y': np.array([0, 1, 0, 1, 0], dtype='float'),
    'S': np.array([0, 1, 1, 0, 0], dtype='float'),
    'W': np.array([1, 0, 0, 1, 0], dtype='float'),
    'K': np.array([0, 0, 1, 1, 0], dtype='float'),
    'M': np.array([1, 1, 0, 0, 0], dtype='float'),
    'D': np.array([1, 0, 1, 1, 0], dtype='float'),
    'H': np.array([1, 1, 0, 1, 0], dtype='float'),
    'B': np.array([0, 1, 1, 1, 0], dtype='float'),
    'V': np.array([1, 1, 1, 0, 0], dtype='float')
    },

'nuc_nogap':{
    'A': np.array([1, 0, 0, 0], dtype='float'),
    'C': np.array([0, 1, 0, 0], dtype='float'),
    'G': np.array([0, 0, 1, 0], dtype='float'),
    'T': np.array([0, 0, 0, 1], dtype='float'),
    '-': np.array([1, 1, 1, 1], dtype='float'), # gaps are completely ignored in distance computations
    'N': np.array([1, 1, 1, 1], dtype='float'),
    'X': np.array([1, 1, 1, 1], dtype='float'),
    'R': np.array([1, 0, 1, 0], dtype='float'),
    'Y': np.array([0, 1, 0, 1], dtype='float'),
    'S': np.array([0, 1, 1, 0], dtype='float'),
    'W': np.array([1, 0, 0, 1], dtype='float'),
    'K': np.array([0, 0, 1, 1], dtype='float'),
    'M': np.array([1, 1, 0, 0], dtype='float'),
    'D': np.array([1, 0, 1, 1], dtype='float'),
    'H': np.array([1, 1, 0, 1], dtype='float'),
    'B': np.array([0, 1, 1, 1], dtype='float'),
    'V': np.array([1, 1, 1, 0], dtype='float')
    },

'aa':{
    'A': np.array([1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Alanine         Ala
    'C': np.array([0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Cysteine        Cys
    'D': np.array([0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Aspartic AciD   Asp
    'E': np.array([0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Glutamic Acid   Glu
    'F': np.array([0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Phenylalanine   Phe
    'G': np.array([0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Glycine         Gly
    'H': np.array([0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Histidine       His
    'I': np.array([0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Isoleucine      Ile
    'K': np.array([0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Lysine          Lys
    'L': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Leucine         Leu
    'M': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Methionine      Met
    'N': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #AsparagiNe      Asn
    'P': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Proline         Pro
    'Q': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Glutamine       Gln
    'R': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #ARginine        Arg
    'S': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0], dtype='float'), #Serine          Ser
    'T': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0], dtype='float'), #Threonine       Thr
    'V': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0], dtype='float'), #Valine          Val
    'W': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0], dtype='float'), #Tryptophan      Trp
    'Y': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0], dtype='float'), #Tyrosine        Tyr
    '*': np.array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 0], dtype='float'), #stop
    '-': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1], dtype='float'), #gap
    'X': np.array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype='float'), #not specified/any
    'B': np.array([0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Asparagine/Aspartic Acid    Asx
    'Z': np.array([0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Glutamine/Glutamic Acid     Glx
    },

'aa_nogap':{
    'A': np.array([1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Alanine         Ala
    'C': np.array([0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Cysteine        Cys
    'D': np.array([0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Aspartic AciD   Asp
    'E': np.array([0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Glutamic Acid   Glu
    'F': np.array([0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Phenylalanine   Phe
    'G': np.array([0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Glycine         Gly
    'H': np.array([0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Histidine       His
    'I': np.array([0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Isoleucine      Ile
    'K': np.array([0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Lysine          Lys
    'L': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Leucine         Leu
    'M': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Methionine      Met
    'N': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #AsparagiNe      Asn
    'P': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Proline         Pro
    'Q': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0], dtype='float'), #Glutamine       Gln
    'R': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0], dtype='float'), #ARginine        Arg
    'S': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0], dtype='float'), #Serine          Ser
    'T': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0], dtype='float'), #Threonine       Thr
    'V': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0], dtype='float'), #Valine          Val
    'W': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0], dtype='float'), #Tryptophan      Trp
    'Y': np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1], dtype='float'), #Tyrosine        Tyr
    'X': np.array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype='float'), #not specified/any
    'B': np.array([0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0], dtype='float'), #Asparagine/Aspartic Acid    Asx
    'Z': np.array([0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0], dtype='float'), #Glutamine/Glutamic Acid     Glx
    }
}

def seq2array(seq, fill_overhangs=True, ambiguous_character='N'):
    """
    Take the raw sequence, substitute the "overhanging" gaps with 'N' (missequenced),
    and convert the sequence to the numpy array of chars.

    Parameters
    ----------
     seq : Biopython.SeqRecord, str, iterable
        Sequence as an object of SeqRecord, string or iterable

     fill_overhangs : bool
        If True, substitute the "overhanging" gaps with ambiguous character symbol

     ambiguous_character : char
        Specify the character for ambiguous state ('N' default for nucleotide)

    Returns
    -------
     sequence : np.array
        Sequence as 1D numpy array of chars

    """
    try:
        sequence = ''.join(seq)
    except TypeError:
        sequence = seq

    sequence = np.array(list(sequence))
    # substitute overhanging unsequenced tails
    if fill_overhangs:
        sequence [:np.where(sequence != '-')[0][0]] = ambiguous_character
        sequence [np.where(sequence != '-')[0][-1]+1:] = ambiguous_character
    return sequence

def seq2prof(seq, profile_map):
    """
    Convert the given character sequence into the profile according to the
    alphabet specified.

    Parameters
    ----------

     seq : numpy.array
        Sequence to be converted to the profile

     profile_map : dic
        Mapping valid characters to profiles

    Returns
    -------

     idx : numpy.array
        Profile for the character. Zero array if the character not found

    """

    return np.array([profile_map[k] for k in seq])

def prof2seq(profile, gtr, sample_from_prof=False, normalize=True):
    """
    Convert profile to sequence and normalize profile across sites.

    Parameters
    ----------

     profile : numpy 2D array
        Profile. Shape of the profile should be (L x a), where L - sequence
        length, a - alphabet size.

     gtr : gtr.GTR
        Instance of the GTR class to supply the sequence alphabet

     collapse_prof : bool
        Whether to convert the profile to the delta-function

    Returns
    -------
     seq : numpy.array
        Sequence as numpy array of length L

     prof_values :  numpy.array
        Values of the profile for the chosen sequence characters (length L)

     idx : numpy.array
        Indices chosen from profile as array of length L
    """

    # normalize profile such that probabilities at each site sum to one
    if normalize:
        tmp_profile, pre=normalize_profile(profile, return_offset=False)
    else:
        tmp_profile = profile

    # sample sequence according to the probabilities in the profile
    # (sampling from cumulative distribution over the different states)
    if sample_from_prof:
        cumdis = tmp_profile.cumsum(axis=1).T
        randnum = np.random.random(size=cumdis.shape[1])
        idx = np.argmax(cumdis>=randnum, axis=0)
    else:
        idx = tmp_profile.argmax(axis=1)
    seq = gtr.alphabet[idx]  # max LH over the alphabet

    prof_values = tmp_profile[np.arange(tmp_profile.shape[0]), idx]

    return seq, prof_values, idx

def normalize_profile(in_profile, log=False, return_offset = True):
    """return a normalized version of a profile matrix

    Parameters
    ----------
    in_profile : np.array
        shape Lxq, will be normalized to one across each row
    log : bool, optional
        treat the input as log probabilities
    return_offset : bool, optional
        return the log of the scale factor for each row

    Returns
    -------
    tuple
        normalized profile (fresh np object) and offset (if return_offset==True)
    """
    if log:
        tmp_prefactor = in_profile.max(axis=1)
        tmp_prof = np.exp(in_profile.T - tmp_prefactor).T
    else:
        tmp_prefactor = 0.0
        tmp_prof = in_profile

    norm_vector = tmp_prof.sum(axis=1)
    return (np.copy(np.einsum('ai,a->ai',tmp_prof,1.0/norm_vector)),
            (np.log(norm_vector) + tmp_prefactor) if return_offset else None)