# pylint: disable=too-many-lines, too-many-function-args, redefined-outer-name
"""Diagnostic functions for ArviZ."""
import warnings
from collections.abc import Sequence

import numpy as np
import packaging
import pandas as pd
import scipy
from scipy import stats

from ..data import convert_to_dataset
from ..utils import Numba, _numba_var, _stack, _var_names
from .density_utils import histogram as _histogram
from .stats_utils import _circular_standard_deviation, _sqrt
from .stats_utils import autocov as _autocov
from .stats_utils import not_valid as _not_valid
from .stats_utils import quantile as _quantile
from .stats_utils import stats_variance_2d as svar
from .stats_utils import wrap_xarray_ufunc as _wrap_xarray_ufunc

__all__ = ["bfmi", "ess", "rhat", "mcse"]


def bfmi(data):
    r"""Calculate the estimated Bayesian fraction of missing information (BFMI).

    BFMI quantifies how well momentum resampling matches the marginal energy distribution. For more
    information on BFMI, see https://arxiv.org/pdf/1604.00695v1.pdf. The current advice is that
    values smaller than 0.3 indicate poor sampling. However, this threshold is
    provisional and may change. See
    `pystan_workflow <http://mc-stan.org/users/documentation/case-studies/pystan_workflow.html>`_
    for more information.

    Parameters
    ----------
    data : obj
        Any object that can be converted to an :class:`arviz.InferenceData` object.
        Refer to documentation of :func:`arviz.convert_to_dataset` for details.
        If InferenceData, energy variable needs to be found.

    Returns
    -------
    z : array
        The Bayesian fraction of missing information of the model and trace. One element per
        chain in the trace.

    See Also
    --------
    plot_energy : Plot energy transition distribution and marginal energy
                  distribution in HMC algorithms.

    Examples
    --------
    Compute the BFMI of an InferenceData object

    .. ipython::

        In [1]: import arviz as az
           ...: data = az.load_arviz_data('radon')
           ...: az.bfmi(data)

    """
    if isinstance(data, np.ndarray):
        return _bfmi(data)

    dataset = convert_to_dataset(data, group="sample_stats")
    if not hasattr(dataset, "energy"):
        raise TypeError("Energy variable was not found.")
    return _bfmi(dataset.energy.transpose("chain", "draw"))


