# pylint: disable=no-member, invalid-name, redefined-outer-name
# pylint: disable=too-many-lines

import os
from collections import namedtuple
from copy import deepcopy
from html import escape
from typing import Dict
from urllib.parse import urlunsplit

import numpy as np
import pytest
import xarray as xr
from xarray.core.options import OPTIONS
from xarray.testing import assert_identical

from ... import (
    InferenceData,
    clear_data_home,
    concat,
    convert_to_dataset,
    convert_to_inference_data,
    from_dict,
    from_json,
    from_netcdf,
    list_datasets,
    load_arviz_data,
    to_netcdf,
)

from ...data.base import dict_to_dataset, generate_dims_coords, infer_stan_dtypes, make_attrs
from ...data.datasets import LOCAL_DATASETS, REMOTE_DATASETS, RemoteFileMetadata
from ..helpers import (  # pylint: disable=unused-import
    chains,
    check_multiple_attrs,
    create_data_random,
    data_random,
    draws,
    eight_schools_params,
    models,
)


@pytest.fixture(autouse=True)
def no_remote_data(monkeypatch, tmpdir):
    """Delete all remote data and replace it with a local dataset."""
    keys = list(REMOTE_DATASETS)
    for key in keys:
        monkeypatch.delitem(REMOTE_DATASETS, key)

    centered = LOCAL_DATASETS["centered_eight"]
    filename = os.path.join(str(tmpdir), os.path.basename(centered.filename))

    url = urlunsplit(("file", "", centered.filename, "", ""))

    monkeypatch.setitem(
        REMOTE_DATASETS,
        "test_remote",
        RemoteFileMetadata(
            filename=filename,
            url=url,
            checksum="2c5501a9f5d7b6998fc7e6a4651030b9765032b2e5a1d7331f5b1f3df6c632a5",
            description=centered.description,
        ),
    )
    monkeypatch.setitem(
        REMOTE_DATASETS,
        "bad_checksum",
        RemoteFileMetadata(
            filename=filename, url=url, checksum="bad!", description=centered.description
        ),
    )
    UnknownFileMetaData = namedtuple(
        "UnknownFileMetaData", ["filename", "url", "checksum", "description"]
    )
    monkeypatch.setitem(
        REMOTE_DATASETS,
        "test_unknown",
        UnknownFileMetaData(
            filename=filename,
            url=url,
            checksum="9ae00c83654b3f061d32c882ec0a270d10838fa36515ecb162b89a290e014849",
            description="Test bad REMOTE_DATASET",
        ),
    )


def test_load_local_arviz_data():
    inference_data = load_arviz_data("centered_eight")
    assert isinstance(inference_data, InferenceData)
    assert set(inference_data.observed_data.obs.coords["school"].values) == {
        "Hotchkiss",
        "Mt. Hermon",
        "Choate",
        "Deerfield",
        "Phillips Andover",
        "St. Paul's",
        "Lawrenceville",
        "Phillips Exeter",
    }
    assert inference_data.posterior["theta"].dims == ("chain", "draw", "school")


def test_clear_data_home():
    resource = REMOTE_DATASETS["test_remote"]
    assert not os.path.exists(resource.filename)
    load_arviz_data("test_remote")
    assert os.path.exists(resource.filename)
    clear_data_home(data_home=os.path.dirname(resource.filename))
    assert not os.path.exists(resource.filename)


def test_load_remote_arviz_data():
    assert load_arviz_data("test_remote")


def test_bad_checksum():
    with pytest.raises(IOError):
        load_arviz_data("bad_checksum")


def test_missing_dataset():
    with pytest.raises(ValueError):
        load_arviz_data("does not exist")


def test_list_datasets():
    dataset_string = list_datasets()
    # make sure all the names of the data sets are in the dataset description
    for key in (
        "centered_eight",
        "non_centered_eight",
        "test_remote",
        "bad_checksum",
        "test_unknown",
    ):
        assert key in dataset_string


def test_dims_coords():
    shape = 4, 20, 5
    var_name = "x"
    dims, coords = generate_dims_coords(shape, var_name)
    assert "x_dim_0" in dims
    assert "x_dim_1" in dims
    assert "x_dim_2" in dims
    assert len(coords["x_dim_0"]) == 4
    assert len(coords["x_dim_1"]) == 20
    assert len(coords["x_dim_2"]) == 5


def test_dims_coords_extra_dims():
    shape = 4, 20
    var_name = "x"
    with pytest.warns(UserWarning):
        dims, coords = generate_dims_coords(shape, var_name, dims=["xx", "xy", "xz"])
    assert "xx" in dims
    assert "xy" in dims
    assert "xz" in dims
    assert len(coords["xx"]) == 4
    assert len(coords["xy"]) == 20


@pytest.mark.parametrize("shape", [(4, 20), (4, 20, 1)])
def test_dims_coords_skip_event_dims(shape):
    coords = {"x": np.arange(4), "y": np.arange(20), "z": np.arange(5)}
    dims, coords = generate_dims_coords(
        shape, "name", dims=["x", "y", "z"], coords=coords, skip_event_dims=True
    )
    assert "x" in dims
    assert "y" in dims
    assert "z" not in dims
    assert len(coords["x"]) == 4
    assert len(coords["y"]) == 20
    assert "z" not in coords


def test_make_attrs():
    extra_attrs = {"key": "Value"}
    attrs = make_attrs(attrs=extra_attrs)
    assert "key" in attrs
    assert attrs["key"] == "Value"


@pytest.mark.parametrize("copy", [True, False])
@pytest.mark.parametrize("inplace", [True, False])
@pytest.mark.parametrize("sequence", [True, False])
def test_concat_group(copy, inplace, sequence):
    idata1 = from_dict(
        posterior={"A": np.random.randn(2, 10, 2), "B": np.random.randn(2, 10, 5, 2)}
    )
    if copy and inplace:
        original_idata1_posterior_id = id(idata1.posterior)
    idata2 = from_dict(prior={"C": np.random.randn(2, 10, 2), "D": np.random.randn(2, 10, 5, 2)})
    idata3 = from_dict(observed_data={"E": np.random.randn(100), "F": np.random.randn(2, 100)})
    # basic case
    assert concat(idata1, idata2, copy=True, inplace=False) is not None
    if sequence:
        new_idata = concat((idata1, idata2, idata3), copy=copy, inplace=inplace)
    else:
        new_idata = concat(idata1, idata2, idata3, copy=copy, inplace=inplace)
    if inplace:
        assert new_idata is None
        new_idata = idata1
    assert new_idata is not None
    test_dict = {"posterior": ["A", "B"], "prior": ["C", "D"], "observed_data": ["E", "F"]}
    fails = check_multiple_attrs(test_dict, new_idata)
    assert not fails
    if copy:
        if inplace:
            assert id(new_idata.posterior) == original_idata1_posterior_id
        else:
            assert id(new_idata.posterior) != id(idata1.posterior)
        assert id(new_idata.prior) != id(idata2.prior)
        assert id(new_idata.observed_data) != id(idata3.observed_data)
    else:
        assert id(new_idata.posterior) == id(idata1.posterior)
        assert id(new_idata.prior) == id(idata2.prior)
        assert id(new_idata.observed_data) == id(idata3.observed_data)


