#  pylint: disable=too-many-instance-attributes,too-many-lines
"""PyStan-specific conversion code."""
import re
import warnings
from collections import OrderedDict
from copy import deepcopy

import numpy as np
import xarray as xr

from .. import _log
from ..rcparams import rcParams
from .base import dict_to_dataset, generate_dims_coords, infer_stan_dtypes, make_attrs, requires
from .inference_data import InferenceData

try:
    import ujson as json
except ImportError:
    # Can't find ujson using json
    # mypy struggles with conditional imports expressed as catching ImportError:
    # https://github.com/python/mypy/issues/1153
    import json  # type: ignore


class PyStanConverter:
    """Encapsulate PyStan specific logic."""

    def __init__(
        self,
        *,
        posterior=None,
        posterior_predictive=None,
        predictions=None,
        prior=None,
        prior_predictive=None,
        observed_data=None,
        constant_data=None,
        predictions_constant_data=None,
        log_likelihood=None,
        coords=None,
        dims=None,
        save_warmup=None,
        dtypes=None,
    ):
        self.posterior = posterior
        self.posterior_predictive = posterior_predictive
        self.predictions = predictions
        self.prior = prior
        self.prior_predictive = prior_predictive
        self.observed_data = observed_data
        self.constant_data = constant_data
        self.predictions_constant_data = predictions_constant_data
        self.log_likelihood = (
            rcParams["data.log_likelihood"] if log_likelihood is None else log_likelihood
        )
        self.coords = coords
        self.dims = dims
        self.save_warmup = rcParams["data.save_warmup"] if save_warmup is None else save_warmup
        self.dtypes = dtypes

        if (
            self.log_likelihood is True
            and self.posterior is not None
            and "log_lik" in self.posterior.sim["pars_oi"]
        ):
            self.log_likelihood = ["log_lik"]
        elif isinstance(self.log_likelihood, bool):
            self.log_likelihood = None

        import pystan  # pylint: disable=import-error

        self.pystan = pystan

    @requires("posterior")
    def posterior_to_xarray(self):
        """Extract posterior samples from fit."""
        posterior = self.posterior
        # filter posterior_predictive and log_likelihood
        posterior_predictive = self.posterior_predictive
        if posterior_predictive is None:
            posterior_predictive = []
        elif isinstance(posterior_predictive, str):
            posterior_predictive = [posterior_predictive]
        predictions = self.predictions
        if predictions is None:
            predictions = []
        elif isinstance(predictions, str):
            predictions = [predictions]
        log_likelihood = self.log_likelihood
        if log_likelihood is None:
            log_likelihood = []
        elif isinstance(log_likelihood, str):
            log_likelihood = [log_likelihood]
        elif isinstance(log_likelihood, dict):
            log_likelihood = list(log_likelihood.values())

        ignore = posterior_predictive + predictions + log_likelihood + ["lp__"]

        data, data_warmup = get_draws(
            posterior, ignore=ignore, warmup=self.save_warmup, dtypes=self.dtypes
        )
        attrs = get_attrs(posterior)
        return (
            dict_to_dataset(
                data, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
            ),
            dict_to_dataset(
                data_warmup, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
            ),
        )

    @requires("posterior")
    def sample_stats_to_xarray(self):
        """Extract sample_stats from posterior."""
        posterior = self.posterior

        data, data_warmup = get_sample_stats(posterior, warmup=self.save_warmup)

        # lp__
        stat_lp, stat_lp_warmup = get_draws(
            posterior, variables="lp__", warmup=self.save_warmup, dtypes=self.dtypes
        )
        data["lp"] = stat_lp["lp__"]
        if stat_lp_warmup:
            data_warmup["lp"] = stat_lp_warmup["lp__"]

        attrs = get_attrs(posterior)
        return (
            dict_to_dataset(
                data, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
            ),
            dict_to_dataset(
                data_warmup, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
            ),
        )

    @requires("posterior")
    @requires("log_likelihood")
    def log_likelihood_to_xarray(self):
        """Store log_likelihood data in log_likelihood group."""
        fit = self.posterior

        # log_likelihood values
        log_likelihood = self.log_likelihood
        if isinstance(log_likelihood, str):
            log_likelihood = [log_likelihood]
        if isinstance(log_likelihood, (list, tuple)):
            log_likelihood = {name: name for name in log_likelihood}
        log_likelihood_draws, log_likelihood_draws_warmup = get_draws(
            fit,
            variables=list(log_likelihood.values()),
            warmup=self.save_warmup,
            dtypes=self.dtypes,
        )
        data = {
            obs_var_name: log_likelihood_draws[log_like_name]
            for obs_var_name, log_like_name in log_likelihood.items()
            if log_like_name in log_likelihood_draws
        }

        data_warmup = {
            obs_var_name: log_likelihood_draws_warmup[log_like_name]
            for obs_var_name, log_like_name in log_likelihood.items()
            if log_like_name in log_likelihood_draws_warmup
        }

        return (
            dict_to_dataset(
                data, library=self.pystan, coords=self.coords, dims=self.dims, skip_event_dims=True
            ),
            dict_to_dataset(
                data_warmup,
                library=self.pystan,
                coords=self.coords,
                dims=self.dims,
                skip_event_dims=True,
            ),
        )

    @requires("posterior")
    @requires("posterior_predictive")
    def posterior_predictive_to_xarray(self):
        """Convert posterior_predictive samples to xarray."""
        posterior = self.posterior
        posterior_predictive = self.posterior_predictive
        data, data_warmup = get_draws(
            posterior, variables=posterior_predictive, warmup=self.save_warmup, dtypes=self.dtypes
        )
        return (
            dict_to_dataset(data, library=self.pystan, coords=self.coords, dims=self.dims),
            dict_to_dataset(data_warmup, library=self.pystan, coords=self.coords, dims=self.dims),
        )

    @requires("posterior")
    @requires("predictions")
    def predictions_to_xarray(self):
        """Convert predictions samples to xarray."""
        posterior = self.posterior
        predictions = self.predictions
        data, data_warmup = get_draws(
            posterior, variables=predictions, warmup=self.save_warmup, dtypes=self.dtypes
        )
        return (
            dict_to_dataset(data, library=self.pystan, coords=self.coords, dims=self.dims),
            dict_to_dataset(data_warmup, library=self.pystan, coords=self.coords, dims=self.dims),
        )

    @requires("prior")
    def prior_to_xarray(self):
        """Convert prior samples to xarray."""
        prior = self.prior
        # filter posterior_predictive and log_likelihood
        prior_predictive = self.prior_predictive
        if prior_predictive is None:
            prior_predictive = []
        elif isinstance(prior_predictive, str):
            prior_predictive = [prior_predictive]

        ignore = prior_predictive + ["lp__"]

        data, _ = get_draws(prior, ignore=ignore, warmup=False, dtypes=self.dtypes)
        attrs = get_attrs(prior)
        return dict_to_dataset(
            data, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
        )

    @requires("prior")
    def sample_stats_prior_to_xarray(self):
        """Extract sample_stats_prior from prior."""
        prior = self.prior
        data, _ = get_sample_stats(prior, warmup=False)

        # lp__
        stat_lp, _ = get_draws(prior, variables="lp__", warmup=False, dtypes=self.dtypes)
        data["lp"] = stat_lp["lp__"]

        attrs = get_attrs(prior)
        return dict_to_dataset(
            data, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
        )

    @requires("prior")
    @requires("prior_predictive")
    def prior_predictive_to_xarray(self):
        """Convert prior_predictive samples to xarray."""
        prior = self.prior
        prior_predictive = self.prior_predictive
        data, _ = get_draws(prior, variables=prior_predictive, warmup=False, dtypes=self.dtypes)
        return dict_to_dataset(data, library=self.pystan, coords=self.coords, dims=self.dims)

    @requires("posterior")
    @requires(["observed_data", "constant_data", "predictions_constant_data"])
    def data_to_xarray(self):
        """Convert observed, constant data and predictions constant data to xarray."""
        posterior = self.posterior
        if self.dims is None:
            dims = {}
        else:
            dims = self.dims
        obs_const_dict = {}
        for group_name in ("observed_data", "constant_data", "predictions_constant_data"):
            names = getattr(self, group_name)
            if names is None:
                continue
            names = [names] if isinstance(names, str) else names
            data = OrderedDict()
            for key in names:
                vals = np.atleast_1d(posterior.data[key])
                val_dims = dims.get(key)
                val_dims, coords = generate_dims_coords(
                    vals.shape, key, dims=val_dims, coords=self.coords
                )
                data[key] = xr.DataArray(vals, dims=val_dims, coords=coords)
            obs_const_dict[group_name] = xr.Dataset(
                data_vars=data, attrs=make_attrs(library=self.pystan)
            )
        return obs_const_dict

    def to_inference_data(self):
        """Convert all available data to an InferenceData object.

        Note that if groups can not be created (i.e., there is no `fit`, so
        the `posterior` and `sample_stats` can not be extracted), then the InferenceData
        will not have those groups.
        """
        data_dict = self.data_to_xarray()
        return InferenceData(
            save_warmup=self.save_warmup,
            **{
                "posterior": self.posterior_to_xarray(),
                "sample_stats": self.sample_stats_to_xarray(),
                "log_likelihood": self.log_likelihood_to_xarray(),
                "posterior_predictive": self.posterior_predictive_to_xarray(),
                "predictions": self.predictions_to_xarray(),
                "prior": self.prior_to_xarray(),
                "sample_stats_prior": self.sample_stats_prior_to_xarray(),
                "prior_predictive": self.prior_predictive_to_xarray(),
                **({} if data_dict is None else data_dict),
            },
        )