def ess(
    data,
    *,
    var_names=None,
    method="bulk",
    relative=False,
    prob=None,
    dask_kwargs=None,
):
    r"""Calculate estimate of the effective sample size (ess).

    Parameters
    ----------
    data : obj
        Any object that can be converted to an :class:`arviz.InferenceData` object.
        Refer to documentation of :func:`arviz.convert_to_dataset` for details.
        For ndarray: shape = (chain, draw).
        For n-dimensional ndarray transform first to dataset with :func:`arviz.convert_to_dataset`.
    var_names : str or list of str
        Names of variables to include in the return value Dataset.
    method : str, optional, default "bulk"
        Select ess method. Valid methods are:

        - "bulk"
        - "tail"     # prob, optional
        - "quantile" # prob
        - "mean" (old ess)
        - "sd"
        - "median"
        - "mad" (mean absolute deviance)
        - "z_scale"
        - "folded"
        - "identity"
        - "local"
    relative : bool
        Return relative ess
        ``ress = ess / n``
    prob : float, or tuple of two floats, optional
        probability value for "tail", "quantile" or "local" ess functions.
    dask_kwargs : dict, optional
        Dask related kwargs passed to :func:`~arviz.wrap_xarray_ufunc`.

    Returns
    -------
    xarray.Dataset
        Return the effective sample size, :math:`\hat{N}_{eff}`

    Notes
    -----
    The basic ess (:math:`N_{\mathit{eff}}`) diagnostic is computed by:

    .. math:: \hat{N}_{\mathit{eff}} = \frac{MN}{\hat{\tau}}

    .. math:: \hat{\tau} = -1 + 2 \sum_{t'=0}^K \hat{P}_{t'}

    where :math:`M` is the number of chains, :math:`N` the number of draws,
    :math:`\hat{\rho}_t` is the estimated _autocorrelation at lag :math:`t`, and
    :math:`K` is the last integer for which :math:`\hat{P}_{K} = \hat{\rho}_{2K} +
    \hat{\rho}_{2K+1}` is still positive.

    The current implementation is similar to Stan, which uses Geyer's initial monotone sequence
    criterion (Geyer, 1992; Geyer, 2011).

    References
    ----------
    * Vehtari et al. (2021). Rank-normalization, folding, and
        localization: An improved Rhat for assessing convergence of
        MCMC. Bayesian analysis, 16(2):667-718.
    * https://mc-stan.org/docs/reference-manual/analysis.html#effective-sample-size.section
    * Gelman et al. BDA3 (2013) Formula 11.8

    See Also
    --------
    arviz.rhat : Compute estimate of rank normalized splitR-hat for a set of traces.
    arviz.mcse : Calculate Markov Chain Standard Error statistic.
    plot_ess : Plot quantile, local or evolution of effective sample sizes (ESS).
    arviz.summary : Create a data frame with summary statistics.

    Examples
    --------
    Calculate the effective_sample_size using the default arguments:

    .. ipython::

        In [1]: import arviz as az
           ...: data = az.load_arviz_data('non_centered_eight')
           ...: az.ess(data)

    Calculate the ress of some of the variables

    .. ipython::

        In [1]: az.ess(data, relative=True, var_names=["mu", "theta_t"])

    Calculate the ess using the "tail" method, leaving the `prob` argument at its default
    value.

    .. ipython::

        In [1]: az.ess(data, method="tail")

    """
    methods = {
        "bulk": _ess_bulk,
        "tail": _ess_tail,
        "quantile": _ess_quantile,
        "mean": _ess_mean,
        "sd": _ess_sd,
        "median": _ess_median,
        "mad": _ess_mad,
        "z_scale": _ess_z_scale,
        "folded": _ess_folded,
        "identity": _ess_identity,
        "local": _ess_local,
    }

    if method not in methods:
        raise TypeError(f"ess method {method} not found. Valid methods are:\n{', '.join(methods)}")
    ess_func = methods[method]

    if (method == "quantile") and prob is None:
        raise TypeError("Quantile (prob) information needs to be defined.")

    if isinstance(data, np.ndarray):
        data = np.atleast_2d(data)
        if len(data.shape) < 3:
            if prob is not None:
                return ess_func(  # pylint: disable=unexpected-keyword-arg
                    data, prob=prob, relative=relative
                )

            return ess_func(data, relative=relative)

        msg = (
            "Only uni-dimensional ndarray variables are supported."
            " Please transform first to dataset with `az.convert_to_dataset`."
        )
        raise TypeError(msg)

    dataset = convert_to_dataset(data, group="posterior")
    var_names = _var_names(var_names, dataset)

    dataset = dataset if var_names is None else dataset[var_names]

    ufunc_kwargs = {"ravel": False}
    func_kwargs = {"relative": relative} if prob is None else {"prob": prob, "relative": relative}
    return _wrap_xarray_ufunc(
        ess_func,
        dataset,
        ufunc_kwargs=ufunc_kwargs,
        func_kwargs=func_kwargs,
        dask_kwargs=dask_kwargs,
    )