@pytest.mark.parametrize("dim", ["chain", "draw"])
@pytest.mark.parametrize("copy", [True, False])
@pytest.mark.parametrize("inplace", [True, False])
@pytest.mark.parametrize("sequence", [True, False])
@pytest.mark.parametrize("reset_dim", [True, False])
def test_concat_dim(dim, copy, inplace, sequence, reset_dim):
    idata1 = from_dict(
        posterior={"A": np.random.randn(2, 10, 2), "B": np.random.randn(2, 10, 5, 2)},
        observed_data={"C": np.random.randn(100), "D": np.random.randn(2, 100)},
    )
    if inplace:
        original_idata1_id = id(idata1)
    idata2 = from_dict(
        posterior={"A": np.random.randn(2, 10, 2), "B": np.random.randn(2, 10, 5, 2)},
        observed_data={"C": np.random.randn(100), "D": np.random.randn(2, 100)},
    )
    idata3 = from_dict(
        posterior={"A": np.random.randn(2, 10, 2), "B": np.random.randn(2, 10, 5, 2)},
        observed_data={"C": np.random.randn(100), "D": np.random.randn(2, 100)},
    )
    # basic case
    assert (
        concat(idata1, idata2, dim=dim, copy=copy, inplace=False, reset_dim=reset_dim) is not None
    )
    if sequence:
        new_idata = concat(
            (idata1, idata2, idata3), copy=copy, dim=dim, inplace=inplace, reset_dim=reset_dim
        )
    else:
        new_idata = concat(
            idata1, idata2, idata3, dim=dim, copy=copy, inplace=inplace, reset_dim=reset_dim
        )
    if inplace:
        assert new_idata is None
        new_idata = idata1
    assert new_idata is not None
    test_dict = {"posterior": ["A", "B"], "observed_data": ["C", "D"]}
    fails = check_multiple_attrs(test_dict, new_idata)
    assert not fails
    if inplace:
        assert id(new_idata) == original_idata1_id
    else:
        assert id(new_idata) != id(idata1)
    assert getattr(new_idata.posterior, dim).size == 6 if dim == "chain" else 30
    if reset_dim:
        assert np.all(
            getattr(new_idata.posterior, dim).values
            == (np.arange(6) if dim == "chain" else np.arange(30))
        )


@pytest.mark.parametrize("copy", [True, False])
@pytest.mark.parametrize("inplace", [True, False])
@pytest.mark.parametrize("sequence", [True, False])
def test_concat_edgecases(copy, inplace, sequence):
    idata = from_dict(posterior={"A": np.random.randn(2, 10, 2), "B": np.random.randn(2, 10, 5, 2)})
    empty = concat()
    assert empty is not None
    if sequence:
        new_idata = concat([idata], copy=copy, inplace=inplace)
    else:
        new_idata = concat(idata, copy=copy, inplace=inplace)
    if inplace:
        assert new_idata is None
        new_idata = idata
    else:
        assert new_idata is not None
    test_dict = {"posterior": ["A", "B"]}
    fails = check_multiple_attrs(test_dict, new_idata)
    assert not fails
    if copy and not inplace:
        assert id(new_idata.posterior) != id(idata.posterior)
    else:
        assert id(new_idata.posterior) == id(idata.posterior)


def test_concat_bad():
    with pytest.raises(TypeError):
        concat("hello", "hello")
    idata = from_dict(posterior={"A": np.random.randn(2, 10, 2), "B": np.random.randn(2, 10, 5, 2)})
    idata2 = from_dict(posterior={"A": np.random.randn(2, 10, 2)})
    idata3 = from_dict(prior={"A": np.random.randn(2, 10, 2)})
    with pytest.raises(TypeError):
        concat(idata, np.array([1, 2, 3, 4, 5]))
    with pytest.raises(TypeError):
        concat(idata, idata, dim=None)
    with pytest.raises(TypeError):
        concat(idata, idata2, dim="chain")
    with pytest.raises(TypeError):
        concat(idata2, idata, dim="chain")
    with pytest.raises(TypeError):
        concat(idata, idata3, dim="chain")
    with pytest.raises(TypeError):
        concat(idata3, idata, dim="chain")


def test_inference_concat_keeps_all_fields():
    """From failures observed in issue #907"""
    idata1 = from_dict(posterior={"A": [1, 2, 3, 4]}, sample_stats={"B": [2, 3, 4, 5]})
    idata2 = from_dict(prior={"C": [1, 2, 3, 4]}, observed_data={"D": [2, 3, 4, 5]})

    idata_c1 = concat(idata1, idata2)
    idata_c2 = concat(idata2, idata1)

    test_dict = {"posterior": ["A"], "sample_stats": ["B"], "prior": ["C"], "observed_data": ["D"]}

    fails_c1 = check_multiple_attrs(test_dict, idata_c1)
    assert not fails_c1
    fails_c2 = check_multiple_attrs(test_dict, idata_c2)
    assert not fails_c2


@pytest.mark.parametrize(
    "model_code,expected",
    [
        ("data {int y;} models {y ~ poisson(3);} generated quantities {int X;}", {"X": "int"}),
        (
            "data {real y;} models {y ~ normal(0,1);} generated quantities {int Y; real G;}",
            {"Y": "int"},
        ),
    ],
)
def test_infer_stan_dtypes(model_code, expected):
    """Test different examples for dtypes in Stan models."""
    res = infer_stan_dtypes(model_code)
    assert res == expected