class PyStan3Converter:
    """Encapsulate PyStan3 specific logic."""

    # pylint: disable=too-many-instance-attributes
    def __init__(
        self,
        *,
        posterior=None,
        posterior_model=None,
        posterior_predictive=None,
        predictions=None,
        prior=None,
        prior_model=None,
        prior_predictive=None,
        observed_data=None,
        constant_data=None,
        predictions_constant_data=None,
        log_likelihood=None,
        coords=None,
        dims=None,
        dtypes=None,
    ):
        self.posterior = posterior
        self.posterior_model = posterior_model
        self.posterior_predictive = posterior_predictive
        self.predictions = predictions
        self.prior = prior
        self.prior_model = prior_model
        self.prior_predictive = prior_predictive
        self.observed_data = observed_data
        self.constant_data = constant_data
        self.predictions_constant_data = predictions_constant_data
        self.log_likelihood = (
            rcParams["data.log_likelihood"] if log_likelihood is None else log_likelihood
        )
        self.coords = coords
        self.dims = dims
        self.dtypes = dtypes

        if (
            self.log_likelihood is True
            and self.posterior is not None
            and "log_lik" in self.posterior.param_names
        ):
            self.log_likelihood = ["log_lik"]
        elif isinstance(self.log_likelihood, bool):
            self.log_likelihood = None

        import stan  # pylint: disable=import-error

        self.stan = stan

    @requires("posterior")
    def posterior_to_xarray(self):
        """Extract posterior samples from fit."""
        posterior = self.posterior
        posterior_model = self.posterior_model
        # filter posterior_predictive and log_likelihood
        posterior_predictive = self.posterior_predictive
        if posterior_predictive is None:
            posterior_predictive = []
        elif isinstance(posterior_predictive, str):
            posterior_predictive = [posterior_predictive]
        predictions = self.predictions
        if predictions is None:
            predictions = []
        elif isinstance(predictions, str):
            predictions = [predictions]
        log_likelihood = self.log_likelihood
        if log_likelihood is None:
            log_likelihood = []
        elif isinstance(log_likelihood, str):
            log_likelihood = [log_likelihood]
        elif isinstance(log_likelihood, dict):
            log_likelihood = list(log_likelihood.values())

        ignore = posterior_predictive + predictions + log_likelihood

        data = get_draws_stan3(posterior, model=posterior_model, ignore=ignore, dtypes=self.dtypes)
        attrs = get_attrs_stan3(posterior, model=posterior_model)
        return dict_to_dataset(
            data, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
        )

    @requires("posterior")
    def sample_stats_to_xarray(self):
        """Extract sample_stats from posterior."""
        posterior = self.posterior
        posterior_model = self.posterior_model
        data = get_sample_stats_stan3(posterior, ignore="lp__", dtypes=self.dtypes)
        data["lp"] = get_sample_stats_stan3(posterior, variables="lp__")["lp"]

        attrs = get_attrs_stan3(posterior, model=posterior_model)
        return dict_to_dataset(
            data, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
        )

    @requires("posterior")
    @requires("log_likelihood")
    def log_likelihood_to_xarray(self):
        """Store log_likelihood data in log_likelihood group."""
        fit = self.posterior

        log_likelihood = self.log_likelihood
        model = self.posterior_model
        if isinstance(log_likelihood, str):
            log_likelihood = [log_likelihood]
        if isinstance(log_likelihood, (list, tuple)):
            log_likelihood = {name: name for name in log_likelihood}
        log_likelihood_draws = get_draws_stan3(
            fit, model=model, variables=list(log_likelihood.values()), dtypes=self.dtypes
        )
        data = {
            obs_var_name: log_likelihood_draws[log_like_name]
            for obs_var_name, log_like_name in log_likelihood.items()
            if log_like_name in log_likelihood_draws
        }

        return dict_to_dataset(data, library=self.stan, coords=self.coords, dims=self.dims)

    @requires("posterior")
    @requires("posterior_predictive")
    def posterior_predictive_to_xarray(self):
        """Convert posterior_predictive samples to xarray."""
        posterior = self.posterior
        posterior_model = self.posterior_model
        posterior_predictive = self.posterior_predictive
        data = get_draws_stan3(
            posterior, model=posterior_model, variables=posterior_predictive, dtypes=self.dtypes
        )
        return dict_to_dataset(data, library=self.stan, coords=self.coords, dims=self.dims)

    @requires("posterior")
    @requires("predictions")
    def predictions_to_xarray(self):
        """Convert predictions samples to xarray."""
        posterior = self.posterior
        posterior_model = self.posterior_model
        predictions = self.predictions
        data = get_draws_stan3(
            posterior, model=posterior_model, variables=predictions, dtypes=self.dtypes
        )
        return dict_to_dataset(data, library=self.stan, coords=self.coords, dims=self.dims)

    @requires("prior")
    def prior_to_xarray(self):
        """Convert prior samples to xarray."""
        prior = self.prior
        prior_model = self.prior_model
        # filter posterior_predictive and log_likelihood
        prior_predictive = self.prior_predictive
        if prior_predictive is None:
            prior_predictive = []
        elif isinstance(prior_predictive, str):
            prior_predictive = [prior_predictive]

        ignore = prior_predictive

        data = get_draws_stan3(prior, model=prior_model, ignore=ignore, dtypes=self.dtypes)
        attrs = get_attrs_stan3(prior, model=prior_model)
        return dict_to_dataset(
            data, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
        )

    @requires("prior")
    def sample_stats_prior_to_xarray(self):
        """Extract sample_stats_prior from prior."""
        prior = self.prior
        prior_model = self.prior_model
        data = get_sample_stats_stan3(prior, dtypes=self.dtypes)
        attrs = get_attrs_stan3(prior, model=prior_model)
        return dict_to_dataset(
            data, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
        )

    @requires("prior")
    @requires("prior_predictive")
    def prior_predictive_to_xarray(self):
        """Convert prior_predictive samples to xarray."""
        prior = self.prior
        prior_model = self.prior_model
        prior_predictive = self.prior_predictive
        data = get_draws_stan3(
            prior, model=prior_model, variables=prior_predictive, dtypes=self.dtypes
        )
        return dict_to_dataset(data, library=self.stan, coords=self.coords, dims=self.dims)

    @requires("posterior_model")
    @requires(["observed_data", "constant_data"])
    def observed_and_constant_data_to_xarray(self):
        """Convert observed data to xarray."""
        posterior_model = self.posterior_model
        if self.dims is None:
            dims = {}
        else:
            dims = self.dims
        obs_const_dict = {}
        for group_name in ("observed_data", "constant_data"):
            names = getattr(self, group_name)
            if names is None:
                continue
            names = [names] if isinstance(names, str) else names
            data = OrderedDict()
            for key in names:
                vals = np.atleast_1d(posterior_model.data[key])
                val_dims = dims.get(key)
                val_dims, coords = generate_dims_coords(
                    vals.shape, key, dims=val_dims, coords=self.coords
                )
                data[key] = xr.DataArray(vals, dims=val_dims, coords=coords)
            obs_const_dict[group_name] = xr.Dataset(
                data_vars=data, attrs=make_attrs(library=self.stan)
            )
        return obs_const_dict

    @requires("posterior_model")
    @requires("predictions_constant_data")
    def predictions_constant_data_to_xarray(self):
        """Convert observed data to xarray."""
        posterior_model = self.posterior_model
        if self.dims is None:
            dims = {}
        else:
            dims = self.dims
        names = self.predictions_constant_data
        names = [names] if isinstance(names, str) else names
        data = OrderedDict()
        for key in names:
            vals = np.atleast_1d(posterior_model.data[key])
            val_dims = dims.get(key)
            val_dims, coords = generate_dims_coords(
                vals.shape, key, dims=val_dims, coords=self.coords
            )
            data[key] = xr.DataArray(vals, dims=val_dims, coords=coords)
        return xr.Dataset(data_vars=data, attrs=make_attrs(library=self.stan))

    def to_inference_data(self):
        """Convert all available data to an InferenceData object.

        Note that if groups can not be created (i.e., there is no `fit`, so
        the `posterior` and `sample_stats` can not be extracted), then the InferenceData
        will not have those groups.
        """
        obs_const_dict = self.observed_and_constant_data_to_xarray()
        predictions_const_data = self.predictions_constant_data_to_xarray()
        return InferenceData(
            **{
                "posterior": self.posterior_to_xarray(),
                "sample_stats": self.sample_stats_to_xarray(),
                "log_likelihood": self.log_likelihood_to_xarray(),
                "posterior_predictive": self.posterior_predictive_to_xarray(),
                "predictions": self.predictions_to_xarray(),
                "prior": self.prior_to_xarray(),
                "sample_stats_prior": self.sample_stats_prior_to_xarray(),
                "prior_predictive": self.prior_predictive_to_xarray(),
                **({} if obs_const_dict is None else obs_const_dict),
                **(
                    {}
                    if predictions_const_data is None
                    else {"predictions_constant_data": predictions_const_data}
                ),
            }
        )