def rhat(data, *, var_names=None, method="rank", dask_kwargs=None):
    r"""Compute estimate of rank normalized splitR-hat for a set of traces.

    The rank normalized R-hat diagnostic tests for lack of convergence by comparing the variance
    between multiple chains to the variance within each chain. If convergence has been achieved,
    the between-chain and within-chain variances should be identical. To be most effective in
    detecting evidence for nonconvergence, each chain should have been initialized to starting
    values that are dispersed relative to the target distribution.

    Parameters
    ----------
    data : obj
        Any object that can be converted to an :class:`arviz.InferenceData` object.
        Refer to documentation of :func:`arviz.convert_to_dataset` for details.
        At least 2 posterior chains are needed to compute this diagnostic of one or more
        stochastic parameters.
        For ndarray: shape = (chain, draw).
        For n-dimensional ndarray transform first to dataset with ``az.convert_to_dataset``.
    var_names : list
        Names of variables to include in the rhat report
    method : str
        Select R-hat method. Valid methods are:
        - "rank"        # recommended by Vehtari et al. (2021)
        - "split"
        - "folded"
        - "z_scale"
        - "identity"
    dask_kwargs : dict, optional
        Dask related kwargs passed to :func:`~arviz.wrap_xarray_ufunc`.

    Returns
    -------
    xarray.Dataset
      Returns dataset of the potential scale reduction factors, :math:`\hat{R}`

    See Also
    --------
    ess : Calculate estimate of the effective sample size (ess).
    mcse : Calculate Markov Chain Standard Error statistic.
    plot_forest : Forest plot to compare HDI intervals from a number of distributions.

    Notes
    -----
    The diagnostic is computed by:

      .. math:: \hat{R} = \sqrt{\frac{\hat{V}}{W}}

    where :math:`W` is the within-chain variance and :math:`\hat{V}` is the posterior variance
    estimate for the pooled rank-traces. This is the potential scale reduction factor, which
    converges to unity when each of the traces is a sample from the target posterior. Values
    greater than one indicate that one or more chains have not yet converged.

    Rank values are calculated over all the chains with ``scipy.stats.rankdata``.
    Each chain is split in two and normalized with the z-transform following
    Vehtari et al. (2021).

    References
    ----------
    * Vehtari et al. (2021).  Rank-normalization, folding, and
      localization: An improved Rhat for assessing convergence of
      MCMC. Bayesian analysis, 16(2):667-718.
    * Gelman et al. BDA3 (2013)
    * Brooks and Gelman (1998)
    * Gelman and Rubin (1992)

    Examples
    --------
    Calculate the R-hat using the default arguments:

    .. ipython::

        In [1]: import arviz as az
           ...: data = az.load_arviz_data("non_centered_eight")
           ...: az.rhat(data)

    Calculate the R-hat of some variables using the folded method:

    .. ipython::

        In [1]: az.rhat(data, var_names=["mu", "theta_t"], method="folded")

    """
    methods = {
        "rank": _rhat_rank,
        "split": _rhat_split,
        "folded": _rhat_folded,
        "z_scale": _rhat_z_scale,
        "identity": _rhat_identity,
    }
    if method not in methods:
        raise TypeError(
            f"R-hat method {method} not found. Valid methods are:\n{', '.join(methods)}"
        )
    rhat_func = methods[method]

    if isinstance(data, np.ndarray):
        data = np.atleast_2d(data)
        if len(data.shape) < 3:
            return rhat_func(data)

        msg = (
            "Only uni-dimensional ndarray variables are supported."
            " Please transform first to dataset with `az.convert_to_dataset`."
        )
        raise TypeError(msg)

    dataset = convert_to_dataset(data, group="posterior")
    var_names = _var_names(var_names, dataset)

    dataset = dataset if var_names is None else dataset[var_names]

    ufunc_kwargs = {"ravel": False}
    func_kwargs = {}
    return _wrap_xarray_ufunc(
        rhat_func,
        dataset,
        ufunc_kwargs=ufunc_kwargs,
        func_kwargs=func_kwargs,
        dask_kwargs=dask_kwargs,
    )


