""" This file holds the parameters of our tHMM in the tHMM class. """

from copy import deepcopy
import numpy as np
import scipy.stats as sp

from .UpwardRecursion import get_Emission_Likelihoods
from .BaumWelch import do_E_step, calculate_log_likelihood, do_M_step, do_M_E_step, do_M_E_step_atonce
from .Viterbi import get_leaf_deltas, get_nonleaf_deltas, Viterbi


class estimate:
    """Estimation class.
    """

    def __init__(self, X, nState: int, fpi=None, fT=None, fE=None):
        self.fpi = fpi
        self.fT = fT
        self.fE = fE
        self.num_states = nState

        if self.fpi is None or self.fpi is True:
            self.pi = np.random.rand(nState)
            self.pi /= np.sum(self.pi)
        else:
            self.pi = self.fpi

        if self.fT is None:
            self.T = np.random.dirichlet(np.random.rand(nState), nState)
        else:
            self.T = self.fT

        if self.fE is None:
            self.E = [deepcopy(X[0].E[0]) for _ in range(self.num_states)]
        else:
            self.E = self.fE


class tHMM:
    """Main tHMM class.
    """

    def __init__(self, X, num_states: int, fpi=None, fT=None, fE=None):
        """Instantiates a tHMM.
        This function uses the following functions and assings them to the cells
        (objects) in the lineage.
        :param X: A list of objects (cells) in a lineage in which
        the NaNs have been removed.
        :type X: list of objects
        :param num_states: The number of hidden states that we want our model have.
        :type num_states: Int
        """
        self.fpi = fpi
        self.fT = fT
        self.fE = fE
        self.X = X  # list containing lineages
        self.num_states = num_states  # number of discrete hidden states, should be integral
        self.estimate = estimate(
            self.X, self.num_states, fpi=self.fpi, fT=self.fT, fE=self.fE)
        self.EL = get_Emission_Likelihoods(self)

    def fit(self, tolerance=1e-9, max_iter=1000):
        """Runs the tHMM function through Baum Welch fitting"""
        MSD_list, NF_list, betas_list, gammas_list, new_LL = fit_list([self], tolerance=tolerance, max_iter=max_iter)

        return self, MSD_list[0], NF_list[0], betas_list[0], gammas_list[0], new_LL

    def predict(self):
        """
        Given a fit model, the model predicts an optimal
        state assignment using the Viterbi algorithm.
        """
        deltas, state_ptrs = get_leaf_deltas(self)
        get_nonleaf_deltas(self, deltas, state_ptrs)
        pred_states_by_lineage = Viterbi(self, deltas, state_ptrs)
        return pred_states_by_lineage

    def get_AIC(self, LL):
        """
        Gets the AIC values. Akaike Information Criterion, used for model selection and deals with the trade off
        between over-fitting and under-fitting.
        :math:`AIC = 2*k - 2 * log(LL)` in which k is the number of free parameters and LL is the maximum of likelihood function.
        Minimum of AIC detremines the relatively better model.
        :param tHMMobj: the tHMM class which has been built.
        :type tHMMobj: object
        :param LL: the likelihood value
        :param AIC_value: containing AIC values relative to 0 for each lineage.
        :type AIC_value: float
        :param AIC_degrees_of_freedom: the degrees of freedom in AIC calculation :math:`(num_{states}^2 + num_{states} * numberOfParameters - 1)` - same for each lineage
        """
        degrees_of_freedom = 0
        # initial prob. matrix
        if self.fpi is None:
            degrees_of_freedom += self.num_states - 1

        # transition matrix
        if self.fT is None:
            degrees_of_freedom += self.num_states * (self.num_states - 1)

        for ii in range(self.num_states):
            degrees_of_freedom += self.estimate.E[ii].dof()

        # the whole population has one AIC value.
        AIC_value = -2 * np.sum(LL) + 2 * degrees_of_freedom

        return AIC_value, degrees_of_freedom

    def log_score(self, X_state_tree_sequence, pi=None, T=None, E=None):
        """
        This function returns the log-likelihood of a possible state assignment
        given the estimated model parameters.
        The user can also provide their own pi, T, or E matrices instead to score
        a possible state assignment.
        :math:`P(x_1,...,x_N,z_1,...,z_N) = P(z_1) * prod_{n=2:N}(P(z_n | z_pn)) * prod_{n=1:N}(P(x_n|z_n))`
        :math:`log{P(x_1,...,x_N,z_1,...,z_N)} = log{P(z_1)} + sum_{n=2:N}(log{P(z_n | z_pn)}) + sum_{n=1:N}(log{P(x_n|z_n)})`
        """
        if pi is None:
            pi = self.estimate.pi
        if T is None:
            T = self.estimate.T
        if E is None:
            E = self.estimate.E

        log_scores = []
        for idx, lineageObj in enumerate(self.X):
            log_score = 0
            # the first term is the value of pi for the state of the first cell
            log_score += np.log(pi[X_state_tree_sequence[idx][0]])
            log_score += log_T_score(T, X_state_tree_sequence[idx], lineageObj)
            log_score += log_E_score(get_Emission_Likelihoods(self, E)
                                     [idx], X_state_tree_sequence[idx])
            assert np.all(np.isfinite(log_score))
            log_scores.append(log_score)
        return log_scores


