# pylint: disable=too-many-lines
"""Statistical functions in ArviZ."""

import warnings
from copy import deepcopy
from typing import List, Optional, Tuple, Union, Mapping, cast, Callable

import numpy as np
import pandas as pd
import scipy.stats as st
from xarray_einstats import stats
import xarray as xr
from scipy.optimize import minimize, LinearConstraint, Bounds
from typing_extensions import Literal

NO_GET_ARGS: bool = False  # pylint: disable=invalid-name
try:
    from typing_extensions import get_args
except ImportError:
    NO_GET_ARGS = True  # pylint: disable=invalid-name

from .. import _log
from ..data import InferenceData, convert_to_dataset, convert_to_inference_data, extract
from ..rcparams import rcParams, ScaleKeyword, ICKeyword
from ..utils import Numba, _numba_var, _var_names, get_coords
from .density_utils import get_bins as _get_bins
from .density_utils import histogram as _histogram
from .density_utils import kde as _kde
from .density_utils import _kde_linear
from .diagnostics import _mc_error, _multichain_statistics, ess
from .stats_utils import ELPDData, _circular_standard_deviation, smooth_data
from .stats_utils import get_log_likelihood as _get_log_likelihood
from .stats_utils import get_log_prior as _get_log_prior
from .stats_utils import logsumexp as _logsumexp
from .stats_utils import make_ufunc as _make_ufunc
from .stats_utils import stats_variance_2d as svar
from .stats_utils import wrap_xarray_ufunc as _wrap_xarray_ufunc
from ..sel_utils import xarray_var_iter
from ..labels import BaseLabeller


__all__ = [
    "apply_test_function",
    "bayes_factor",
    "compare",
    "hdi",
    "loo",
    "loo_pit",
    "psislw",
    "r2_samples",
    "r2_score",
    "summary",
    "waic",
    "weight_predictions",
    "_calculate_ics",
    "psens",
]


def compare(
    compare_dict: Mapping[str, InferenceData],
    ic: Optional[ICKeyword] = None,
    method: Literal["stacking", "BB-pseudo-BMA", "pseudo-BMA"] = "stacking",
    b_samples: int = 1000,
    alpha: float = 1,
    seed=None,
    scale: Optional[ScaleKeyword] = None,
    var_name: Optional[str] = None,
):
    r"""Compare models based on  their expected log pointwise predictive density (ELPD).

    The ELPD is estimated either by Pareto smoothed importance sampling leave-one-out
    cross-validation (LOO) or using the widely applicable information criterion (WAIC).
    We recommend loo. Read more theory here - in a paper by some of the
    leading authorities on model comparison dx.doi.org/10.1111/1467-9868.00353

    Parameters
    ----------
    compare_dict: dict of {str: InferenceData or ELPDData}
        A dictionary of model names and :class:`arviz.InferenceData` or ``ELPDData``.
    ic: str, optional
        Method to estimate the ELPD, available options are "loo" or "waic". Defaults to
        ``rcParams["stats.information_criterion"]``.
    method: str, optional
        Method used to estimate the weights for each model. Available options are:

        - 'stacking' : stacking of predictive distributions.
        - 'BB-pseudo-BMA' : pseudo-Bayesian Model averaging using Akaike-type
          weighting. The weights are stabilized using the Bayesian bootstrap.
        - 'pseudo-BMA': pseudo-Bayesian Model averaging using Akaike-type
          weighting, without Bootstrap stabilization (not recommended).

        For more information read https://arxiv.org/abs/1704.02030
    b_samples: int, optional default = 1000
        Number of samples taken by the Bayesian bootstrap estimation.
        Only useful when method = 'BB-pseudo-BMA'.
        Defaults to ``rcParams["stats.ic_compare_method"]``.
    alpha: float, optional
        The shape parameter in the Dirichlet distribution used for the Bayesian bootstrap. Only
        useful when method = 'BB-pseudo-BMA'. When alpha=1 (default), the distribution is uniform
        on the simplex. A smaller alpha will keeps the final weights more away from 0 and 1.
    seed: int or np.random.RandomState instance, optional
        If int or RandomState, use it for seeding Bayesian bootstrap. Only
        useful when method = 'BB-pseudo-BMA'. Default None the global
        :mod:`numpy.random` state is used.
    scale: str, optional
        Output scale for IC. Available options are:

        - `log` : (default) log-score (after Vehtari et al. (2017))
        - `negative_log` : -1 * (log-score)
        - `deviance` : -2 * (log-score)

        A higher log-score (or a lower deviance) indicates a model with better predictive
        accuracy.
    var_name: str, optional
        If there is more than a single observed variable in the ``InferenceData``, which
        should be used as the basis for comparison.

    Returns
    -------
    A DataFrame, ordered from best to worst model (measured by the ELPD).
    The index reflects the key with which the models are passed to this function. The columns are:
    rank: The rank-order of the models. 0 is the best.
    elpd: ELPD estimated either using (PSIS-LOO-CV `elpd_loo` or WAIC `elpd_waic`).
        Higher ELPD indicates higher out-of-sample predictive fit ("better" model).
        If `scale` is `deviance` or `negative_log` smaller values indicates
        higher out-of-sample predictive fit ("better" model).
    pIC: Estimated effective number of parameters.
    elpd_diff: The difference in ELPD between two models.
        If more than two models are compared, the difference is computed relative to the
        top-ranked model, that always has a elpd_diff of 0.
    weight: Relative weight for each model.
        This can be loosely interpreted as the probability of each model (among the compared model)
        given the data. By default the uncertainty in the weights estimation is considered using
        Bayesian bootstrap.
    SE: Standard error of the ELPD estimate.
        If method = BB-pseudo-BMA these values are estimated using Bayesian bootstrap.
    dSE: Standard error of the difference in ELPD between each model and the top-ranked model.
        It's always 0 for the top-ranked model.
    warning: A value of 1 indicates that the computation of the ELPD may not be reliable.
        This could be indication of WAIC/LOO starting to fail see
        http://arxiv.org/abs/1507.04544 for details.
    scale: Scale used for the ELPD.

    Examples
    --------
    Compare the centered and non centered models of the eight school problem:

    .. ipython::
       :okwarning:

        In [1]: import arviz as az
           ...: data1 = az.load_arviz_data("non_centered_eight")
           ...: data2 = az.load_arviz_data("centered_eight")
           ...: compare_dict = {"non centered": data1, "centered": data2}
           ...: az.compare(compare_dict)

    Compare the models using PSIS-LOO-CV, returning the ELPD in log scale and calculating the
    weights using the stacking method.

    .. ipython::
       :okwarning:

        In [1]: az.compare(compare_dict, ic="loo", method="stacking", scale="log")

    See Also
    --------
    loo :
        Compute the ELPD using the Pareto smoothed importance sampling Leave-one-out
        cross-validation method.
    waic : Compute the ELPD using the widely applicable information criterion.
    plot_compare : Summary plot for model comparison.

    References
    ----------
    .. [1] Vehtari, A., Gelman, A. & Gabry, J. Practical Bayesian model evaluation using
        leave-one-out cross-validation and WAIC. Stat Comput 27, 1413–1432 (2017)
        see https://doi.org/10.1007/s11222-016-9696-4

    """
    try:
        (ics_dict, scale, ic) = _calculate_ics(compare_dict, scale=scale, ic=ic, var_name=var_name)
    except Exception as e:
        raise e.__class__("Encountered error in ELPD computation of compare.") from e
    names = list(ics_dict.keys())
    if ic in {"loo", "waic"}:
        df_comp = pd.DataFrame(
            {
                "rank": pd.Series(index=names, dtype="int"),
                f"elpd_{ic}": pd.Series(index=names, dtype="float"),
                f"p_{ic}": pd.Series(index=names, dtype="float"),
                "elpd_diff": pd.Series(index=names, dtype="float"),
                "weight": pd.Series(index=names, dtype="float"),
                "se": pd.Series(index=names, dtype="float"),
                "dse": pd.Series(index=names, dtype="float"),
                "warning": pd.Series(index=names, dtype="boolean"),
                "scale": pd.Series(index=names, dtype="str"),
            }
        )
    else:
        raise NotImplementedError(f"The information criterion {ic} is not supported.")

    if scale == "log":
        scale_value = 1
        ascending = False
    else:
        if scale == "negative_log":
            scale_value = -1
        else:
            scale_value = -2
        ascending = True

    method = rcParams["stats.ic_compare_method"] if method is None else method
    if method.lower() not in ["stacking", "bb-pseudo-bma", "pseudo-bma"]:
        raise ValueError(f"The method {method}, to compute weights, is not supported.")

    p_ic = f"p_{ic}"
    ic_i = f"{ic}_i"

    ics = pd.DataFrame.from_dict(ics_dict, orient="index")
    ics.sort_values(by=f"elpd_{ic}", inplace=True, ascending=ascending)
    ics[ic_i] = ics[ic_i].apply(lambda x: x.values.flatten())

    if method.lower() == "stacking":
        rows, cols, ic_i_val = _ic_matrix(ics, ic_i)
        exp_ic_i = np.exp(ic_i_val / scale_value)

        def log_score(weights):
            return -np.sum(np.log(exp_ic_i @ weights))

        def gradient(weights):
            denominator = exp_ic_i @ weights
            return -np.sum(exp_ic_i / denominator[:, np.newaxis], axis=0)

        theta = np.full(cols, 1.0 / cols)
        bounds = Bounds(lb=np.zeros(cols), ub=np.ones(cols))
        constraints = LinearConstraint(np.ones(cols), lb=1.0, ub=1.0)

        minimize_result = minimize(
            fun=log_score, x0=theta, jac=gradient, bounds=bounds, constraints=constraints
        )

        weights = minimize_result["x"]
        ses = ics["se"]

    elif method.lower() == "bb-pseudo-bma":
        rows, cols, ic_i_val = _ic_matrix(ics, ic_i)
        ic_i_val = ic_i_val * rows

        b_weighting = st.dirichlet.rvs(alpha=[alpha] * rows, size=b_samples, random_state=seed)
        weights = np.zeros((b_samples, cols))
        z_bs = np.zeros_like(weights)
        for i in range(b_samples):
            z_b = np.dot(b_weighting[i], ic_i_val)
            u_weights = np.exp((z_b - np.max(z_b)) / scale_value)
            z_bs[i] = z_b  # pylint: disable=unsupported-assignment-operation
            weights[i] = u_weights / np.sum(u_weights)

        weights = weights.mean(axis=0)
        ses = pd.Series(z_bs.std(axis=0), index=ics.index)  # pylint: disable=no-member

    elif method.lower() == "pseudo-bma":
        min_ic = ics.iloc[0][f"elpd_{ic}"]
        z_rv = np.exp((ics[f"elpd_{ic}"] - min_ic) / scale_value)
        weights = (z_rv / np.sum(z_rv)).to_numpy()
        ses = ics["se"]

    if np.any(weights):
        min_ic_i_val = ics[ic_i].iloc[0]
        for idx, val in enumerate(ics.index):
            res = ics.loc[val]
            if scale_value < 0:
                diff = res[ic_i] - min_ic_i_val
            else:
                diff = min_ic_i_val - res[ic_i]
            d_ic = np.sum(diff)
            d_std_err = np.sqrt(len(diff) * np.var(diff))
            std_err = ses.loc[val]
            weight = weights[idx]
            df_comp.loc[val] = (
                idx,
                res[f"elpd_{ic}"],
                res[p_ic],
                d_ic,
                weight,
                std_err,
                d_std_err,
                res["warning"],
                res["scale"],
            )

    df_comp["rank"] = df_comp["rank"].astype(int)
    df_comp["warning"] = df_comp["warning"].astype(bool)
    return df_comp.sort_values(by=f"elpd_{ic}", ascending=ascending)