def mcse(data, *, var_names=None, method="mean", prob=None, dask_kwargs=None):
    """Calculate Markov Chain Standard Error statistic.

    Parameters
    ----------
    data : obj
        Any object that can be converted to an :class:`arviz.InferenceData` object
        Refer to documentation of :func:`arviz.convert_to_dataset` for details
        For ndarray: shape = (chain, draw).
        For n-dimensional ndarray transform first to dataset with ``az.convert_to_dataset``.
    var_names : list
        Names of variables to include in the rhat report
    method : str
        Select mcse method. Valid methods are:
        - "mean"
        - "sd"
        - "median"
        - "quantile"

    prob : float
        Quantile information.
    dask_kwargs : dict, optional
        Dask related kwargs passed to :func:`~arviz.wrap_xarray_ufunc`.

    Returns
    -------
    xarray.Dataset
        Return the msce dataset

    See Also
    --------
    ess : Compute autocovariance estimates for every lag for the input array.
    summary : Create a data frame with summary statistics.
    plot_mcse : Plot quantile or local Monte Carlo Standard Error.

    Examples
    --------
    Calculate the Markov Chain Standard Error using the default arguments:

    .. ipython::

        In [1]: import arviz as az
           ...: data = az.load_arviz_data("non_centered_eight")
           ...: az.mcse(data)

    Calculate the Markov Chain Standard Error using the quantile method:

    .. ipython::

        In [1]: az.mcse(data, method="quantile", prob=0.7)

    """
    methods = {
        "mean": _mcse_mean,
        "sd": _mcse_sd,
        "median": _mcse_median,
        "quantile": _mcse_quantile,
    }
    if method not in methods:
        raise TypeError(
            "mcse method {} not found. Valid methods are:\n{}".format(
                method, "\n    ".join(methods)
            )
        )
    mcse_func = methods[method]

    if method == "quantile" and prob is None:
        raise TypeError("Quantile (prob) information needs to be defined.")

    if isinstance(data, np.ndarray):
        data = np.atleast_2d(data)
        if len(data.shape) < 3:
            if prob is not None:
                return mcse_func(data, prob=prob)  # pylint: disable=unexpected-keyword-arg

            return mcse_func(data)

        msg = (
            "Only uni-dimensional ndarray variables are supported."
            " Please transform first to dataset with `az.convert_to_dataset`."
        )
        raise TypeError(msg)

    dataset = convert_to_dataset(data, group="posterior")
    var_names = _var_names(var_names, dataset)

    dataset = dataset if var_names is None else dataset[var_names]

    ufunc_kwargs = {"ravel": False}
    func_kwargs = {} if prob is None else {"prob": prob}
    return _wrap_xarray_ufunc(
        mcse_func,
        dataset,
        ufunc_kwargs=ufunc_kwargs,
        func_kwargs=func_kwargs,
        dask_kwargs=dask_kwargs,
    )


def ks_summary(pareto_tail_indices):
    """Display a summary of Pareto tail indices.

    Parameters
    ----------
    pareto_tail_indices : array
      Pareto tail indices.

    Returns
    -------
    df_k : dataframe
      Dataframe containing k diagnostic values.
    """
    _numba_flag = Numba.numba_flag
    if _numba_flag:
        bins = np.asarray([-np.inf, 0.5, 0.7, 1, np.inf])
        kcounts, *_ = _histogram(pareto_tail_indices, bins)
    else:
        kcounts, *_ = _histogram(pareto_tail_indices, bins=[-np.inf, 0.5, 0.7, 1, np.inf])
    kprop = kcounts / len(pareto_tail_indices) * 100
    df_k = pd.DataFrame(
        dict(_=["(good)", "(ok)", "(bad)", "(very bad)"], Count=kcounts, Pct=kprop)
    ).rename(index={0: "(-Inf, 0.5]", 1: " (0.5, 0.7]", 2: "   (0.7, 1]", 3: "   (1, Inf)"})

    if np.sum(kcounts[1:]) == 0:
        warnings.warn("All Pareto k estimates are good (k < 0.5)")
    elif np.sum(kcounts[2:]) == 0:
        warnings.warn("All Pareto k estimates are ok (k < 0.7)")

    return df_k


def _bfmi(energy):
    r"""Calculate the estimated Bayesian fraction of missing information (BFMI).

    BFMI quantifies how well momentum resampling matches the marginal energy distribution. For more
    information on BFMI, see https://arxiv.org/pdf/1604.00695v1.pdf. The current advice is that
    values smaller than 0.3 indicate poor sampling. However, this threshold is provisional and may
    change. See http://mc-stan.org/users/documentation/case-studies/pystan_workflow.html for more
    information.

    Parameters
    ----------
    energy : NumPy array
        Should be extracted from a gradient based sampler, such as in Stan or PyMC3. Typically,
        after converting a trace or fit to InferenceData, the energy will be in
        `data.sample_stats.energy`.

    Returns
    -------
    z : array
        The Bayesian fraction of missing information of the model and trace. One element per
        chain in the trace.
    """
    energy_mat = np.atleast_2d(energy)
    num = np.square(np.diff(energy_mat, axis=1)).mean(axis=1)  # pylint: disable=no-member
    if energy_mat.ndim == 2:
        den = _numba_var(svar, np.var, energy_mat, axis=1, ddof=1)
    else:
        den = np.var(energy, axis=1, ddof=1)
    return num / den


