r"""Expected transition counts for Markov chain in equilibrium.

    Parameters
    ----------
    T : (M, M) ndarray or sparse matrix
        Transition matrix.
    N : int
        Number of steps for chain.
    mu : (M,) ndarray (optional)
        Stationary distribution for T. If mu is not specified it will be
        computed from T.

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

    Notes
    -----
    Since :math:`\mu` is stationary for :math:`T` we have

    .. math::

        \mathbb{E}[C^{(N)}]=N D_{\mu}T.

    :math:`D_{\mu}` is a diagonal matrix. Elements on the diagonal are
    given by the stationary vector :math:`\mu`

    Examples
    --------

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

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

    >>> EC
    array([[ 40.90909091,   4.54545455,   0.        ],
           [  4.54545455,   0.        ,   4.54545455],
           [  0.        ,   4.54545455,  40.90909091]])

def expected_counts_stationary(T, N, mu=None):
    r"""Expected transition counts for Markov chain in equilibrium.

    Parameters
    ----------
    T : (M, M) ndarray or sparse matrix
        Transition matrix.
    N : int
        Number of steps for chain.
    mu : (M,) ndarray (optional)
        Stationary distribution for T. If mu is not specified it will be
        computed from T.

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

    Notes
    -----
    Since :math:`\mu` is stationary for :math:`T` we have

    .. math::

        \mathbb{E}[C^{(N)}]=N D_{\mu}T.

    :math:`D_{\mu}` is a diagonal matrix. Elements on the diagonal are
    given by the stationary vector :math:`\mu`

    Examples
    --------

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

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

    >>> EC
    array([[ 40.90909091,   4.54545455,   0.        ],
           [  4.54545455,   0.        ,   4.54545455],
           [  0.        ,   4.54545455,  40.90909091]])

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