def _ic_matrix(ics, ic_i):
    """Store the previously computed pointwise predictive accuracy values (ics) in a 2D matrix."""
    cols, _ = ics.shape
    rows = len(ics[ic_i].iloc[0])
    ic_i_val = np.zeros((rows, cols))

    for idx, val in enumerate(ics.index):
        ic = ics.loc[val][ic_i]

        if len(ic) != rows:
            raise ValueError("The number of observations should be the same across all models")

        ic_i_val[:, idx] = ic

    return rows, cols, ic_i_val


def _calculate_ics(
    compare_dict,
    scale: Optional[ScaleKeyword] = None,
    ic: Optional[ICKeyword] = None,
    var_name: Optional[str] = None,
):
    """Calculate LOO or WAIC only if necessary.

    It always calls the ic function with ``pointwise=True``.

    Parameters
    ----------
    compare_dict :  dict of {str : InferenceData or ELPDData}
        A dictionary of model names and InferenceData or ELPDData objects
    scale : str, optional
        Output scale for IC. Available options are:

        - `log` : (default) log-score (after Vehtari et al. (2017))
        - `negative_log` : -1 * (log-score)
        - `deviance` : -2 * (log-score)

        A higher log-score (or a lower deviance) indicates a model with better predictive accuracy.
    ic : str, optional
        Information Criterion (PSIS-LOO `loo` or WAIC `waic`) used to compare models.
        Defaults to ``rcParams["stats.information_criterion"]``.
    var_name : str, optional
        Name of the variable storing pointwise log likelihood values in ``log_likelihood`` group.


    Returns
    -------
    compare_dict : dict of ELPDData
    scale : str
    ic : str

    """
    precomputed_elpds = {
        name: elpd_data
        for name, elpd_data in compare_dict.items()
        if isinstance(elpd_data, ELPDData)
    }
    precomputed_ic = None
    precomputed_scale = None
    if precomputed_elpds:
        _, arbitrary_elpd = precomputed_elpds.popitem()
        precomputed_ic = arbitrary_elpd.index[0].split("_")[1]
        precomputed_scale = arbitrary_elpd["scale"]
        raise_non_pointwise = f"{precomputed_ic}_i" not in arbitrary_elpd
        if any(
            elpd_data.index[0].split("_")[1] != precomputed_ic
            for elpd_data in precomputed_elpds.values()
        ):
            raise ValueError(
                "All information criteria to be compared must be the same "
                "but found both loo and waic."
            )
        if any(elpd_data["scale"] != precomputed_scale for elpd_data in precomputed_elpds.values()):
            raise ValueError("All information criteria to be compared must use the same scale")
        if (
            any(f"{precomputed_ic}_i" not in elpd_data for elpd_data in precomputed_elpds.values())
            or raise_non_pointwise
        ):
            raise ValueError("Not all provided ELPDData have been calculated with pointwise=True")
        if ic is not None and ic.lower() != precomputed_ic:
            warnings.warn(
                "Provided ic argument is incompatible with precomputed elpd data. "
                f"Using ic from precomputed elpddata: {precomputed_ic}"
            )
            ic = precomputed_ic
        if scale is not None and scale.lower() != precomputed_scale:
            warnings.warn(
                "Provided scale argument is incompatible with precomputed elpd data. "
                f"Using scale from precomputed elpddata: {precomputed_scale}"
            )
            scale = precomputed_scale

    if ic is None and precomputed_ic is None:
        ic = cast(ICKeyword, rcParams["stats.information_criterion"])
    elif ic is None:
        ic = precomputed_ic
    else:
        ic = cast(ICKeyword, ic.lower())
    allowable = ["loo", "waic"] if NO_GET_ARGS else get_args(ICKeyword)
    if ic not in allowable:
        raise ValueError(f"{ic} is not a valid value for ic: must be in {allowable}")

    if scale is None and precomputed_scale is None:
        scale = cast(ScaleKeyword, rcParams["stats.ic_scale"])
    elif scale is None:
        scale = precomputed_scale
    else:
        scale = cast(ScaleKeyword, scale.lower())
    allowable = ["log", "negative_log", "deviance"] if NO_GET_ARGS else get_args(ScaleKeyword)
    if scale not in allowable:
        raise ValueError(f"{scale} is not a valid value for scale: must be in {allowable}")

    if ic == "loo":
        ic_func: Callable = loo
    elif ic == "waic":
        ic_func = waic
    else:
        raise NotImplementedError(f"The information criterion {ic} is not supported.")

    compare_dict = deepcopy(compare_dict)
    for name, dataset in compare_dict.items():
        if not isinstance(dataset, ELPDData):
            try:
                compare_dict[name] = ic_func(
                    convert_to_inference_data(dataset),
                    pointwise=True,
                    scale=scale,
                    var_name=var_name,
                )
            except Exception as e:
                raise e.__class__(
                    f"Encountered error trying to compute {ic} from model {name}."
                ) from e
    return (compare_dict, scale, ic)


def hdi(
    ary,
    hdi_prob=None,
    circular=False,
    multimodal=False,
    skipna=False,
    group="posterior",
    var_names=None,
    filter_vars=None,
    coords=None,
    max_modes=10,
    dask_kwargs=None,
    **kwargs,
):
    """
    Calculate highest density interval (HDI) of array for given probability.

    The HDI is the minimum width Bayesian credible interval (BCI).

    Parameters
    ----------
    ary: obj
        object containing posterior samples.
        Any object that can be converted to an :class:`arviz.InferenceData` object.
        Refer to documentation of :func:`arviz.convert_to_dataset` for details.
    hdi_prob: float, optional
        Prob for which the highest density interval will be computed. Defaults to
        ``stats.ci_prob`` rcParam.
    circular: bool, optional
        Whether to compute the hdi taking into account `x` is a circular variable
        (in the range [-np.pi, np.pi]) or not. Defaults to False (i.e non-circular variables).
        Only works if multimodal is False.
    multimodal: bool, optional
        If true it may compute more than one hdi if the distribution is multimodal and the
        modes are well separated.
    skipna: bool, optional
        If true ignores nan values when computing the hdi. Defaults to false.
    group: str, optional
        Specifies which InferenceData group should be used to calculate hdi.
        Defaults to 'posterior'
    var_names: list, optional
        Names of variables to include in the hdi report. Prefix the variables by ``~``
        when you want to exclude them from the report: `["~beta"]` instead of `["beta"]`
        (see :func:`arviz.summary` for more details).
    filter_vars: {None, "like", "regex"}, optional, default=None
        If `None` (default), interpret var_names as the real variables names. If "like",
        interpret var_names as substrings of the real variables names. If "regex",
        interpret var_names as regular expressions on the real variables names. A la
        ``pandas.filter``.
    coords: mapping, optional
        Specifies the subset over to calculate hdi.
    max_modes: int, optional
        Specifies the maximum number of modes for multimodal case.
    dask_kwargs : dict, optional
        Dask related kwargs passed to :func:`~arviz.wrap_xarray_ufunc`.
    kwargs: dict, optional
        Additional keywords passed to :func:`~arviz.wrap_xarray_ufunc`.

    Returns
    -------
    np.ndarray or xarray.Dataset, depending upon input
        lower(s) and upper(s) values of the interval(s).

    See Also
    --------
    plot_hdi : Plot highest density intervals for regression data.
    xarray.Dataset.quantile : Calculate quantiles of array for given probabilities.

    Examples
    --------
    Calculate the HDI of a Normal random variable:

    .. ipython::

        In [1]: import arviz as az
           ...: import numpy as np
           ...: data = np.random.normal(size=2000)
           ...: az.hdi(data, hdi_prob=.68)

    Calculate the HDI of a dataset:

    .. ipython::

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

    We can also calculate the HDI of some of the variables of dataset:

    .. ipython::

        In [1]: az.hdi(data, var_names=["mu", "theta"])

    By default, ``hdi`` is calculated over the ``chain`` and ``draw`` dimensions. We can use the
    ``input_core_dims`` argument of :func:`~arviz.wrap_xarray_ufunc` to change this. In this example
    we calculate the HDI also over the ``school`` dimension:

    .. ipython::

        In [1]: az.hdi(data, var_names="theta", input_core_dims = [["chain","draw", "school"]])

    We can also calculate the hdi over a particular selection:

    .. ipython::

        In [1]: az.hdi(data, coords={"chain":[0, 1, 3]}, input_core_dims = [["draw"]])

    """
    if hdi_prob is None:
        hdi_prob = rcParams["stats.ci_prob"]
    elif not 1 >= hdi_prob > 0:
        raise ValueError("The value of hdi_prob should be in the interval (0, 1]")

    func_kwargs = {
        "hdi_prob": hdi_prob,
        "skipna": skipna,
        "out_shape": (max_modes, 2) if multimodal else (2,),
    }
    kwargs.setdefault("output_core_dims", [["mode", "hdi"] if multimodal else ["hdi"]])
    if not multimodal:
        func_kwargs["circular"] = circular
    else:
        func_kwargs["max_modes"] = max_modes

    func = _hdi_multimodal if multimodal else _hdi

    isarray = isinstance(ary, np.ndarray)
    if isarray and ary.ndim <= 1:
        func_kwargs.pop("out_shape")
        hdi_data = func(ary, **func_kwargs)  # pylint: disable=unexpected-keyword-arg
        return hdi_data[~np.isnan(hdi_data).all(axis=1), :] if multimodal else hdi_data

    if isarray and ary.ndim == 2:
        warnings.warn(
            "hdi currently interprets 2d data as (draw, shape) but this will change in "
            "a future release to (chain, draw) for coherence with other functions",
            FutureWarning,
            stacklevel=2,
        )
        ary = np.expand_dims(ary, 0)

    ary = convert_to_dataset(ary, group=group)
    if coords is not None:
        ary = get_coords(ary, coords)
    var_names = _var_names(var_names, ary, filter_vars)
    ary = ary[var_names] if var_names else ary

    hdi_coord = xr.DataArray(["lower", "higher"], dims=["hdi"], attrs=dict(hdi_prob=hdi_prob))
    hdi_data = _wrap_xarray_ufunc(
        func, ary, func_kwargs=func_kwargs, dask_kwargs=dask_kwargs, **kwargs
    ).assign_coords({"hdi": hdi_coord})
    hdi_data = hdi_data.dropna("mode", how="all") if multimodal else hdi_data
    return hdi_data.x.values if isarray else hdi_data