def _backtransform_ranks(arr, c=3 / 8):  # pylint: disable=invalid-name
    """Backtransformation of ranks.

    Parameters
    ----------
    arr : np.ndarray
        Ranks array
    c : float
        Fractional offset. Defaults to c = 3/8 as recommended by Blom (1958).

    Returns
    -------
    np.ndarray

    References
    ----------
    Blom, G. (1958). Statistical Estimates and Transformed Beta-Variables. Wiley; New York.
    """
    arr = np.asarray(arr)
    size = arr.size
    return (arr - c) / (size - 2 * c + 1)


def _z_scale(ary):
    """Calculate z_scale.

    Parameters
    ----------
    ary : np.ndarray

    Returns
    -------
    np.ndarray
    """
    ary = np.asarray(ary)
    if packaging.version.parse(scipy.__version__) < packaging.version.parse("1.10.0.dev0"):
        rank = stats.rankdata(ary, method="average")
    else:
        # the .ravel part is only needed to overcom a bug in scipy 1.10.0.rc1
        rank = stats.rankdata(  # pylint: disable=unexpected-keyword-arg
            ary, method="average", nan_policy="omit"
        )
    rank = _backtransform_ranks(rank)
    z = stats.norm.ppf(rank)
    z = z.reshape(ary.shape)
    return z


def _split_chains(ary):
    """Split and stack chains."""
    ary = np.asarray(ary)
    if len(ary.shape) <= 1:
        ary = np.atleast_2d(ary)
    _, n_draw = ary.shape
    half = n_draw // 2
    return _stack(ary[:, :half], ary[:, -half:])


def _z_fold(ary):
    """Fold and z-scale values."""
    ary = np.asarray(ary)
    ary = abs(ary - np.median(ary))
    ary = _z_scale(ary)
    return ary


def _rhat(ary):
    """Compute the rhat for a 2d array."""
    _numba_flag = Numba.numba_flag
    ary = np.asarray(ary, dtype=float)
    if _not_valid(ary, check_shape=False):
        return np.nan
    _, num_samples = ary.shape

    # Calculate chain mean
    chain_mean = np.mean(ary, axis=1)
    # Calculate chain variance
    chain_var = _numba_var(svar, np.var, ary, axis=1, ddof=1)
    # Calculate between-chain variance
    between_chain_variance = num_samples * _numba_var(svar, np.var, chain_mean, axis=None, ddof=1)
    # Calculate within-chain variance
    within_chain_variance = np.mean(chain_var)
    # Estimate of marginal posterior variance
    rhat_value = np.sqrt(
        (between_chain_variance / within_chain_variance + num_samples - 1) / (num_samples)
    )
    return rhat_value


def _rhat_rank(ary):
    """Compute the rank normalized rhat for 2d array.

    Computation follows https://arxiv.org/abs/1903.08008
    """
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=2)):
        return np.nan
    split_ary = _split_chains(ary)
    rhat_bulk = _rhat(_z_scale(split_ary))

    split_ary_folded = abs(split_ary - np.median(split_ary))
    rhat_tail = _rhat(_z_scale(split_ary_folded))

    rhat_rank = max(rhat_bulk, rhat_tail)
    return rhat_rank


def _rhat_folded(ary):
    """Calculate split-Rhat for folded z-values."""
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=2)):
        return np.nan
    ary = _z_fold(_split_chains(ary))
    return _rhat(ary)


def _rhat_z_scale(ary):
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=2)):
        return np.nan
    return _rhat(_z_scale(_split_chains(ary)))


