r"""Mean first passage times to a set of target states.

    Parameters
    ----------
    T : scipy.sparse matrix
        Transition matrix.
    target : int or list of int
        Target states for mfpt calculation.

    Returns
    -------
    m_t : ndarray, shape=(n,)
         Vector of mean first passage times to target states.

    Notes
    -----
    The mean first passage time :math:`\mathbf{E}_x[T_Y]` is the expected
    hitting time of one state :math:`y` in :math:`Y` when starting in state :math:`x`.

    For a fixed target state :math:`y` it is given by

    .. math :: \mathbb{E}_x[T_y] = \left \{  \begin{array}{cc}
                                             0 & x=y \\
                                             1+\sum_{z} T_{x,z} \mathbb{E}_z[T_y] & x \neq y
                                             \end{array}  \right.

    For a set of target states :math:`Y` it is given by

    .. math :: \mathbb{E}_x[T_Y] = \left \{  \begin{array}{cc}
                                             0 & x \in Y \\
                                             1+\sum_{z} T_{x,z} \mathbb{E}_z[T_Y] & x \notin Y
                                             \end{array}  \right.

    References
    ----------
    .. [1] Hoel, P G and S C Port and C J Stone. 1972. Introduction to
        Stochastic Processes.

    Examples
    --------

    >>> from msmtools.analysis import mfpt

    >>> T = np.array([[0.9, 0.1, 0.0], [0.5, 0.0, 0.5], [0.0, 0.1, 0.9]])
    >>> m_t = mfpt(T, 0)
    >>> m_t
    array([  0.,  12.,  22.])

def mfpt(T, target):
    r"""Mean first passage times to a set of target states.

    Parameters
    ----------
    T : scipy.sparse matrix
        Transition matrix.
    target : int or list of int
        Target states for mfpt calculation.

    Returns
    -------
    m_t : ndarray, shape=(n,)
         Vector of mean first passage times to target states.

    Notes
    -----
    The mean first passage time :math:`\mathbf{E}_x[T_Y]` is the expected
    hitting time of one state :math:`y` in :math:`Y` when starting in state :math:`x`.

    For a fixed target state :math:`y` it is given by

    .. math :: \mathbb{E}_x[T_y] = \left \{  \begin{array}{cc}
                                             0 & x=y \\
                                             1+\sum_{z} T_{x,z} \mathbb{E}_z[T_y] & x \neq y
                                             \end{array}  \right.

    For a set of target states :math:`Y` it is given by

    .. math :: \mathbb{E}_x[T_Y] = \left \{  \begin{array}{cc}
                                             0 & x \in Y \\
                                             1+\sum_{z} T_{x,z} \mathbb{E}_z[T_Y] & x \notin Y
                                             \end{array}  \right.

    References
    ----------
    .. [1] Hoel, P G and S C Port and C J Stone. 1972. Introduction to
        Stochastic Processes.

    Examples
    --------

    >>> from msmtools.analysis import mfpt

    >>> T = np.array([[0.9, 0.1, 0.0], [0.5, 0.0, 0.5], [0.0, 0.1, 0.9]])
    >>> m_t = mfpt(T, 0)
    >>> m_t
    array([  0.,  12.,  22.])

    """
    dim = T.shape[0]
    A = eye(dim, dim) - T

    """Convert to DOK (dictionary of keys) matrix to enable
    row-slicing and assignement"""
    A = A.todok()
    D = A.diagonal()
    A[target, :] = 0.0
    D[target] = 1.0
    A.setdiag(D)
    """Convert back to CSR-format for fast sparse linear algebra"""
    A = A.tocsr()
    b = np.ones(dim)
    b[target] = 0.0
    m_t = spsolve(A, b)
    return m_t