import contextlib
import copy
import itertools
import os
import glob
import pickle
import hashlib
import logging
import multiprocessing
import warnings

import tqdm

import numpy as np
import pandas as pd
import xarray

import arviz as az

try:
    import pystan

    pystan_success = True
except:
    pystan_success = False

try:
    import cmdstanpy

    cmdstanpy_success = True
except:
    cmdstanpy_success = False

if pystan_success and cmdstanpy_success:
    warnings.warn(
        "Both pystan and cmdstanpy are importable in this environment. As per the cmdstanpy documentation, this is not advised."
    )

if not pystan_success and not cmdstanpy_success:
    raise RuntimeError("Neither PyStan nor CmdStanPy could be imported.")

import bokeh.plotting

from . import viz
from . import az_utils


def StanModel(
    file=None,
    charset="utf-8",
    model_name="anon_model",
    model_code=None,
    force_compile=False,
    **kwargs,
):
    """"Utility to load/save cached compiled Stan models using PyStan.

    Parameters
    ----------
    file : str or open file
        File from which Stan model is to be read. Cannot be specified
        if `model_code` is not None.
    charset : str, default utf-8
        Character set to be used to decode model.
    model_name : str, 'default anon_model'
        The name of the model.
    model_code : str, default None
        The Stan code to be compiled. If not None, `file` must be None.
    force_compile : bool, deafult False
        If True, compile, even if a cached version exists.
    kwargs : dict
        And kwargs to pass to pystem.StanModel().

    Returns
    -------
    output : pystan.model.StanModel
        A compiled Stan model.

    Notes
    -----
    .. If a Stan model does not exist in the pwd that matches the code
       provided in either `file` or `model_code`, a new Stan model is
       built and compiled and then pickled and stored in pwd. If such a
       model does exist, the pickled model is loaded.
    """

    logger = logging.getLogger("pystan")

    if file and model_code:
        raise ValueError("Specify stan model with `file` or `model_code`, " "not both.")
    if file is None and model_code is None:
        raise ValueError("Model file missing and empty model_code.")
    if file is not None:
        if isinstance(file, str):
            try:
                with open(file, "rt", encoding=charset) as f:
                    model_code = f.read()
            except:
                logger.critical("Unable to read file specified by `file`.")
                raise
        else:
            model_code = file.read()

    # Make a code_hash to use for file name
    code_hash = hashlib.md5(model_code.encode("ascii")).hexdigest()

    if model_name is None:
        cache_fname = "cached-model-{}.pkl".format(code_hash)
    else:
        cache_fname = "cached-{}-{}.pkl".format(model_name, code_hash)

    if force_compile:
        sm = pystan.StanModel(model_code=model_code, **kwargs)

        with open(cache_fname, "wb") as f:
            pickle.dump(sm, f)
    else:
        try:
            sm = pickle.load(open(cache_fname, "rb"))
        except:
            sm = pystan.StanModel(model_code=model_code, **kwargs)

            with open(cache_fname, "wb") as f:
                pickle.dump(sm, f)
        else:
            print("Using cached StanModel.")

    return sm


def clean_cmdstan(path="./", prefix=None, delete_sampling_output=False):
    """Remove all .hpp, .o, and executable files resulting from
    compilation of Stan models using CmdStanPy.

    Parameters
    ----------
    path : str, default './'
        Path to directory containing files to delete.
    prefix : str, default None
        String of prefix of model name. This name of the Stan file from
        which the model was generated is <prefix>.stan. If None, then
        all .hpp, .o, and executable files that match a prefix of both
        a .hpp and .o file are deleted.
    delete_sampling_output: bool, default False
        If True, also delete all output generated by CmdStan.

    Notes
    -----
    .. If your files are stored in a temporary directory, as for the
       CmdStanPy default, use `cmdstanpy.cleanup_tmpdir()` instead.

    """
    if path == "/":
        raise RuntimeError(
            "No way you can delete stuff from root. You are making a big mistake."
        )

    if prefix is None:
        prefix = ""

    hpp_files = glob.glob(os.path.join(path, prefix, "*.hpp"))
    o_files = glob.glob(os.path.join(path, prefix, "*.o"))
    prefixes = set([fname[:-4] for fname in hpp_files]) & set(
        [fname[:-2] for fname in o_files]
    )

    for fname in hpp_files + o_files:
        if os.path.isfile(fname):
            os.remove(fname)

    for fname in prefixes:
        if os.path.isfile(fname) and os.access(fname, os.X_OK):
            os.remove(fname)

    if delete_sampling_output:
        for pre in prefixes:
            out_files = glob.glob(pre + "*.csv")
            out_files += glob.glob(pre + "*.txt")
            for fname in out_files:
                if os.path.isfile(fname):
                    os.remove(fname)