def _rhat_split(ary):
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=2)):
        return np.nan
    return _rhat(_split_chains(ary))


def _rhat_identity(ary):
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=2)):
        return np.nan
    return _rhat(ary)


def _ess(ary, relative=False):
    """Compute the effective sample size for a 2D array."""
    _numba_flag = Numba.numba_flag
    ary = np.asarray(ary, dtype=float)
    if _not_valid(ary, check_shape=False):
        return np.nan
    if (np.max(ary) - np.min(ary)) < np.finfo(float).resolution:  # pylint: disable=no-member
        return ary.size
    if len(ary.shape) < 2:
        ary = np.atleast_2d(ary)
    n_chain, n_draw = ary.shape
    acov = _autocov(ary, axis=1)
    chain_mean = ary.mean(axis=1)
    mean_var = np.mean(acov[:, 0]) * n_draw / (n_draw - 1.0)
    var_plus = mean_var * (n_draw - 1.0) / n_draw
    if n_chain > 1:
        var_plus += _numba_var(svar, np.var, chain_mean, axis=None, ddof=1)

    rho_hat_t = np.zeros(n_draw)
    rho_hat_even = 1.0
    rho_hat_t[0] = rho_hat_even
    rho_hat_odd = 1.0 - (mean_var - np.mean(acov[:, 1])) / var_plus
    rho_hat_t[1] = rho_hat_odd

    # Geyer's initial positive sequence
    t = 1
    while t < (n_draw - 3) and (rho_hat_even + rho_hat_odd) > 0.0:
        rho_hat_even = 1.0 - (mean_var - np.mean(acov[:, t + 1])) / var_plus
        rho_hat_odd = 1.0 - (mean_var - np.mean(acov[:, t + 2])) / var_plus
        if (rho_hat_even + rho_hat_odd) >= 0:
            rho_hat_t[t + 1] = rho_hat_even
            rho_hat_t[t + 2] = rho_hat_odd
        t += 2

    max_t = t - 2
    # improve estimation
    if rho_hat_even > 0:
        rho_hat_t[max_t + 1] = rho_hat_even
    # Geyer's initial monotone sequence
    t = 1
    while t <= max_t - 2:
        if (rho_hat_t[t + 1] + rho_hat_t[t + 2]) > (rho_hat_t[t - 1] + rho_hat_t[t]):
            rho_hat_t[t + 1] = (rho_hat_t[t - 1] + rho_hat_t[t]) / 2.0
            rho_hat_t[t + 2] = rho_hat_t[t + 1]
        t += 2

    ess = n_chain * n_draw
    tau_hat = -1.0 + 2.0 * np.sum(rho_hat_t[: max_t + 1]) + np.sum(rho_hat_t[max_t + 1 : max_t + 2])
    tau_hat = max(tau_hat, 1 / np.log10(ess))
    ess = (1 if relative else ess) / tau_hat
    if np.isnan(rho_hat_t).any():
        ess = np.nan
    return ess


def _ess_bulk(ary, relative=False):
    """Compute the effective sample size for the bulk."""
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan
    z_scaled = _z_scale(_split_chains(ary))
    ess_bulk = _ess(z_scaled, relative=relative)
    return ess_bulk


def _ess_tail(ary, prob=None, relative=False):
    """Compute the effective sample size for the tail.

    If `prob` defined, ess = min(qess(prob), qess(1-prob))
    """
    if prob is None:
        prob = (0.05, 0.95)
    elif not isinstance(prob, Sequence):
        prob = (prob, 1 - prob)

    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan

    prob_low, prob_high = prob
    quantile_low_ess = _ess_quantile(ary, prob_low, relative=relative)
    quantile_high_ess = _ess_quantile(ary, prob_high, relative=relative)
    return min(quantile_low_ess, quantile_high_ess)


def _ess_mean(ary, relative=False):
    """Compute the effective sample size for the mean."""
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan
    return _ess(_split_chains(ary), relative=relative)