def log_T_score(T, state_tree_sequence, lineageObj):
    """
    To calculate the joint probability of state and observations.
    This function calculates the second term.
    :math:`P(x_1,...,x_N,z_1,...,z_N) = P(z_1) * prod_{n=2:N}(P(z_n | z_pn)) * prod_{n=1:N}(P(x_n|z_n))`
    :math:`log{P(x_1,...,x_N,z_1,...,z_N)} = log{P(z_1)} + sum_{n=2:N}(log{P(z_n | z_pn)}) + sum_{n=1:N}(log{P(x_n|z_n)})`
    """
    log_T_score_holder = 0
    log_T = np.log(T)
    # we start with the first transition, from the root cell
    for level in lineageObj.output_list_of_gens[1:]:
        for cell in level:
            cell_idx = lineageObj.output_lineage.index(cell)
            cell_state = state_tree_sequence[cell_idx]
            if not cell.isLeaf():
                for daughter in cell.get_daughters():
                    child_idx = lineageObj.output_lineage.index(daughter)
                    daughter_state = state_tree_sequence[child_idx]
                    log_T_score_holder += log_T[cell_state, daughter_state]
    return log_T_score_holder


def log_E_score(EL_array, state_tree_sequence):
    """
    To calculate the joint probability of state and observations.
    This function calculates the thid term.
    :math:`P(x_1,...,x_N,z_1,...,z_N) = P(z_1) * prod_{n=2:N}(P(z_n | z_pn)) * prod_{n=1:N}(P(x_n|z_n))`
    :math:`log{P(x_1,...,x_N,z_1,...,z_N)} = log{P(z_1)} + sum_{n=2:N}(log{P(z_n | z_pn)}) + sum_{n=1:N}(log{P(x_n|z_n)})`
    """
    log_EL_array = np.log(EL_array)
    log_E_score_holder = 0
    for idx, row in enumerate(log_EL_array):
        log_E_score_holder += row[state_tree_sequence[idx]]
    return log_E_score_holder


def fit_list(tHMMobj_list, tolerance=1e-9, max_iter=1000):
    """Runs the tHMM function through Baum Welch fitting for a list containing a set of data for different concentrations"""

    # Step 0: initialize with random assignments and do an M step
    # when there are no fixed emissions, we need to randomize the start
    init_all_gammas = [[sp.multinomial.rvs(n=1, p=[1. / tHMMobj.num_states] * tHMMobj.num_states, size=len(lineage))
                        for lineage in tHMMobj.X] for tHMMobj in tHMMobj_list]

    if len(tHMMobj_list) > 1:  # it means we are fitting several concentrations at once.
        do_M_E_step_atonce(tHMMobj_list, init_all_gammas)
    else:  # means we are fitting one condition at a time.
        do_M_E_step(tHMMobj_list[0], init_all_gammas[0])

    # Step 1: first E step
    MSD_list, NF_list, betas_list, gammas_list = map(list, zip(*[do_E_step(tHMM) for tHMM in tHMMobj_list]))
    old_LL = np.sum([np.sum(calculate_log_likelihood(NF)) for NF in NF_list])

    # first stopping condition check
    for _ in range(max_iter):
        do_M_step(tHMMobj_list, MSD_list, betas_list, gammas_list)
        MSD_list, NF_list, betas_list, gammas_list = map(list, zip(*[do_E_step(tHMM) for tHMM in tHMMobj_list]))
        new_LL = np.sum([np.sum(calculate_log_likelihood(NF)) for NF in NF_list])
        if new_LL - old_LL < tolerance:
            break

        old_LL = new_LL

    return MSD_list, NF_list, betas_list, gammas_list, new_LL

86e147bcc01d9ef28c7eadbcab8985b0e5c65aa7
<mask token>