def df_to_datadict_hier(
    df=None, level_cols=None, data_cols=None, sort_cols=[], cowardly=False
):
    """Convert a tidy data frame to a data dictionary for a hierarchical
    Stan model.

    Parameters
    ----------
    df : DataFrame
        A tidy Pandas data frame.
    level_cols : list
        A list of column names containing variables that specify the
        level of the hierarchical model. These must be given in order
        of the hierarchy of levels, with the first entry being the
        farthest from the data.
    data_cols : list
        A list of column names containing the data.
    sort_cols : list, default []]
        List of columns names to use in sorting the data. They will be
        first sorted by the level indices, and the subsequently sorted
        according to sort_cols.
    cowardly : bool, default False
        If True, refuse to generate new columns if they already exist
        in the data frame. If you run this function using a data frame
        that was outputted previously by this function, you will get an
        error if `cowardly` is True. Otherwise, the columns may be
        overwritten.

    Returns
    -------
    data : dict
        A dictionary that can be passed to into a Stan program. The
        dictionary contains keys/entries:
          'N': Total number of data points
          'J_1': Number of hyper parameters for hierarchical level 1.
          'J_2': Number of hyper parameters for hierarchical level 2.
            ... and so on with 'J_3', 'J_4', ...
          'index_1': Set of `J_2` indices defining which level 1
            parameters condition the level 2 parameters.
          'index_2': Set of `J_3` indices defining which level 2
            parameters condition the level 3 parameters.
            ...and so on for 'index_3', etc.
          'index_k': Set of `N` indices defining which of the level k
            parameters condition the data, for a k-level hierarchical
            model.
          `data_col[0]` : Data from first data_col
          `data_col[1]` : Data from second data_col ...and so on.
    df : DataFrame
        Updated input data frame with added columnes with names given by
        `level_col[0] + '_stan'`, `level_col[1] + '_stan'`, etc. These
        contain the integer indices that correspond to the possibly
        non-integer values in the `level_col`s of the original data
        frame. This enables interpretation of Stan results, which have
        everything integer indexed.

    Notes
    -----
    .. Assumes no missing data.
    .. The ordering of data sets is not guaranteed. So, e.g., if you
       have time series data, you should use caution.

    Example
    -------
    >>> import io
    >>> import pandas as pd
    >>> import bebi103
    >>> df = pd.read_csv(io.StringIO('''
        day,batch,colony,x
        monday,1,1,9.31
        monday,1,1,8.35
        monday,1,1,10.48
        monday,1,1,9.91
        monday,1,1,10.43
        monday,1,2,9.98
        monday,1,2,9.76
        monday,1,3,9.30
        monday,2,1,10.56
        monday,2,1,11.40
        monday,2,2,10.36
        monday,2,2,12.04
        monday,2,2,9.92
        monday,2,2,10.10
        monday,2,2,8.72
        monday,2,2,10.36
        monday,2,2,11.56
        monday,2,2,10.87
        monday,2,2,10.43
        monday,2,2,10.67
        monday,2,2,9.05
        monday,3,1,10.32
        monday,3,1,9.07
        monday,4,1,9.86
        monday,4,1,9.21
        monday,4,1,11.36
        monday,4,2,8.60
        monday,4,2,10.54
        monday,4,2,8.93
        monday,4,2,9.43
        monday,4,2,9.23
        monday,4,2,9.66
        monday,4,2,11.26
        monday,4,2,9.61
        monday,4,2,11.99
        monday,4,2,10.27
        monday,4,2,9.97
        monday,4,2,9.37
        monday,4,2,10.10
        monday,4,3,10.39
        monday,4,3,8.79
        wednesday,1,1,10.76
        wednesday,1,2,10.72
        wednesday,1,2,8.97
        wednesday,1,2,9.14
        wednesday,1,2,11.31
        wednesday,1,2,9.49
        wednesday,1,2,10.21
        wednesday,1,2,10.04
        wednesday,2,1,13.16
        wednesday,2,1,7.07
        wednesday,2,1,12.74
        wednesday,3,1,9.45
        wednesday,3,1,9.62
        wednesday,3,1,10.46
        wednesday,3,1,11.11
        wednesday,3,1,10.56
        wednesday,3,1,9.93
        thursday,1,1,8.60
        thursday,1,2,11.24
        thursday,1,2,9.10
        thursday,1,2,9.10
        thursday,1,2,11.30
        thursday,1,2,10.65
        thursday,1,2,9.98
        thursday,1,2,9.85
        thursday,1,2,12.41
        thursday,1,3,10.03
        thursday,1,3,10.53
        thursday,1,4,10.85
        '''), skipinitialspace=True)
    >>> data, df = bebi103.stan.df_to_datadict_hier(df,
                                level_cols=['day', 'batch', 'colony'],
                                data_cols=['x'])
    >>> data
    {'N': 70,
     'J_1': 3,
     'J_2': 8,
     'J_3': 17,
     'index_1': array([1, 1, 1, 1, 2, 3, 3, 3]),
     'index_2': array([1, 1, 1, 2, 2, 3, 4, 4, 4, 5, 5, 5, 5, 6, 6, 7, 8]),
     'index_3': array([ 1,  1,  1,  1,  1,  2,  2,  3,  4,  4,  5,  5,  5,  5,  5,  5,  5,
             5,  5,  5,  5,  6,  6,  7,  7,  7,  8,  8,  8,  8,  8,  8,  8,  8,
             8,  8,  8,  8,  8,  9,  9, 10, 11, 11, 11, 11, 11, 11, 11, 11, 12,
            12, 13, 14, 15, 15, 15, 15, 15, 15, 15, 16, 16, 16, 17, 17, 17, 17,
            17, 17]),
     'x': array([ 9.31,  8.35, 10.48,  9.91, 10.43,  9.98,  9.76,  9.3 , 10.56,
            11.4 , 10.36, 12.04,  9.92, 10.1 ,  8.72, 10.36, 11.56, 10.87,
            10.43, 10.67,  9.05, 10.32,  9.07,  9.86,  9.21, 11.36,  8.6 ,
            10.54,  8.93,  9.43,  9.23,  9.66, 11.26,  9.61, 11.99, 10.27,
             9.97,  9.37, 10.1 , 10.39,  8.79,  8.6 , 11.24,  9.1 ,  9.1 ,
            11.3 , 10.65,  9.98,  9.85, 12.41, 10.03, 10.53, 10.85, 10.76,
            10.72,  8.97,  9.14, 11.31,  9.49, 10.21, 10.04, 13.16,  7.07,
            12.74,  9.45,  9.62, 10.46, 11.11, 10.56,  9.93])}

    >>> df.head(10)
          day  batch  colony      x  day_stan  batch_stan  colony_stan
    0  monday      1       1   9.31         1           1            1
    1  monday      1       1   8.35         1           1            1
    2  monday      1       1  10.48         1           1            1
    3  monday      1       1   9.91         1           1            1
    4  monday      1       1  10.43         1           1            1
    5  monday      1       2   9.98         1           1            2
    6  monday      1       2   9.76         1           1            2
    7  monday      1       3   9.30         1           1            3
    8  monday      2       1  10.56         1           2            4
    9  monday      2       1  11.40         1           2            4
    """
    if df is None or level_cols is None or data_cols is None:
        raise RuntimeError("`df`, `level_cols`, and `data_cols` must all be specified.")

    if type(sort_cols) != list:
        raise RuntimeError("`sort_cols` must be a list.")

    # Get a copy so we don't overwrite
    new_df = df.copy(deep=True)

    if type(level_cols) not in [list, tuple]:
        level_cols = [level_cols]

    if type(data_cols) not in [list, tuple]:
        data_cols = [data_cols]

    level_cols_stan = [col + "_stan" for col in level_cols]

    if cowardly:
        for col in level_cols_stan:
            if col in df:
                raise RuntimeError(
                    "column "
                    + col
                    + " already in data frame. Cowardly deciding not to overwrite."
                )

    for col_ind, col in enumerate(level_cols):
        new_df[str(col) + "_stan"] = df.groupby(level_cols[: col_ind + 1]).ngroup() + 1

    new_df = new_df.sort_values(by=level_cols_stan + sort_cols)

    data = dict()
    data["N"] = len(new_df)
    for i, col in enumerate(level_cols_stan):
        data["J_" + str(i + 1)] = len(new_df[col].unique())
    for i, _ in enumerate(level_cols_stan[1:]):
        data["index_" + str(i + 1)] = np.array(
            [key[i] for key in new_df.groupby(level_cols_stan[: i + 2]).groups]
        ).astype(int)
    data["index_" + str(len(level_cols_stan))] = new_df[
        level_cols_stan[-1]
    ].values.astype(int)
    for col in data_cols:
        # Check string naming
        new_col = str(col)
        try:
            bytes(new_col, "ascii")
        except UnicodeEncodeError:
            raise RuntimeError("Column names must be ASCII.")
        if new_col[0].isdigit():
            raise RuntimeError("Column name cannot start with a number.")
        for char in new_col:
            if char in "`~!@#$%^&*()- =+[]{}\\|:;\"',<.>/?":
                raise RuntimeError("Invalid column name for Stan variable.")

        data[new_col] = new_df[col].values

    return data, new_df