def _ess_sd(ary, relative=False):
    """Compute the effective sample size for the sd."""
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan
    ary = (ary - ary.mean()) ** 2
    return _ess(_split_chains(ary), relative=relative)


def _ess_quantile(ary, prob, relative=False):
    """Compute the effective sample size for the specific residual."""
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan
    if prob is None:
        raise TypeError("Prob not defined.")
    (quantile,) = _quantile(ary, prob)
    iquantile = ary <= quantile
    return _ess(_split_chains(iquantile), relative=relative)


def _ess_local(ary, prob, relative=False):
    """Compute the effective sample size for the specific residual."""
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan
    if prob is None:
        raise TypeError("Prob not defined.")
    if len(prob) != 2:
        raise ValueError("Prob argument in ess local must be upper and lower bound")
    quantile = _quantile(ary, prob)
    iquantile = (quantile[0] <= ary) & (ary <= quantile[1])
    return _ess(_split_chains(iquantile), relative=relative)


def _ess_z_scale(ary, relative=False):
    """Calculate ess for z-scaLe."""
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan
    return _ess(_z_scale(_split_chains(ary)), relative=relative)


def _ess_folded(ary, relative=False):
    """Calculate split-ess for folded data."""
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan
    return _ess(_z_fold(_split_chains(ary)), relative=relative)


def _ess_median(ary, relative=False):
    """Calculate split-ess for median."""
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan
    return _ess_quantile(ary, 0.5, relative=relative)


def _ess_mad(ary, relative=False):
    """Calculate split-ess for mean absolute deviance."""
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan
    ary = abs(ary - np.median(ary))
    ary = ary <= np.median(ary)
    ary = _z_scale(_split_chains(ary))
    return _ess(ary, relative=relative)


def _ess_identity(ary, relative=False):
    """Calculate ess."""
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan
    return _ess(ary, relative=relative)


def _mcse_mean(ary):
    """Compute the Markov Chain mean error."""
    _numba_flag = Numba.numba_flag
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan
    ess = _ess_mean(ary)
    if _numba_flag:
        sd = _sqrt(svar(np.ravel(ary), ddof=1), np.zeros(1))
    else:
        sd = np.std(ary, ddof=1)
    mcse_mean_value = sd / np.sqrt(ess)
    return mcse_mean_value


def _mcse_sd(ary):
    """Compute the Markov Chain sd error."""
    _numba_flag = Numba.numba_flag
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan
    sims_c2 = (ary - ary.mean()) ** 2
    ess = _ess_mean(sims_c2)
    evar = (sims_c2).mean()
    varvar = ((sims_c2**2).mean() - evar**2) / ess
    varsd = varvar / evar / 4
    if _numba_flag:
        mcse_sd_value = float(_sqrt(np.ravel(varsd), np.zeros(1)))
    else:
        mcse_sd_value = np.sqrt(varsd)
    return mcse_sd_value


def _mcse_median(ary):
    """Compute the Markov Chain median error."""
    return _mcse_quantile(ary, 0.5)