class TestInferenceData:  # pylint: disable=too-many-public-methods
    def test_addition(self):
        idata1 = from_dict(
            posterior={"A": np.random.randn(2, 10, 2), "B": np.random.randn(2, 10, 5, 2)}
        )
        idata2 = from_dict(
            prior={"C": np.random.randn(2, 10, 2), "D": np.random.randn(2, 10, 5, 2)}
        )
        new_idata = idata1 + idata2
        assert new_idata is not None
        test_dict = {"posterior": ["A", "B"], "prior": ["C", "D"]}
        fails = check_multiple_attrs(test_dict, new_idata)
        assert not fails

    def test_iter(self, models):
        idata = models.model_1
        for group in idata:
            assert group in idata._groups_all  # pylint: disable=protected-access

    def test_groups(self, models):
        idata = models.model_1
        for group in idata.groups():
            assert group in idata._groups_all  # pylint: disable=protected-access

    def test_values(self, models):
        idata = models.model_1
        datasets = idata.values()
        for group in idata.groups():
            assert group in idata._groups_all  # pylint: disable=protected-access
            dataset = getattr(idata, group)
            assert dataset in datasets

    def test_items(self, models):
        idata = models.model_1
        for group, dataset in idata.items():
            assert group in idata._groups_all  # pylint: disable=protected-access
            assert dataset.equals(getattr(idata, group))

    @pytest.mark.parametrize("inplace", [True, False])
    def test_extend_xr_method(self, data_random, inplace):
        idata = data_random
        idata_copy = deepcopy(idata)
        kwargs = {"groups": "posterior_groups"}
        if inplace:
            idata_copy.sum(dim="draw", inplace=inplace, **kwargs)
        else:
            idata2 = idata_copy.sum(dim="draw", inplace=inplace, **kwargs)
            assert idata2 is not idata_copy
            idata_copy = idata2
        assert_identical(idata_copy.posterior, idata.posterior.sum(dim="draw"))
        assert_identical(
            idata_copy.posterior_predictive, idata.posterior_predictive.sum(dim="draw")
        )
        assert_identical(idata_copy.observed_data, idata.observed_data)

    @pytest.mark.parametrize("inplace", [False, True])
    def test_sel(self, data_random, inplace):
        idata = data_random
        original_groups = getattr(idata, "_groups")
        ndraws = idata.posterior.draw.values.size
        kwargs = {"draw": slice(200, None), "chain": slice(None, None, 2), "b_dim_0": [1, 2, 7]}
        if inplace:
            idata.sel(inplace=inplace, **kwargs)
        else:
            idata2 = idata.sel(inplace=inplace, **kwargs)
            assert idata2 is not idata
            idata = idata2
        groups = getattr(idata, "_groups")
        assert np.all(np.isin(groups, original_groups))
        for group in groups:
            dataset = getattr(idata, group)
            assert "b_dim_0" in dataset.dims
            assert np.all(dataset.b_dim_0.values == np.array(kwargs["b_dim_0"]))
            if group != "observed_data":
                assert np.all(np.isin(["chain", "draw"], dataset.dims))
                assert np.all(dataset.chain.values == np.arange(0, 4, 2))
                assert np.all(dataset.draw.values == np.arange(200, ndraws))

    def test_sel_chain_prior(self):
        idata = load_arviz_data("centered_eight")
        original_groups = getattr(idata, "_groups")
        idata_subset = idata.sel(inplace=False, chain_prior=False, chain=[0, 1, 3])
        groups = getattr(idata_subset, "_groups")
        assert np.all(np.isin(groups, original_groups))
        for group in groups:
            dataset_subset = getattr(idata_subset, group)
            dataset = getattr(idata, group)
            if "chain" in dataset.dims:
                assert "chain" in dataset_subset.dims
                if "prior" not in group:
                    assert np.all(dataset_subset.chain.values == np.array([0, 1, 3]))
            else:
                assert "chain" not in dataset_subset.dims
        with pytest.raises(KeyError):
            idata.sel(inplace=False, chain_prior=True, chain=[0, 1, 3])

    @pytest.mark.parametrize("use", ("del", "delattr"))
    def test_del(self, use):
        # create inference data object
        data = np.random.normal(size=(4, 500, 8))
        idata = from_dict(
            posterior={"a": data[..., 0], "b": data},
            sample_stats={"a": data[..., 0], "b": data},
            observed_data={"b": data[0, 0, :]},
            posterior_predictive={"a": data[..., 0], "b": data},
        )

        # assert inference data object has all attributes
        test_dict = {
            "posterior": ("a", "b"),
            "sample_stats": ("a", "b"),
            "observed_data": ["b"],
            "posterior_predictive": ("a", "b"),
        }
        fails = check_multiple_attrs(test_dict, idata)
        assert not fails
        # assert _groups attribute contains all groups
        groups = getattr(idata, "_groups")
        assert all((group in groups for group in test_dict))

        # Use del method
        if use == "del":
            del idata.sample_stats
        else:
            delattr(idata, "sample_stats")

        # assert attribute has been removed
        test_dict.pop("sample_stats")
        fails = check_multiple_attrs(test_dict, idata)
        assert not fails
        assert not hasattr(idata, "sample_stats")
        # assert _groups attribute has been updated
        assert "sample_stats" not in getattr(idata, "_groups")

    @pytest.mark.parametrize(
        "args_res",
        (
            ([("posterior", "sample_stats")], ("posterior", "sample_stats")),
            (["posterior", "like"], ("posterior", "warmup_posterior", "posterior_predictive")),
            (["^posterior", "regex"], ("posterior", "posterior_predictive")),
            (
                [("~^warmup", "~^obs"), "regex"],
                ("posterior", "sample_stats", "posterior_predictive"),
            ),
            (
                ["~observed_vars"],
                ("posterior", "sample_stats", "warmup_posterior", "warmup_sample_stats"),
            ),
        ),
    )
    def test_group_names(self, args_res):
        args, result = args_res
        ds = dict_to_dataset({"a": np.random.normal(size=(3, 10))})
        idata = InferenceData(
            posterior=(ds, ds),
            sample_stats=(ds, ds),
            observed_data=ds,
            posterior_predictive=ds,
        )
        group_names = idata._group_names(*args)  # pylint: disable=protected-access
        assert np.all([name in result for name in group_names])

    def test_group_names_invalid_args(self):
        ds = dict_to_dataset({"a": np.random.normal(size=(3, 10))})
        idata = InferenceData(posterior=(ds, ds))
        msg = r"^\'filter_groups\' can only be None, \'like\', or \'regex\', got: 'foo'$"
        with pytest.raises(ValueError, match=msg):
            idata._group_names(  # pylint: disable=protected-access
                ("posterior",), filter_groups="foo"
            )

    @pytest.mark.parametrize("inplace", [False, True])
    def test_isel(self, data_random, inplace):
        idata = data_random
        original_groups = getattr(idata, "_groups")
        ndraws = idata.posterior.draw.values.size
        kwargs = {"draw": slice(200, None), "chain": slice(None, None, 2), "b_dim_0": [1, 2, 7]}
        if inplace:
            idata.isel(inplace=inplace, **kwargs)
        else:
            idata2 = idata.isel(inplace=inplace, **kwargs)
            assert idata2 is not idata
            idata = idata2
        groups = getattr(idata, "_groups")
        assert np.all(np.isin(groups, original_groups))
        for group in groups:
            dataset = getattr(idata, group)
            assert "b_dim_0" in dataset.dims
            assert np.all(dataset.b_dim_0.values == np.array(kwargs["b_dim_0"]))
            if group != "observed_data":
                assert np.all(np.isin(["chain", "draw"], dataset.dims))
                assert np.all(dataset.chain.values == np.arange(0, 4, 2))
                assert np.all(dataset.draw.values == np.arange(200, ndraws))

    def test_rename(self, data_random):
        idata = data_random
        original_groups = getattr(idata, "_groups")
        renamed_idata = idata.rename({"b": "b_new"})
        for group in original_groups:
            xr_data = getattr(renamed_idata, group)
            assert "b_new" in list(xr_data.data_vars)
            assert "b" not in list(xr_data.data_vars)

        renamed_idata = idata.rename({"b_dim_0": "b_new"})
        for group in original_groups:
            xr_data = getattr(renamed_idata, group)
            assert "b_new" in list(xr_data.dims)
            assert "b_dim_0" not in list(xr_data.dims)

    def test_rename_vars(self, data_random):
        idata = data_random
        original_groups = getattr(idata, "_groups")
        renamed_idata = idata.rename_vars({"b": "b_new"})
        for group in original_groups:
            xr_data = getattr(renamed_idata, group)
            assert "b_new" in list(xr_data.data_vars)
            assert "b" not in list(xr_data.data_vars)

        renamed_idata = idata.rename_vars({"b_dim_0": "b_new"})
        for group in original_groups:
            xr_data = getattr(renamed_idata, group)
            assert "b_new" not in list(xr_data.dims)
            assert "b_dim_0" in list(xr_data.dims)

    def test_rename_dims(self, data_random):
        idata = data_random
        original_groups = getattr(idata, "_groups")
        renamed_idata = idata.rename_dims({"b_dim_0": "b_new"})
        for group in original_groups:
            xr_data = getattr(renamed_idata, group)
            assert "b_new" in list(xr_data.dims)
            assert "b_dim_0" not in list(xr_data.dims)

        renamed_idata = idata.rename_dims({"b": "b_new"})
        for group in original_groups:
            xr_data = getattr(renamed_idata, group)
            assert "b_new" not in list(xr_data.data_vars)
            assert "b" in list(xr_data.data_vars)

    def test_stack_unstack(self):
        datadict = {
            "a": np.random.randn(100),
            "b": np.random.randn(1, 100, 10),
            "c": np.random.randn(1, 100, 3, 4),
        }
        coords = {
            "c1": np.arange(3),
            "c99": np.arange(4),
            "b1": np.arange(10),
        }
        dims = {"c": ["c1", "c99"], "b": ["b1"]}
        dataset = from_dict(posterior=datadict, coords=coords, dims=dims)
        assert_identical(
            dataset.stack(z=["c1", "c99"]).posterior, dataset.posterior.stack(z=["c1", "c99"])
        )
        assert_identical(dataset.stack(z=["c1", "c99"]).unstack().posterior, dataset.posterior)
        assert_identical(
            dataset.stack(z=["c1", "c99"]).unstack(dim="z").posterior, dataset.posterior
        )

    def test_to_dict(self, models):
        idata = models.model_1
        test_data = from_dict(**idata.to_dict())
        assert test_data
        for group in idata._groups_all:  # pylint: disable=protected-access
            xr_data = getattr(idata, group)
            test_xr_data = getattr(test_data, group)
            assert xr_data.equals(test_xr_data)

    def test_to_dict_warmup(self):
        idata = create_data_random(
            groups=[
                "posterior",
                "sample_stats",
                "observed_data",
                "warmup_posterior",
                "warmup_posterior_predictive",
            ]
        )
        test_data = from_dict(**idata.to_dict(), save_warmup=True)
        assert test_data
        for group in idata._groups_all:  # pylint: disable=protected-access
            xr_data = getattr(idata, group)
            test_xr_data = getattr(test_data, group)
            assert xr_data.equals(test_xr_data)

    @pytest.mark.parametrize(
        "kwargs",
        (
            {
                "groups": "posterior",
                "include_coords": True,
                "include_index": True,
                "index_origin": 0,
            },
            {
                "groups": ["posterior", "sample_stats"],
                "include_coords": False,
                "include_index": True,
                "index_origin": 0,
            },
            {
                "groups": "posterior_groups",
                "include_coords": True,
                "include_index": False,
                "index_origin": 1,
            },
        ),
    )
    def test_to_dataframe(self, kwargs):
        idata = from_dict(
            posterior={"a": np.random.randn(4, 100, 3, 4, 5), "b": np.random.randn(4, 100)},
            sample_stats={"a": np.random.randn(4, 100, 3, 4, 5), "b": np.random.randn(4, 100)},
            observed_data={"a": np.random.randn(3, 4, 5), "b": np.random.randn(4)},
        )
        test_data = idata.to_dataframe(**kwargs)
        assert not test_data.empty
        groups = kwargs.get("groups", idata._groups_all)  # pylint: disable=protected-access
        for group in idata._groups_all:  # pylint: disable=protected-access
            if "data" in group:
                continue
            assert test_data.shape == (
                (4 * 100, 3 * 4 * 5 + 1 + 2)
                if groups == "posterior"
                else (4 * 100, (3 * 4 * 5 + 1) * 2 + 2)
            )
            if groups == "posterior":
                if kwargs.get("include_coords", True) and kwargs.get("include_index", True):
                    assert any(
                        f"[{kwargs.get('index_origin', 0)}," in item[0]
                        for item in test_data.columns
                        if isinstance(item, tuple)
                    )
                if kwargs.get("include_coords", True):
                    assert any(isinstance(item, tuple) for item in test_data.columns)
                else:
                    assert not any(isinstance(item, tuple) for item in test_data.columns)
            else:
                if not kwargs.get("include_index", True):
                    assert all(
                        item in test_data.columns
                        for item in (("posterior", "a", 1, 1, 1), ("posterior", "b"))
                    )
            assert all(item in test_data.columns for item in ("chain", "draw"))

    def test_to_dataframe_bad(self):
        idata = from_dict(
            posterior={"a": np.random.randn(4, 100, 3, 4, 5), "b": np.random.randn(4, 100)},
            sample_stats={"a": np.random.randn(4, 100, 3, 4, 5), "b": np.random.randn(4, 100)},
            observed_data={"a": np.random.randn(3, 4, 5), "b": np.random.randn(4)},
        )
        with pytest.raises(TypeError):
            idata.to_dataframe(index_origin=2)

        with pytest.raises(TypeError):
            idata.to_dataframe(include_coords=False, include_index=False)

        with pytest.raises(TypeError):
            idata.to_dataframe(groups=["observed_data"])

        with pytest.raises(KeyError):
            idata.to_dataframe(groups=["invalid_group"])

    @pytest.mark.parametrize("use", (None, "args", "kwargs"))
    def test_map(self, use):
        idata = load_arviz_data("centered_eight")
        args = []
        kwargs = {}
        if use is None:
            fun = lambda x: x + 3
        elif use == "args":
            fun = lambda x, a: x + a
            args = [3]
        else:
            fun = lambda x, a: x + a
            kwargs = {"a": 3}
        groups = ("observed_data", "posterior_predictive")
        idata_map = idata.map(fun, groups, args=args, **kwargs)
        groups_map = idata_map._groups  # pylint: disable=protected-access
        assert groups_map == idata._groups  # pylint: disable=protected-access
        assert np.allclose(
            idata_map.observed_data.obs, fun(idata.observed_data.obs, *args, **kwargs)
        )
        assert np.allclose(
            idata_map.posterior_predictive.obs, fun(idata.posterior_predictive.obs, *args, **kwargs)
        )
        assert np.allclose(idata_map.posterior.mu, idata.posterior.mu)

    def test_repr_html(self):
        """Test if the function _repr_html is generating html."""
        idata = load_arviz_data("centered_eight")
        display_style = OPTIONS["display_style"]
        xr.set_options(display_style="html")
        html = idata._repr_html_()  # pylint: disable=protected-access

        assert html is not None
        assert "<div" in html
        for group in idata._groups:  # pylint: disable=protected-access
            assert group in html
            xr_data = getattr(idata, group)
            for item, _ in xr_data.items():
                assert item in html
        specific_style = ".xr-wrap{width:700px!important;}"
        assert specific_style in html

        xr.set_options(display_style="text")
        html = idata._repr_html_()  # pylint: disable=protected-access
        assert escape(idata.__repr__()) in html
        xr.set_options(display_style=display_style)

    def test_add_groups(self, data_random):
        data = np.random.normal(size=(4, 500, 8))
        idata = data_random
        idata.add_groups({"prior": {"a": data[..., 0], "b": data}})
        assert "prior" in idata._groups  # pylint: disable=protected-access
        assert isinstance(idata.prior, xr.Dataset)
        assert hasattr(idata, "prior")

        idata.add_groups(warmup_posterior={"a": data[..., 0], "b": data})
        assert "warmup_posterior" in idata._groups_all  # pylint: disable=protected-access
        assert isinstance(idata.warmup_posterior, xr.Dataset)
        assert hasattr(idata, "warmup_posterior")

    def test_add_groups_warning(self, data_random):
        data = np.random.normal(size=(4, 500, 8))
        idata = data_random
        with pytest.warns(UserWarning, match="The group.+not defined in the InferenceData scheme"):
            idata.add_groups({"new_group": idata.posterior})
        with pytest.warns(UserWarning, match="the default dims.+will be added automatically"):
            idata.add_groups(constant_data={"a": data[..., 0], "b": data})
        assert idata.new_group.equals(idata.posterior)

    def test_add_groups_error(self, data_random):
        idata = data_random
        with pytest.raises(ValueError, match="One of.+must be provided."):
            idata.add_groups()
        with pytest.raises(ValueError, match="Arguments.+xr.Dataset, xr.Dataarray or dicts"):
            idata.add_groups({"new_group": "new_group"})
        with pytest.raises(ValueError, match="group.+already exists"):
            idata.add_groups({"posterior": idata.posterior})

    def test_extend(self, data_random):
        idata = data_random
        idata2 = create_data_random(
            groups=["prior", "prior_predictive", "observed_data", "warmup_posterior"], seed=7
        )
        idata.extend(idata2)
        assert "prior" in idata._groups_all  # pylint: disable=protected-access
        assert "warmup_posterior" in idata._groups_all  # pylint: disable=protected-access
        assert hasattr(idata, "prior")
        assert hasattr(idata, "prior_predictive")
        assert idata.prior.equals(idata2.prior)
        assert not idata.observed_data.equals(idata2.observed_data)
        assert idata.prior_predictive.equals(idata2.prior_predictive)

        idata.extend(idata2, join="right")
        assert idata.prior.equals(idata2.prior)
        assert idata.observed_data.equals(idata2.observed_data)
        assert idata.prior_predictive.equals(idata2.prior_predictive)

    def test_extend_errors_warnings(self, data_random):
        idata = data_random
        idata2 = create_data_random(groups=["prior", "prior_predictive", "observed_data"], seed=7)
        with pytest.raises(ValueError, match="Extending.+InferenceData objects only."):
            idata.extend("something")
        with pytest.raises(ValueError, match="join must be either"):
            idata.extend(idata2, join="outer")
        idata2.add_groups(new_group=idata2.prior)
        with pytest.warns(UserWarning):
            idata.extend(idata2)