def _hdi(ary, hdi_prob, circular, skipna):
    """Compute hpi over the flattened array."""
    ary = ary.flatten()
    if skipna:
        nans = np.isnan(ary)
        if not nans.all():
            ary = ary[~nans]
    n = len(ary)

    if circular:
        mean = st.circmean(ary, high=np.pi, low=-np.pi)
        ary = ary - mean
        ary = np.arctan2(np.sin(ary), np.cos(ary))

    ary = np.sort(ary)
    interval_idx_inc = int(np.floor(hdi_prob * n))
    n_intervals = n - interval_idx_inc
    interval_width = np.subtract(ary[interval_idx_inc:], ary[:n_intervals], dtype=np.float64)

    if len(interval_width) == 0:
        raise ValueError("Too few elements for interval calculation. ")

    min_idx = np.argmin(interval_width)
    hdi_min = ary[min_idx]
    hdi_max = ary[min_idx + interval_idx_inc]

    if circular:
        hdi_min = hdi_min + mean
        hdi_max = hdi_max + mean
        hdi_min = np.arctan2(np.sin(hdi_min), np.cos(hdi_min))
        hdi_max = np.arctan2(np.sin(hdi_max), np.cos(hdi_max))

    hdi_interval = np.array([hdi_min, hdi_max])

    return hdi_interval


def _hdi_multimodal(ary, hdi_prob, skipna, max_modes):
    """Compute HDI if the distribution is multimodal."""
    ary = ary.flatten()
    if skipna:
        ary = ary[~np.isnan(ary)]

    if ary.dtype.kind == "f":
        bins, density = _kde(ary)
        lower, upper = bins[0], bins[-1]
        range_x = upper - lower
        dx = range_x / len(density)
    else:
        bins = _get_bins(ary)
        _, density, _ = _histogram(ary, bins=bins)
        dx = np.diff(bins)[0]

    density *= dx

    idx = np.argsort(-density)
    intervals = bins[idx][density[idx].cumsum() <= hdi_prob]
    intervals.sort()

    intervals_splitted = np.split(intervals, np.where(np.diff(intervals) >= dx * 1.1)[0] + 1)

    hdi_intervals = np.full((max_modes, 2), np.nan)
    for i, interval in enumerate(intervals_splitted):
        if i == max_modes:
            warnings.warn(
                f"found more modes than {max_modes}, returning only the first {max_modes} modes"
            )
            break
        if interval.size == 0:
            hdi_intervals[i] = np.asarray([bins[0], bins[0]])
        else:
            hdi_intervals[i] = np.asarray([interval[0], interval[-1]])

    return np.array(hdi_intervals)


def loo(data, pointwise=None, var_name=None, reff=None, scale=None):
    """Compute Pareto-smoothed importance sampling leave-one-out cross-validation (PSIS-LOO-CV).

    Estimates the expected log pointwise predictive density (elpd) using Pareto-smoothed
    importance sampling leave-one-out cross-validation (PSIS-LOO-CV). Also calculates LOO's
    standard error and the effective number of parameters. Read more theory here
    https://arxiv.org/abs/1507.04544 and here https://arxiv.org/abs/1507.02646

    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.
    pointwise: bool, optional
        If True the pointwise predictive accuracy will be returned. Defaults to
        ``stats.ic_pointwise`` rcParam.
    var_name : str, optional
        The name of the variable in log_likelihood groups storing the pointwise log
        likelihood data to use for loo computation.
    reff: float, optional
        Relative MCMC efficiency, ``ess / n`` i.e. number of effective samples divided by the number
        of actual samples. Computed from trace by default.
    scale: str
        Output scale for loo. Available options are:

        - ``log`` : (default) log-score
        - ``negative_log`` : -1 * log-score
        - ``deviance`` : -2 * log-score

        A higher log-score (or a lower deviance or negative log_score) indicates a model with
        better predictive accuracy.

    Returns
    -------
    ELPDData object (inherits from :class:`pandas.Series`) with the following row/attributes:
    elpd_loo: approximated expected log pointwise predictive density (elpd)
    se: standard error of the elpd
    p_loo: effective number of parameters
    n_samples: number of samples
    n_data_points: number of data points
    warning: bool
        True if the estimated shape parameter of Pareto distribution is greater than
        ``good_k``.
    loo_i: :class:`~xarray.DataArray` with the pointwise predictive accuracy,
            only if pointwise=True
    pareto_k: array of Pareto shape values, only if pointwise True
    scale: scale of the elpd
    good_k: For a sample size S, the thresold is compute as min(1 - 1/log10(S), 0.7)

        The returned object has a custom print method that overrides pd.Series method.

    See Also
    --------
    compare : Compare models based on PSIS-LOO loo or WAIC waic cross-validation.
    waic : Compute the widely applicable information criterion.
    plot_compare : Summary plot for model comparison.
    plot_elpd : Plot pointwise elpd differences between two or more models.
    plot_khat : Plot Pareto tail indices for diagnosing convergence.

    Examples
    --------
    Calculate LOO of a model:

    .. ipython::

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

    Calculate LOO of a model and return the pointwise values:

    .. ipython::

        In [2]: data_loo = az.loo(data, pointwise=True)
           ...: data_loo.loo_i
    """
    inference_data = convert_to_inference_data(data)
    log_likelihood = _get_log_likelihood(inference_data, var_name=var_name)
    pointwise = rcParams["stats.ic_pointwise"] if pointwise is None else pointwise

    log_likelihood = log_likelihood.stack(__sample__=("chain", "draw"))
    shape = log_likelihood.shape
    n_samples = shape[-1]
    n_data_points = np.prod(shape[:-1])
    scale = rcParams["stats.ic_scale"] if scale is None else scale.lower()

    if scale == "deviance":
        scale_value = -2
    elif scale == "log":
        scale_value = 1
    elif scale == "negative_log":
        scale_value = -1
    else:
        raise TypeError('Valid scale values are "deviance", "log", "negative_log"')

    if reff is None:
        if not hasattr(inference_data, "posterior"):
            raise TypeError("Must be able to extract a posterior group from data.")
        posterior = inference_data.posterior
        n_chains = len(posterior.chain)
        if n_chains == 1:
            reff = 1.0
        else:
            ess_p = ess(posterior, method="mean")
            # this mean is over all data variables
            reff = (
                np.hstack([ess_p[v].values.flatten() for v in ess_p.data_vars]).mean() / n_samples
            )

    log_weights, pareto_shape = psislw(-log_likelihood, reff)
    log_weights += log_likelihood

    warn_mg = False
    good_k = min(1 - 1 / np.log10(n_samples), 0.7)

    if np.any(pareto_shape > good_k):
        warnings.warn(
            f"Estimated shape parameter of Pareto distribution is greater than {good_k:.2f} "
            "for one or more samples. You should consider using a more robust model, this is "
            "because importance sampling is less likely to work well if the marginal posterior "
            "and LOO posterior are very different. This is more likely to happen with a "
            "non-robust model and highly influential observations."
        )
        warn_mg = True

    ufunc_kwargs = {"n_dims": 1, "ravel": False}
    kwargs = {"input_core_dims": [["__sample__"]]}
    loo_lppd_i = scale_value * _wrap_xarray_ufunc(
        _logsumexp, log_weights, ufunc_kwargs=ufunc_kwargs, **kwargs
    )
    loo_lppd = loo_lppd_i.values.sum()
    loo_lppd_se = (n_data_points * np.var(loo_lppd_i.values)) ** 0.5

    lppd = np.sum(
        _wrap_xarray_ufunc(
            _logsumexp,
            log_likelihood,
            func_kwargs={"b_inv": n_samples},
            ufunc_kwargs=ufunc_kwargs,
            **kwargs,
        ).values
    )
    p_loo = lppd - loo_lppd / scale_value

    if not pointwise:
        return ELPDData(
            data=[loo_lppd, loo_lppd_se, p_loo, n_samples, n_data_points, warn_mg, scale, good_k],
            index=[
                "elpd_loo",
                "se",
                "p_loo",
                "n_samples",
                "n_data_points",
                "warning",
                "scale",
                "good_k",
            ],
        )
    if np.equal(loo_lppd, loo_lppd_i).all():  # pylint: disable=no-member
        warnings.warn(
            "The point-wise LOO is the same with the sum LOO, please double check "
            "the Observed RV in your model to make sure it returns element-wise logp."
        )
    return ELPDData(
        data=[
            loo_lppd,
            loo_lppd_se,
            p_loo,
            n_samples,
            n_data_points,
            warn_mg,
            loo_lppd_i.rename("loo_i"),
            pareto_shape,
            scale,
            good_k,
        ],
        index=[
            "elpd_loo",
            "se",
            "p_loo",
            "n_samples",
            "n_data_points",
            "warning",
            "loo_i",
            "pareto_k",
            "scale",
            "good_k",
        ],
    )


