r"""Compute expected transition counts for Markov chain with n steps.

    Parameters
    ----------
    T : (M, M) ndarray or sparse matrix
        Transition matrix
    p0 : (M,) ndarray
        Initial (probability) vector
    N : int
        Number of steps to take

    Returns
    --------
    EC : (M, M) ndarray or sparse matrix
        Expected value for transition counts after N steps

    Notes
    -----
    Expected counts can be computed via the following expression

    .. math::

        \mathbb{E}[C^{(N)}]=\sum_{k=0}^{N-1} \text{diag}(p^{T} T^{k}) T

    Examples
    --------

    >>> import numpy as np
    >>> from msmtools.analysis import expected_counts

    >>> T = np.array([[0.9, 0.1, 0.0], [0.5, 0.0, 0.5], [0.0, 0.1, 0.9]])
    >>> p0 = np.array([1.0, 0.0, 0.0])
    >>> N = 100
    >>> EC = expected_counts(T, p0, N)

    >>> EC
    array([[ 45.44616147,   5.0495735 ,   0.        ],
           [  4.50413223,   0.        ,   4.50413223],
           [  0.        ,   4.04960006,  36.44640052]])

def expected_counts(T, p0, N):
    r"""Compute expected transition counts for Markov chain with n steps.

    Parameters
    ----------
    T : (M, M) ndarray or sparse matrix
        Transition matrix
    p0 : (M,) ndarray
        Initial (probability) vector
    N : int
        Number of steps to take

    Returns
    --------
    EC : (M, M) ndarray or sparse matrix
        Expected value for transition counts after N steps

    Notes
    -----
    Expected counts can be computed via the following expression

    .. math::

        \mathbb{E}[C^{(N)}]=\sum_{k=0}^{N-1} \text{diag}(p^{T} T^{k}) T

    Examples
    --------

    >>> import numpy as np
    >>> from msmtools.analysis import expected_counts

    >>> T = np.array([[0.9, 0.1, 0.0], [0.5, 0.0, 0.5], [0.0, 0.1, 0.9]])
    >>> p0 = np.array([1.0, 0.0, 0.0])
    >>> N = 100
    >>> EC = expected_counts(T, p0, N)

    >>> EC
    array([[ 45.44616147,   5.0495735 ,   0.        ],
           [  4.50413223,   0.        ,   4.50413223],
           [  0.        ,   4.04960006,  36.44640052]])

    """
    # check input
    T = _types.ensure_ndarray_or_sparse(T, ndim=2, uniform=True, kind='numeric')
    p0 = _types.ensure_float_vector(p0, require_order=True)
    # go
    if _issparse(T):
        return sparse.expectations.expected_counts(p0, T, N)
    else:
        return dense.expectations.expected_counts(p0, T, N)