class TestNumpyToDataArray:
    def test_1d_dataset(self):
        size = 100
        dataset = convert_to_dataset(np.random.randn(size))
        assert len(dataset.data_vars) == 1

        assert set(dataset.coords) == {"chain", "draw"}
        assert dataset.chain.shape == (1,)
        assert dataset.draw.shape == (size,)

    def test_warns_bad_shape(self):
        # Shape should be (chain, draw, *shape)
        with pytest.warns(UserWarning):
            convert_to_dataset(np.random.randn(100, 4))

    def test_nd_to_dataset(self):
        shape = (1, 2, 3, 4, 5)
        dataset = convert_to_dataset(np.random.randn(*shape))
        assert len(dataset.data_vars) == 1
        var_name = list(dataset.data_vars)[0]

        assert len(dataset.coords) == len(shape)
        assert dataset.chain.shape == shape[:1]
        assert dataset.draw.shape == shape[1:2]
        assert dataset[var_name].shape == shape

    def test_nd_to_inference_data(self):
        shape = (1, 2, 3, 4, 5)
        inference_data = convert_to_inference_data(np.random.randn(*shape), group="prior")
        assert hasattr(inference_data, "prior")
        assert len(inference_data.prior.data_vars) == 1
        var_name = list(inference_data.prior.data_vars)[0]

        assert len(inference_data.prior.coords) == len(shape)
        assert inference_data.prior.chain.shape == shape[:1]
        assert inference_data.prior.draw.shape == shape[1:2]
        assert inference_data.prior[var_name].shape == shape
        assert repr(inference_data).startswith("Inference data with groups")

    def test_more_chains_than_draws(self):
        shape = (10, 4)
        with pytest.warns(UserWarning):
            inference_data = convert_to_inference_data(np.random.randn(*shape), group="prior")
        assert hasattr(inference_data, "prior")
        assert len(inference_data.prior.data_vars) == 1
        var_name = list(inference_data.prior.data_vars)[0]

        assert len(inference_data.prior.coords) == len(shape)
        assert inference_data.prior.chain.shape == shape[:1]
        assert inference_data.prior.draw.shape == shape[1:2]
        assert inference_data.prior[var_name].shape == shape