def posterior_to_dataframe(data, var_names=None):
    """Convert ArviZ InferenceData to a Pandas data frame.

    Any multi-dimensional parameters are converted to one-dimensional
    equivalents. For example, a 2x2 matrix A is converted to columns
    in the data frame with headings `'A[0,0]'`, `'A[0,1]'`, `'A[1,0]'`,
    and `'A[1,1]'`.

    Only divergence information, chain ID, and draw number (not
    additional sampling stats) are stored in the data frame.

    Parameters
    ----------
    data : ArviZ InferenceData instance
        Results from MCMC sampling to convert to a data frame. It must
        have a `posterior` attribute, and it should have a
        `sample_stats` attribute if divergence information is to be
        stored.

    var_names : list of strings or None
        List of variables to store in the data frame. If None, all
        variables are stored.

    Returns
    -------
    output : Pandas DataFrame
        DataFrame with samples.

    """
    if type(data) == xarray.Dataset:
        raise RuntimeError(
            "You need to pass in an ArviZ InferenceData instance, not an xarray Dataset. Maybe you passed in the posterior attribute of the InferenceData instance?"
        )
    if hasattr(data, "sample_stats") and hasattr(data.sample_stats, "diverging"):
        diverging = np.ravel(data.sample_stats["diverging"])
    else:
        diverging = False

    if not hasattr(data, "posterior"):
        raise RuntimeError("No posterior contained in `data`.")

    cols, data_as_ndarray = xarray_to_ndarray(data.posterior, var_names=var_names)

    chain = np.concatenate(
        [[i] * data.posterior.dims["draw"] for i in range(data.posterior.dims["chain"])]
    )

    draw = np.concatenate(
        [data.posterior["draw"].values for i in range(data.posterior.dims["chain"])]
    )

    df = pd.DataFrame(data=data_as_ndarray.T, columns=cols)
    df["chain__"] = chain
    df["draw__"] = draw
    df["divergent__"] = diverging

    return df


def check_divergences(samples, quiet=False, return_diagnostics=False):
    """Check transitions that ended with a divergence.

    Parameters
    ----------
    samples : ArviZ InferenceData instance
        Contains samples to be checked. Must contain both `posterior`
        and `samples_stats`.
    quiet : bool, default False
        If True, do not print diagnostic result to the screen.
    return_diagnostics : bool, default False
        If True, return both a Boolean about whether the diagnostic
        passed and the number of samples where the tree depth was too
        deep. Otherwise, only return Boolean if the test passed.

    Returns
    -------
    passed : bool
        Return True if tree depth test passed. Return False otherwise.
    n_divergent : int, optional
        Number of divergent samples.
    """
    n_divergent = samples.sample_stats["diverging"].values.sum()
    n_total = samples.sample_stats.dims["chain"] * samples.sample_stats.dims["draw"]

    if not quiet:
        msg = "{} of {} ({}%) iterations ended with a divergence.".format(
            n_divergent, n_total, 100 * n_divergent / n_total
        )
        print(msg)

    pass_check = n_divergent == 0

    if not pass_check and not quiet:
        print("  Try running with larger adapt_delta to remove divergences.")

    if return_diagnostics:
        return pass_check, n_divergent
    return pass_check


def check_treedepth(samples, max_treedepth=10, quiet=False, return_diagnostics=False):
    """Check transitions that ended prematurely due to maximum tree depth limit.

    Parameters
    ----------
    samples : ArviZ InferenceData instance
        Contains samples to be checked. Must contain both `posterior`
        and `samples_stats`.
    max_treedepth: int, default 10
        Maximum tree depth used in the calculation.
    quiet : bool, default False
        If True, do not print diagnostic result to the screen.
    return_diagnostics : bool, default False
        If True, return both a Boolean about whether the diagnostic
        passed and the number of samples where the tree depth was too
        deep. Otherwise, only return Boolean if the test passed.

    Returns
    -------
    passed : bool
        Return True if tree depth test passed. Return False otherwise.
    n_too_deep : int, optional
        Number of samplers wherein the tree depth was greater than
        `max_treedepth`.
    """
    n_too_deep = (samples.sample_stats.treedepth.values >= max_treedepth).sum()
    n_total = samples.sample_stats.dims["chain"] * samples.sample_stats.dims["draw"]

    if not quiet:
        msg = "{} of {} ({}%) iterations saturated".format(
            n_too_deep, n_total, 100 * n_too_deep / n_total
        )
        msg += " the maximum tree depth of {}.".format(max_treedepth)
        print(msg)

    pass_check = n_too_deep == 0

    if not pass_check and not quiet:
        print(
            "  Try running again with max_treedepth set to a larger value"
            + " to avoid saturation."
        )

    if return_diagnostics:
        return pass_check, n_too_deep
    return pass_check