class tHMM:
    """Main tHMM class.
    """

    def __init__(self, X, num_states: int, fpi=None, fT=None, fE=None):
        """Instantiates a tHMM.
        This function uses the following functions and assings them to the cells
        (objects) in the lineage.
        :param X: A list of objects (cells) in a lineage in which
        the NaNs have been removed.
        :type X: list of objects
        :param num_states: The number of hidden states that we want our model have.
        :type num_states: Int
        """
        self.fpi = fpi
        self.fT = fT
        self.fE = fE
        self.X = X
        self.num_states = num_states
        self.estimate = estimate(self.X, self.num_states, fpi=self.fpi, fT=
            self.fT, fE=self.fE)
        self.EL = get_Emission_Likelihoods(self)

    def fit(self, tolerance=1e-09, max_iter=1000):
        """Runs the tHMM function through Baum Welch fitting"""
        MSD_list, NF_list, betas_list, gammas_list, new_LL = fit_list([self
            ], tolerance=tolerance, max_iter=max_iter)
        return self, MSD_list[0], NF_list[0], betas_list[0], gammas_list[0
            ], new_LL

    def predict(self):
        """
        Given a fit model, the model predicts an optimal
        state assignment using the Viterbi algorithm.
        """
        deltas, state_ptrs = get_leaf_deltas(self)
        get_nonleaf_deltas(self, deltas, state_ptrs)
        pred_states_by_lineage = Viterbi(self, deltas, state_ptrs)
        return pred_states_by_lineage

    def get_AIC(self, LL):
        """
        Gets the AIC values. Akaike Information Criterion, used for model selection and deals with the trade off
        between over-fitting and under-fitting.
        :math:`AIC = 2*k - 2 * log(LL)` in which k is the number of free parameters and LL is the maximum of likelihood function.
        Minimum of AIC detremines the relatively better model.
        :param tHMMobj: the tHMM class which has been built.
        :type tHMMobj: object
        :param LL: the likelihood value
        :param AIC_value: containing AIC values relative to 0 for each lineage.
        :type AIC_value: float
        :param AIC_degrees_of_freedom: the degrees of freedom in AIC calculation :math:`(num_{states}^2 + num_{states} * numberOfParameters - 1)` - same for each lineage
        """
        degrees_of_freedom = 0
        if self.fpi is None:
            degrees_of_freedom += self.num_states - 1
        if self.fT is None:
            degrees_of_freedom += self.num_states * (self.num_states - 1)
        for ii in range(self.num_states):
            degrees_of_freedom += self.estimate.E[ii].dof()
        AIC_value = -2 * np.sum(LL) + 2 * degrees_of_freedom
        return AIC_value, degrees_of_freedom

    def log_score(self, X_state_tree_sequence, pi=None, T=None, E=None):
        """
        This function returns the log-likelihood of a possible state assignment
        given the estimated model parameters.
        The user can also provide their own pi, T, or E matrices instead to score
        a possible state assignment.
        :math:`P(x_1,...,x_N,z_1,...,z_N) = P(z_1) * prod_{n=2:N}(P(z_n | z_pn)) * prod_{n=1:N}(P(x_n|z_n))`
        :math:`log{P(x_1,...,x_N,z_1,...,z_N)} = log{P(z_1)} + sum_{n=2:N}(log{P(z_n | z_pn)}) + sum_{n=1:N}(log{P(x_n|z_n)})`
        """
        if pi is None:
            pi = self.estimate.pi
        if T is None:
            T = self.estimate.T
        if E is None:
            E = self.estimate.E
        log_scores = []
        for idx, lineageObj in enumerate(self.X):
            log_score = 0
            log_score += np.log(pi[X_state_tree_sequence[idx][0]])
            log_score += log_T_score(T, X_state_tree_sequence[idx], lineageObj)
            log_score += log_E_score(get_Emission_Likelihoods(self, E)[idx],
                X_state_tree_sequence[idx])
            assert np.all(np.isfinite(log_score))
            log_scores.append(log_score)
        return log_scores


<mask token>

<mask token>


class estimate:
    <mask token>
    <mask token>