class TestConvertToDataset:
    @pytest.fixture(scope="class")
    def data(self):
        # pylint: disable=attribute-defined-outside-init
        class Data:
            datadict = {
                "a": np.random.randn(100),
                "b": np.random.randn(1, 100, 10),
                "c": np.random.randn(1, 100, 3, 4),
            }
            coords = {"c1": np.arange(3), "c2": np.arange(4), "b1": np.arange(10)}
            dims = {"b": ["b1"], "c": ["c1", "c2"]}

        return Data

    def test_use_all(self, data):
        dataset = convert_to_dataset(data.datadict, coords=data.coords, dims=data.dims)
        assert set(dataset.data_vars) == {"a", "b", "c"}
        assert set(dataset.coords) == {"chain", "draw", "c1", "c2", "b1"}

        assert set(dataset.a.coords) == {"chain", "draw"}
        assert set(dataset.b.coords) == {"chain", "draw", "b1"}
        assert set(dataset.c.coords) == {"chain", "draw", "c1", "c2"}

    def test_missing_coords(self, data):
        dataset = convert_to_dataset(data.datadict, coords=None, dims=data.dims)
        assert set(dataset.data_vars) == {"a", "b", "c"}
        assert set(dataset.coords) == {"chain", "draw", "c1", "c2", "b1"}

        assert set(dataset.a.coords) == {"chain", "draw"}
        assert set(dataset.b.coords) == {"chain", "draw", "b1"}
        assert set(dataset.c.coords) == {"chain", "draw", "c1", "c2"}

    def test_missing_dims(self, data):
        # missing dims
        coords = {"c_dim_0": np.arange(3), "c_dim_1": np.arange(4), "b_dim_0": np.arange(10)}
        dataset = convert_to_dataset(data.datadict, coords=coords, dims=None)
        assert set(dataset.data_vars) == {"a", "b", "c"}
        assert set(dataset.coords) == {"chain", "draw", "c_dim_0", "c_dim_1", "b_dim_0"}

        assert set(dataset.a.coords) == {"chain", "draw"}
        assert set(dataset.b.coords) == {"chain", "draw", "b_dim_0"}
        assert set(dataset.c.coords) == {"chain", "draw", "c_dim_0", "c_dim_1"}

    def test_skip_dim_0(self, data):
        dims = {"c": [None, "c2"]}
        coords = {"c_dim_0": np.arange(3), "c2": np.arange(4), "b_dim_0": np.arange(10)}
        dataset = convert_to_dataset(data.datadict, coords=coords, dims=dims)
        assert set(dataset.data_vars) == {"a", "b", "c"}
        assert set(dataset.coords) == {"chain", "draw", "c_dim_0", "c2", "b_dim_0"}

        assert set(dataset.a.coords) == {"chain", "draw"}
        assert set(dataset.b.coords) == {"chain", "draw", "b_dim_0"}
        assert set(dataset.c.coords) == {"chain", "draw", "c_dim_0", "c2"}