def check_energy(
    samples, e_bfmi_rule_of_thumb=0.3, quiet=False, return_diagnostics=False
):
    """Checks the energy-Bayes fraction of missing information (E-BFMI)

    Parameters
    ----------
    samples : ArviZ InferenceData instance
        Contains samples to be checked. Must contain both `posterior`
        and `samples_stats`.
    e_bfmi_rule_of_thumb : float, default 0.3 (as per cmdstan)
        Rule of thumb value for E-BFMI. If below this value, there may
        be cause for concern.
    quiet : bool, default False
        If True, do no print diagnostic result to the screen.
    return_diagnostics : bool, default False
        If True, return both a Boolean about whether the diagnostic
        passed and a data frame containing results about the E-BMFI
        tests. Otherwise, only return Boolean if the test passed.

    Returns
    -------
    passed : bool
        Return True if test passed. Return False otherwise.
    e_bfmi_diagnostics : DataFrame
        DataFrame with information about which chains had problematic
        E-BFMIs.
    """
    # Alternatively:
    # samples.sample_stats["energy"].groupby("chain").apply(_ebfmi).values
    ebfmi = az.bfmi(samples)

    problematic = ebfmi < e_bfmi_rule_of_thumb

    pass_check = (~problematic).all()

    if not quiet:
        if pass_check:
            print("E-BFMI indicated no pathological behavior.")
        else:
            for i, (problem, ebfmi_val) in enumerate(zip(problematic, ebfmi)):
                print("Chain {}: E-BFMI = {}".format(i, ebfmi_val))
            print(
                "  E-BFMI below 0.3 indicates you may need to "
                + "reparametrize your model."
            )

    if return_diagnostics:
        return (
            pass_check,
            pd.DataFrame(
                {
                    "chain": np.arange(len(ebfmi)),
                    "E-BFMI": ebfmi,
                    "problematic": ebfmi < e_bfmi_rule_of_thumb,
                }
            ),
        )
    return pass_check


def check_ess(
    samples,
    var_names=None,
    total_ess_rule_of_thumb=100,
    fractional_ess_rule_of_thumb=0.001,
    quiet=False,
    return_diagnostics=False,
):
    """Checks the effective sample size (ESS).

    Parameters
    ----------
    samples : ArviZ InferenceData instance
        Contains samples to be checked. Must contain both `posterior`
        and `samples_stats`.
    var_names : list of str, or None (default)
        Names of parameters to use in doing check for ESS. If None,
        use all parameters.
    quiet : bool, default False
        If True, do not print diagnostic result to the screen.
    total_ess_rule_of_thumb : int, default 100
        Rule of thumb for number of effective samples per chain. The
        default of 100 is based on the suggestion of Vehtari, et al., to
        have 50 effective samples per split chain.
    fractional_ess_rule_of_thumb : float, default 0.001
        Rule of thumb value for fractional number of effective samples.
        The default of 0.001 is based on CmdStan's defaults.
    return_diagnostics : bool, default False
        If True, return both a Boolean about whether the diagnostic
        passed and a data frame containing results about the number
        of effective samples tests. Otherwise, only return Boolean if
        the test passed.

    Returns
    -------
    passed : bool
        Return True if test passed. Return False otherwise.
    ess_diagnostics : DataFrame
        Data frame with information about problematic n_eff.

    Notes
    -----
    .. Parameters with n_eff given as NaN are not checked.

    References
    ----------
    .. Vehtari, et al., 2019, https://arxiv.org/abs/1903.08008.
    """
    # For convenience
    N = samples.posterior.dims["draw"]
    M = samples.posterior.dims["chain"]
    total_ess_rule_of_thumb *= M

    # Compute ESS
    if var_names is None:
        var_names = [var for var in samples.posterior.data_vars]
    ess = az.ess(samples, var_names=var_names, method="bulk")
    tail_ess = az.ess(samples, var_names=var_names, method="tail")

    # Convert to list of names and numpy arrays
    names, ess = xarray_to_ndarray(ess)
    _, tail_ess = xarray_to_ndarray(tail_ess)
    if ess.shape == (1, 1):
        ess = np.array([ess[0, 0]])
        tail_ess = np.array([tail_ess[0, 0]])
    else:
        ess = ess.squeeze()
        tail_ess = tail_ess.squeeze()

    # Fractional ESS
    frac_ess = ess / M / N
    frac_tail_ess = tail_ess / M / N

    pass_check = (
        (ess[~np.isnan(ess)] > total_ess_rule_of_thumb).all()
        and (frac_ess[~np.isnan(ess)] > fractional_ess_rule_of_thumb).all()
        and (tail_ess[~np.isnan(tail_ess)] > total_ess_rule_of_thumb).all()
        and (frac_tail_ess[~np.isnan(tail_ess)] > fractional_ess_rule_of_thumb).all()
    )

    if not quiet:
        if not pass_check:
            n_e = 0
            n_f = 0
            for name, e, f, te, tf in zip(
                names, ess, frac_ess, tail_ess, frac_tail_ess
            ):
                if e < total_ess_rule_of_thumb:
                    print("ESS for parameter {} is {}.".format(name, e))
                    n_e += 1
                if f < fractional_ess_rule_of_thumb:
                    print("ESS / iter for parameter {} is {}.".format(name, f))
                    n_f += 1
                if te < total_ess_rule_of_thumb:
                    print("tail-ESS for parameter {} is {}.".format(name, te))
                    n_e += 1
                if tf < fractional_ess_rule_of_thumb:
                    print("ESS / iter for parameter {} is {}.".format(name, tf))
                    n_f += 1
            if n_e > 0:
                print(
                    """  ESS or tail-ESS below 100 per chain indicates that expectation values
  computed from samples are unlikely to be good approximations of the
  true expectation values."""
                )
            if n_f > 0:
                print(
                    """  ESS / iter or tail-ESS / iter below 0.001 indicates that the effective
  sample size has likely been overestimated."""
                )
        else:
            print("Effective sample size looks reasonable for all parameters.")

    if return_diagnostics:
        return (
            pass_check,
            pd.DataFrame(data={"parameter": names, "ESS": ess, "ESS/iter": frac_ess}),
        )
    return pass_check