class tHMM:
    """Main tHMM class.
    """

    def __init__(self, X, num_states: int, fpi=None, fT=None, fE=None):
        """Instantiates a tHMM.
        This function uses the following functions and assings them to the cells
        (objects) in the lineage.
        :param X: A list of objects (cells) in a lineage in which
        the NaNs have been removed.
        :type X: list of objects
        :param num_states: The number of hidden states that we want our model have.
        :type num_states: Int
        """
        self.fpi = fpi
        self.fT = fT
        self.fE = fE
        self.X = X
        self.num_states = num_states
        self.estimate = estimate(self.X, self.num_states, fpi=self.fpi, fT=
            self.fT, fE=self.fE)
        self.EL = get_Emission_Likelihoods(self)

    def fit(self, tolerance=1e-09, max_iter=1000):
        """Runs the tHMM function through Baum Welch fitting"""
        MSD_list, NF_list, betas_list, gammas_list, new_LL = fit_list([self
            ], tolerance=tolerance, max_iter=max_iter)
        return self, MSD_list[0], NF_list[0], betas_list[0], gammas_list[0
            ], new_LL

    def predict(self):
        """
        Given a fit model, the model predicts an optimal
        state assignment using the Viterbi algorithm.
        """
        deltas, state_ptrs = get_leaf_deltas(self)
        get_nonleaf_deltas(self, deltas, state_ptrs)
        pred_states_by_lineage = Viterbi(self, deltas, state_ptrs)
        return pred_states_by_lineage

    def get_AIC(self, LL):
        """
        Gets the AIC values. Akaike Information Criterion, used for model selection and deals with the trade off
        between over-fitting and under-fitting.
        :math:`AIC = 2*k - 2 * log(LL)` in which k is the number of free parameters and LL is the maximum of likelihood function.
        Minimum of AIC detremines the relatively better model.
        :param tHMMobj: the tHMM class which has been built.
        :type tHMMobj: object
        :param LL: the likelihood value
        :param AIC_value: containing AIC values relative to 0 for each lineage.
        :type AIC_value: float
        :param AIC_degrees_of_freedom: the degrees of freedom in AIC calculation :math:`(num_{states}^2 + num_{states} * numberOfParameters - 1)` - same for each lineage
        """
        degrees_of_freedom = 0
        if self.fpi is None:
            degrees_of_freedom += self.num_states - 1
        if self.fT is None:
            degrees_of_freedom += self.num_states * (self.num_states - 1)
        for ii in range(self.num_states):
            degrees_of_freedom += self.estimate.E[ii].dof()
        AIC_value = -2 * np.sum(LL) + 2 * degrees_of_freedom
        return AIC_value, degrees_of_freedom

    def log_score(self, X_state_tree_sequence, pi=None, T=None, E=None):
        """
        This function returns the log-likelihood of a possible state assignment
        given the estimated model parameters.
        The user can also provide their own pi, T, or E matrices instead to score
        a possible state assignment.
        :math:`P(x_1,...,x_N,z_1,...,z_N) = P(z_1) * prod_{n=2:N}(P(z_n | z_pn)) * prod_{n=1:N}(P(x_n|z_n))`
        :math:`log{P(x_1,...,x_N,z_1,...,z_N)} = log{P(z_1)} + sum_{n=2:N}(log{P(z_n | z_pn)}) + sum_{n=1:N}(log{P(x_n|z_n)})`
        """
        if pi is None:
            pi = self.estimate.pi
        if T is None:
            T = self.estimate.T
        if E is None:
            E = self.estimate.E
        log_scores = []
        for idx, lineageObj in enumerate(self.X):
            log_score = 0
            log_score += np.log(pi[X_state_tree_sequence[idx][0]])
            log_score += log_T_score(T, X_state_tree_sequence[idx], lineageObj)
            log_score += log_E_score(get_Emission_Likelihoods(self, E)[idx],
                X_state_tree_sequence[idx])
            assert np.all(np.isfinite(log_score))
            log_scores.append(log_score)
        return log_scores


<mask token>

<mask token>


class estimate:
    <mask token>

    def __init__(self, X, nState: int, fpi=None, fT=None, fE=None):
        self.fpi = fpi
        self.fT = fT
        self.fE = fE
        self.num_states = nState
        if self.fpi is None or self.fpi is True:
            self.pi = np.random.rand(nState)
            self.pi /= np.sum(self.pi)
        else:
            self.pi = self.fpi
        if self.fT is None:
            self.T = np.random.dirichlet(np.random.rand(nState), nState)
        else:
            self.T = self.fT
        if self.fE is None:
            self.E = [deepcopy(X[0].E[0]) for _ in range(self.num_states)]
        else:
            self.E = self.fE