def _mcse_quantile(ary, prob):
    """Compute the Markov Chain quantile error at quantile=prob."""
    ary = np.asarray(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        return np.nan
    ess = _ess_quantile(ary, prob)
    probability = [0.1586553, 0.8413447]
    with np.errstate(invalid="ignore"):
        ppf = stats.beta.ppf(probability, ess * prob + 1, ess * (1 - prob) + 1)
    sorted_ary = np.sort(ary.ravel())
    size = sorted_ary.size
    ppf_size = ppf * size - 1
    th1 = sorted_ary[int(np.floor(np.nanmax((ppf_size[0], 0))))]
    th2 = sorted_ary[int(np.ceil(np.nanmin((ppf_size[1], size - 1))))]
    return (th2 - th1) / 2


def _mc_error(ary, batches=5, circular=False):
    """Calculate the simulation standard error, accounting for non-independent samples.

    The trace is divided into batches, and the standard deviation of the batch
    means is calculated.

    Parameters
    ----------
    ary : Numpy array
        An array containing MCMC samples
    batches : integer
        Number of batches
    circular : bool
        Whether to compute the error taking into account `ary` is a circular variable
        (in the range [-np.pi, np.pi]) or not. Defaults to False (i.e non-circular variables).

    Returns
    -------
    mc_error : float
        Simulation standard error
    """
    _numba_flag = Numba.numba_flag
    if ary.ndim > 1:
        dims = np.shape(ary)
        trace = np.transpose([t.ravel() for t in ary])

        return np.reshape([_mc_error(t, batches) for t in trace], dims[1:])

    else:
        if _not_valid(ary, check_shape=False):
            return np.nan
        if batches == 1:
            if circular:
                if _numba_flag:
                    std = _circular_standard_deviation(ary, high=np.pi, low=-np.pi)
                else:
                    std = stats.circstd(ary, high=np.pi, low=-np.pi)
            elif _numba_flag:
                std = float(_sqrt(svar(ary), np.zeros(1)).item())
            else:
                std = np.std(ary)
            return std / np.sqrt(len(ary))

        batched_traces = np.resize(ary, (batches, int(len(ary) / batches)))

        if circular:
            means = stats.circmean(batched_traces, high=np.pi, low=-np.pi, axis=1)
            if _numba_flag:
                std = _circular_standard_deviation(means, high=np.pi, low=-np.pi)
            else:
                std = stats.circstd(means, high=np.pi, low=-np.pi)
        else:
            means = np.mean(batched_traces, 1)
            std = _sqrt(svar(means), np.zeros(1)) if _numba_flag else np.std(means)
        return std / np.sqrt(batches)


def _multichain_statistics(ary, focus="mean"):
    """Calculate efficiently multichain statistics for summary.

    Parameters
    ----------
    ary : numpy.ndarray
    focus : select focus for the statistics. Deafault is mean.

    Returns
    -------
    tuple
        Order of return parameters is
            If focus equals "mean"
                - mcse_mean, mcse_sd, ess_bulk, ess_tail, r_hat
            Else if focus equals "median"
                - mcse_median, ess_median, ess_tail, r_hat
    """
    ary = np.atleast_2d(ary)
    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
        if focus == "mean":
            return np.nan, np.nan, np.nan, np.nan, np.nan
        return np.nan, np.nan, np.nan, np.nan

    z_split = _z_scale(_split_chains(ary))

    # ess tail
    quantile05, quantile95 = _quantile(ary, [0.05, 0.95])
    iquantile05 = ary <= quantile05
    quantile05_ess = _ess(_split_chains(iquantile05))
    iquantile95 = ary <= quantile95
    quantile95_ess = _ess(_split_chains(iquantile95))
    ess_tail_value = min(quantile05_ess, quantile95_ess)

    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=2)):
        rhat_value = np.nan
    else:
        # r_hat
        rhat_bulk = _rhat(z_split)
        ary_folded = np.abs(ary - np.median(ary))
        rhat_tail = _rhat(_z_scale(_split_chains(ary_folded)))
        rhat_value = max(rhat_bulk, rhat_tail)

    if focus == "mean":
        # ess mean
        ess_mean_value = _ess_mean(ary)

        # mcse_mean
        sims_c2 = (ary - ary.mean()) ** 2
        sims_c2_sum = sims_c2.sum()
        var = sims_c2_sum / (sims_c2.size - 1)
        mcse_mean_value = np.sqrt(var / ess_mean_value)

        # ess bulk
        ess_bulk_value = _ess(z_split)

        # mcse_sd
        evar = sims_c2_sum / sims_c2.size
        ess_mean_sims = _ess_mean(sims_c2)
        varvar = ((sims_c2**2).mean() - evar**2) / ess_mean_sims
        varsd = varvar / evar / 4
        mcse_sd_value = np.sqrt(varsd)

        return (
            mcse_mean_value,
            mcse_sd_value,
            ess_bulk_value,
            ess_tail_value,
            rhat_value,
        )

    # ess median
    ess_median_value = _ess_median(ary)

    # mcse_median
    mcse_median_value = _mcse_median(ary)

    return (
        mcse_median_value,
        ess_median_value,
        ess_tail_value,
        rhat_value,
    )