def check_rhat(
    samples,
    var_names=None,
    rhat_rule_of_thumb=1.01,
    known_rhat_nans=None,
    quiet=False,
    return_diagnostics=False,
):
    """Checks the potential issues with scale reduction factors. Rhat
    is computed with rank-normalization and folding.

    Parameters
    ----------
    samples : ArviZ InferenceData instance
        Contains samples to be checked. Must contain both `posterior`
        and `samples_stats`.
    var_names : list of str, or None (default)
        Names of parameters to use in doing diagnostic checks. If None,
        use all parameters.
    quiet : bool, default False
        If True, do no print diagnostic result to the screen.
    rhat_rule_of_thumb : float, default 1.01
        Rule of thumb value for maximum allowed R-hat, as per Vehtari,
        et al.
    known_rhat_nans : list, default None
        List of parameter names which are known to have R-hat be NaN.
        These are typically parameters that are deterministic. If a
        parameter is multidimensional, the specific parameter should
        be listed. For example, if A is a 2x2 symmetric matrix and the
        entry in the lower left corner has a known NaN for R-hat, then
        known_rhat_nan = ['A[1,1]']. Parameters in `known_rhat_nans` are
        ignored.
    return_diagnostics : bool, default False
        If True, return both a Boolean about whether the diagnostic
        passed and a data frame containing results about the number
        of effective samples tests. Otherwise, only return Boolean if
        the test passed.

    Results
    -------
    passed : bool
        Return True if test passed. Return False otherwise.
    rhat_diagnostics : DataFrame
        Data frame with information about problematic R-hat values.

    References
    ----------
    .. Vehtari, et al., 2019, https://arxiv.org/abs/1903.08008.

    """
    if known_rhat_nans is None:
        known_rhat_nans = []
    elif type(known_rhat_nans) == str:
        known_rhat_nans = [known_rhat_nans]

    if var_names is None:
        var_names = [var for var in samples.posterior.data_vars]

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        rhat = az.rhat(samples, var_names=var_names, method="rank")

    # Convert to list of names and numpy arrays
    names, rhat = xarray_to_ndarray(rhat)
    if rhat.shape == (1, 1):
        rhat = np.array([rhat[0, 0]])
    else:
        rhat = rhat.squeeze()
    names = [name.replace("\n", "[") + "]" if "\n" in name else name for name in names]

    known_nan = [True if name in known_rhat_nans else False for name in names]

    pass_check = np.isnan(rhat[~np.array(known_nan)]).sum() == 0 and np.all(
        rhat[~np.array(known_nan)] < rhat_rule_of_thumb
    )

    if not quiet:
        if not pass_check:
            non_nan_fail = 0
            for name, r, nan in zip(names, rhat, known_nan):
                if np.isnan(r) and not nan:
                    print("Rhat for parameter {} is {}.".format(name, r))
                elif r > rhat_rule_of_thumb:
                    non_nan_fail += 1
                    print("Rhat for parameter {} is {}.".format(name, r))
            if non_nan_fail > 0:
                print(
                    "  Rank-normalized Rhat above 1.01 indicates that the chains very likely"
                    + " have not mixed"
                )
        else:
            print("Rhat looks reasonable for all parameters.")

    if return_diagnostics:
        return pass_check, pd.DataFrame(data={"parameter": names, "Rhat": rhat})
    return pass_check


def check_all_diagnostics(
    samples,
    var_names=None,
    e_bfmi_rule_of_thumb=0.3,
    total_ess_rule_of_thumb=100,
    fractional_ess_rule_of_thumb=0.001,
    rhat_rule_of_thumb=1.01,
    known_rhat_nans=None,
    max_treedepth=10,
    quiet=False,
    return_diagnostics=False,
):
    """Checks all MCMC diagnostics

    Parameters
    ----------
    samples : ArviZ InferenceData instance
        Contains samples to be checked. Must contain both `posterior`
        and `samples_stats`.
    var_names : list of str, or None (default)
        Names of parameters to use in doing diagnostic checks. If None,
        use all parameters.
    e_bfmi_rule_of_thumb : float, default 0.3
        Rule of thumb value for E-BFMI. If below this value, there may
        be cause for concern.
    total_ess_rule_of_thumb : int, default 100
        Rule of thumb for number of effective samples per chain. The
        default of 100 is based on the suggestion of Vehtari, et al., to
        have 50 effective samples per split chain.
    fractional_ess_rule_of_thumb : float, default 0.001
        Rule of thumb value for fractional number of effective samples.
        The default of 0.001 is based on CmdStan's defaults.
    rhat_rule_of_thumb : float, default 1.1
        Rule of thumb value for maximum allowed R-hat.
    known_rhat_nans : list, default []
        List of parameter names which are known to have R-hat be NaN.
        These are typically parameters that are deterministic.
        Parameters in this list are ignored.
    max_treedepth: int, default 'infer'
        If int, specification of maximum treedepth allowed. If 'infer',
        inferred from `fit`.
    quiet : bool, default False
        If True, do no print diagnostic result to the screen.
    return_diagnostics : bool, default False
        If True, return a dictionary containing the results of each test.

    Results
    -------
    warning_code : int
        When converted to binary, each digit in the code stands for
        whether or not a test passed. A digit of zero indicates the test
        passed. The ordering of the tests goes:
            ess
            r_hat
            divergences
            tree depth
            E-BFMI
        For example, a warning code of 12 has a binary representation
        of 01100, which means that R-hat and divergences tests failed.
    return_dict : dict
        Returned if `return_dict` is True. A dictionary with the result
        of each diagnostic test.
    """
    warning_code = 0
    diag_results = {}

    pass_check, diag_results["ess"] = check_ess(
        samples,
        var_names=var_names,
        total_ess_rule_of_thumb=total_ess_rule_of_thumb,
        fractional_ess_rule_of_thumb=fractional_ess_rule_of_thumb,
        quiet=quiet,
        return_diagnostics=True,
    )
    if not pass_check:
        warning_code = warning_code | (1 << 0)

    pass_check, diag_results["rhat"] = check_rhat(
        samples,
        var_names=var_names,
        rhat_rule_of_thumb=rhat_rule_of_thumb,
        known_rhat_nans=known_rhat_nans,
        quiet=quiet,
        return_diagnostics=True,
    )
    if not pass_check:
        warning_code = warning_code | (1 << 1)

    pass_check, diag_results["n_divergences"] = check_divergences(
        samples, quiet=quiet, return_diagnostics=True
    )
    if not pass_check:
        warning_code = warning_code | (1 << 2)

    pass_check, diag_results["n_max_treedepth"] = check_treedepth(
        samples, max_treedepth=max_treedepth, quiet=quiet, return_diagnostics=True
    )
    if not pass_check:
        warning_code = warning_code | (1 << 3)

    pass_check, diag_results["e_bfmi"] = check_energy(
        samples,
        e_bfmi_rule_of_thumb=e_bfmi_rule_of_thumb,
        quiet=quiet,
        return_diagnostics=True,
    )
    if not pass_check:
        warning_code = warning_code | (1 << 4)

    if return_diagnostics:
        return warning_code, diag_results

    return warning_code