def test_dict_to_dataset():
    datadict = {"a": np.random.randn(100), "b": np.random.randn(1, 100, 10)}
    dataset = convert_to_dataset(datadict, coords={"c": np.arange(10)}, dims={"b": ["c"]})
    assert set(dataset.data_vars) == {"a", "b"}
    assert set(dataset.coords) == {"chain", "draw", "c"}

    assert set(dataset.a.coords) == {"chain", "draw"}
    assert set(dataset.b.coords) == {"chain", "draw", "c"}


def test_dict_to_dataset_event_dims_error():
    datadict = {"a": np.random.randn(1, 100, 10)}
    coords = {"b": np.arange(10), "c": ["x", "y", "z"]}
    msg = "different number of dimensions on data and dims"
    with pytest.raises(ValueError, match=msg):
        convert_to_dataset(datadict, coords=coords, dims={"a": ["b", "c"]})


def test_convert_to_dataset_idempotent():
    first = convert_to_dataset(np.random.randn(100))
    second = convert_to_dataset(first)
    assert first.equals(second)


def test_convert_to_inference_data_idempotent():
    first = convert_to_inference_data(np.random.randn(100), group="prior")
    second = convert_to_inference_data(first)
    assert first.prior is second.prior


def test_convert_to_inference_data_from_file(tmpdir):
    first = convert_to_inference_data(np.random.randn(100), group="prior")
    filename = str(tmpdir.join("test_file.nc"))
    first.to_netcdf(filename)
    second = convert_to_inference_data(filename)
    assert first.prior.equals(second.prior)


def test_convert_to_inference_data_bad():
    with pytest.raises(ValueError):
        convert_to_inference_data(1)


def test_convert_to_dataset_bad(tmpdir):
    first = convert_to_inference_data(np.random.randn(100), group="prior")
    filename = str(tmpdir.join("test_file.nc"))
    first.to_netcdf(filename)
    with pytest.raises(ValueError):
        convert_to_dataset(filename, group="bar")


def test_bad_inference_data():
    with pytest.raises(ValueError):
        InferenceData(posterior=[1, 2, 3])


def test_inference_data_other_groups():
    datadict = {"a": np.random.randn(100), "b": np.random.randn(1, 100, 10)}
    dataset = convert_to_dataset(datadict, coords={"c": np.arange(10)}, dims={"b": ["c"]})
    with pytest.warns(UserWarning, match="not.+in.+InferenceData scheme"):
        idata = InferenceData(other_group=dataset)
    fails = check_multiple_attrs({"other_group": ["a", "b"]}, idata)
    assert not fails


class TestDataConvert:
    @pytest.fixture(scope="class")
    def data(self, draws, chains):
        class Data:
            # fake 8-school output
            obj = {}
            for key, shape in {"mu": [], "tau": [], "eta": [8], "theta": [8]}.items():
                obj[key] = np.random.randn(chains, draws, *shape)

        return Data

    def get_inference_data(self, data):
        return convert_to_inference_data(
            data.obj,
            group="posterior",
            coords={"school": np.arange(8)},
            dims={"theta": ["school"], "eta": ["school"]},
        )

    def check_var_names_coords_dims(self, dataset):
        assert set(dataset.data_vars) == {"mu", "tau", "eta", "theta"}
        assert set(dataset.coords) == {"chain", "draw", "school"}

    def test_convert_to_inference_data(self, data):
        inference_data = self.get_inference_data(data)
        assert hasattr(inference_data, "posterior")
        self.check_var_names_coords_dims(inference_data.posterior)

    def test_convert_to_dataset(self, draws, chains, data):
        dataset = convert_to_dataset(
            data.obj,
            group="posterior",
            coords={"school": np.arange(8)},
            dims={"theta": ["school"], "eta": ["school"]},
        )
        assert dataset.draw.shape == (draws,)
        assert dataset.chain.shape == (chains,)
        assert dataset.school.shape == (8,)
        assert dataset.theta.shape == (chains, draws, 8)