def get_draws(fit, variables=None, ignore=None, warmup=False, dtypes=None):
    """Extract draws from PyStan fit."""
    if ignore is None:
        ignore = []
    if fit.mode == 1:
        msg = "Model in mode 'test_grad'. Sampling is not conducted."
        raise AttributeError(msg)

    if fit.mode == 2 or fit.sim.get("samples") is None:
        msg = "Fit doesn't contain samples."
        raise AttributeError(msg)

    if dtypes is None:
        dtypes = {}

    dtypes = {**infer_dtypes(fit), **dtypes}

    if variables is None:
        variables = fit.sim["pars_oi"]
    elif isinstance(variables, str):
        variables = [variables]
    variables = list(variables)

    for var, dim in zip(fit.sim["pars_oi"], fit.sim["dims_oi"]):
        if var in variables and np.prod(dim) == 0:
            del variables[variables.index(var)]

    ndraws_warmup = fit.sim["warmup2"]
    if max(ndraws_warmup) == 0:
        warmup = False
    ndraws = [s - w for s, w in zip(fit.sim["n_save"], ndraws_warmup)]
    nchain = len(fit.sim["samples"])

    # check if the values are in 0-based (<=2.17) or 1-based indexing (>=2.18)
    shift = 1
    if any(dim and np.prod(dim) != 0 for dim in fit.sim["dims_oi"]):
        # choose variable with lowest number of dims > 1
        par_idx = min(
            (dim, i) for i, dim in enumerate(fit.sim["dims_oi"]) if (dim and np.prod(dim) != 0)
        )[1]
        offset = int(sum(map(np.product, fit.sim["dims_oi"][:par_idx])))
        par_offset = int(np.product(fit.sim["dims_oi"][par_idx]))
        par_keys = fit.sim["fnames_oi"][offset : offset + par_offset]
        shift = len(par_keys)
        for item in par_keys:
            _, shape = item.replace("]", "").split("[")
            shape_idx_min = min(int(shape_value) for shape_value in shape.split(","))
            if shape_idx_min < shift:
                shift = shape_idx_min
        # If shift is higher than 1, this will probably mean that Stan
        # has implemented sparse structure (saves only non-zero parts),
        # but let's hope that dims are still corresponding to the full shape
        shift = int(min(shift, 1))

    var_keys = OrderedDict((var, []) for var in fit.sim["pars_oi"])
    for key in fit.sim["fnames_oi"]:
        var, *tails = key.split("[")
        loc = [Ellipsis]
        for tail in tails:
            loc = []
            for i in tail[:-1].split(","):
                loc.append(int(i) - shift)
        var_keys[var].append((key, loc))

    shapes = dict(zip(fit.sim["pars_oi"], fit.sim["dims_oi"]))

    variables = [var for var in variables if var not in ignore]

    data = OrderedDict()
    data_warmup = OrderedDict()

    for var in variables:
        if var in data:
            continue
        keys_locs = var_keys.get(var, [(var, [Ellipsis])])
        shape = shapes.get(var, [])
        dtype = dtypes.get(var)

        ndraw = max(ndraws)
        ary_shape = [nchain, ndraw] + shape
        ary = np.empty(ary_shape, dtype=dtype, order="F")

        if warmup:
            nwarmup = max(ndraws_warmup)
            ary_warmup_shape = [nchain, nwarmup] + shape
            ary_warmup = np.empty(ary_warmup_shape, dtype=dtype, order="F")

        for chain, (pyholder, ndraw, ndraw_warmup) in enumerate(
            zip(fit.sim["samples"], ndraws, ndraws_warmup)
        ):
            axes = [chain, slice(None)]
            for key, loc in keys_locs:
                ary_slice = tuple(axes + loc)
                ary[ary_slice] = pyholder.chains[key][-ndraw:]
                if warmup:
                    ary_warmup[ary_slice] = pyholder.chains[key][:ndraw_warmup]
        data[var] = ary
        if warmup:
            data_warmup[var] = ary_warmup
    return data, data_warmup


def get_sample_stats(fit, warmup=False, dtypes=None):
    """Extract sample stats from PyStan fit."""
    if dtypes is None:
        dtypes = {}
    dtypes = {"divergent__": bool, "n_leapfrog__": np.int64, "treedepth__": np.int64, **dtypes}

    rename_dict = {
        "divergent": "diverging",
        "n_leapfrog": "n_steps",
        "treedepth": "tree_depth",
        "stepsize": "step_size",
        "accept_stat": "acceptance_rate",
    }

    ndraws_warmup = fit.sim["warmup2"]
    if max(ndraws_warmup) == 0:
        warmup = False
    ndraws = [s - w for s, w in zip(fit.sim["n_save"], ndraws_warmup)]

    extraction = OrderedDict()
    extraction_warmup = OrderedDict()
    for chain, (pyholder, ndraw, ndraw_warmup) in enumerate(
        zip(fit.sim["samples"], ndraws, ndraws_warmup)
    ):
        if chain == 0:
            for key in pyholder["sampler_param_names"]:
                extraction[key] = []
                if warmup:
                    extraction_warmup[key] = []
        for key, values in zip(pyholder["sampler_param_names"], pyholder["sampler_params"]):
            extraction[key].append(values[-ndraw:])
            if warmup:
                extraction_warmup[key].append(values[:ndraw_warmup])

    data = OrderedDict()
    for key, values in extraction.items():
        values = np.stack(values, axis=0)
        dtype = dtypes.get(key)
        values = values.astype(dtype)
        name = re.sub("__$", "", key)
        name = rename_dict.get(name, name)
        data[name] = values

    data_warmup = OrderedDict()
    if warmup:
        for key, values in extraction_warmup.items():
            values = np.stack(values, axis=0)
            values = values.astype(dtypes.get(key))
            name = re.sub("__$", "", key)
            name = rename_dict.get(name, name)
            data_warmup[name] = values

    return data, data_warmup


def get_attrs(fit):
    """Get attributes from PyStan fit object."""
    attrs = {}

    try:
        attrs["args"] = [deepcopy(holder.args) for holder in fit.sim["samples"]]
    except Exception as exp:  # pylint: disable=broad-except
        _log.warning("Failed to fetch args from fit: %s", exp)
    if "args" in attrs:
        for arg in attrs["args"]:
            if isinstance(arg["init"], bytes):
                arg["init"] = arg["init"].decode("utf-8")
        attrs["args"] = json.dumps(attrs["args"])
    try:
        attrs["inits"] = [holder.inits for holder in fit.sim["samples"]]
    except Exception as exp:  # pylint: disable=broad-except
        _log.warning("Failed to fetch `args` from fit: %s", exp)
    else:
        attrs["inits"] = json.dumps(attrs["inits"])

    attrs["step_size"] = []
    attrs["metric"] = []
    attrs["inv_metric"] = []
    for holder in fit.sim["samples"]:
        try:
            step_size = float(
                re.search(
                    r"step\s*size\s*=\s*([0-9]+.?[0-9]+)\s*",
                    holder.adaptation_info,
                    flags=re.IGNORECASE,
                ).group(1)
            )
        except AttributeError:
            step_size = np.nan
        attrs["step_size"].append(step_size)

        inv_metric_match = re.search(
            r"mass matrix:\s*(.*)\s*$", holder.adaptation_info, flags=re.DOTALL
        )
        if inv_metric_match:
            inv_metric_str = inv_metric_match.group(1)
            if "Diagonal elements of inverse mass matrix" in holder.adaptation_info:
                metric = "diag_e"
                inv_metric = [float(item) for item in inv_metric_str.strip(" #\n").split(",")]
            else:
                metric = "dense_e"
                inv_metric = [
                    list(map(float, item.split(",")))
                    for item in re.sub(r"#\s", "", inv_metric_str).splitlines()
                ]
        else:
            metric = "unit_e"
            inv_metric = None

        attrs["metric"].append(metric)
        attrs["inv_metric"].append(inv_metric)
    attrs["inv_metric"] = json.dumps(attrs["inv_metric"])

    if not attrs["step_size"]:
        del attrs["step_size"]

    attrs["adaptation_info"] = fit.get_adaptation_info()
    attrs["stan_code"] = fit.get_stancode()

    return attrs


def get_draws_stan3(fit, model=None, variables=None, ignore=None, dtypes=None):
    """Extract draws from PyStan3 fit."""
    if ignore is None:
        ignore = []

    if dtypes is None:
        dtypes = {}

    if model is not None:
        dtypes = {**infer_dtypes(fit, model), **dtypes}

    if variables is None:
        variables = fit.param_names
    elif isinstance(variables, str):
        variables = [variables]
    variables = list(variables)

    data = OrderedDict()

    for var in variables:
        if var in data:
            continue
        dtype = dtypes.get(var)

        # in future fix the correct number of draws if fit.save_warmup is True
        new_shape = (*fit.dims[fit.param_names.index(var)], -1, fit.num_chains)
        if 0 in new_shape:
            continue
        values = fit._draws[fit._parameter_indexes(var), :]  # pylint: disable=protected-access
        values = values.reshape(new_shape, order="F")
        values = np.moveaxis(values, [-2, -1], [1, 0])
        values = values.astype(dtype)
        data[var] = values

    return data