class tHMM:
    """Main tHMM class.
    """

    def __init__(self, X, num_states: int, fpi=None, fT=None, fE=None):
        """Instantiates a tHMM.
        This function uses the following functions and assings them to the cells
        (objects) in the lineage.
        :param X: A list of objects (cells) in a lineage in which
        the NaNs have been removed.
        :type X: list of objects
        :param num_states: The number of hidden states that we want our model have.
        :type num_states: Int
        """
        self.fpi = fpi
        self.fT = fT
        self.fE = fE
        self.X = X
        self.num_states = num_states
        self.estimate = estimate(self.X, self.num_states, fpi=self.fpi, fT=
            self.fT, fE=self.fE)
        self.EL = get_Emission_Likelihoods(self)

    def fit(self, tolerance=1e-09, max_iter=1000):
        """Runs the tHMM function through Baum Welch fitting"""
        MSD_list, NF_list, betas_list, gammas_list, new_LL = fit_list([self
            ], tolerance=tolerance, max_iter=max_iter)
        return self, MSD_list[0], NF_list[0], betas_list[0], gammas_list[0
            ], new_LL

    def predict(self):
        """
        Given a fit model, the model predicts an optimal
        state assignment using the Viterbi algorithm.
        """
        deltas, state_ptrs = get_leaf_deltas(self)
        get_nonleaf_deltas(self, deltas, state_ptrs)
        pred_states_by_lineage = Viterbi(self, deltas, state_ptrs)
        return pred_states_by_lineage

    def get_AIC(self, LL):
        """
        Gets the AIC values. Akaike Information Criterion, used for model selection and deals with the trade off
        between over-fitting and under-fitting.
        :math:`AIC = 2*k - 2 * log(LL)` in which k is the number of free parameters and LL is the maximum of likelihood function.
        Minimum of AIC detremines the relatively better model.
        :param tHMMobj: the tHMM class which has been built.
        :type tHMMobj: object
        :param LL: the likelihood value
        :param AIC_value: containing AIC values relative to 0 for each lineage.
        :type AIC_value: float
        :param AIC_degrees_of_freedom: the degrees of freedom in AIC calculation :math:`(num_{states}^2 + num_{states} * numberOfParameters - 1)` - same for each lineage
        """
        degrees_of_freedom = 0
        if self.fpi is None:
            degrees_of_freedom += self.num_states - 1
        if self.fT is None:
            degrees_of_freedom += self.num_states * (self.num_states - 1)
        for ii in range(self.num_states):
            degrees_of_freedom += self.estimate.E[ii].dof()
        AIC_value = -2 * np.sum(LL) + 2 * degrees_of_freedom
        return AIC_value, degrees_of_freedom

    def log_score(self, X_state_tree_sequence, pi=None, T=None, E=None):
        """
        This function returns the log-likelihood of a possible state assignment
        given the estimated model parameters.
        The user can also provide their own pi, T, or E matrices instead to score
        a possible state assignment.
        :math:`P(x_1,...,x_N,z_1,...,z_N) = P(z_1) * prod_{n=2:N}(P(z_n | z_pn)) * prod_{n=1:N}(P(x_n|z_n))`
        :math:`log{P(x_1,...,x_N,z_1,...,z_N)} = log{P(z_1)} + sum_{n=2:N}(log{P(z_n | z_pn)}) + sum_{n=1:N}(log{P(x_n|z_n)})`
        """
        if pi is None:
            pi = self.estimate.pi
        if T is None:
            T = self.estimate.T
        if E is None:
            E = self.estimate.E
        log_scores = []
        for idx, lineageObj in enumerate(self.X):
            log_score = 0
            log_score += np.log(pi[X_state_tree_sequence[idx][0]])
            log_score += log_T_score(T, X_state_tree_sequence[idx], lineageObj)
            log_score += log_E_score(get_Emission_Likelihoods(self, E)[idx],
                X_state_tree_sequence[idx])
            assert np.all(np.isfinite(log_score))
            log_scores.append(log_score)
        return log_scores


<mask token>

<mask token>


class estimate:
    """Estimation class.
    """

    def __init__(self, X, nState: int, fpi=None, fT=None, fE=None):
        self.fpi = fpi
        self.fT = fT
        self.fE = fE
        self.num_states = nState
        if self.fpi is None or self.fpi is True:
            self.pi = np.random.rand(nState)
            self.pi /= np.sum(self.pi)
        else:
            self.pi = self.fpi
        if self.fT is None:
            self.T = np.random.dirichlet(np.random.rand(nState), nState)
        else:
            self.T = self.fT
        if self.fE is None:
            self.E = [deepcopy(X[0].E[0]) for _ in range(self.num_states)]
        else:
            self.E = self.fE