class TestDataDict:
    @pytest.fixture(scope="class")
    def data(self, draws, chains):
        class Data:
            # fake 8-school output
            obj = {}
            for key, shape in {"mu": [], "tau": [], "eta": [8], "theta": [8]}.items():
                obj[key] = np.random.randn(chains, draws, *shape)

        return Data

    def check_var_names_coords_dims(self, dataset):
        assert set(dataset.data_vars) == {"mu", "tau", "eta", "theta"}
        assert set(dataset.coords) == {"chain", "draw", "school"}

    def get_inference_data(self, data, eight_schools_params, save_warmup=False):
        return from_dict(
            posterior=data.obj,
            posterior_predictive=data.obj,
            sample_stats=data.obj,
            prior=data.obj,
            prior_predictive=data.obj,
            sample_stats_prior=data.obj,
            warmup_posterior=data.obj,
            warmup_posterior_predictive=data.obj,
            predictions=data.obj,
            observed_data=eight_schools_params,
            coords={
                "school": np.arange(8),
            },
            pred_coords={
                "school_pred": np.arange(8),
            },
            dims={"theta": ["school"], "eta": ["school"]},
            pred_dims={"theta": ["school_pred"], "eta": ["school_pred"]},
            save_warmup=save_warmup,
        )

    def test_inference_data(self, data, eight_schools_params):
        inference_data = self.get_inference_data(data, eight_schools_params)
        test_dict = {
            "posterior": [],
            "prior": [],
            "sample_stats": [],
            "posterior_predictive": [],
            "prior_predictive": [],
            "sample_stats_prior": [],
            "observed_data": ["J", "y", "sigma"],
        }
        fails = check_multiple_attrs(test_dict, inference_data)
        assert not fails
        self.check_var_names_coords_dims(inference_data.posterior)
        self.check_var_names_coords_dims(inference_data.posterior_predictive)
        self.check_var_names_coords_dims(inference_data.sample_stats)
        self.check_var_names_coords_dims(inference_data.prior)
        self.check_var_names_coords_dims(inference_data.prior_predictive)
        self.check_var_names_coords_dims(inference_data.sample_stats_prior)

        pred_dims = inference_data.predictions.dims["school_pred"]
        assert pred_dims == 8

    def test_inference_data_warmup(self, data, eight_schools_params):
        inference_data = self.get_inference_data(data, eight_schools_params, save_warmup=True)
        test_dict = {
            "posterior": [],
            "prior": [],
            "sample_stats": [],
            "posterior_predictive": [],
            "prior_predictive": [],
            "sample_stats_prior": [],
            "observed_data": ["J", "y", "sigma"],
            "warmup_posterior_predictive": [],
            "warmup_posterior": [],
        }
        fails = check_multiple_attrs(test_dict, inference_data)
        assert not fails
        self.check_var_names_coords_dims(inference_data.posterior)
        self.check_var_names_coords_dims(inference_data.posterior_predictive)
        self.check_var_names_coords_dims(inference_data.sample_stats)
        self.check_var_names_coords_dims(inference_data.prior)
        self.check_var_names_coords_dims(inference_data.prior_predictive)
        self.check_var_names_coords_dims(inference_data.sample_stats_prior)
        self.check_var_names_coords_dims(inference_data.warmup_posterior)
        self.check_var_names_coords_dims(inference_data.warmup_posterior_predictive)

    def test_inference_data_edge_cases(self):
        # create data
        log_likelihood = {
            "y": np.random.randn(4, 100),
            "log_likelihood": np.random.randn(4, 100, 8),
        }

        # log_likelihood to posterior
        with pytest.warns(UserWarning, match="log_likelihood.+in posterior"):
            assert from_dict(posterior=log_likelihood) is not None

        # dims == None
        assert from_dict(observed_data=log_likelihood, dims=None) is not None

    def test_inference_data_bad(self):
        # create data
        x = np.random.randn(4, 100)

        # input ndarray
        with pytest.raises(TypeError):
            from_dict(posterior=x)
        with pytest.raises(TypeError):
            from_dict(posterior_predictive=x)
        with pytest.raises(TypeError):
            from_dict(sample_stats=x)
        with pytest.raises(TypeError):
            from_dict(prior=x)
        with pytest.raises(TypeError):
            from_dict(prior_predictive=x)
        with pytest.raises(TypeError):
            from_dict(sample_stats_prior=x)
        with pytest.raises(TypeError):
            from_dict(observed_data=x)

    def test_from_dict_warning(self):
        bad_posterior_dict = {"log_likelihood": np.ones((5, 1000, 2))}
        with pytest.warns(UserWarning):
            from_dict(posterior=bad_posterior_dict)