def get_sample_stats_stan3(fit, variables=None, ignore=None, dtypes=None):
    """Extract sample stats from PyStan3 fit."""
    if dtypes is None:
        dtypes = {}
    dtypes = {"divergent__": bool, "n_leapfrog__": np.int64, "treedepth__": np.int64, **dtypes}

    rename_dict = {
        "divergent": "diverging",
        "n_leapfrog": "n_steps",
        "treedepth": "tree_depth",
        "stepsize": "step_size",
        "accept_stat": "acceptance_rate",
    }

    if isinstance(variables, str):
        variables = [variables]
    if isinstance(ignore, str):
        ignore = [ignore]

    data = OrderedDict()
    for key in fit.sample_and_sampler_param_names:
        if (variables and key not in variables) or (ignore and key in ignore):
            continue
        new_shape = -1, fit.num_chains
        values = fit._draws[fit._parameter_indexes(key)]  # pylint: disable=protected-access
        values = values.reshape(new_shape, order="F")
        values = np.moveaxis(values, [-2, -1], [1, 0])
        dtype = dtypes.get(key)
        values = values.astype(dtype)
        name = re.sub("__$", "", key)
        name = rename_dict.get(name, name)
        data[name] = values

    return data


def get_attrs_stan3(fit, model=None):
    """Get attributes from PyStan3 fit and model object."""
    attrs = {}
    for key in ["num_chains", "num_samples", "num_thin", "num_warmup", "save_warmup"]:
        try:
            attrs[key] = getattr(fit, key)
        except AttributeError as exp:
            _log.warning("Failed to access attribute %s in fit object %s", key, exp)

    if model is not None:
        for key in ["model_name", "program_code", "random_seed"]:
            try:
                attrs[key] = getattr(model, key)
            except AttributeError as exp:
                _log.warning("Failed to access attribute %s in model object %s", key, exp)

    return attrs


def infer_dtypes(fit, model=None):
    """Infer dtypes from Stan model code.

    Function strips out generated quantities block and searches for `int`
    dtypes after stripping out comments inside the block.
    """
    if model is None:
        stan_code = fit.get_stancode()
        model_pars = fit.model_pars
    else:
        stan_code = model.program_code
        model_pars = fit.param_names

    dtypes = {key: item for key, item in infer_stan_dtypes(stan_code).items() if key in model_pars}
    return dtypes


# pylint disable=too-many-instance-attributes
def from_pystan(
    posterior=None,
    *,
    posterior_predictive=None,
    predictions=None,
    prior=None,
    prior_predictive=None,
    observed_data=None,
    constant_data=None,
    predictions_constant_data=None,
    log_likelihood=None,
    coords=None,
    dims=None,
    posterior_model=None,
    prior_model=None,
    save_warmup=None,
    dtypes=None,
):
    """Convert PyStan data into an InferenceData object.

    For a usage example read the
    :ref:`Creating InferenceData section on from_pystan <creating_InferenceData>`

    Parameters
    ----------
    posterior : StanFit4Model or stan.fit.Fit
        PyStan fit object for posterior.
    posterior_predictive : str, a list of str
        Posterior predictive samples for the posterior.
    predictions : str, a list of str
        Out-of-sample predictions for the posterior.
    prior : StanFit4Model or stan.fit.Fit
        PyStan fit object for prior.
    prior_predictive : str, a list of str
        Posterior predictive samples for the prior.
    observed_data : str or a list of str
        observed data used in the sampling.
        Observed data is extracted from the `posterior.data`.
        PyStan3 needs model object for the extraction.
        See `posterior_model`.
    constant_data : str or list of str
        Constants relevant to the model (i.e. x values in a linear
        regression).
    predictions_constant_data : str or list of str
        Constants relevant to the model predictions (i.e. new x values in a linear
        regression).
    log_likelihood : dict of {str: str}, list of str or str, optional
        Pointwise log_likelihood for the data. log_likelihood is extracted from the
        posterior. It is recommended to use this argument as a dictionary whose keys
        are observed variable names and its values are the variables storing log
        likelihood arrays in the Stan code. In other cases, a dictionary with keys
        equal to its values is used. By default, if a variable ``log_lik`` is
        present in the Stan model, it will be retrieved as pointwise log
        likelihood values. Use ``False`` or set ``data.log_likelihood`` to
        false to avoid this behaviour.
    coords : dict[str, iterable]
        A dictionary containing the values that are used as index. The key
        is the name of the dimension, the values are the index values.
    dims : dict[str, List(str)]
        A mapping from variables to a list of coordinate names for the variable.
    posterior_model : stan.model.Model
        PyStan3 specific model object. Needed for automatic dtype parsing
        and for the extraction of observed data.
    prior_model : stan.model.Model
        PyStan3 specific model object. Needed for automatic dtype parsing.
    save_warmup : bool
        Save warmup iterations into InferenceData object. If not defined, use default
        defined by the rcParams. Not supported in PyStan3.
    dtypes: dict
        A dictionary containing dtype information (int, float) for parameters.
        By default dtype information is extracted from the model code.
        Model code is extracted from fit object in PyStan 2 and from model object
        in PyStan 3.

    Returns
    -------
    InferenceData object
    """
    check_posterior = (posterior is not None) and (type(posterior).__module__ == "stan.fit")
    check_prior = (prior is not None) and (type(prior).__module__ == "stan.fit")
    if check_posterior or check_prior:
        if save_warmup:
            warnings.warn(
                "save_warmup is not currently supported for PyStan3",
                UserWarning,
            )
        return PyStan3Converter(
            posterior=posterior,
            posterior_model=posterior_model,
            posterior_predictive=posterior_predictive,
            predictions=predictions,
            prior=prior,
            prior_model=prior_model,
            prior_predictive=prior_predictive,
            observed_data=observed_data,
            constant_data=constant_data,
            predictions_constant_data=predictions_constant_data,
            log_likelihood=log_likelihood,
            coords=coords,
            dims=dims,
            dtypes=dtypes,
        ).to_inference_data()
    else:
        return PyStanConverter(
            posterior=posterior,
            posterior_predictive=posterior_predictive,
            predictions=predictions,
            prior=prior,
            prior_predictive=prior_predictive,
            observed_data=observed_data,
            constant_data=constant_data,
            predictions_constant_data=predictions_constant_data,
            log_likelihood=log_likelihood,
            coords=coords,
            dims=dims,
            save_warmup=save_warmup,
            dtypes=dtypes,
        ).to_inference_data()

arviz/data/io_pystan.py
#  pylint: disable=too-many-instance-attributes,too-many-lines
"""PyStan-specific conversion code."""
import re
import warnings
from collections import OrderedDict
from copy import deepcopy

import numpy as np
import xarray as xr

from .. import _log
from ..rcparams import rcParams
from .base import dict_to_dataset, generate_dims_coords, infer_stan_dtypes, make_attrs, requires
from .inference_data import InferenceData

try:
    import ujson as json
except ImportError:
    # Can't find ujson using json
    # mypy struggles with conditional imports expressed as catching ImportError:
    # https://github.com/python/mypy/issues/1153
    import json  # type: ignore