class tHMM:
    """Main tHMM class.
    """

    def __init__(self, X, num_states: int, fpi=None, fT=None, fE=None):
        """Instantiates a tHMM.
        This function uses the following functions and assings them to the cells
        (objects) in the lineage.
        :param X: A list of objects (cells) in a lineage in which
        the NaNs have been removed.
        :type X: list of objects
        :param num_states: The number of hidden states that we want our model have.
        :type num_states: Int
        """
        self.fpi = fpi
        self.fT = fT
        self.fE = fE
        self.X = X
        self.num_states = num_states
        self.estimate = estimate(self.X, self.num_states, fpi=self.fpi, fT=
            self.fT, fE=self.fE)
        self.EL = get_Emission_Likelihoods(self)

    def fit(self, tolerance=1e-09, max_iter=1000):
        """Runs the tHMM function through Baum Welch fitting"""
        MSD_list, NF_list, betas_list, gammas_list, new_LL = fit_list([self
            ], tolerance=tolerance, max_iter=max_iter)
        return self, MSD_list[0], NF_list[0], betas_list[0], gammas_list[0
            ], new_LL

    def predict(self):
        """
        Given a fit model, the model predicts an optimal
        state assignment using the Viterbi algorithm.
        """
        deltas, state_ptrs = get_leaf_deltas(self)
        get_nonleaf_deltas(self, deltas, state_ptrs)
        pred_states_by_lineage = Viterbi(self, deltas, state_ptrs)
        return pred_states_by_lineage

    def get_AIC(self, LL):
        """
        Gets the AIC values. Akaike Information Criterion, used for model selection and deals with the trade off
        between over-fitting and under-fitting.
        :math:`AIC = 2*k - 2 * log(LL)` in which k is the number of free parameters and LL is the maximum of likelihood function.
        Minimum of AIC detremines the relatively better model.
        :param tHMMobj: the tHMM class which has been built.
        :type tHMMobj: object
        :param LL: the likelihood value
        :param AIC_value: containing AIC values relative to 0 for each lineage.
        :type AIC_value: float
        :param AIC_degrees_of_freedom: the degrees of freedom in AIC calculation :math:`(num_{states}^2 + num_{states} * numberOfParameters - 1)` - same for each lineage
        """
        degrees_of_freedom = 0
        if self.fpi is None:
            degrees_of_freedom += self.num_states - 1
        if self.fT is None:
            degrees_of_freedom += self.num_states * (self.num_states - 1)
        for ii in range(self.num_states):
            degrees_of_freedom += self.estimate.E[ii].dof()
        AIC_value = -2 * np.sum(LL) + 2 * degrees_of_freedom
        return AIC_value, degrees_of_freedom

    def log_score(self, X_state_tree_sequence, pi=None, T=None, E=None):
        """
        This function returns the log-likelihood of a possible state assignment
        given the estimated model parameters.
        The user can also provide their own pi, T, or E matrices instead to score
        a possible state assignment.
        :math:`P(x_1,...,x_N,z_1,...,z_N) = P(z_1) * prod_{n=2:N}(P(z_n | z_pn)) * prod_{n=1:N}(P(x_n|z_n))`
        :math:`log{P(x_1,...,x_N,z_1,...,z_N)} = log{P(z_1)} + sum_{n=2:N}(log{P(z_n | z_pn)}) + sum_{n=1:N}(log{P(x_n|z_n)})`
        """
        if pi is None:
            pi = self.estimate.pi
        if T is None:
            T = self.estimate.T
        if E is None:
            E = self.estimate.E
        log_scores = []
        for idx, lineageObj in enumerate(self.X):
            log_score = 0
            log_score += np.log(pi[X_state_tree_sequence[idx][0]])
            log_score += log_T_score(T, X_state_tree_sequence[idx], lineageObj)
            log_score += log_E_score(get_Emission_Likelihoods(self, E)[idx],
                X_state_tree_sequence[idx])
            assert np.all(np.isfinite(log_score))
            log_scores.append(log_score)
        return log_scores


def log_T_score(T, state_tree_sequence, lineageObj):
    """
    To calculate the joint probability of state and observations.
    This function calculates the second term.
    :math:`P(x_1,...,x_N,z_1,...,z_N) = P(z_1) * prod_{n=2:N}(P(z_n | z_pn)) * prod_{n=1:N}(P(x_n|z_n))`
    :math:`log{P(x_1,...,x_N,z_1,...,z_N)} = log{P(z_1)} + sum_{n=2:N}(log{P(z_n | z_pn)}) + sum_{n=1:N}(log{P(x_n|z_n)})`
    """
    log_T_score_holder = 0
    log_T = np.log(T)
    for level in lineageObj.output_list_of_gens[1:]:
        for cell in level:
            cell_idx = lineageObj.output_lineage.index(cell)
            cell_state = state_tree_sequence[cell_idx]
            if not cell.isLeaf():
                for daughter in cell.get_daughters():
                    child_idx = lineageObj.output_lineage.index(daughter)
                    daughter_state = state_tree_sequence[child_idx]
                    log_T_score_holder += log_T[cell_state, daughter_state]
    return log_T_score_holder


<mask token>


def fit_list(tHMMobj_list, tolerance=1e-09, max_iter=1000):
    """Runs the tHMM function through Baum Welch fitting for a list containing a set of data for different concentrations"""
    init_all_gammas = [[sp.multinomial.rvs(n=1, p=[1.0 / tHMMobj.num_states
        ] * tHMMobj.num_states, size=len(lineage)) for lineage in tHMMobj.X
        ] for tHMMobj in tHMMobj_list]
    if len(tHMMobj_list) > 1:
        do_M_E_step_atonce(tHMMobj_list, init_all_gammas)
    else:
        do_M_E_step(tHMMobj_list[0], init_all_gammas[0])
    MSD_list, NF_list, betas_list, gammas_list = map(list, zip(*[do_E_step(
        tHMM) for tHMM in tHMMobj_list]))
    old_LL = np.sum([np.sum(calculate_log_likelihood(NF)) for NF in NF_list])
    for _ in range(max_iter):
        do_M_step(tHMMobj_list, MSD_list, betas_list, gammas_list)
        MSD_list, NF_list, betas_list, gammas_list = map(list, zip(*[
            do_E_step(tHMM) for tHMM in tHMMobj_list]))
        new_LL = np.sum([np.sum(calculate_log_likelihood(NF)) for NF in
            NF_list])
        if new_LL - old_LL < tolerance:
            break
        old_LL = new_LL
    return MSD_list, NF_list, betas_list, gammas_list, new_LL