class TestDataNetCDF:
    @pytest.fixture(scope="class")
    def data(self, draws, chains):
        class Data:
            # fake 8-school output
            obj = {}
            for key, shape in {"mu": [], "tau": [], "eta": [8], "theta": [8]}.items():
                obj[key] = np.random.randn(chains, draws, *shape)

        return Data

    def get_inference_data(self, data, eight_schools_params):
        return from_dict(
            posterior=data.obj,
            posterior_predictive=data.obj,
            sample_stats=data.obj,
            prior=data.obj,
            prior_predictive=data.obj,
            sample_stats_prior=data.obj,
            observed_data=eight_schools_params,
            coords={"school": np.arange(8)},
            dims={"theta": ["school"], "eta": ["school"]},
        )

    def test_io_function(self, data, eight_schools_params):
        # create inference data and assert all attributes are present
        inference_data = self.get_inference_data(  # pylint: disable=W0612
            data, eight_schools_params
        )
        test_dict = {
            "posterior": ["eta", "theta", "mu", "tau"],
            "posterior_predictive": ["eta", "theta", "mu", "tau"],
            "sample_stats": ["eta", "theta", "mu", "tau"],
            "prior": ["eta", "theta", "mu", "tau"],
            "prior_predictive": ["eta", "theta", "mu", "tau"],
            "sample_stats_prior": ["eta", "theta", "mu", "tau"],
            "observed_data": ["J", "y", "sigma"],
        }
        fails = check_multiple_attrs(test_dict, inference_data)
        assert not fails

        # check filename does not exist and save InferenceData
        here = os.path.dirname(os.path.abspath(__file__))
        data_directory = os.path.join(here, "..", "saved_models")
        filepath = os.path.join(data_directory, "io_function_testfile.nc")
        # az -function
        to_netcdf(inference_data, filepath)

        # Assert InferenceData has been saved correctly
        assert os.path.exists(filepath)
        assert os.path.getsize(filepath) > 0
        inference_data2 = from_netcdf(filepath)
        fails = check_multiple_attrs(test_dict, inference_data2)
        assert not fails
        os.remove(filepath)
        assert not os.path.exists(filepath)

    @pytest.mark.parametrize("groups_arg", [False, True])
    def test_io_method(self, data, eight_schools_params, groups_arg):
        # create InferenceData and check it has been properly created
        inference_data = self.get_inference_data(  # pylint: disable=W0612
            data, eight_schools_params
        )
        test_dict = {
            "posterior": ["eta", "theta", "mu", "tau"],
            "posterior_predictive": ["eta", "theta", "mu", "tau"],
            "sample_stats": ["eta", "theta", "mu", "tau"],
            "prior": ["eta", "theta", "mu", "tau"],
            "prior_predictive": ["eta", "theta", "mu", "tau"],
            "sample_stats_prior": ["eta", "theta", "mu", "tau"],
            "observed_data": ["J", "y", "sigma"],
        }
        fails = check_multiple_attrs(test_dict, inference_data)
        assert not fails

        # check filename does not exist and use to_netcdf method
        here = os.path.dirname(os.path.abspath(__file__))
        data_directory = os.path.join(here, "..", "saved_models")
        filepath = os.path.join(data_directory, "io_method_testfile.nc")
        assert not os.path.exists(filepath)
        # InferenceData method
        inference_data.to_netcdf(
            filepath, groups=("posterior", "observed_data") if groups_arg else None
        )

        # assert file has been saved correctly
        assert os.path.exists(filepath)
        assert os.path.getsize(filepath) > 0
        inference_data2 = InferenceData.from_netcdf(filepath)
        if groups_arg:  # if groups arg, update test dict to contain only saved groups
            test_dict = {
                "posterior": ["eta", "theta", "mu", "tau"],
                "observed_data": ["J", "y", "sigma"],
            }
            assert not hasattr(inference_data2, "sample_stats")
        fails = check_multiple_attrs(test_dict, inference_data2)
        assert not fails

        os.remove(filepath)
        assert not os.path.exists(filepath)

    def test_empty_inference_data_object(self):
        inference_data = InferenceData()
        here = os.path.dirname(os.path.abspath(__file__))
        data_directory = os.path.join(here, "..", "saved_models")
        filepath = os.path.join(data_directory, "empty_test_file.nc")
        assert not os.path.exists(filepath)
        inference_data.to_netcdf(filepath)
        assert os.path.exists(filepath)
        os.remove(filepath)
        assert not os.path.exists(filepath)


class TestJSON:
    def test_json_converters(self, models):
        idata = models.model_1

        filepath = os.path.realpath("test.json")
        idata.to_json(filepath)

        idata_copy = from_json(filepath)
        for group in idata._groups_all:  # pylint: disable=protected-access
            xr_data = getattr(idata, group)
            test_xr_data = getattr(idata_copy, group)
            assert xr_data.equals(test_xr_data)

        os.remove(filepath)
        assert not os.path.exists(filepath)


class TestConversions:
    def test_id_conversion_idempotent(self):
        stored = load_arviz_data("centered_eight")
        inference_data = convert_to_inference_data(stored)
        assert isinstance(inference_data, InferenceData)
        assert set(inference_data.observed_data.obs.coords["school"].values) == {
            "Hotchkiss",
            "Mt. Hermon",
            "Choate",
            "Deerfield",
            "Phillips Andover",
            "St. Paul's",
            "Lawrenceville",
            "Phillips Exeter",
        }
        assert inference_data.posterior["theta"].dims == ("chain", "draw", "school")

    def test_dataset_conversion_idempotent(self):
        inference_data = load_arviz_data("centered_eight")
        data_set = convert_to_dataset(inference_data.posterior)
        assert isinstance(data_set, xr.Dataset)
        assert set(data_set.coords["school"].values) == {
            "Hotchkiss",
            "Mt. Hermon",
            "Choate",
            "Deerfield",
            "Phillips Andover",
            "St. Paul's",
            "Lawrenceville",
            "Phillips Exeter",
        }
        assert data_set["theta"].dims == ("chain", "draw", "school")

    def test_id_conversion_args(self):
        stored = load_arviz_data("centered_eight")
        IVIES = ["Yale", "Harvard", "MIT", "Princeton", "Cornell", "Dartmouth", "Columbia", "Brown"]
        # test dictionary argument...
        # I reverse engineered a dictionary out of the centered_eight
        # data. That's what this block of code does.
        d = stored.posterior.to_dict()
        d = d["data_vars"]
        test_dict = {}  # type: Dict[str, np.ndarray]
        for var_name in d:
            data = d[var_name]["data"]
            # this is a list of chains that is a list of samples...
            chain_arrs = []
            for chain in data:  # list of samples
                chain_arrs.append(np.array(chain))
            data_arr = np.stack(chain_arrs)
            test_dict[var_name] = data_arr

        inference_data = convert_to_inference_data(
            test_dict, dims={"theta": ["Ivies"]}, coords={"Ivies": IVIES}
        )

        assert isinstance(inference_data, InferenceData)
        assert set(inference_data.posterior.coords["Ivies"].values) == set(IVIES)
        assert inference_data.posterior["theta"].dims == ("chain", "draw", "Ivies")


class TestDataArrayToDataset:
    def test_1d_dataset(self):
        size = 100
        dataset = convert_to_dataset(
            xr.DataArray(np.random.randn(1, size), name="plot", dims=("chain", "draw"))
        )
        assert len(dataset.data_vars) == 1
        assert "plot" in dataset.data_vars
        assert dataset.chain.shape == (1,)
        assert dataset.draw.shape == (size,)

    def test_nd_to_dataset(self):
        shape = (1, 2, 3, 4, 5)
        dataset = convert_to_dataset(
            xr.DataArray(np.random.randn(*shape), dims=("chain", "draw", "dim_0", "dim_1", "dim_2"))
        )
        var_name = list(dataset.data_vars)[0]

        assert len(dataset.data_vars) == 1
        assert dataset.chain.shape == shape[:1]
        assert dataset.draw.shape == shape[1:2]
        assert dataset[var_name].shape == shape

    def test_nd_to_inference_data(self):
        shape = (1, 2, 3, 4, 5)
        inference_data = convert_to_inference_data(
            xr.DataArray(
                np.random.randn(*shape), dims=("chain", "draw", "dim_0", "dim_1", "dim_2")
            ),
            group="prior",
        )
        var_name = list(inference_data.prior.data_vars)[0]

        assert hasattr(inference_data, "prior")
        assert len(inference_data.prior.data_vars) == 1
        assert inference_data.prior.chain.shape == shape[:1]
        assert inference_data.prior.draw.shape == shape[1:2]
        assert inference_data.prior[var_name].shape == shape