class PyStanConverter:
    """Encapsulate PyStan specific logic."""

    def __init__(
        self,
        *,
        posterior=None,
        posterior_predictive=None,
        predictions=None,
        prior=None,
        prior_predictive=None,
        observed_data=None,
        constant_data=None,
        predictions_constant_data=None,
        log_likelihood=None,
        coords=None,
        dims=None,
        save_warmup=None,
        dtypes=None,
    ):
        self.posterior = posterior
        self.posterior_predictive = posterior_predictive
        self.predictions = predictions
        self.prior = prior
        self.prior_predictive = prior_predictive
        self.observed_data = observed_data
        self.constant_data = constant_data
        self.predictions_constant_data = predictions_constant_data
        self.log_likelihood = (
            rcParams["data.log_likelihood"] if log_likelihood is None else log_likelihood
        )
        self.coords = coords
        self.dims = dims
        self.save_warmup = rcParams["data.save_warmup"] if save_warmup is None else save_warmup
        self.dtypes = dtypes

        if (
            self.log_likelihood is True
            and self.posterior is not None
            and "log_lik" in self.posterior.sim["pars_oi"]
        ):
            self.log_likelihood = ["log_lik"]
        elif isinstance(self.log_likelihood, bool):
            self.log_likelihood = None

        import pystan  # pylint: disable=import-error

        self.pystan = pystan

    @requires("posterior")
    def posterior_to_xarray(self):
        """Extract posterior samples from fit."""
        posterior = self.posterior
        # filter posterior_predictive and log_likelihood
        posterior_predictive = self.posterior_predictive
        if posterior_predictive is None:
            posterior_predictive = []
        elif isinstance(posterior_predictive, str):
            posterior_predictive = [posterior_predictive]
        predictions = self.predictions
        if predictions is None:
            predictions = []
        elif isinstance(predictions, str):
            predictions = [predictions]
        log_likelihood = self.log_likelihood
        if log_likelihood is None:
            log_likelihood = []
        elif isinstance(log_likelihood, str):
            log_likelihood = [log_likelihood]
        elif isinstance(log_likelihood, dict):
            log_likelihood = list(log_likelihood.values())

        ignore = posterior_predictive + predictions + log_likelihood + ["lp__"]

        data, data_warmup = get_draws(
            posterior, ignore=ignore, warmup=self.save_warmup, dtypes=self.dtypes
        )
        attrs = get_attrs(posterior)
        return (
            dict_to_dataset(
                data, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
            ),
            dict_to_dataset(
                data_warmup, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
            ),
        )

    @requires("posterior")
    def sample_stats_to_xarray(self):
        """Extract sample_stats from posterior."""
        posterior = self.posterior

        data, data_warmup = get_sample_stats(posterior, warmup=self.save_warmup)

        # lp__
        stat_lp, stat_lp_warmup = get_draws(
            posterior, variables="lp__", warmup=self.save_warmup, dtypes=self.dtypes
        )
        data["lp"] = stat_lp["lp__"]
        if stat_lp_warmup:
            data_warmup["lp"] = stat_lp_warmup["lp__"]

        attrs = get_attrs(posterior)
        return (
            dict_to_dataset(
                data, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
            ),
            dict_to_dataset(
                data_warmup, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
            ),
        )

    @requires("posterior")
    @requires("log_likelihood")
    def log_likelihood_to_xarray(self):
        """Store log_likelihood data in log_likelihood group."""
        fit = self.posterior

        # log_likelihood values
        log_likelihood = self.log_likelihood
        if isinstance(log_likelihood, str):
            log_likelihood = [log_likelihood]
        if isinstance(log_likelihood, (list, tuple)):
            log_likelihood = {name: name for name in log_likelihood}
        log_likelihood_draws, log_likelihood_draws_warmup = get_draws(
            fit,
            variables=list(log_likelihood.values()),
            warmup=self.save_warmup,
            dtypes=self.dtypes,
        )
        data = {
            obs_var_name: log_likelihood_draws[log_like_name]
            for obs_var_name, log_like_name in log_likelihood.items()
            if log_like_name in log_likelihood_draws
        }

        data_warmup = {
            obs_var_name: log_likelihood_draws_warmup[log_like_name]
            for obs_var_name, log_like_name in log_likelihood.items()
            if log_like_name in log_likelihood_draws_warmup
        }

        return (
            dict_to_dataset(
                data, library=self.pystan, coords=self.coords, dims=self.dims, skip_event_dims=True
            ),
            dict_to_dataset(
                data_warmup,
                library=self.pystan,
                coords=self.coords,
                dims=self.dims,
                skip_event_dims=True,
            ),
        )

    @requires("posterior")
    @requires("posterior_predictive")
    def posterior_predictive_to_xarray(self):
        """Convert posterior_predictive samples to xarray."""
        posterior = self.posterior
        posterior_predictive = self.posterior_predictive
        data, data_warmup = get_draws(
            posterior, variables=posterior_predictive, warmup=self.save_warmup, dtypes=self.dtypes
        )
        return (
            dict_to_dataset(data, library=self.pystan, coords=self.coords, dims=self.dims),
            dict_to_dataset(data_warmup, library=self.pystan, coords=self.coords, dims=self.dims),
        )

    @requires("posterior")
    @requires("predictions")
    def predictions_to_xarray(self):
        """Convert predictions samples to xarray."""
        posterior = self.posterior
        predictions = self.predictions
        data, data_warmup = get_draws(
            posterior, variables=predictions, warmup=self.save_warmup, dtypes=self.dtypes
        )
        return (
            dict_to_dataset(data, library=self.pystan, coords=self.coords, dims=self.dims),
            dict_to_dataset(data_warmup, library=self.pystan, coords=self.coords, dims=self.dims),
        )

    @requires("prior")
    def prior_to_xarray(self):
        """Convert prior samples to xarray."""
        prior = self.prior
        # filter posterior_predictive and log_likelihood
        prior_predictive = self.prior_predictive
        if prior_predictive is None:
            prior_predictive = []
        elif isinstance(prior_predictive, str):
            prior_predictive = [prior_predictive]

        ignore = prior_predictive + ["lp__"]

        data, _ = get_draws(prior, ignore=ignore, warmup=False, dtypes=self.dtypes)
        attrs = get_attrs(prior)
        return dict_to_dataset(
            data, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
        )

    @requires("prior")
    def sample_stats_prior_to_xarray(self):
        """Extract sample_stats_prior from prior."""
        prior = self.prior
        data, _ = get_sample_stats(prior, warmup=False)

        # lp__
        stat_lp, _ = get_draws(prior, variables="lp__", warmup=False, dtypes=self.dtypes)
        data["lp"] = stat_lp["lp__"]

        attrs = get_attrs(prior)
        return dict_to_dataset(
            data, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
        )

    @requires("prior")
    @requires("prior_predictive")
    def prior_predictive_to_xarray(self):
        """Convert prior_predictive samples to xarray."""
        prior = self.prior
        prior_predictive = self.prior_predictive
        data, _ = get_draws(prior, variables=prior_predictive, warmup=False, dtypes=self.dtypes)
        return dict_to_dataset(data, library=self.pystan, coords=self.coords, dims=self.dims)

    @requires("posterior")
    @requires(["observed_data", "constant_data", "predictions_constant_data"])
    def data_to_xarray(self):
        """Convert observed, constant data and predictions constant data to xarray."""
        posterior = self.posterior
        if self.dims is None:
            dims = {}
        else:
            dims = self.dims
        obs_const_dict = {}
        for group_name in ("observed_data", "constant_data", "predictions_constant_data"):
            names = getattr(self, group_name)
            if names is None:
                continue
            names = [names] if isinstance(names, str) else names
            data = OrderedDict()
            for key in names:
                vals = np.atleast_1d(posterior.data[key])
                val_dims = dims.get(key)
                val_dims, coords = generate_dims_coords(
                    vals.shape, key, dims=val_dims, coords=self.coords
                )
                data[key] = xr.DataArray(vals, dims=val_dims, coords=coords)
            obs_const_dict[group_name] = xr.Dataset(
                data_vars=data, attrs=make_attrs(library=self.pystan)
            )
        return obs_const_dict

    def to_inference_data(self):
        """Convert all available data to an InferenceData object.

        Note that if groups can not be created (i.e., there is no `fit`, so
        the `posterior` and `sample_stats` can not be extracted), then the InferenceData
        will not have those groups.
        """
        data_dict = self.data_to_xarray()
        return InferenceData(
            save_warmup=self.save_warmup,
            **{
                "posterior": self.posterior_to_xarray(),
                "sample_stats": self.sample_stats_to_xarray(),
                "log_likelihood": self.log_likelihood_to_xarray(),
                "posterior_predictive": self.posterior_predictive_to_xarray(),
                "predictions": self.predictions_to_xarray(),
                "prior": self.prior_to_xarray(),
                "sample_stats_prior": self.sample_stats_prior_to_xarray(),
                "prior_predictive": self.prior_predictive_to_xarray(),
                **({} if data_dict is None else data_dict),
            },
        )