def parse_warning_code(warning_code, quiet=False, return_dict=False):
    """Parses warning code from `check_all_diagnostics()` into
    individual failures and prints results.

    Parameters
    ----------
    warning_code : int
        When converted to binary, each digit in the code stands for
        whether or not a test passed. A digit of zero indicates the test
        passed. The ordering of the tests goes:
            ESS
            r_hat
            divergences
            tree depth
            E-BFMI
        For example, a warning code of 12 has a binary representation
        of 01100, which means that R-hat and divergences tests failed.
    quiet : bool, default False
        If True, suppress print results to the screen.
    return_dict : bool, default False
        If True, return a dictionary of containing test passage info.

    Returns
    -------
    output : dict
        If `return_dict` is True, returns a dictionary where each entry
        is True is the respective diagnostic check was passed and False
        if it was not.
    """

    if quiet and not return_dict:
        raise RuntimeError(
            "`quiet` is True and `return_dict` is False, "
            + "so there is nothing to do."
        )

    passed_tests = dict(
        neff=True, rhat=True, divergence=True, treedepth=True, energy=True
    )

    if warning_code & (1 << 0):
        passed_tests["neff"] = False
        if not quiet:
            print("n_eff / iteration warning")
    if warning_code & (1 << 1):
        passed_tests["rhat"] = False
        if not quiet:
            print("Rhat warning")
    if warning_code & (1 << 2):
        passed_tests["divergence"] = False
        if not quiet:
            print("divergence warning")
    if warning_code & (1 << 3):
        passed_tests["treedepth"] = False
        if not quiet:
            print("treedepth warning")
    if warning_code & (1 << 4):
        passed_tests["energy"] = False
        if not quiet:
            print("energy warning")
    if warning_code == 0:
        if not quiet:
            print("No diagnostic warnings")

    if return_dict:
        return passed_tests


def sbc(
    prior_predictive_model=None,
    posterior_model=None,
    prior_predictive_model_data=None,
    posterior_model_data=None,
    measured_data=None,
    parameters=None,
    measured_data_dtypes={},
    sampling_kwargs={},
    posterior_predictive_var_names=None,
    cores=1,
    N=400,
    n_prior_draws_for_sd=1000,
    known_rhat_nans=None,
    progress_bar=False,
):
    """Perform simulation-based calibration on a Stan Model.

    Parameters
    ----------
    prior_predictive_model : pystan.model.StanModel
        A Stan model for generating prior predictive data sets.
    posterior_model : pystan.model.StanModel
        A Stan model of the posterior that allows sampling.
    prior_predictive_model_data : dict
        Dictionary with entries specified by the data block of the prior
        predictive Stan model.
    posterior_model_data : dict
        Dictionary with entries specified by the data block of the prior
        predictive Stan model. Measured data in this dictionary will be
        replaced in each simulation by what was generated by the prior
        predictive model.
    measured_data : list
        A list of strings containing the variable names of measured
        data. Each entry in `measured_data` must be a key in
        `posterior_model_data`.
    parameters : list of strings
        A list of strings containing parameter names to be considered
        in the SBC analysis. Not all parameters of the model need be
        considered; only those in `parameters` have rank statistics
        calculated.
    posterior_predictive_var_names : list of strings, default None
        List of variables that are posterior predictive. These are
        ignored in the SBC analysis.
    measured_data_dtypes : dict, default {}
        The key in the dtypes dict is a string representing the data
        name, and the corresponding item is its dtype, almost always
        either `int` or `float`.
    sampling_kwargs : dict, default {}
        kwargs to be passed to `sm.sample()` for a CmdStanPy model `sm`
        or to `sm.sampling()` for a PyStan model `sm`.
    cores : int, default 1
        Number of cores to use in the SBC calculation.
    N : int, 400
        Number of simulations to run.
    n_prior_draws_for_sd : int, default 1000
        Number of prior draws to compute the prior standard deviation
        for a parameter in the prior distribution. This standard
        deviation is used in the shrinkage calculation.
    known_rhat_nans : list, default []
        List of parameter names which are known to have R-hat be NaN.
        These are typically parameters that are deterministic.
        Parameters in this list are ignored.
    progress_bar : bool, default False
        If True, display a progress bar for the calculation using tqdm.

    Returns
    -------
    output : Pandas DataFrame
        A Pandas DataFrame with the output of the SBC analysis. It has
        the following columns.
        - trial : Unique trial number for the simulation.
        - warning_code : Warning code based on diagnostic checks
            outputted by `check_all_diagnostics()`.
        - parameter: The name of the scalar parameter.
        - prior: Value of the parameter used in the simulation. This
            value was drawn out of the prior distribution.
        - mean : mean parameter value based on sampling out of the
            posterior in the simulation.
        - sd : standard deviation of the parameter value based on
            sampling out of the posterior in the simulation.
        - L : The number of bins used in computing the rank statistic.
            The rank statistic should be uniform on the integers [0, L].
        - rank_statistic : Value of the rank statistic for the parameter
            for the trial.
        - shrinkage : The shrinkage for the parameter for the given
            trial. This is computed as 1 - sd / sd_prior, where sd_prior
            is the standard deviation of the parameters as determined
            from drawing out of the prior.
        - z_score : The z-score for the parameter for the given trial.
            This is computed as |mean - prior| / sd.

    Notes
    -----
    .. Each simulation is done by sampling a parameter set out of the
       prior distribution, using those parameters to generate data from
       the likelihood, and then performing posterior sampling based on
       the generated data. A rank statistic for each simulation is
       computed. This rank statistic should be uniformly distributed
       over its L possible values. See https://arxiv.org/abs/1804.06788,
       by Talts, et al., for details.

    """
    if prior_predictive_model is None:
        raise RuntimeError("`prior_predictive_model` must be specified.")
    if posterior_model is None:
        raise RuntimeError("`posterior_model` must be specified.")
    if prior_predictive_model_data is None:
        raise RuntimeError("`prior_predictive_model_data` must be specified.")
    if posterior_model_data is None:
        raise RuntimeError("`posterior_model_data` must be specified.")
    if measured_data is None:
        raise RuntimeError("`measured_data` must be specified.")

    # We parallelize simulations, not chains within each simulation
    if "n_jobs" in sampling_kwargs:
        del sampling_kwargs["n_jobs"]
    if "cores" in sampling_kwargs:
        del sampling_kwargs["cores"]

    # Take a prior sample to infer data types
    if "pystan" in str(type(prior_predictive_model)):
        with disable_logging():
            prior_sample = prior_predictive_model.sampling(
                data=prior_predictive_model_data,
                algorithm="Fixed_param",
                iter=1,
                chains=1,
                warmup=0,
            )
        if az.__version__ > "0.6.1":
            prior_sample = az.from_pystan(
                prior=prior_sample, prior_predictive=measured_data
            )
        else:
            prior_sample = az.from_pystan(
                posterior=prior_sample,
                prior=prior_sample,
                prior_predictive=measured_data,
            )
    elif "cmdstanpy" in str(type(prior_predictive_model)):
        with disable_logging():
            prior_sample = prior_predictive_model.sample(
                data=prior_predictive_model_data, fixed_param=True, sampling_iters=1
            )
        if az.__version__ > "0.6.1":
            prior_sample = az.from_cmdstanpy(
                prior=prior_sample, prior_predictive=measured_data
            )
        else:
            prior_sample = az.from_cmdstanpy(
                posterior=prior_sample,
                prior=prior_sample,
                prior_predictive=measured_data,
            )
    elif callable(prior_predictive_model):
        raise NotImplementedError("Python-based prior models not yet supported.")
    else:
        raise RuntimeError("Improper type for `prior_predictive_model`.")

    # Infer dtypes of measured data
    for data in measured_data:
        ar = prior_sample.prior_predictive[data].squeeze()
        if data not in measured_data_dtypes:
            if np.sum(ar != ar.astype(int)) == 0:
                warnings.warn(f"Inferring int dtype for {data}.")
                measured_data_dtypes[data] = int
            else:
                measured_data_dtypes[data] = float

    # Check parameters
    if parameters is None:
        parameters = [
            param
            for param in prior_sample.prior.data_vars
            if len(param) < 2 or param[-2:] != "__"
        ]
    else:
        for param in parameters:
            if param not in prior_sample.prior.data_vars:
                raise RuntimeError(
                    f"parameter `{param}` not in prior sample generated from `prior_predictive_model`."
                )

    # Determine prior SDs for parameters of interest
    prior_sd = _get_prior_sds(
        prior_predictive_model,
        prior_predictive_model_data,
        parameters,
        measured_data,
        n_prior_draws_for_sd,
    )

    def arg_input_generator():
        counter = 0
        while counter < N:
            counter += 1
            yield (
                prior_predictive_model,
                posterior_model,
                prior_predictive_model_data,
                posterior_model_data,
                measured_data,
                parameters,
                measured_data_dtypes,
                sampling_kwargs,
                posterior_predictive_var_names,
                prior_sd,
                known_rhat_nans,
            )

    with multiprocessing.Pool(cores) as pool:
        if progress_bar == "notebook":
            output = list(
                tqdm.tqdm_notebook(
                    pool.imap(_perform_sbc, arg_input_generator()), total=N
                )
            )
        elif progress_bar == True:
            output = list(
                tqdm.tqdm(pool.imap(_perform_sbc, arg_input_generator()), total=N)
            )
        elif progress_bar == False:
            output = pool.map(_perform_sbc, arg_input_generator())
        else:
            raise RuntimeError("Invalid `progress_bar`.")

    output = pd.DataFrame(output)

    # Determine number of iterations
    thin = sampling_kwargs["thin"] if "thin" in sampling_kwargs else 1
    chains = sampling_kwargs["chains"] if "chains" in sampling_kwargs else 4
    if "pystan" in str(type(posterior_model)):
        if "iter" in sampling_kwargs:
            if "warmup" in sampling_kwargs:
                sampling_iters = sampling_kwargs["iter"] - sampling_kwargs["warmup"]
            else:
                sampling_iters = sampling_kwargs["iter"] - sampling_kwargs["iter"] // 2
        elif "warmup" in sampling_kwargs:
            sampling_iters = 2000 - sampling_kwargs["warmup"]
        else:
            sampling_iters = 1000
    elif "cmdstanpy" in str(type(prior_predictive_model)):
        if "sampling_iters" in sampling_kwargs:
            sampling_iters = sampling_kwargs["sampling_iters"]
        else:
            sampling_iters = 1000

    output["L"] = sampling_iters * chains // thin

    return _tidy_sbc_output(output)