""" This file holds the parameters of our tHMM in the tHMM class. """

from copy import deepcopy
import numpy as np
import scipy.stats as sp

from .UpwardRecursion import get_Emission_Likelihoods
from .BaumWelch import do_E_step, calculate_log_likelihood, do_M_step, do_M_E_step, do_M_E_step_atonce
from .Viterbi import get_leaf_deltas, get_nonleaf_deltas, Viterbi


class estimate:
    """Estimation class.
    """

    def __init__(self, X, nState: int, fpi=None, fT=None, fE=None):
        self.fpi = fpi
        self.fT = fT
        self.fE = fE
        self.num_states = nState

        if self.fpi is None or self.fpi is True:
            self.pi = np.random.rand(nState)
            self.pi /= np.sum(self.pi)
        else:
            self.pi = self.fpi

        if self.fT is None:
            self.T = np.random.dirichlet(np.random.rand(nState), nState)
        else:
            self.T = self.fT

        if self.fE is None:
            self.E = [deepcopy(X[0].E[0]) for _ in range(self.num_states)]
        else:
            self.E = self.fE


class tHMM:
    """Main tHMM class.
    """

    def __init__(self, X, num_states: int, fpi=None, fT=None, fE=None):
        """Instantiates a tHMM.
        This function uses the following functions and assings them to the cells
        (objects) in the lineage.
        :param X: A list of objects (cells) in a lineage in which
        the NaNs have been removed.
        :type X: list of objects
        :param num_states: The number of hidden states that we want our model have.
        :type num_states: Int
        """
        self.fpi = fpi
        self.fT = fT
        self.fE = fE
        self.X = X  # list containing lineages
        self.num_states = num_states  # number of discrete hidden states, should be integral
        self.estimate = estimate(
            self.X, self.num_states, fpi=self.fpi, fT=self.fT, fE=self.fE)
        self.EL = get_Emission_Likelihoods(self)

    def fit(self, tolerance=1e-9, max_iter=1000):
        """Runs the tHMM function through Baum Welch fitting"""
        MSD_list, NF_list, betas_list, gammas_list, new_LL = fit_list([self], tolerance=tolerance, max_iter=max_iter)

        return self, MSD_list[0], NF_list[0], betas_list[0], gammas_list[0], new_LL

    def predict(self):
        """
        Given a fit model, the model predicts an optimal
        state assignment using the Viterbi algorithm.
        """
        deltas, state_ptrs = get_leaf_deltas(self)
        get_nonleaf_deltas(self, deltas, state_ptrs)
        pred_states_by_lineage = Viterbi(self, deltas, state_ptrs)
        return pred_states_by_lineage

    def get_AIC(self, LL):
        """
        Gets the AIC values. Akaike Information Criterion, used for model selection and deals with the trade off
        between over-fitting and under-fitting.
        :math:`AIC = 2*k - 2 * log(LL)` in which k is the number of free parameters and LL is the maximum of likelihood function.
        Minimum of AIC detremines the relatively better model.
        :param tHMMobj: the tHMM class which has been built.
        :type tHMMobj: object
        :param LL: the likelihood value
        :param AIC_value: containing AIC values relative to 0 for each lineage.
        :type AIC_value: float
        :param AIC_degrees_of_freedom: the degrees of freedom in AIC calculation :math:`(num_{states}^2 + num_{states} * numberOfParameters - 1)` - same for each lineage
        """
        degrees_of_freedom = 0
        # initial prob. matrix
        if self.fpi is None:
            degrees_of_freedom += self.num_states - 1

        # transition matrix
        if self.fT is None:
            degrees_of_freedom += self.num_states * (self.num_states - 1)

        for ii in range(self.num_states):
            degrees_of_freedom += self.estimate.E[ii].dof()

        # the whole population has one AIC value.
        AIC_value = -2 * np.sum(LL) + 2 * degrees_of_freedom

        return AIC_value, degrees_of_freedom

    def log_score(self, X_state_tree_sequence, pi=None, T=None, E=None):
        """
        This function returns the log-likelihood of a possible state assignment
        given the estimated model parameters.
        The user can also provide their own pi, T, or E matrices instead to score
        a possible state assignment.
        :math:`P(x_1,...,x_N,z_1,...,z_N) = P(z_1) * prod_{n=2:N}(P(z_n | z_pn)) * prod_{n=1:N}(P(x_n|z_n))`
        :math:`log{P(x_1,...,x_N,z_1,...,z_N)} = log{P(z_1)} + sum_{n=2:N}(log{P(z_n | z_pn)}) + sum_{n=1:N}(log{P(x_n|z_n)})`
        """
        if pi is None:
            pi = self.estimate.pi
        if T is None:
            T = self.estimate.T
        if E is None:
            E = self.estimate.E

        log_scores = []
        for idx, lineageObj in enumerate(self.X):
            log_score = 0
            # the first term is the value of pi for the state of the first cell
            log_score += np.log(pi[X_state_tree_sequence[idx][0]])
            log_score += log_T_score(T, X_state_tree_sequence[idx], lineageObj)
            log_score += log_E_score(get_Emission_Likelihoods(self, E)
                                     [idx], X_state_tree_sequence[idx])
            assert np.all(np.isfinite(log_score))
            log_scores.append(log_score)
        return log_scores