def psislw(log_weights, reff=1.0, normalize=True):
    """
    Pareto smoothed importance sampling (PSIS).

    Notes
    -----
    If the ``log_weights`` input is an :class:`~xarray.DataArray` with a dimension
    named ``__sample__`` (recommended) ``psislw`` will interpret this dimension as samples,
    and all other dimensions as dimensions of the observed data, looping over them to
    calculate the psislw of each observation. If no ``__sample__`` dimension is present or
    the input is a numpy array, the last dimension will be interpreted as ``__sample__``.

    Parameters
    ----------
    log_weights : DataArray or (..., N) array-like
        Array of size (n_observations, n_samples)
    reff : float, default 1
        relative MCMC efficiency, ``ess / n``
    normalize : bool, default True
        return normalized log weights

    Returns
    -------
    lw_out : DataArray or (..., N) ndarray
        Smoothed, truncated and possibly normalized log weights.
    kss : DataArray or (...) ndarray
        Estimates of the shape parameter *k* of the generalized Pareto
        distribution.

    References
    ----------
    * Vehtari et al. (2024). Pareto smoothed importance sampling. Journal of Machine
      Learning Research, 25(72):1-58.

    See Also
    --------
    loo : Compute Pareto-smoothed importance sampling leave-one-out cross-validation (PSIS-LOO-CV).

    Examples
    --------
    Get Pareto smoothed importance sampling (PSIS) log weights:

    .. ipython::

        In [1]: import arviz as az
           ...: data = az.load_arviz_data("non_centered_eight")
           ...: log_likelihood = data.log_likelihood["obs"].stack(
           ...:     __sample__=["chain", "draw"]
           ...: )
           ...: az.psislw(-log_likelihood, reff=0.8)

    """
    log_weights = deepcopy(log_weights)
    if hasattr(log_weights, "__sample__"):
        n_samples = len(log_weights.__sample__)
        shape = [
            size for size, dim in zip(log_weights.shape, log_weights.dims) if dim != "__sample__"
        ]
    else:
        n_samples = log_weights.shape[-1]
        shape = log_weights.shape[:-1]
    # precalculate constants
    cutoff_ind = -int(np.ceil(min(n_samples / 5.0, 3 * (n_samples / reff) ** 0.5))) - 1
    cutoffmin = np.log(np.finfo(float).tiny)  # pylint: disable=no-member, assignment-from-no-return

    # create output array with proper dimensions
    out = np.empty_like(log_weights), np.empty(shape)

    # define kwargs
    func_kwargs = {
        "cutoff_ind": cutoff_ind,
        "cutoffmin": cutoffmin,
        "out": out,
        "normalize": normalize,
    }
    ufunc_kwargs = {"n_dims": 1, "n_output": 2, "ravel": False, "check_shape": False}
    kwargs = {"input_core_dims": [["__sample__"]], "output_core_dims": [["__sample__"], []]}
    log_weights, pareto_shape = _wrap_xarray_ufunc(
        _psislw,
        log_weights,
        ufunc_kwargs=ufunc_kwargs,
        func_kwargs=func_kwargs,
        **kwargs,
    )
    if isinstance(log_weights, xr.DataArray):
        log_weights = log_weights.rename("log_weights")
    if isinstance(pareto_shape, xr.DataArray):
        pareto_shape = pareto_shape.rename("pareto_shape")
    return log_weights, pareto_shape


def _psislw(log_weights, cutoff_ind, cutoffmin, normalize):
    """
    Pareto smoothed importance sampling (PSIS) for a 1D vector.

    Parameters
    ----------
    log_weights: array
        Array of length n_observations
    cutoff_ind: int
    cutoffmin: float
    normalize: bool

    Returns
    -------
    lw_out: array
        Smoothed log weights
    kss: float
        Pareto tail index
    """
    x = np.asarray(log_weights)

    # improve numerical accuracy
    max_x = np.max(x)
    x -= max_x
    # sort the array
    x_sort_ind = np.argsort(x)
    # divide log weights into body and right tail
    xcutoff = max(x[x_sort_ind[cutoff_ind]], cutoffmin)

    expxcutoff = np.exp(xcutoff)
    (tailinds,) = np.where(x > xcutoff)  # pylint: disable=unbalanced-tuple-unpacking
    x_tail = x[tailinds]
    tail_len = len(x_tail)
    if tail_len <= 4:
        # not enough tail samples for gpdfit
        k = np.inf
    else:
        # order of tail samples
        x_tail_si = np.argsort(x_tail)
        # fit generalized Pareto distribution to the right tail samples
        x_tail = np.exp(x_tail) - expxcutoff
        k, sigma = _gpdfit(x_tail[x_tail_si])

        if np.isfinite(k):
            # no smoothing if GPD fit failed
            # compute ordered statistic for the fit
            sti = np.arange(0.5, tail_len) / tail_len
            smoothed_tail = _gpinv(sti, k, sigma)
            smoothed_tail = np.log(  # pylint: disable=assignment-from-no-return
                smoothed_tail + expxcutoff
            )
            # place the smoothed tail into the output array
            x[tailinds[x_tail_si]] = smoothed_tail
            # truncate smoothed values to the largest raw weight 0
            x[x > 0] = 0

    # renormalize weights
    if normalize:
        x -= _logsumexp(x)
    else:
        x += max_x

    return x, k


def _gpdfit(ary):
    """Estimate the parameters for the Generalized Pareto Distribution (GPD).

    Empirical Bayes estimate for the parameters of the generalized Pareto
    distribution given the data.

    Parameters
    ----------
    ary: array
        sorted 1D data array

    Returns
    -------
    k: float
        estimated shape parameter
    sigma: float
        estimated scale parameter
    """
    prior_bs = 3
    prior_k = 10
    n = len(ary)
    m_est = 30 + int(n**0.5)

    b_ary = 1 - np.sqrt(m_est / (np.arange(1, m_est + 1, dtype=float) - 0.5))
    b_ary /= prior_bs * ary[int(n / 4 + 0.5) - 1]
    b_ary += 1 / ary[-1]

    k_ary = np.log1p(-b_ary[:, None] * ary).mean(axis=1)  # pylint: disable=no-member
    len_scale = n * (np.log(-(b_ary / k_ary)) - k_ary - 1)
    weights = 1 / np.exp(len_scale - len_scale[:, None]).sum(axis=1)

    # remove negligible weights
    real_idxs = weights >= 10 * np.finfo(float).eps
    if not np.all(real_idxs):
        weights = weights[real_idxs]
        b_ary = b_ary[real_idxs]
    # normalise weights
    weights /= weights.sum()

    # posterior mean for b
    b_post = np.sum(b_ary * weights)
    # estimate for k
    k_post = np.log1p(-b_post * ary).mean()  # pylint: disable=invalid-unary-operand-type,no-member
    # add prior for k_post
    sigma = -k_post / b_post
    k_post = (n * k_post + prior_k * 0.5) / (n + prior_k)

    return k_post, sigma


def _gpinv(probs, kappa, sigma):
    """Inverse Generalized Pareto distribution function."""
    # pylint: disable=unsupported-assignment-operation, invalid-unary-operand-type
    x = np.full_like(probs, np.nan)
    if sigma <= 0:
        return x
    ok = (probs > 0) & (probs < 1)
    if np.all(ok):
        if np.abs(kappa) < np.finfo(float).eps:
            x = -np.log1p(-probs)
        else:
            x = np.expm1(-kappa * np.log1p(-probs)) / kappa
        x *= sigma
    else:
        if np.abs(kappa) < np.finfo(float).eps:
            x[ok] = -np.log1p(-probs[ok])
        else:
            x[ok] = np.expm1(-kappa * np.log1p(-probs[ok])) / kappa
        x *= sigma
        x[probs == 0] = 0
        x[probs == 1] = np.inf if kappa >= 0 else -sigma / kappa
    return x


def r2_samples(y_true, y_pred):
    """R² samples for Bayesian regression models. Only valid for linear models.

    Parameters
    ----------
    y_true: array-like of shape = (n_outputs,)
        Ground truth (correct) target values.
    y_pred: array-like of shape = (n_posterior_samples, n_outputs)
        Estimated target values.

    Returns
    -------
    Pandas Series with the following indices:
    Bayesian R² samples.

    See Also
    --------
    plot_lm : Posterior predictive and mean plots for regression-like data.

    Examples
    --------
    Calculate R² samples for Bayesian regression models :

    .. ipython::

        In [1]: import arviz as az
           ...: data = az.load_arviz_data('regression1d')
           ...: y_true = data.observed_data["y"].values
           ...: y_pred = data.posterior_predictive.stack(sample=("chain", "draw"))["y"].values.T
           ...: az.r2_samples(y_true, y_pred)

    """
    _numba_flag = Numba.numba_flag
    if y_pred.ndim == 1:
        var_y_est = _numba_var(svar, np.var, y_pred)
        var_e = _numba_var(svar, np.var, (y_true - y_pred))
    else:
        var_y_est = _numba_var(svar, np.var, y_pred, axis=1)
        var_e = _numba_var(svar, np.var, (y_true - y_pred), axis=1)
    r_squared = var_y_est / (var_y_est + var_e)

    return r_squared


def r2_score(y_true, y_pred):
    """R² for Bayesian regression models. Only valid for linear models.

    Parameters
    ----------
    y_true: array-like of shape = (n_outputs,)
        Ground truth (correct) target values.
    y_pred: array-like of shape = (n_posterior_samples, n_outputs)
        Estimated target values.

    Returns
    -------
    Pandas Series with the following indices:
    r2: Bayesian R²
    r2_std: standard deviation of the Bayesian R².

    See Also
    --------
    plot_lm : Posterior predictive and mean plots for regression-like data.

    Examples
    --------
    Calculate R² for Bayesian regression models :

    .. ipython::

        In [1]: import arviz as az
           ...: data = az.load_arviz_data('regression1d')
           ...: y_true = data.observed_data["y"].values
           ...: y_pred = data.posterior_predictive.stack(sample=("chain", "draw"))["y"].values.T
           ...: az.r2_score(y_true, y_pred)

    """
    r_squared = r2_samples(y_true=y_true, y_pred=y_pred)
    return pd.Series([np.mean(r_squared), np.std(r_squared)], index=["r2", "r2_std"])