def _perform_sbc(args):
    """Perform an SBC analysis"""
    logging.getLogger("pystan").setLevel(logging.CRITICAL)

    (
        prior_predictive_model,
        posterior_model,
        prior_predictive_model_data,
        posterior_model_data,
        measured_data,
        parameters,
        measured_data_dtypes,
        sampling_kwargs,
        posterior_predictive_var_names,
        prior_sd,
        known_rhat_nans,
    ) = args

    posterior_model_data = copy.deepcopy(posterior_model_data)

    if "pystan" in str(type(prior_predictive_model)):
        with disable_logging():
            prior_sample = prior_predictive_model.sampling(
                data=prior_predictive_model_data,
                algorithm="Fixed_param",
                iter=1,
                chains=1,
                warmup=0,
            )
        if az.__version__ > "0.6.1":
            prior_sample = az.from_pystan(
                prior=prior_sample, prior_predictive=measured_data
            )
        else:
            prior_sample = az.from_pystan(
                posterior=prior_sample,
                prior=prior_sample,
                prior_predictive=measured_data,
            )
    elif "cmdstanpy" in str(type(prior_predictive_model)):
        with disable_logging():
            prior_sample = prior_predictive_model.sample(
                data=prior_predictive_model_data, fixed_param=True, sampling_iters=1
            )
        if az.__version__ > "0.6.1":
            prior_sample = az.from_cmdstanpy(
                prior=prior_sample, prior_predictive=measured_data
            )
        else:
            prior_sample = az.from_cmdstanpy(
                posterior=prior_sample,
                prior=prior_sample,
                prior_predictive=measured_data,
            )

    # Extract data generated from the prior predictive calculation
    for data in measured_data:
        ar = prior_sample.prior_predictive[data].squeeze().values
        if len(ar.shape) == 0:
            posterior_model_data[data] = ar.astype(measured_data_dtypes[data]).item()
        else:
            posterior_model_data[data] = ar.astype(measured_data_dtypes[data])

    # Store what the parameters were to generate prior predictive data
    names, vals = xarray_to_ndarray(prior_sample.prior)
    param_priors = {name: val.item() for name, val in zip(names, vals)}

    # Generate posterior samples
    if "pystan" in str(type(posterior_model)):
        with disable_logging():
            posterior_samples = posterior_model.sampling(
                data=posterior_model_data, n_jobs=1, **sampling_kwargs
            )
        posterior_samples = az.from_pystan(
            posterior=posterior_samples,
            posterior_predictive=posterior_predictive_var_names,
        )

    elif "cmdstanpy" in str(type(posterior_model)):
        with disable_logging():
            posterior_samples = posterior_model.sample(
                data=posterior_model_data, cores=1, **sampling_kwargs
            )
        posterior_samples = az.from_cmdstanpy(
            posterior=posterior_samples,
            posterior_predictive=posterior_predictive_var_names,
        )

    # Check diagnostics
    warning_code, diagnostics = check_all_diagnostics(
        posterior_samples,
        var_names=parameters,
        quiet=True,
        known_rhat_nans=known_rhat_nans,
        return_diagnostics=True,
    )

    # Convert output to Numpy array
    names, vals = xarray_to_ndarray(posterior_samples.posterior, var_names=parameters)

    # Generate output dictionary
    output = {
        name + "_rank_statistic": (vals[i] < param_priors[name]).sum()
        for i, name in enumerate(names)
    }
    for name, p_prior in param_priors.items():
        output[name + "_ground_truth"] = p_prior

    # Compute posterior sensitivities
    for name, val in zip(names, vals):
        output[name + "_mean"] = np.mean(val)
        output[name + "_sd"] = np.std(val)
        output[name + "_z_score"] = (
            output[name + "_mean"] - output[name + "_ground_truth"]
        ) / output[name + "_sd"]
        output[name + "_shrinkage"] = 1 - (output[name + "_sd"] / prior_sd[name]) ** 2
        output[name + "_ESS"] = (
            diagnostics["ess"]
            .loc[diagnostics["ess"]["parameter"] == name, "ESS"]
            .values[0]
        )
        output[name + "_ESS/iter"] = (
            diagnostics["ess"]
            .loc[diagnostics["ess"]["parameter"] == name, "ESS/iter"]
            .values[0]
        )
        output[name + "_Rhat"] = (
            diagnostics["rhat"]
            .loc[diagnostics["rhat"]["parameter"] == name, "Rhat"]
            .values[0]
        )

    output["n_bad_ebfmi"] = diagnostics["e_bfmi"]["problematic"].sum()
    output["n_divergences"] = int(diagnostics["n_divergences"])
    output["n_max_treedepth"] = int(diagnostics["n_max_treedepth"])
    output["warning_code"] = warning_code

    return output