def log_T_score(T, state_tree_sequence, lineageObj):
    """
    To calculate the joint probability of state and observations.
    This function calculates the second term.
    :math:`P(x_1,...,x_N,z_1,...,z_N) = P(z_1) * prod_{n=2:N}(P(z_n | z_pn)) * prod_{n=1:N}(P(x_n|z_n))`
    :math:`log{P(x_1,...,x_N,z_1,...,z_N)} = log{P(z_1)} + sum_{n=2:N}(log{P(z_n | z_pn)}) + sum_{n=1:N}(log{P(x_n|z_n)})`
    """
    log_T_score_holder = 0
    log_T = np.log(T)
    # we start with the first transition, from the root cell
    for level in lineageObj.output_list_of_gens[1:]:
        for cell in level:
            cell_idx = lineageObj.output_lineage.index(cell)
            cell_state = state_tree_sequence[cell_idx]
            if not cell.isLeaf():
                for daughter in cell.get_daughters():
                    child_idx = lineageObj.output_lineage.index(daughter)
                    daughter_state = state_tree_sequence[child_idx]
                    log_T_score_holder += log_T[cell_state, daughter_state]
    return log_T_score_holder


def log_E_score(EL_array, state_tree_sequence):
    """
    To calculate the joint probability of state and observations.
    This function calculates the thid term.
    :math:`P(x_1,...,x_N,z_1,...,z_N) = P(z_1) * prod_{n=2:N}(P(z_n | z_pn)) * prod_{n=1:N}(P(x_n|z_n))`
    :math:`log{P(x_1,...,x_N,z_1,...,z_N)} = log{P(z_1)} + sum_{n=2:N}(log{P(z_n | z_pn)}) + sum_{n=1:N}(log{P(x_n|z_n)})`
    """
    log_EL_array = np.log(EL_array)
    log_E_score_holder = 0
    for idx, row in enumerate(log_EL_array):
        log_E_score_holder += row[state_tree_sequence[idx]]
    return log_E_score_holder


def fit_list(tHMMobj_list, tolerance=1e-9, max_iter=1000):
    """Runs the tHMM function through Baum Welch fitting for a list containing a set of data for different concentrations"""

    # Step 0: initialize with random assignments and do an M step
    # when there are no fixed emissions, we need to randomize the start
    init_all_gammas = [[sp.multinomial.rvs(n=1, p=[1. / tHMMobj.num_states] * tHMMobj.num_states, size=len(lineage))
                        for lineage in tHMMobj.X] for tHMMobj in tHMMobj_list]

    if len(tHMMobj_list) > 1:  # it means we are fitting several concentrations at once.
        do_M_E_step_atonce(tHMMobj_list, init_all_gammas)
    else:  # means we are fitting one condition at a time.
        do_M_E_step(tHMMobj_list[0], init_all_gammas[0])

    # Step 1: first E step
    MSD_list, NF_list, betas_list, gammas_list = map(list, zip(*[do_E_step(tHMM) for tHMM in tHMMobj_list]))
    old_LL = np.sum([np.sum(calculate_log_likelihood(NF)) for NF in NF_list])

    # first stopping condition check
    for _ in range(max_iter):
        do_M_step(tHMMobj_list, MSD_list, betas_list, gammas_list)
        MSD_list, NF_list, betas_list, gammas_list = map(list, zip(*[do_E_step(tHMM) for tHMM in tHMMobj_list]))
        new_LL = np.sum([np.sum(calculate_log_likelihood(NF)) for NF in NF_list])
        if new_LL - old_LL < tolerance:
            break

        old_LL = new_LL

    return MSD_list, NF_list, betas_list, gammas_list, new_LL