def summary(
    data,
    var_names: Optional[List[str]] = None,
    filter_vars=None,
    group=None,
    fmt: "Literal['wide', 'long', 'xarray']" = "wide",
    kind: "Literal['all', 'stats', 'diagnostics']" = "all",
    round_to=None,
    circ_var_names=None,
    stat_focus="mean",
    stat_funcs=None,
    extend=True,
    hdi_prob=None,
    skipna=False,
    labeller=None,
    coords=None,
    index_origin=None,
    order=None,
) -> Union[pd.DataFrame, xr.Dataset]:
    """Create a data frame with summary statistics.

    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
    var_names: list
        Names of variables to include in summary. Prefix the variables by ``~`` when you
        want to exclude them from the summary: `["~beta"]` instead of `["beta"]` (see
        examples below).
    filter_vars: {None, "like", "regex"}, optional, default=None
        If `None` (default), interpret var_names as the real variables names. If "like",
        interpret var_names as substrings of the real variables names. If "regex",
        interpret var_names as regular expressions on the real variables names. A la
        ``pandas.filter``.
    coords: Dict[str, List[Any]], optional
        Coordinate subset for which to calculate the summary.
    group: str
        Select a group for summary. Defaults to "posterior", "prior" or first group
        in that order, depending what groups exists.
    fmt: {'wide', 'long', 'xarray'}
        Return format is either pandas.DataFrame {'wide', 'long'} or xarray.Dataset {'xarray'}.
    kind: {'all', 'stats', 'diagnostics'}
        Whether to include the `stats`: `mean`, `sd`, `hdi_3%`, `hdi_97%`, or the `diagnostics`:
        `mcse_mean`, `mcse_sd`, `ess_bulk`, `ess_tail`, and `r_hat`. Default to include `all` of
        them.
    round_to: int
        Number of decimals used to round results. Defaults to 2. Use "none" to return raw numbers.
    circ_var_names: list
        A list of circular variables to compute circular stats for
    stat_focus : str, default "mean"
        Select the focus for summary.
    stat_funcs: dict
        A list of functions or a dict of functions with function names as keys used to calculate
        statistics. By default, the mean, standard deviation, simulation standard error, and
        highest posterior density intervals are included.

        The functions will be given one argument, the samples for a variable as an nD array,
        The functions should be in the style of a ufunc and return a single number. For example,
        :func:`numpy.mean`, or ``scipy.stats.var`` would both work.
    extend: boolean
        If True, use the statistics returned by ``stat_funcs`` in addition to, rather than in place
        of, the default statistics. This is only meaningful when ``stat_funcs`` is not None.
    hdi_prob: float, optional
        Highest density interval to compute. Defaults to 0.94. This is only meaningful when
        ``stat_funcs`` is None.
    skipna: bool
        If true ignores nan values when computing the summary statistics, it does not affect the
        behaviour of the functions passed to ``stat_funcs``. Defaults to false.
    labeller : labeller instance, optional
        Class providing the method `make_label_flat` to generate the labels in the plot titles.
        For more details on ``labeller`` usage see :ref:`label_guide`
    credible_interval: float, optional
        deprecated: Please see hdi_prob
    order
        deprecated: order is now ignored.
    index_origin
        deprecated: index_origin is now ignored, modify the coordinate values to change the
        value used in summary.

    Returns
    -------
    pandas.DataFrame or xarray.Dataset
        Return type dicated by `fmt` argument.

        Return value will contain summary statistics for each variable. Default statistics depend on
        the value of ``stat_focus``:

        ``stat_focus="mean"``: `mean`, `sd`, `hdi_3%`, `hdi_97%`, `mcse_mean`, `mcse_sd`,
        `ess_bulk`, `ess_tail`, and `r_hat`

        ``stat_focus="median"``: `median`, `mad`, `eti_3%`, `eti_97%`, `mcse_median`, `ess_median`,
        `ess_tail`, and `r_hat`

        `r_hat` is only computed for traces with 2 or more chains.

    See Also
    --------
    waic : Compute the widely applicable information criterion.
    loo : Compute Pareto-smoothed importance sampling leave-one-out
          cross-validation (PSIS-LOO-CV).
    ess : Calculate estimate of the effective sample size (ess).
    rhat : Compute estimate of rank normalized splitR-hat for a set of traces.
    mcse : Calculate Markov Chain Standard Error statistic.

    Examples
    --------
    .. ipython::

        In [1]: import arviz as az
           ...: data = az.load_arviz_data("centered_eight")
           ...: az.summary(data, var_names=["mu", "tau"])

    You can use ``filter_vars`` to select variables without having to specify all the exact
    names. Use ``filter_vars="like"`` to select based on partial naming:

    .. ipython::

        In [1]: az.summary(data, var_names=["the"], filter_vars="like")

    Use ``filter_vars="regex"`` to select based on regular expressions, and prefix the variables
    you want to exclude by ``~``. Here, we exclude from the summary all the variables
    starting with the letter t:

    .. ipython::

        In [1]: az.summary(data, var_names=["~^t"], filter_vars="regex")

    Other statistics can be calculated by passing a list of functions
    or a dictionary with key, function pairs.

    .. ipython::

        In [1]: import numpy as np
           ...: def median_sd(x):
           ...:     median = np.percentile(x, 50)
           ...:     sd = np.sqrt(np.mean((x-median)**2))
           ...:     return sd
           ...:
           ...: func_dict = {
           ...:     "std": np.std,
           ...:     "median_std": median_sd,
           ...:     "5%": lambda x: np.percentile(x, 5),
           ...:     "median": lambda x: np.percentile(x, 50),
           ...:     "95%": lambda x: np.percentile(x, 95),
           ...: }
           ...: az.summary(
           ...:     data,
           ...:     var_names=["mu", "tau"],
           ...:     stat_funcs=func_dict,
           ...:     extend=False
           ...: )

    Use ``stat_focus`` to change the focus of summary statistics obatined to median:

    .. ipython::

        In [1]: az.summary(data, stat_focus="median")

    """
    _log.cache = []

    if coords is None:
        coords = {}

    if index_origin is not None:
        warnings.warn(
            "index_origin has been deprecated. summary now shows coordinate values, "
            "to change the label shown, modify the coordinate values before calling summary",
            DeprecationWarning,
        )
        index_origin = rcParams["data.index_origin"]
    if labeller is None:
        labeller = BaseLabeller()
    if hdi_prob is None:
        hdi_prob = rcParams["stats.ci_prob"]
    elif not 1 >= hdi_prob > 0:
        raise ValueError("The value of hdi_prob should be in the interval (0, 1]")

    if isinstance(data, InferenceData):
        if group is None:
            if not data.groups():
                raise TypeError("InferenceData does not contain any groups")
            if "posterior" in data:
                dataset = data["posterior"]
            elif "prior" in data:
                dataset = data["prior"]
            else:
                warnings.warn(f"Selecting first found group: {data.groups()[0]}")
                dataset = data[data.groups()[0]]
        elif group in data.groups():
            dataset = data[group]
        else:
            raise TypeError(f"InferenceData does not contain group: {group}")
    else:
        dataset = convert_to_dataset(data, group="posterior")
    var_names = _var_names(var_names, dataset, filter_vars)
    dataset = dataset if var_names is None else dataset[var_names]
    dataset = get_coords(dataset, coords)

    fmt_group = ("wide", "long", "xarray")
    if not isinstance(fmt, str) or (fmt.lower() not in fmt_group):
        raise TypeError(f"Invalid format: '{fmt}'. Formatting options are: {fmt_group}")

    kind_group = ("all", "stats", "diagnostics")
    if not isinstance(kind, str) or kind not in kind_group:
        raise TypeError(f"Invalid kind: '{kind}'. Kind options are: {kind_group}")

    focus_group = ("mean", "median")
    if not isinstance(stat_focus, str) or (stat_focus not in focus_group):
        raise TypeError(f"Invalid format: '{stat_focus}'. Focus options are: {focus_group}")

    if stat_focus != "mean" and circ_var_names is not None:
        raise TypeError(f"Invalid format: Circular stats not supported for '{stat_focus}'")

    if order is not None:
        warnings.warn(
            "order has been deprecated. summary now shows coordinate values.", DeprecationWarning
        )

    alpha = 1 - hdi_prob

    extra_metrics = []
    extra_metric_names = []

    if stat_funcs is not None:
        if isinstance(stat_funcs, dict):
            for stat_func_name, stat_func in stat_funcs.items():
                extra_metrics.append(
                    xr.apply_ufunc(
                        _make_ufunc(stat_func), dataset, input_core_dims=(("chain", "draw"),)
                    )
                )
                extra_metric_names.append(stat_func_name)
        else:
            for stat_func in stat_funcs:
                extra_metrics.append(
                    xr.apply_ufunc(
                        _make_ufunc(stat_func), dataset, input_core_dims=(("chain", "draw"),)
                    )
                )
                extra_metric_names.append(stat_func.__name__)

    metrics: List[xr.Dataset] = []
    metric_names: List[str] = []
    if extend and kind in ["all", "stats"]:
        if stat_focus == "mean":
            mean = dataset.mean(dim=("chain", "draw"), skipna=skipna)

            sd = dataset.std(dim=("chain", "draw"), ddof=1, skipna=skipna)

            hdi_post = hdi(dataset, hdi_prob=hdi_prob, multimodal=False, skipna=skipna)
            hdi_lower = hdi_post.sel(hdi="lower", drop=True)
            hdi_higher = hdi_post.sel(hdi="higher", drop=True)
            metrics.extend((mean, sd, hdi_lower, hdi_higher))
            metric_names.extend(
                ("mean", "sd", f"hdi_{100 * alpha / 2:g}%", f"hdi_{100 * (1 - alpha / 2):g}%")
            )
        elif stat_focus == "median":
            median = dataset.median(dim=("chain", "draw"), skipna=skipna)

            mad = stats.median_abs_deviation(dataset, dims=("chain", "draw"))
            eti_post = dataset.quantile(
                (alpha / 2, 1 - alpha / 2), dim=("chain", "draw"), skipna=skipna
            )
            eti_lower = eti_post.isel(quantile=0, drop=True)
            eti_higher = eti_post.isel(quantile=1, drop=True)
            metrics.extend((median, mad, eti_lower, eti_higher))
            metric_names.extend(
                ("median", "mad", f"eti_{100 * alpha / 2:g}%", f"eti_{100 * (1 - alpha / 2):g}%")
            )

    if circ_var_names:
        nan_policy = "omit" if skipna else "propagate"
        circ_mean = stats.circmean(
            dataset, dims=["chain", "draw"], high=np.pi, low=-np.pi, nan_policy=nan_policy
        )
        _numba_flag = Numba.numba_flag
        if _numba_flag:
            circ_sd = xr.apply_ufunc(
                _make_ufunc(_circular_standard_deviation),
                dataset,
                kwargs=dict(high=np.pi, low=-np.pi, skipna=skipna),
                input_core_dims=(("chain", "draw"),),
            )
        else:
            circ_sd = stats.circstd(
                dataset, dims=["chain", "draw"], high=np.pi, low=-np.pi, nan_policy=nan_policy
            )
        circ_mcse = xr.apply_ufunc(
            _make_ufunc(_mc_error),
            dataset,
            kwargs=dict(circular=True),
            input_core_dims=(("chain", "draw"),),
        )

        circ_hdi = hdi(dataset, hdi_prob=hdi_prob, circular=True, skipna=skipna)
        circ_hdi_lower = circ_hdi.sel(hdi="lower", drop=True)
        circ_hdi_higher = circ_hdi.sel(hdi="higher", drop=True)

    if kind in ["all", "diagnostics"] and extend:
        diagnostics_names: Tuple[str, ...]
        if stat_focus == "mean":
            diagnostics = xr.apply_ufunc(
                _make_ufunc(_multichain_statistics, n_output=5, ravel=False),
                dataset,
                input_core_dims=(("chain", "draw"),),
                output_core_dims=tuple([] for _ in range(5)),
            )
            diagnostics_names = (
                "mcse_mean",
                "mcse_sd",
                "ess_bulk",
                "ess_tail",
                "r_hat",
            )

        elif stat_focus == "median":
            diagnostics = xr.apply_ufunc(
                _make_ufunc(_multichain_statistics, n_output=4, ravel=False),
                dataset,
                kwargs=dict(focus="median"),
                input_core_dims=(("chain", "draw"),),
                output_core_dims=tuple([] for _ in range(4)),
            )
            diagnostics_names = (
                "mcse_median",
                "ess_median",
                "ess_tail",
                "r_hat",
            )
        metrics.extend(diagnostics)
        metric_names.extend(diagnostics_names)

    if circ_var_names and kind != "diagnostics" and stat_focus == "mean":
        for metric, circ_stat in zip(
            # Replace only the first 5 statistics for their circular equivalent
            metrics[:5],
            (circ_mean, circ_sd, circ_hdi_lower, circ_hdi_higher, circ_mcse),
        ):
            for circ_var in circ_var_names:
                metric[circ_var] = circ_stat[circ_var]

    metrics.extend(extra_metrics)
    metric_names.extend(extra_metric_names)
    joined = (
        xr.concat(metrics, dim="metric").assign_coords(metric=metric_names).reset_coords(drop=True)
    )
    n_metrics = len(metric_names)
    n_vars = np.sum([joined[var].size // n_metrics for var in joined.data_vars])

    if fmt.lower() == "wide":
        summary_df = pd.DataFrame(
            (np.full((cast(int, n_vars), n_metrics), np.nan)), columns=metric_names
        )
        indices = []
        for i, (var_name, sel, isel, values) in enumerate(
            xarray_var_iter(joined, skip_dims={"metric"})
        ):
            summary_df.iloc[i] = values
            indices.append(labeller.make_label_flat(var_name, sel, isel))
        summary_df.index = indices
    elif fmt.lower() == "long":
        df = joined.to_dataframe().reset_index().set_index("metric")
        df.index = list(df.index)
        summary_df = df
    else:
        # format is 'xarray'
        summary_df = joined
    if (round_to is not None) and (round_to not in ("None", "none")):
        summary_df = summary_df.round(round_to)
    elif round_to not in ("None", "none") and (fmt.lower() in ("long", "wide")):
        # Don't round xarray object by default (even with "none")
        decimals = {
            col: 3 if col not in {"ess_bulk", "ess_tail", "r_hat"} else 2 if col == "r_hat" else 0
            for col in summary_df.columns
        }
        summary_df = summary_df.round(decimals)

    return summary_df


def waic(data, pointwise=None, var_name=None, scale=None, dask_kwargs=None):
    """Compute the widely applicable information criterion.

    Estimates the expected log pointwise predictive density (elpd) using WAIC. Also calculates the
    WAIC's standard error and the effective number of parameters.
    Read more theory here https://arxiv.org/abs/1507.04544 and here https://arxiv.org/abs/1004.2316

    Parameters
    ----------
    data: obj
        Any object that can be converted to an :class:`arviz.InferenceData` object.
        Refer to documentation of :func:`arviz.convert_to_inference_data` for details.
    pointwise: bool
        If True the pointwise predictive accuracy will be returned. Defaults to
        ``stats.ic_pointwise`` rcParam.
    var_name : str, optional
        The name of the variable in log_likelihood groups storing the pointwise log
        likelihood data to use for waic computation.
    scale: str
        Output scale for WAIC. Available options are:

        - `log` : (default) log-score
        - `negative_log` : -1 * log-score
        - `deviance` : -2 * log-score

        A higher log-score (or a lower deviance or negative log_score) indicates a model with
        better predictive accuracy.
    dask_kwargs : dict, optional
        Dask related kwargs passed to :func:`~arviz.wrap_xarray_ufunc`.

    Returns
    -------
    ELPDData object (inherits from :class:`pandas.Series`) with the following row/attributes:
    elpd_waic: approximated expected log pointwise predictive density (elpd)
    se: standard error of the elpd
    p_waic: effective number parameters
    n_samples: number of samples
    n_data_points: number of data points
    warning: bool
        True if posterior variance of the log predictive densities exceeds 0.4
    waic_i: :class:`~xarray.DataArray` with the pointwise predictive accuracy,
            only if pointwise=True
    scale: scale of the elpd

        The returned object has a custom print method that overrides pd.Series method.

    See Also
    --------
    loo : Compute Pareto-smoothed importance sampling leave-one-out cross-validation (PSIS-LOO-CV).
    compare : Compare models based on PSIS-LOO-CV or WAIC.
    plot_compare : Summary plot for model comparison.

    Examples
    --------
    Calculate WAIC of a model:

    .. ipython::

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

    Calculate WAIC of a model and return the pointwise values:

    .. ipython::

        In [2]: data_waic = az.waic(data, pointwise=True)
           ...: data_waic.waic_i
    """
    inference_data = convert_to_inference_data(data)
    log_likelihood = _get_log_likelihood(inference_data, var_name=var_name)
    scale = rcParams["stats.ic_scale"] if scale is None else scale.lower()
    pointwise = rcParams["stats.ic_pointwise"] if pointwise is None else pointwise

    if scale == "deviance":
        scale_value = -2
    elif scale == "log":
        scale_value = 1
    elif scale == "negative_log":
        scale_value = -1
    else:
        raise TypeError('Valid scale values are "deviance", "log", "negative_log"')

    log_likelihood = log_likelihood.stack(__sample__=("chain", "draw"))
    shape = log_likelihood.shape
    n_samples = shape[-1]
    n_data_points = np.prod(shape[:-1])

    ufunc_kwargs = {"n_dims": 1, "ravel": False}
    kwargs = {"input_core_dims": [["__sample__"]]}
    lppd_i = _wrap_xarray_ufunc(
        _logsumexp,
        log_likelihood,
        func_kwargs={"b_inv": n_samples},
        ufunc_kwargs=ufunc_kwargs,
        dask_kwargs=dask_kwargs,
        **kwargs,
    )

    vars_lpd = log_likelihood.var(dim="__sample__")
    warn_mg = False
    if np.any(vars_lpd > 0.4):
        warnings.warn(
            (
                "For one or more samples the posterior variance of the log predictive "
                "densities exceeds 0.4. This could be indication of WAIC starting to fail. \n"
                "See http://arxiv.org/abs/1507.04544 for details"
            )
        )
        warn_mg = True

    waic_i = scale_value * (lppd_i - vars_lpd)
    waic_se = (n_data_points * np.var(waic_i.values)) ** 0.5
    waic_sum = np.sum(waic_i.values)
    p_waic = np.sum(vars_lpd.values)

    if not pointwise:
        return ELPDData(
            data=[waic_sum, waic_se, p_waic, n_samples, n_data_points, warn_mg, scale],
            index=[
                "waic",
                "se",
                "p_waic",
                "n_samples",
                "n_data_points",
                "warning",
                "scale",
            ],
        )
    if np.equal(waic_sum, waic_i).all():  # pylint: disable=no-member
        warnings.warn(
            """The point-wise WAIC is the same with the sum WAIC, please double check
            the Observed RV in your model to make sure it returns element-wise logp.
            """
        )
    return ELPDData(
        data=[
            waic_sum,
            waic_se,
            p_waic,
            n_samples,
            n_data_points,
            warn_mg,
            waic_i.rename("waic_i"),
            scale,
        ],
        index=[
            "elpd_waic",
            "se",
            "p_waic",
            "n_samples",
            "n_data_points",
            "warning",
            "waic_i",
            "scale",
        ],
    )


def loo_pit(idata=None, *, y=None, y_hat=None, log_weights=None):
    """Compute leave one out (PSIS-LOO) probability integral transform (PIT) values.

    Parameters
    ----------
    idata: InferenceData
        :class:`arviz.InferenceData` object.
    y: array, DataArray or str
        Observed data. If str, ``idata`` must be present and contain the observed data group
    y_hat: array, DataArray or str
        Posterior predictive samples for ``y``. It must have the same shape as y plus an
        extra dimension at the end of size n_samples (chains and draws stacked). If str or
        None, ``idata`` must contain the posterior predictive group. If None, y_hat is taken
        equal to y, thus, y must be str too.
    log_weights: array or DataArray
        Smoothed log_weights. It must have the same shape as ``y_hat``
    dask_kwargs : dict, optional
        Dask related kwargs passed to :func:`~arviz.wrap_xarray_ufunc`.

    Returns
    -------
    loo_pit: array or DataArray
        Value of the LOO-PIT at each observed data point.

    See Also
    --------
    plot_loo_pit : Plot Leave-One-Out probability integral transformation (PIT) predictive checks.
    loo : Compute Pareto-smoothed importance sampling leave-one-out
          cross-validation (PSIS-LOO-CV).
    plot_elpd : Plot pointwise elpd differences between two or more models.
    plot_khat : Plot Pareto tail indices for diagnosing convergence.

    Examples
    --------
    Calculate LOO-PIT values using as test quantity the observed values themselves.

    .. ipython::

        In [1]: import arviz as az
           ...: data = az.load_arviz_data("centered_eight")
           ...: az.loo_pit(idata=data, y="obs")

    Calculate LOO-PIT values using as test quantity the square of the difference between
    each observation and `mu`. Both ``y`` and ``y_hat`` inputs will be array-like,
    but ``idata`` will still be passed in order to calculate the ``log_weights`` from
    there.

    .. ipython::

        In [1]: T = data.observed_data.obs - data.posterior.mu.median(dim=("chain", "draw"))
           ...: T_hat = data.posterior_predictive.obs - data.posterior.mu
           ...: T_hat = T_hat.stack(__sample__=("chain", "draw"))
           ...: az.loo_pit(idata=data, y=T**2, y_hat=T_hat**2)

    """
    y_str = ""
    if idata is not None and not isinstance(idata, InferenceData):
        raise ValueError("idata must be of type InferenceData or None")

    if idata is None:
        if not all(isinstance(arg, (np.ndarray, xr.DataArray)) for arg in (y, y_hat, log_weights)):
            raise ValueError(
                "all 3 y, y_hat and log_weights must be array or DataArray when idata is None "
                f"but they are of types {[type(arg) for arg in (y, y_hat, log_weights)]}"
            )

    else:
        if y_hat is None and isinstance(y, str):
            y_hat = y
        elif y_hat is None:
            raise ValueError("y_hat cannot be None if y is not a str")
        if isinstance(y, str):
            y_str = y
            y = idata.observed_data[y].values
        elif not isinstance(y, (np.ndarray, xr.DataArray)):
            raise ValueError(f"y must be of types array, DataArray or str, not {type(y)}")
        if isinstance(y_hat, str):
            y_hat = idata.posterior_predictive[y_hat].stack(__sample__=("chain", "draw")).values
        elif not isinstance(y_hat, (np.ndarray, xr.DataArray)):
            raise ValueError(f"y_hat must be of types array, DataArray or str, not {type(y_hat)}")
        if log_weights is None:
            if y_str:
                try:
                    log_likelihood = _get_log_likelihood(idata, var_name=y_str)
                except TypeError:
                    log_likelihood = _get_log_likelihood(idata)
            else:
                log_likelihood = _get_log_likelihood(idata)
            log_likelihood = log_likelihood.stack(__sample__=("chain", "draw"))
            posterior = convert_to_dataset(idata, group="posterior")
            n_chains = len(posterior.chain)
            n_samples = len(log_likelihood.__sample__)
            ess_p = ess(posterior, method="mean")
            # this mean is over all data variables
            reff = (
                (np.hstack([ess_p[v].values.flatten() for v in ess_p.data_vars]).mean() / n_samples)
                if n_chains > 1
                else 1
            )
            log_weights = psislw(-log_likelihood, reff=reff)[0].values
        elif not isinstance(log_weights, (np.ndarray, xr.DataArray)):
            raise ValueError(
                f"log_weights must be None or of types array or DataArray, not {type(log_weights)}"
            )

    if len(y.shape) + 1 != len(y_hat.shape):
        raise ValueError(
            f"y_hat must have 1 more dimension than y, but y_hat has {len(y_hat.shape)} dims and "
            f"y has {len(y.shape)} dims"
        )

    if y.shape != y_hat.shape[:-1]:
        raise ValueError(
            f"y has shape: {y.shape} which should be equal to y_hat shape (omitting the last "
            f"dimension): {y_hat.shape}"
        )

    if y_hat.shape != log_weights.shape:
        raise ValueError(
            "y_hat and log_weights must have the same shape but have shapes "
            f"{y_hat.shape,} and {log_weights.shape}"
        )

    kwargs = {
        "input_core_dims": [[], ["__sample__"], ["__sample__"]],
        "output_core_dims": [[]],
        "join": "left",
    }
    ufunc_kwargs = {"n_dims": 1}

    if y.dtype.kind == "i" or y_hat.dtype.kind == "i":
        y, y_hat = smooth_data(y, y_hat)

    return _wrap_xarray_ufunc(
        _loo_pit,
        y,
        y_hat,
        log_weights,
        ufunc_kwargs=ufunc_kwargs,
        **kwargs,
    )


def _loo_pit(y, y_hat, log_weights):
    """Compute LOO-PIT values."""
    sel = y_hat <= y
    if np.sum(sel) > 0:
        value = np.exp(_logsumexp(log_weights[sel]))
        return min(1, value)
    else:
        return 0


def apply_test_function(
    idata,
    func,
    group="both",
    var_names=None,
    pointwise=False,
    out_data_shape=None,
    out_pp_shape=None,
    out_name_data="T",
    out_name_pp=None,
    func_args=None,
    func_kwargs=None,
    ufunc_kwargs=None,
    wrap_data_kwargs=None,
    wrap_pp_kwargs=None,
    inplace=True,
    overwrite=None,
):
    """Apply a Bayesian test function to an InferenceData object.

    Parameters
    ----------
    idata: InferenceData
        :class:`arviz.InferenceData` object on which to apply the test function.
        This function will add new variables to the InferenceData object
        to store the result without modifying the existing ones.
    func: callable
        Callable that calculates the test function. It must have the following call signature
        ``func(y, theta, *args, **kwargs)`` (where ``y`` is the observed data or posterior
        predictive and ``theta`` the model parameters) even if not all the arguments are
        used.
    group: str, optional
        Group on which to apply the test function. Can be observed_data, posterior_predictive
        or both.
    var_names: dict group -> var_names, optional
        Mapping from group name to the variables to be passed to func. It can be a dict of
        strings or lists of strings. There is also the option of using ``both`` as key,
        in which case, the same variables are used in observed data and posterior predictive
        groups
    pointwise: bool, optional
        If True, apply the test function to each observation and sample, otherwise, apply
        test function to each sample.
    out_data_shape, out_pp_shape: tuple, optional
        Output shape of the test function applied to the observed/posterior predictive data.
        If None, the default depends on the value of pointwise.
    out_name_data, out_name_pp: str, optional
        Name of the variables to add to the observed_data and posterior_predictive datasets
        respectively. ``out_name_pp`` can be ``None``, in which case will be taken equal to
        ``out_name_data``.
    func_args: sequence, optional
        Passed as is to ``func``
    func_kwargs: mapping, optional
        Passed as is to ``func``
    wrap_data_kwargs, wrap_pp_kwargs: mapping, optional
        kwargs passed to :func:`~arviz.wrap_xarray_ufunc`. By default, some suitable input_core_dims
        are used.
    inplace: bool, optional
        If True, add the variables inplace, otherwise, return a copy of idata with the variables
        added.
    overwrite: bool, optional
        Overwrite data in case ``out_name_data`` or ``out_name_pp`` are already variables in
        dataset. If ``None`` it will be the opposite of inplace.

    Returns
    -------
    idata: InferenceData
        Output InferenceData object. If ``inplace=True``, it is the same input object modified
        inplace.

    See Also
    --------
    plot_bpv :  Plot Bayesian p-value for observed data and Posterior/Prior predictive.

    Notes
    -----
    This function is provided for convenience to wrap scalar or functions working on low
    dims to inference data object. It is not optimized to be faster nor as fast as vectorized
    computations.

    Examples
    --------
    Use ``apply_test_function`` to wrap ``numpy.min`` for illustration purposes. And plot the
    results.

    .. plot::
        :context: close-figs

        >>> import arviz as az
        >>> idata = az.load_arviz_data("centered_eight")
        >>> az.apply_test_function(idata, lambda y, theta: np.min(y))
        >>> T = idata.observed_data.T.item()
        >>> az.plot_posterior(idata, var_names=["T"], group="posterior_predictive", ref_val=T)

    """
    out = idata if inplace else deepcopy(idata)

    valid_groups = ("observed_data", "posterior_predictive", "both")
    if group not in valid_groups:
        raise ValueError(f"Invalid group argument. Must be one of {valid_groups} not {group}.")
    if overwrite is None:
        overwrite = not inplace

    if out_name_pp is None:
        out_name_pp = out_name_data

    if func_args is None:
        func_args = tuple()

    if func_kwargs is None:
        func_kwargs = {}

    if ufunc_kwargs is None:
        ufunc_kwargs = {}
    ufunc_kwargs.setdefault("check_shape", False)
    ufunc_kwargs.setdefault("ravel", False)

    if wrap_data_kwargs is None:
        wrap_data_kwargs = {}
    if wrap_pp_kwargs is None:
        wrap_pp_kwargs = {}
    if var_names is None:
        var_names = {}

    both_var_names = var_names.pop("both", None)
    var_names.setdefault("posterior", list(out.posterior.data_vars))

    in_posterior = out.posterior[var_names["posterior"]]
    if isinstance(in_posterior, xr.Dataset):
        in_posterior = in_posterior.to_array().squeeze()

    groups = ("posterior_predictive", "observed_data") if group == "both" else [group]
    for grp in groups:
        out_group_shape = out_data_shape if grp == "observed_data" else out_pp_shape
        out_name_group = out_name_data if grp == "observed_data" else out_name_pp
        wrap_group_kwargs = wrap_data_kwargs if grp == "observed_data" else wrap_pp_kwargs
        if not hasattr(out, grp):
            raise ValueError(f"InferenceData object must have {grp} group")
        if not overwrite and out_name_group in getattr(out, grp).data_vars:
            raise ValueError(
                f"Should overwrite: {out_name_group} variable present in group {grp},"
                " but overwrite is False"
            )
        var_names.setdefault(
            grp, list(getattr(out, grp).data_vars) if both_var_names is None else both_var_names
        )
        in_group = getattr(out, grp)[var_names[grp]]
        if isinstance(in_group, xr.Dataset):
            in_group = in_group.to_array(dim=f"{grp}_var").squeeze()

        if pointwise:
            out_group_shape = in_group.shape if out_group_shape is None else out_group_shape
        elif grp == "observed_data":
            out_group_shape = () if out_group_shape is None else out_group_shape
        elif grp == "posterior_predictive":
            out_group_shape = in_group.shape[:2] if out_group_shape is None else out_group_shape
        loop_dims = in_group.dims[: len(out_group_shape)]

        wrap_group_kwargs.setdefault(
            "input_core_dims",
            [
                [dim for dim in dataset.dims if dim not in loop_dims]
                for dataset in [in_group, in_posterior]
            ],
        )
        func_kwargs["out"] = np.empty(out_group_shape)

        out_group = getattr(out, grp)
        try:
            out_group[out_name_group] = _wrap_xarray_ufunc(
                func,
                in_group.values,
                in_posterior.values,
                func_args=func_args,
                func_kwargs=func_kwargs,
                ufunc_kwargs=ufunc_kwargs,
                **wrap_group_kwargs,
            )
        except IndexError:
            excluded_dims = set(
                wrap_group_kwargs["input_core_dims"][0] + wrap_group_kwargs["input_core_dims"][1]
            )
            out_group[out_name_group] = _wrap_xarray_ufunc(
                func,
                *xr.broadcast(in_group, in_posterior, exclude=excluded_dims),
                func_args=func_args,
                func_kwargs=func_kwargs,
                ufunc_kwargs=ufunc_kwargs,
                **wrap_group_kwargs,
            )
        setattr(out, grp, out_group)

    return out


def weight_predictions(idatas, weights=None):
    """
    Generate weighted posterior predictive samples from a list of InferenceData
    and a set of weights.

    Parameters
    ---------
    idatas : list[InferenceData]
        List of :class:`arviz.InferenceData` objects containing the groups `posterior_predictive`
        and `observed_data`. Observations should be the same for all InferenceData objects.
    weights : array-like, optional
        Individual weights for each model. Weights should be positive. If they do not sum up to 1,
        they will be normalized. Default, same weight for each model.
        Weights can be computed using many different methods including those in
        :func:`arviz.compare`.

    Returns
    -------
    idata: InferenceData
        Output InferenceData object with the groups `posterior_predictive` and `observed_data`.

    See Also
    --------
    compare :  Compare models based on PSIS-LOO `loo` or WAIC `waic` cross-validation
    """
    if len(idatas) < 2:
        raise ValueError("You should provide a list with at least two InferenceData objects")

    if not all("posterior_predictive" in idata.groups() for idata in idatas):
        raise ValueError(
            "All the InferenceData objects must contain the `posterior_predictive` group"
        )

    if not all(idatas[0].observed_data.equals(idata.observed_data) for idata in idatas[1:]):
        raise ValueError("The observed data should be the same for all InferenceData objects")

    if weights is None:
        weights = np.ones(len(idatas)) / len(idatas)
    elif len(idatas) != len(weights):
        raise ValueError(
            "The number of weights should be the same as the number of InferenceData objects"
        )

    weights = np.array(weights, dtype=float)
    weights /= weights.sum()

    len_idatas = [
        idata.posterior_predictive.sizes["chain"] * idata.posterior_predictive.sizes["draw"]
        for idata in idatas
    ]

    if not all(len_idatas):
        raise ValueError("At least one of your idatas has 0 samples")

    new_samples = (np.min(len_idatas) * weights).astype(int)

    new_idatas = [
        extract(idata, group="posterior_predictive", num_samples=samples).reset_coords()
        for samples, idata in zip(new_samples, idatas)
    ]

    weighted_samples = InferenceData(
        posterior_predictive=xr.concat(new_idatas, dim="sample"),
        observed_data=idatas[0].observed_data,
    )

    return weighted_samples


def psens(
    data,
    *,
    component="prior",
    component_var_names=None,
    component_coords=None,
    var_names=None,
    coords=None,
    filter_vars=None,
    delta=0.01,
    dask_kwargs=None,
):
    """Compute power-scaling sensitivity diagnostic.

    Power-scales the prior or likelihood and calculates how much the posterior is affected.

    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``.
    component : {"prior", "likelihood"}, default "prior"
        When `component` is "likelihood", the log likelihood values are retrieved
        from the ``log_likelihood`` group as pointwise log likelihood and added
        together. With "prior", the log prior values are retrieved from the
        ``log_prior`` group.
    component_var_names : str, optional
        Name of the prior or log likelihood variables to use
    component_coords : dict, optional
        Coordinates defining a subset over the component element for which to
        compute the prior sensitivity diagnostic.
    var_names : list of str, optional
        Names of posterior variables to include in the power scaling sensitivity diagnostic
    coords : dict, optional
        Coordinates defining a subset over the posterior. Only these variables will
        be used when computing the prior sensitivity.
    filter_vars: {None, "like", "regex"}, default None
        If ``None`` (default), interpret var_names as the real variables names.
        If "like", interpret var_names as substrings of the real variables names.
        If "regex", interpret var_names as regular expressions on the real variables names.
    delta : float
        Value for finite difference derivative calculation.
    dask_kwargs : dict, optional
        Dask related kwargs passed to :func:`~arviz.wrap_xarray_ufunc`.

    Returns
    -------
    xarray.Dataset
        Returns dataset of power-scaling sensitivity diagnostic values.
        Higher sensitivity values indicate greater sensitivity.
        Prior sensitivity above 0.05 indicates informative prior.
        Likelihood sensitivity below 0.05 indicates weak or nonin-formative likelihood.

    Examples
    --------
    Compute the likelihood sensitivity for the non centered eight model:

    .. ipython::

        In [1]: import arviz as az
           ...: data = az.load_arviz_data("non_centered_eight")
           ...: az.psens(data, component="likelihood")

    To compute the prior sensitivity, we need to first compute the log prior
    at each posterior sample. In our case, we know mu has a normal prior :math:`N(0, 5)`,
    tau is a half cauchy prior with scale/beta parameter 5,
    and theta has a standard normal as prior.
    We add this information to the ``log_prior`` group before computing powerscaling
    check with ``psens``

    .. ipython::

        In [1]: from xarray_einstats.stats import XrContinuousRV
           ...: from scipy.stats import norm, halfcauchy
           ...: post = data.posterior
           ...: log_prior = {
           ...:     "mu": XrContinuousRV(norm, 0, 5).logpdf(post["mu"]),
           ...:     "tau": XrContinuousRV(halfcauchy, scale=5).logpdf(post["tau"]),
           ...:     "theta_t": XrContinuousRV(norm, 0, 1).logpdf(post["theta_t"]),
           ...: }
           ...: data.add_groups({"log_prior": log_prior})
           ...: az.psens(data, component="prior")

    Notes
    -----
    The diagnostic is computed by power-scaling the specified component (prior or likelihood)
    and determining the degree to which the posterior changes as described in [1]_.
    It uses Pareto-smoothed importance sampling to avoid refitting the model.

    References
    ----------
    .. [1] Kallioinen et al, *Detecting and diagnosing prior and likelihood sensitivity with
       power-scaling*, 2022, https://arxiv.org/abs/2107.14054

    """
    dataset = extract(data, var_names=var_names, filter_vars=filter_vars, group="posterior")
    if coords is None:
        dataset = dataset.sel(coords)

    if component == "likelihood":
        component_draws = _get_log_likelihood(data, var_name=component_var_names, single_var=False)
    elif component == "prior":
        component_draws = _get_log_prior(data, var_names=component_var_names)
    else:
        raise ValueError("Value for `component` argument not recognized")

    component_draws = component_draws.stack(__sample__=("chain", "draw"))
    if component_coords is None:
        component_draws = component_draws.sel(component_coords)
    if isinstance(component_draws, xr.DataArray):
        component_draws = component_draws.to_dataset()
    if len(component_draws.dims):
        component_draws = component_draws.to_stacked_array(
            "latent-obs_var", sample_dims=("__sample__",)
        ).sum("latent-obs_var")
    # from here component_draws is a 1d object with dimensions (sample,)

    # calculate lower and upper alpha values
    lower_alpha = 1 / (1 + delta)
    upper_alpha = 1 + delta

    # calculate importance sampling weights for lower and upper alpha power-scaling
    lower_w = np.exp(_powerscale_lw(component_draws=component_draws, alpha=lower_alpha))
    lower_w = lower_w / np.sum(lower_w)

    upper_w = np.exp(_powerscale_lw(component_draws=component_draws, alpha=upper_alpha))
    upper_w = upper_w / np.sum(upper_w)

    ufunc_kwargs = {"n_dims": 1, "ravel": False}
    func_kwargs = {"lower_weights": lower_w.values, "upper_weights": upper_w.values, "delta": delta}

    # calculate the sensitivity diagnostic based on the importance weights and draws
    return _wrap_xarray_ufunc(
        _powerscale_sens,
        dataset,
        ufunc_kwargs=ufunc_kwargs,
        func_kwargs=func_kwargs,
        dask_kwargs=dask_kwargs,
        input_core_dims=[["sample"]],
    )


def _powerscale_sens(draws, *, lower_weights=None, upper_weights=None, delta=0.01):
    """
    Calculate power-scaling sensitivity by finite difference
    second derivative of CJS
    """
    lower_cjs = max(
        _cjs_dist(draws=draws, weights=lower_weights),
        _cjs_dist(draws=-1 * draws, weights=lower_weights),
    )
    upper_cjs = max(
        _cjs_dist(draws=draws, weights=upper_weights),
        _cjs_dist(draws=-1 * draws, weights=upper_weights),
    )
    logdiffsquare = 2 * np.log2(1 + delta)
    grad = (lower_cjs + upper_cjs) / logdiffsquare

    return grad


def _powerscale_lw(alpha, component_draws):
    """
    Calculate log weights for power-scaling component by alpha.
    """
    log_weights = (alpha - 1) * component_draws
    log_weights = psislw(log_weights)[0]

    return log_weights


def _cjs_dist(draws, weights):
    """
    Calculate the cumulative Jensen-Shannon distance between original draws and weighted draws.
    """

    # sort draws and weights
    order = np.argsort(draws)
    draws = draws[order]
    weights = weights[order]

    binwidth = np.diff(draws)

    # ecdfs
    cdf_p = np.linspace(1 / len(draws), 1 - 1 / len(draws), len(draws) - 1)
    cdf_q = np.cumsum(weights / np.sum(weights))[:-1]

    # integrals of ecdfs
    cdf_p_int = np.dot(cdf_p, binwidth)
    cdf_q_int = np.dot(cdf_q, binwidth)

    # cjs calculation
    pq_numer = np.log2(cdf_p, out=np.zeros_like(cdf_p), where=cdf_p != 0)
    qp_numer = np.log2(cdf_q, out=np.zeros_like(cdf_q), where=cdf_q != 0)

    denom = 0.5 * (cdf_p + cdf_q)
    denom = np.log2(denom, out=np.zeros_like(denom), where=denom != 0)

    cjs_pq = np.sum(binwidth * (cdf_p * (pq_numer - denom))) + 0.5 / np.log(2) * (
        cdf_q_int - cdf_p_int
    )

    cjs_qp = np.sum(binwidth * (cdf_q * (qp_numer - denom))) + 0.5 / np.log(2) * (
        cdf_p_int - cdf_q_int
    )

    cjs_pq = max(0, cjs_pq)
    cjs_qp = max(0, cjs_qp)

    bound = cdf_p_int + cdf_q_int

    return np.sqrt((cjs_pq + cjs_qp) / bound)


def bayes_factor(idata, var_name, ref_val=0, prior=None, return_ref_vals=False):
    r"""Approximated Bayes Factor for comparing hypothesis of two nested models.

    The Bayes factor is estimated by comparing a model (H1) against a model in which the
    parameter of interest has been restricted to be a point-null (H0). This computation
    assumes the models are nested and thus H0 is a special case of H1.

    Notes
    -----
    The bayes Factor is approximated as the Savage-Dickey density ratio
    algorithm presented in [1]_.

    Parameters
    ----------
    idata : InferenceData
        Any object that can be converted to an :class:`arviz.InferenceData` object
        Refer to documentation of :func:`arviz.convert_to_dataset` for details.
    var_name : str, optional
        Name of variable we want to test.
    ref_val : int, default 0
        Point-null for Bayes factor estimation.
    prior : numpy.array, optional
        In case we want to use different prior, for example for sensitivity analysis.
    return_ref_vals : bool, optional
        Whether to return the values of the prior and posterior at the reference value.
        Used by :func:`arviz.plot_bf` to display the distribution comparison.


    Returns
    -------
    dict : A dictionary with BF10 (Bayes Factor 10 (H1/H0 ratio), and BF01 (H0/H1 ratio).

    References
    ----------
    .. [1] Heck, D., 2019. A caveat on the Savage-Dickey density ratio:
       The case of computing Bayes factors for regression parameters.

    Examples
    --------
    Moderate evidence indicating that the parameter "a" is different from zero.

    .. ipython::

        In [1]: import numpy as np
           ...: import arviz as az
           ...: idata = az.from_dict(posterior={"a":np.random.normal(1, 0.5, 5000)},
           ...: prior={"a":np.random.normal(0, 1, 5000)})
           ...: az.bayes_factor(idata, var_name="a", ref_val=0)

    """

    posterior = extract(idata, var_names=var_name).values

    if ref_val > posterior.max() or ref_val < posterior.min():
        _log.warning(
            "The reference value is outside of the posterior. "
            "This translate into infinite support for H1, which is most likely an overstatement."
        )

    if posterior.ndim > 1:
        _log.warning("Posterior distribution has {posterior.ndim} dimensions")

    if prior is None:
        prior = extract(idata, var_names=var_name, group="prior").values

    if posterior.dtype.kind == "f":
        posterior_grid, posterior_pdf, *_ = _kde_linear(posterior)
        prior_grid, prior_pdf, *_ = _kde_linear(prior)
        posterior_at_ref_val = np.interp(ref_val, posterior_grid, posterior_pdf)
        prior_at_ref_val = np.interp(ref_val, prior_grid, prior_pdf)

    elif posterior.dtype.kind == "i":
        posterior_at_ref_val = (posterior == ref_val).mean()
        prior_at_ref_val = (prior == ref_val).mean()

    bf_10 = prior_at_ref_val / posterior_at_ref_val
    bf = {"BF10": bf_10, "BF01": 1 / bf_10}

    if return_ref_vals:
        return (bf, {"prior": prior_at_ref_val, "posterior": posterior_at_ref_val})
    else:
        return bf