def _get_prior_sds(
    prior_predictive_model,
    prior_predictive_model_data,
    parameters,
    measured_data,
    n_prior_draws_for_sd,
):
    """Compute standard deviations of prior parameters."""
    if "pystan" in str(type(prior_predictive_model)):
        with disable_logging():
            prior_samples = prior_predictive_model.sampling(
                data=prior_predictive_model_data,
                algorithm="Fixed_param",
                iter=n_prior_draws_for_sd,
                chains=1,
                warmup=0,
            )
        if az.__version__ > "0.6.1":
            prior_samples = az.from_pystan(
                prior=prior_samples, prior_predictive=measured_data
            )
        else:
            prior_samples = az.from_pystan(
                posterior=prior_samples,
                prior=prior_samples,
                prior_predictive=measured_data,
            )
    elif "cmdstanpy" in str(type(prior_predictive_model)):
        with disable_logging():
            prior_samples = prior_predictive_model.sample(
                data=prior_predictive_model_data,
                fixed_param=True,
                sampling_iters=n_prior_draws_for_sd,
            )
        if az.__version__ > "0.6.1":
            prior_samples = az.from_cmdstanpy(
                prior=prior_samples, prior_predictive=measured_data
            )
        else:
            prior_samples = az.from_cmdstanpy(
                posterior=prior_samples,
                prior=prior_samples,
                prior_predictive=measured_data,
            )

    # Compute prior sd's
    names, vals = xarray_to_ndarray(prior_samples.prior)
    prior_sd = {
        name: np.std(val) for name, val in zip(names, vals) if name in parameters
    }

    return prior_sd


def _tidy_sbc_output(sbc_output):
    """Tidy output from sbc().

    Returns
    -------
    output : DataFrame
        Tidy data frame with SBC results.
    """
    df = sbc_output.copy()
    df["trial"] = df.index.values

    rank_stat_cols = list(df.columns[df.columns.str.contains("_rank_statistic")])
    params = [col[: col.rfind("_rank_statistic")] for col in rank_stat_cols]

    dfs = []
    stats = [
        "ground_truth",
        "rank_statistic",
        "mean",
        "sd",
        "shrinkage",
        "z_score",
        "Rhat",
        "ESS",
        "ESS/iter",
    ]

    aux_cols = [
        "n_divergences",
        "n_bad_ebfmi",
        "n_max_treedepth",
        "warning_code",
        "L",
        "trial",
    ]
    for param in params:
        cols = [param + "_" + stat for stat in stats]
        sub_df = df[cols + aux_cols].rename(
            columns={old_col: new_col for old_col, new_col in zip(cols, stats)}
        )
        sub_df["parameter"] = param
        dfs.append(sub_df)

    return pd.concat(dfs, ignore_index=True)


def xarray_to_ndarray(ds, var_names=None, omit_dunders=True):
    """Convert xarray data set with coordinates `chain` and `draw` to a
    Numpy array and a list of row labels for the Numpy array.

    Parameters
    ----------
    ds : xarray Dataset
        Dataset stored in an ArviZ InferenceData instance. This is
        usually contained in the `prior` or `posterior` fields of the
        InferenceData instance.
    var_names : list of strings, default None
        Names of variables to include. If None, include all.
    omit_dunders : bool, default True
        If True, omit any variable that ends in '__' unless it is
        explicitly included in `var_names`. These are not variables, but
        sampling information computed by Stan.

    Output
    ------
    names : list of strings
        List of names of headings. For a mulitdimensional xarray
        DataArray, each entry in the array is the variable name,
        followed by a newline character and indexing information. As
        an example, if `A` is a 2x2 matrix, then there will be names
        'A[0,0]', 'A[0,1]', 'A[1,0]', and 'A[1,1]'.
    vals : 2D Numpy array
        Each row has the samples for a given variable. If combined is
        True, each row is a concatenation of samples from all chains.
    """

    names, vals = az_utils.xarray_to_ndarray(ds, var_names=var_names, combined=True)

    names = [
        name.replace("\n", "[").replace(", ", ",") + "]" if "\n" in name else name
        for name in names
    ]

    if var_names is None:
        var_names = []

    inds = [
        i
        for i, name in enumerate(names)
        if name in var_names or len(name) < 2 or name[-2:] != "__"
    ]

    return [name for i, name in enumerate(names) if i in inds], vals[inds, :]


def _ebfmi(energy):
    """Compute energy-Bayes fraction of missing information"""
    return np.sum(np.diff(energy) ** 2) / (len(energy) - 1) / np.var(energy)


@contextlib.contextmanager
def disable_logging(level=logging.CRITICAL):
    """Context manager for disabling logging."""
    previous_level = logging.root.manager.disable

    logging.disable(level)

    try:
        yield
    finally:
        logging.disable(previous_level)