class PyStan3Converter:
    """Encapsulate PyStan3 specific logic."""

    # pylint: disable=too-many-instance-attributes
    def __init__(
        self,
        *,
        posterior=None,
        posterior_model=None,
        posterior_predictive=None,
        predictions=None,
        prior=None,
        prior_model=None,
        prior_predictive=None,
        observed_data=None,
        constant_data=None,
        predictions_constant_data=None,
        log_likelihood=None,
        coords=None,
        dims=None,
        dtypes=None,
    ):
        self.posterior = posterior
        self.posterior_model = posterior_model
        self.posterior_predictive = posterior_predictive
        self.predictions = predictions
        self.prior = prior
        self.prior_model = prior_model
        self.prior_predictive = prior_predictive
        self.observed_data = observed_data
        self.constant_data = constant_data
        self.predictions_constant_data = predictions_constant_data
        self.log_likelihood = (
            rcParams["data.log_likelihood"] if log_likelihood is None else log_likelihood
        )
        self.coords = coords
        self.dims = dims
        self.dtypes = dtypes

        if (
            self.log_likelihood is True
            and self.posterior is not None
            and "log_lik" in self.posterior.param_names
        ):
            self.log_likelihood = ["log_lik"]
        elif isinstance(self.log_likelihood, bool):
            self.log_likelihood = None

        import stan  # pylint: disable=import-error

        self.stan = stan

    @requires("posterior")
    def posterior_to_xarray(self):
        """Extract posterior samples from fit."""
        posterior = self.posterior
        posterior_model = self.posterior_model
        # filter posterior_predictive and log_likelihood
        posterior_predictive = self.posterior_predictive
        if posterior_predictive is None:
            posterior_predictive = []
        elif isinstance(posterior_predictive, str):
            posterior_predictive = [posterior_predictive]
        predictions = self.predictions
        if predictions is None:
            predictions = []
        elif isinstance(predictions, str):
            predictions = [predictions]
        log_likelihood = self.log_likelihood
        if log_likelihood is None:
            log_likelihood = []
        elif isinstance(log_likelihood, str):
            log_likelihood = [log_likelihood]
        elif isinstance(log_likelihood, dict):
            log_likelihood = list(log_likelihood.values())

        ignore = posterior_predictive + predictions + log_likelihood

        data = get_draws_stan3(posterior, model=posterior_model, ignore=ignore, dtypes=self.dtypes)
        attrs = get_attrs_stan3(posterior, model=posterior_model)
        return dict_to_dataset(
            data, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
        )

    @requires("posterior")
    def sample_stats_to_xarray(self):
        """Extract sample_stats from posterior."""
        posterior = self.posterior
        posterior_model = self.posterior_model
        data = get_sample_stats_stan3(posterior, ignore="lp__", dtypes=self.dtypes)
        data["lp"] = get_sample_stats_stan3(posterior, variables="lp__")["lp"]

        attrs = get_attrs_stan3(posterior, model=posterior_model)
        return dict_to_dataset(
            data, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
        )

    @requires("posterior")
    @requires("log_likelihood")
    def log_likelihood_to_xarray(self):
        """Store log_likelihood data in log_likelihood group."""
        fit = self.posterior

        log_likelihood = self.log_likelihood
        model = self.posterior_model
        if isinstance(log_likelihood, str):
            log_likelihood = [log_likelihood]
        if isinstance(log_likelihood, (list, tuple)):
            log_likelihood = {name: name for name in log_likelihood}
        log_likelihood_draws = get_draws_stan3(
            fit, model=model, variables=list(log_likelihood.values()), dtypes=self.dtypes
        )
        data = {
            obs_var_name: log_likelihood_draws[log_like_name]
            for obs_var_name, log_like_name in log_likelihood.items()
            if log_like_name in log_likelihood_draws
        }

        return dict_to_dataset(data, library=self.stan, coords=self.coords, dims=self.dims)

    @requires("posterior")
    @requires("posterior_predictive")
    def posterior_predictive_to_xarray(self):
        """Convert posterior_predictive samples to xarray."""
        posterior = self.posterior
        posterior_model = self.posterior_model
        posterior_predictive = self.posterior_predictive
        data = get_draws_stan3(
            posterior, model=posterior_model, variables=posterior_predictive, dtypes=self.dtypes
        )
        return dict_to_dataset(data, library=self.stan, coords=self.coords, dims=self.dims)

    @requires("posterior")
    @requires("predictions")
    def predictions_to_xarray(self):
        """Convert predictions samples to xarray."""
        posterior = self.posterior
        posterior_model = self.posterior_model
        predictions = self.predictions
        data = get_draws_stan3(
            posterior, model=posterior_model, variables=predictions, dtypes=self.dtypes
        )
        return dict_to_dataset(data, library=self.stan, coords=self.coords, dims=self.dims)

    @requires("prior")
    def prior_to_xarray(self):
        """Convert prior samples to xarray."""
        prior = self.prior
        prior_model = self.prior_model
        # filter posterior_predictive and log_likelihood
        prior_predictive = self.prior_predictive
        if prior_predictive is None:
            prior_predictive = []
        elif isinstance(prior_predictive, str):
            prior_predictive = [prior_predictive]

        ignore = prior_predictive

        data = get_draws_stan3(prior, model=prior_model, ignore=ignore, dtypes=self.dtypes)
        attrs = get_attrs_stan3(prior, model=prior_model)
        return dict_to_dataset(
            data, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
        )

    @requires("prior")
    def sample_stats_prior_to_xarray(self):
        """Extract sample_stats_prior from prior."""
        prior = self.prior
        prior_model = self.prior_model
        data = get_sample_stats_stan3(prior, dtypes=self.dtypes)
        attrs = get_attrs_stan3(prior, model=prior_model)
        return dict_to_dataset(
            data, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
        )

    @requires("prior")
    @requires("prior_predictive")
    def prior_predictive_to_xarray(self):
        """Convert prior_predictive samples to xarray."""
        prior = self.prior
        prior_model = self.prior_model
        prior_predictive = self.prior_predictive
        data = get_draws_stan3(
            prior, model=prior_model, variables=prior_predictive, dtypes=self.dtypes
        )
        return dict_to_dataset(data, library=self.stan, coords=self.coords, dims=self.dims)

    @requires("posterior_model")
    @requires(["observed_data", "constant_data"])
    def observed_and_constant_data_to_xarray(self):
        """Convert observed data to xarray."""
        posterior_model = self.posterior_model
        if self.dims is None:
            dims = {}
        else:
            dims = self.dims
        obs_const_dict = {}
        for group_name in ("observed_data", "constant_data"):
            names = getattr(self, group_name)
            if names is None:
                continue
            names = [names] if isinstance(names, str) else names
            data = OrderedDict()
            for key in names:
                vals = np.atleast_1d(posterior_model.data[key])
                val_dims = dims.get(key)
                val_dims, coords = generate_dims_coords(
                    vals.shape, key, dims=val_dims, coords=self.coords
                )
                data[key] = xr.DataArray(vals, dims=val_dims, coords=coords)
            obs_const_dict[group_name] = xr.Dataset(
                data_vars=data, attrs=make_attrs(library=self.stan)
            )
        return obs_const_dict

    @requires("posterior_model")
    @requires("predictions_constant_data")
    def predictions_constant_data_to_xarray(self):
        """Convert observed data to xarray."""
        posterior_model = self.posterior_model
        if self.dims is None:
            dims = {}
        else:
            dims = self.dims
        names = self.predictions_constant_data
        names = [names] if isinstance(names, str) else names
        data = OrderedDict()
        for key in names:
            vals = np.atleast_1d(posterior_model.data[key])
            val_dims = dims.get(key)
            val_dims, coords = generate_dims_coords(
                vals.shape, key, dims=val_dims, coords=self.coords
            )
            data[key] = xr.DataArray(vals, dims=val_dims, coords=coords)
        return xr.Dataset(data_vars=data, attrs=make_attrs(library=self.stan))

    def to_inference_data(self):
        """Convert all available data to an InferenceData object.

        Note that if groups can not be created (i.e., there is no `fit`, so
        the `posterior` and `sample_stats` can not be extracted), then the InferenceData
        will not have those groups.
        """
        obs_const_dict = self.observed_and_constant_data_to_xarray()
        predictions_const_data = self.predictions_constant_data_to_xarray()
        return InferenceData(
            **{
                "posterior": self.posterior_to_xarray(),
                "sample_stats": self.sample_stats_to_xarray(),
                "log_likelihood": self.log_likelihood_to_xarray(),
                "posterior_predictive": self.posterior_predictive_to_xarray(),
                "predictions": self.predictions_to_xarray(),
                "prior": self.prior_to_xarray(),
                "sample_stats_prior": self.sample_stats_prior_to_xarray(),
                "prior_predictive": self.prior_predictive_to_xarray(),
                **({} if obs_const_dict is None else obs_const_dict),
                **(
                    {}
                    if predictions_const_data is None
                    else {"predictions_constant_data": predictions_const_data}
                ),
            }
        )


def get_draws(fit, variables=None, ignore=None, warmup=False, dtypes=None):
    """Extract draws from PyStan fit."""
    if ignore is None:
        ignore = []
    if fit.mode == 1:
        msg = "Model in mode 'test_grad'. Sampling is not conducted."
        raise AttributeError(msg)

    if fit.mode == 2 or fit.sim.get("samples") is None:
        msg = "Fit doesn't contain samples."
        raise AttributeError(msg)

    if dtypes is None:
        dtypes = {}

    dtypes = {**infer_dtypes(fit), **dtypes}

    if variables is None:
        variables = fit.sim["pars_oi"]
    elif isinstance(variables, str):
        variables = [variables]
    variables = list(variables)

    for var, dim in zip(fit.sim["pars_oi"], fit.sim["dims_oi"]):
        if var in variables and np.prod(dim) == 0:
            del variables[variables.index(var)]

    ndraws_warmup = fit.sim["warmup2"]
    if max(ndraws_warmup) == 0:
        warmup = False
    ndraws = [s - w for s, w in zip(fit.sim["n_save"], ndraws_warmup)]
    nchain = len(fit.sim["samples"])

    # check if the values are in 0-based (<=2.17) or 1-based indexing (>=2.18)
    shift = 1
    if any(dim and np.prod(dim) != 0 for dim in fit.sim["dims_oi"]):
        # choose variable with lowest number of dims > 1
        par_idx = min(
            (dim, i) for i, dim in enumerate(fit.sim["dims_oi"]) if (dim and np.prod(dim) != 0)
        )[1]
        offset = int(sum(map(np.product, fit.sim["dims_oi"][:par_idx])))
        par_offset = int(np.product(fit.sim["dims_oi"][par_idx]))
        par_keys = fit.sim["fnames_oi"][offset : offset + par_offset]
        shift = len(par_keys)
        for item in par_keys:
            _, shape = item.replace("]", "").split("[")
            shape_idx_min = min(int(shape_value) for shape_value in shape.split(","))
            if shape_idx_min < shift:
                shift = shape_idx_min
        # If shift is higher than 1, this will probably mean that Stan
        # has implemented sparse structure (saves only non-zero parts),
        # but let's hope that dims are still corresponding to the full shape
        shift = int(min(shift, 1))

    var_keys = OrderedDict((var, []) for var in fit.sim["pars_oi"])
    for key in fit.sim["fnames_oi"]:
        var, *tails = key.split("[")
        loc = [Ellipsis]
        for tail in tails:
            loc = []
            for i in tail[:-1].split(","):
                loc.append(int(i) - shift)
        var_keys[var].append((key, loc))

    shapes = dict(zip(fit.sim["pars_oi"], fit.sim["dims_oi"]))

    variables = [var for var in variables if var not in ignore]

    data = OrderedDict()
    data_warmup = OrderedDict()

    for var in variables:
        if var in data:
            continue
        keys_locs = var_keys.get(var, [(var, [Ellipsis])])
        shape = shapes.get(var, [])
        dtype = dtypes.get(var)

        ndraw = max(ndraws)
        ary_shape = [nchain, ndraw] + shape
        ary = np.empty(ary_shape, dtype=dtype, order="F")

        if warmup:
            nwarmup = max(ndraws_warmup)
            ary_warmup_shape = [nchain, nwarmup] + shape
            ary_warmup = np.empty(ary_warmup_shape, dtype=dtype, order="F")

        for chain, (pyholder, ndraw, ndraw_warmup) in enumerate(
            zip(fit.sim["samples"], ndraws, ndraws_warmup)
        ):
            axes = [chain, slice(None)]
            for key, loc in keys_locs:
                ary_slice = tuple(axes + loc)
                ary[ary_slice] = pyholder.chains[key][-ndraw:]
                if warmup:
                    ary_warmup[ary_slice] = pyholder.chains[key][:ndraw_warmup]
        data[var] = ary
        if warmup:
            data_warmup[var] = ary_warmup
    return data, data_warmup


def get_sample_stats(fit, warmup=False, dtypes=None):
    """Extract sample stats from PyStan fit."""
    if dtypes is None:
        dtypes = {}
    dtypes = {"divergent__": bool, "n_leapfrog__": np.int64, "treedepth__": np.int64, **dtypes}

    rename_dict = {
        "divergent": "diverging",
        "n_leapfrog": "n_steps",
        "treedepth": "tree_depth",
        "stepsize": "step_size",
        "accept_stat": "acceptance_rate",
    }

    ndraws_warmup = fit.sim["warmup2"]
    if max(ndraws_warmup) == 0:
        warmup = False
    ndraws = [s - w for s, w in zip(fit.sim["n_save"], ndraws_warmup)]

    extraction = OrderedDict()
    extraction_warmup = OrderedDict()
    for chain, (pyholder, ndraw, ndraw_warmup) in enumerate(
        zip(fit.sim["samples"], ndraws, ndraws_warmup)
    ):
        if chain == 0:
            for key in pyholder["sampler_param_names"]:
                extraction[key] = []
                if warmup:
                    extraction_warmup[key] = []
        for key, values in zip(pyholder["sampler_param_names"], pyholder["sampler_params"]):
            extraction[key].append(values[-ndraw:])
            if warmup:
                extraction_warmup[key].append(values[:ndraw_warmup])

    data = OrderedDict()
    for key, values in extraction.items():
        values = np.stack(values, axis=0)
        dtype = dtypes.get(key)
        values = values.astype(dtype)
        name = re.sub("__$", "", key)
        name = rename_dict.get(name, name)
        data[name] = values

    data_warmup = OrderedDict()
    if warmup:
        for key, values in extraction_warmup.items():
            values = np.stack(values, axis=0)
            values = values.astype(dtypes.get(key))
            name = re.sub("__$", "", key)
            name = rename_dict.get(name, name)
            data_warmup[name] = values

    return data, data_warmup


def get_attrs(fit):
    """Get attributes from PyStan fit object."""
    attrs = {}

    try:
        attrs["args"] = [deepcopy(holder.args) for holder in fit.sim["samples"]]
    except Exception as exp:  # pylint: disable=broad-except
        _log.warning("Failed to fetch args from fit: %s", exp)
    if "args" in attrs:
        for arg in attrs["args"]:
            if isinstance(arg["init"], bytes):
                arg["init"] = arg["init"].decode("utf-8")
        attrs["args"] = json.dumps(attrs["args"])
    try:
        attrs["inits"] = [holder.inits for holder in fit.sim["samples"]]
    except Exception as exp:  # pylint: disable=broad-except
        _log.warning("Failed to fetch `args` from fit: %s", exp)
    else:
        attrs["inits"] = json.dumps(attrs["inits"])

    attrs["step_size"] = []
    attrs["metric"] = []
    attrs["inv_metric"] = []
    for holder in fit.sim["samples"]:
        try:
            step_size = float(
                re.search(
                    r"step\s*size\s*=\s*([0-9]+.?[0-9]+)\s*",
                    holder.adaptation_info,
                    flags=re.IGNORECASE,
                ).group(1)
            )
        except AttributeError:
            step_size = np.nan
        attrs["step_size"].append(step_size)

        inv_metric_match = re.search(
            r"mass matrix:\s*(.*)\s*$", holder.adaptation_info, flags=re.DOTALL
        )
        if inv_metric_match:
            inv_metric_str = inv_metric_match.group(1)
            if "Diagonal elements of inverse mass matrix" in holder.adaptation_info:
                metric = "diag_e"
                inv_metric = [float(item) for item in inv_metric_str.strip(" #\n").split(",")]
            else:
                metric = "dense_e"
                inv_metric = [
                    list(map(float, item.split(",")))
                    for item in re.sub(r"#\s", "", inv_metric_str).splitlines()
                ]
        else:
            metric = "unit_e"
            inv_metric = None

        attrs["metric"].append(metric)
        attrs["inv_metric"].append(inv_metric)
    attrs["inv_metric"] = json.dumps(attrs["inv_metric"])

    if not attrs["step_size"]:
        del attrs["step_size"]

    attrs["adaptation_info"] = fit.get_adaptation_info()
    attrs["stan_code"] = fit.get_stancode()

    return attrs


def get_draws_stan3(fit, model=None, variables=None, ignore=None, dtypes=None):
    """Extract draws from PyStan3 fit."""
    if ignore is None:
        ignore = []

    if dtypes is None:
        dtypes = {}

    if model is not None:
        dtypes = {**infer_dtypes(fit, model), **dtypes}

    if variables is None:
        variables = fit.param_names
    elif isinstance(variables, str):
        variables = [variables]
    variables = list(variables)

    data = OrderedDict()

    for var in variables:
        if var in ignore:
            continue
        if var in data:
            continue
        dtype = dtypes.get(var)

        # in future fix the correct number of draws if fit.save_warmup is True
        new_shape = (*fit.dims[fit.param_names.index(var)], -1, fit.num_chains)
        if 0 in new_shape:
            continue
        values = fit._draws[fit._parameter_indexes(var), :]  # pylint: disable=protected-access
        values = values.reshape(new_shape, order="F")
        values = np.moveaxis(values, [-2, -1], [1, 0])
        values = values.astype(dtype)
        data[var] = values

    return data


def get_sample_stats_stan3(fit, variables=None, ignore=None, dtypes=None):
    """Extract sample stats from PyStan3 fit."""
    if dtypes is None:
        dtypes = {}
    dtypes = {"divergent__": bool, "n_leapfrog__": np.int64, "treedepth__": np.int64, **dtypes}

    rename_dict = {
        "divergent": "diverging",
        "n_leapfrog": "n_steps",
        "treedepth": "tree_depth",
        "stepsize": "step_size",
        "accept_stat": "acceptance_rate",
    }

    if isinstance(variables, str):
        variables = [variables]
    if isinstance(ignore, str):
        ignore = [ignore]

    data = OrderedDict()
    for key in fit.sample_and_sampler_param_names:
        if (variables and key not in variables) or (ignore and key in ignore):
            continue
        new_shape = -1, fit.num_chains
        values = fit._draws[fit._parameter_indexes(key)]  # pylint: disable=protected-access
        values = values.reshape(new_shape, order="F")
        values = np.moveaxis(values, [-2, -1], [1, 0])
        dtype = dtypes.get(key)
        values = values.astype(dtype)
        name = re.sub("__$", "", key)
        name = rename_dict.get(name, name)
        data[name] = values

    return data


def get_attrs_stan3(fit, model=None):
    """Get attributes from PyStan3 fit and model object."""
    attrs = {}
    for key in ["num_chains", "num_samples", "num_thin", "num_warmup", "save_warmup"]:
        try:
            attrs[key] = getattr(fit, key)
        except AttributeError as exp:
            _log.warning("Failed to access attribute %s in fit object %s", key, exp)

    if model is not None:
        for key in ["model_name", "program_code", "random_seed"]:
            try:
                attrs[key] = getattr(model, key)
            except AttributeError as exp:
                _log.warning("Failed to access attribute %s in model object %s", key, exp)

    return attrs


def infer_dtypes(fit, model=None):
    """Infer dtypes from Stan model code.

    Function strips out generated quantities block and searches for `int`
    dtypes after stripping out comments inside the block.
    """
    if model is None:
        stan_code = fit.get_stancode()
        model_pars = fit.model_pars
    else:
        stan_code = model.program_code
        model_pars = fit.param_names

    dtypes = {key: item for key, item in infer_stan_dtypes(stan_code).items() if key in model_pars}
    return dtypes


# pylint disable=too-many-instance-attributes
def from_pystan(
    posterior=None,
    *,
    posterior_predictive=None,
    predictions=None,
    prior=None,
    prior_predictive=None,
    observed_data=None,
    constant_data=None,
    predictions_constant_data=None,
    log_likelihood=None,
    coords=None,
    dims=None,
    posterior_model=None,
    prior_model=None,
    save_warmup=None,
    dtypes=None,
):
    """Convert PyStan data into an InferenceData object.

    For a usage example read the
    :ref:`Creating InferenceData section on from_pystan <creating_InferenceData>`

    Parameters
    ----------
    posterior : StanFit4Model or stan.fit.Fit
        PyStan fit object for posterior.
    posterior_predictive : str, a list of str
        Posterior predictive samples for the posterior.
    predictions : str, a list of str
        Out-of-sample predictions for the posterior.
    prior : StanFit4Model or stan.fit.Fit
        PyStan fit object for prior.
    prior_predictive : str, a list of str
        Posterior predictive samples for the prior.
    observed_data : str or a list of str
        observed data used in the sampling.
        Observed data is extracted from the `posterior.data`.
        PyStan3 needs model object for the extraction.
        See `posterior_model`.
    constant_data : str or list of str
        Constants relevant to the model (i.e. x values in a linear
        regression).
    predictions_constant_data : str or list of str
        Constants relevant to the model predictions (i.e. new x values in a linear
        regression).
    log_likelihood : dict of {str: str}, list of str or str, optional
        Pointwise log_likelihood for the data. log_likelihood is extracted from the
        posterior. It is recommended to use this argument as a dictionary whose keys
        are observed variable names and its values are the variables storing log
        likelihood arrays in the Stan code. In other cases, a dictionary with keys
        equal to its values is used. By default, if a variable ``log_lik`` is
        present in the Stan model, it will be retrieved as pointwise log
        likelihood values. Use ``False`` or set ``data.log_likelihood`` to
        false to avoid this behaviour.
    coords : dict[str, iterable]
        A dictionary containing the values that are used as index. The key
        is the name of the dimension, the values are the index values.
    dims : dict[str, List(str)]
        A mapping from variables to a list of coordinate names for the variable.
    posterior_model : stan.model.Model
        PyStan3 specific model object. Needed for automatic dtype parsing
        and for the extraction of observed data.
    prior_model : stan.model.Model
        PyStan3 specific model object. Needed for automatic dtype parsing.
    save_warmup : bool
        Save warmup iterations into InferenceData object. If not defined, use default
        defined by the rcParams. Not supported in PyStan3.
    dtypes: dict
        A dictionary containing dtype information (int, float) for parameters.
        By default dtype information is extracted from the model code.
        Model code is extracted from fit object in PyStan 2 and from model object
        in PyStan 3.

    Returns
    -------
    InferenceData object
    """
    check_posterior = (posterior is not None) and (type(posterior).__module__ == "stan.fit")
    check_prior = (prior is not None) and (type(prior).__module__ == "stan.fit")
    if check_posterior or check_prior:
        if save_warmup:
            warnings.warn(
                "save_warmup is not currently supported for PyStan3",
                UserWarning,
            )
        return PyStan3Converter(
            posterior=posterior,
            posterior_model=posterior_model,
            posterior_predictive=posterior_predictive,
            predictions=predictions,
            prior=prior,
            prior_model=prior_model,
            prior_predictive=prior_predictive,
            observed_data=observed_data,
            constant_data=constant_data,
            predictions_constant_data=predictions_constant_data,
            log_likelihood=log_likelihood,
            coords=coords,
            dims=dims,
            dtypes=dtypes,
        ).to_inference_data()
    else:
        return PyStanConverter(
            posterior=posterior,
            posterior_predictive=posterior_predictive,
            predictions=predictions,
            prior=prior,
            prior_predictive=prior_predictive,
            observed_data=observed_data,
            constant_data=constant_data,
            predictions_constant_data=predictions_constant_data,
            log_likelihood=log_likelihood,
            coords=coords,
            dims=dims,
            save_warmup=save_warmup,
            dtypes=dtypes,
        ).to_inference_data()

arviz/data/io_pystan.py
diff --git a/arviz/data/io_pystan.py b/arviz/data/io_pystan.py
index 0de5db3b1d..2401c7084e 100644
--- a/arviz/data/io_pystan.py
+++ b/arviz/data/io_pystan.py
@@ -809,6 +809,8 @@ def get_draws_stan3(fit, model=None, variables=None, ignore=None, dtypes=None):
     data = OrderedDict()
 
     for var in variables:
+        if var in ignore:
+            continue
         if var in data:
             continue
         dtype = dtypes.get(var)
diff --git a/arviz/tests/external_tests/test_data_pystan.py b/arviz/tests/external_tests/test_data_pystan.py
index feb82d1a9a..c2aa9322d9 100644
--- a/arviz/tests/external_tests/test_data_pystan.py
+++ b/arviz/tests/external_tests/test_data_pystan.py
@@ -153,7 +153,7 @@ def test_inference_data(self, data, eight_schools_params):
         inference_data5 = self.get_inference_data5(data)
         # inference_data 1
         test_dict = {
-            "posterior": ["theta"],
+            "posterior": ["theta", "~log_lik"],
             "posterior_predictive": ["y_hat"],
             "predictions": ["y_hat"],
             "observed_data": ["y"],

from_pystan puts log likelihood and posterior predictions in posterior data
## Short Description

I would like to convert a posterior sample from a PyStan `stan.fit.Fit` object to an `az.InferenceData` object, but the log likelihood and posterior prediction variables end up in the `posterior` section. (They also appear in their respective sections of the `InferenceData` object). I was wondering if there is a way to have them not be in the `posterior` section. This would make using them with other ArviZ functions much easier because then they will *not* be included in tables or plots of the *posterior distributions* such as `az.plot_trace(data)`.

## Code Example or link

The example is from the PyStan section of the *Creating InferenceData* guide: https://arviz-devs.github.io/arviz/getting_started/CreatingInferenceData.html#from-pystan

> Note that there are some changes to reflect changes in the PyStan API (such as importing `stan` instead of `pystan`.

```python
import arviz as az
import numpy as np
import stan

schools_code = """
data {
    int<lower=0> J;
    real y[J];
    real<lower=0> sigma[J];
}

parameters {
    real mu;
    real<lower=0> tau;
    real theta_tilde[J];
}

transformed parameters {
    real theta[J];
    for (j in 1:J)
        theta[j] = mu + tau * theta_tilde[j];
}

model {
    mu ~ normal(0, 5);
    tau ~ cauchy(0, 5);
    theta_tilde ~ normal(0, 1);
    y ~ normal(theta, sigma);
}

generated quantities {
    vector[J] log_lik;
    vector[J] y_hat;
    for (j in 1:J) {
        log_lik[j] = normal_lpdf(y[j] | theta[j], sigma[j]);
        y_hat[j] = normal_rng(theta[j], sigma[j]);
    }
}
"""

eight_school_data = {
    "J": 8,
    "y": np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0]),
    "sigma": np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0]),
}

stan_model = stan.build(schools_code, data=eight_school_data)
fit = stan_model.sample()

stan_data = az.from_pystan(
    posterior=fit,
    posterior_predictive="y_hat",
    posterior_model= stan_model,
    observed_data=["y"],
    log_likelihood={"y": "log_lik"},
    coords={"school": np.arange(eight_school_data["J"])},
    dims={
        "theta": ["school"],
        "y": ["school"],
        "log_lik": ["school"],
        "y_hat": ["school"],
        "theta_tilde": ["school"],
    },
)

stan_data
#> Inference data with groups:
#>         > posterior
#>         > posterior_predictive
#>         > log_likelihood
#>         > sample_stats

stan_data.posterior
#> <xarray.Dataset>
#> Dimensions:      (chain: 4, draw: 1000, school: 8)
#> Coordinates:
#>   * chain        (chain) int64 0 1 2 3
#>   * draw         (draw) int64 0 1 2 3 4 5 6 7 ... 993 994 995 996 997 998 999
#>   * school       (school) int64 0 1 2 3 4 5 6 7
#> Data variables:
#>     mu           (chain, draw) float64 5.583 13.74 -1.643 ... 7.135 1.818 2.276
#>     tau          (chain, draw) float64 3.375 2.563 2.734 ... 1.585 0.2655 8.359
#>     theta_tilde  (chain, draw, school) float64 0.9058 0.2812 ... 1.286 -0.5058
#>     theta        (chain, draw, school) float64 8.64 6.531 6.381 ... 13.03 -1.952
#>     log_lik      (chain, draw, school) float64 -4.46 -3.232 ... -3.345 -4.11
#>     y_hat        (chain, draw, school) float64 -15.09 7.846 ... 19.34 8.089
#> Attributes:
#>     created_at:                 2022-01-13T22:30:54.799215
#>     arviz_version:              0.11.4
#>     inference_library:          stan
#>     inference_library_version:  3.3.0
#>     num_chains:                 4
#>     num_samples:                1000
#>     num_thin:                   1
#>     num_warmup:                 1000
#>     save_warmup:                0
#>     model_name:                 models/6kxc6kz6
#>     program_code:               \ndata {\n    int<lower=0> J;\n    real y[J];...
#>     random_seed:                None
```

PyStan version: 3.3.0
ArviZ version: 0.11.4

## Relevant documentation or public examples

- Guide on creating InferenceData objects from PyStan: https://arviz-devs.github.io/arviz/getting_started/CreatingInferenceData.html#from-pystan
- documentation for `from_pystan()`: https://arviz-devs.github.io/arviz/api/generated/arviz.from_pystan.html#arviz-from-pystan
