"""Dictionary specific conversion code."""

import warnings
from typing import Optional

from ..rcparams import rcParams
from .base import dict_to_dataset, requires
from .inference_data import WARMUP_TAG, InferenceData


# pylint: disable=too-many-instance-attributes
class DictConverter:
    """Encapsulate Dictionary specific logic."""

    def __init__(
        self,
        *,
        posterior=None,
        posterior_predictive=None,
        predictions=None,
        sample_stats=None,
        log_likelihood=None,
        prior=None,
        prior_predictive=None,
        sample_stats_prior=None,
        observed_data=None,
        constant_data=None,
        predictions_constant_data=None,
        warmup_posterior=None,
        warmup_posterior_predictive=None,
        warmup_predictions=None,
        warmup_log_likelihood=None,
        warmup_sample_stats=None,
        save_warmup=None,
        index_origin=None,
        coords=None,
        dims=None,
        pred_dims=None,
        pred_coords=None,
        attrs=None,
        **kwargs,
    ):
        self.posterior = posterior
        self.posterior_predictive = posterior_predictive
        self.predictions = predictions
        self.sample_stats = sample_stats
        self.log_likelihood = log_likelihood
        self.prior = prior
        self.prior_predictive = prior_predictive
        self.sample_stats_prior = sample_stats_prior
        self.observed_data = observed_data
        self.constant_data = constant_data
        self.predictions_constant_data = predictions_constant_data
        self.warmup_posterior = warmup_posterior
        self.warmup_posterior_predictive = warmup_posterior_predictive
        self.warmup_predictions = warmup_predictions
        self.warmup_log_likelihood = warmup_log_likelihood
        self.warmup_sample_stats = warmup_sample_stats
        self.save_warmup = rcParams["data.save_warmup"] if save_warmup is None else save_warmup
        self.coords = (
            coords
            if pred_coords is None
            else pred_coords if coords is None else {**coords, **pred_coords}
        )
        self.index_origin = index_origin
        self.coords = coords
        self.dims = dims
        self.pred_dims = dims if pred_dims is None else pred_dims
        self.attrs = {} if attrs is None else attrs
        self.attrs.pop("created_at", None)
        self.attrs.pop("arviz_version", None)
        self._kwargs = kwargs

    def _init_dict(self, attr_name):
        dict_or_none = getattr(self, attr_name, {})
        return {} if dict_or_none is None else dict_or_none

    @requires(["posterior", f"{WARMUP_TAG}posterior"])
    def posterior_to_xarray(self):
        """Convert posterior samples to xarray."""
        data = self._init_dict("posterior")
        data_warmup = self._init_dict(f"{WARMUP_TAG}posterior")
        if not isinstance(data, dict):
            raise TypeError("DictConverter.posterior is not a dictionary")
        if not isinstance(data_warmup, dict):
            raise TypeError("DictConverter.warmup_posterior is not a dictionary")

        if "log_likelihood" in data:
            warnings.warn(
                "log_likelihood variable found in posterior group."
                " For stats functions log likelihood data needs to be in log_likelihood group.",
                UserWarning,
            )
        posterior_attrs = self._kwargs.get("posterior_attrs")
        posterior_warmup_attrs = self._kwargs.get("posterior_warmup_attrs")
        return (
            dict_to_dataset(
                data,
                library=None,
                coords=self.coords,
                dims=self.dims,
                attrs=posterior_attrs,
                index_origin=self.index_origin,
            ),
            dict_to_dataset(
                data_warmup,
                library=None,
                coords=self.coords,
                dims=self.dims,
                attrs=posterior_warmup_attrs,
                index_origin=self.index_origin,
            ),
        )

    @requires(["sample_stats", f"{WARMUP_TAG}sample_stats"])
    def sample_stats_to_xarray(self):
        """Convert sample_stats samples to xarray."""
        data = self._init_dict("sample_stats")
        data_warmup = self._init_dict(f"{WARMUP_TAG}sample_stats")
        if not isinstance(data, dict):
            raise TypeError("DictConverter.sample_stats is not a dictionary")
        if not isinstance(data_warmup, dict):
            raise TypeError("DictConverter.warmup_sample_stats is not a dictionary")

        if "log_likelihood" in data:
            warnings.warn(
                "log_likelihood variable found in sample_stats."
                " Storing log_likelihood data in sample_stats group will be deprecated in "
                "favour of storing them in the log_likelihood group.",
                PendingDeprecationWarning,
            )
        sample_stats_attrs = self._kwargs.get("sample_stats_attrs")
        sample_stats_warmup_attrs = self._kwargs.get("sample_stats_warmup_attrs")
        return (
            dict_to_dataset(
                data,
                library=None,
                coords=self.coords,
                dims=self.dims,
                attrs=sample_stats_attrs,
                index_origin=self.index_origin,
            ),
            dict_to_dataset(
                data_warmup,
                library=None,
                coords=self.coords,
                dims=self.dims,
                attrs=sample_stats_warmup_attrs,
                index_origin=self.index_origin,
            ),
        )

    @requires(["log_likelihood", f"{WARMUP_TAG}log_likelihood"])
    def log_likelihood_to_xarray(self):
        """Convert log_likelihood samples to xarray."""
        data = self._init_dict("log_likelihood")
        data_warmup = self._init_dict(f"{WARMUP_TAG}log_likelihood")
        if not isinstance(data, dict):
            raise TypeError("DictConverter.log_likelihood is not a dictionary")
        if not isinstance(data_warmup, dict):
            raise TypeError("DictConverter.warmup_log_likelihood is not a dictionary")
        log_likelihood_attrs = self._kwargs.get("log_likelihood_attrs")
        log_likelihood_warmup_attrs = self._kwargs.get("log_likelihood_warmup_attrs")
        return (
            dict_to_dataset(
                data,
                library=None,
                coords=self.coords,
                dims=self.dims,
                attrs=log_likelihood_attrs,
                index_origin=self.index_origin,
                skip_event_dims=True,
            ),
            dict_to_dataset(
                data_warmup,
                library=None,
                coords=self.coords,
                dims=self.dims,
                attrs=log_likelihood_warmup_attrs,
                index_origin=self.index_origin,
                skip_event_dims=True,
            ),
        )

    @requires(["posterior_predictive", f"{WARMUP_TAG}posterior_predictive"])
    def posterior_predictive_to_xarray(self):
        """Convert posterior_predictive samples to xarray."""
        data = self._init_dict("posterior_predictive")
        data_warmup = self._init_dict(f"{WARMUP_TAG}posterior_predictive")
        if not isinstance(data, dict):
            raise TypeError("DictConverter.posterior_predictive is not a dictionary")
        if not isinstance(data_warmup, dict):
            raise TypeError("DictConverter.warmup_posterior_predictive is not a dictionary")
        posterior_predictive_attrs = self._kwargs.get("posterior_predictive_attrs")
        posterior_predictive_warmup_attrs = self._kwargs.get("posterior_predictive_warmup_attrs")
        return (
            dict_to_dataset(
                data,
                library=None,
                coords=self.coords,
                dims=self.dims,
                attrs=posterior_predictive_attrs,
                index_origin=self.index_origin,
            ),
            dict_to_dataset(
                data_warmup,
                library=None,
                coords=self.coords,
                dims=self.dims,
                attrs=posterior_predictive_warmup_attrs,
                index_origin=self.index_origin,
            ),
        )

    @requires(["predictions", f"{WARMUP_TAG}predictions"])
    def predictions_to_xarray(self):
        """Convert predictions to xarray."""
        data = self._init_dict("predictions")
        data_warmup = self._init_dict(f"{WARMUP_TAG}predictions")
        if not isinstance(data, dict):
            raise TypeError("DictConverter.predictions is not a dictionary")
        if not isinstance(data_warmup, dict):
            raise TypeError("DictConverter.warmup_predictions is not a dictionary")
        predictions_attrs = self._kwargs.get("predictions_attrs")
        predictions_warmup_attrs = self._kwargs.get("predictions_warmup_attrs")
        return (
            dict_to_dataset(
                data,
                library=None,
                coords=self.coords,
                dims=self.pred_dims,
                attrs=predictions_attrs,
                index_origin=self.index_origin,
            ),
            dict_to_dataset(
                data_warmup,
                library=None,
                coords=self.coords,
                dims=self.pred_dims,
                attrs=predictions_warmup_attrs,
                index_origin=self.index_origin,
            ),
        )

    @requires("prior")
    def prior_to_xarray(self):
        """Convert prior samples to xarray."""
        data = self.prior
        if not isinstance(data, dict):
            raise TypeError("DictConverter.prior is not a dictionary")
        prior_attrs = self._kwargs.get("prior_attrs")
        return dict_to_dataset(
            data,
            library=None,
            coords=self.coords,
            dims=self.dims,
            attrs=prior_attrs,
            index_origin=self.index_origin,
        )

    @requires("sample_stats_prior")
    def sample_stats_prior_to_xarray(self):
        """Convert sample_stats_prior samples to xarray."""
        data = self.sample_stats_prior
        if not isinstance(data, dict):
            raise TypeError("DictConverter.sample_stats_prior is not a dictionary")
        sample_stats_prior_attrs = self._kwargs.get("sample_stats_prior_attrs")
        return dict_to_dataset(
            data,
            library=None,
            coords=self.coords,
            dims=self.dims,
            attrs=sample_stats_prior_attrs,
            index_origin=self.index_origin,
        )

    @requires("prior_predictive")
    def prior_predictive_to_xarray(self):
        """Convert prior_predictive samples to xarray."""
        data = self.prior_predictive
        if not isinstance(data, dict):
            raise TypeError("DictConverter.prior_predictive is not a dictionary")
        prior_predictive_attrs = self._kwargs.get("prior_predictive_attrs")
        return dict_to_dataset(
            data,
            library=None,
            coords=self.coords,
            dims=self.dims,
            attrs=prior_predictive_attrs,
            index_origin=self.index_origin,
        )

    def data_to_xarray(self, data, group, dims=None):
        """Convert data to xarray."""
        if not isinstance(data, dict):
            raise TypeError(f"DictConverter.{group} is not a dictionary")
        if dims is None:
            dims = {} if self.dims is None else self.dims
        return dict_to_dataset(
            data,
            library=None,
            coords=self.coords,
            dims=self.dims,
            default_dims=[],
            attrs=self.attrs,
            index_origin=self.index_origin,
        )

    @requires("observed_data")
    def observed_data_to_xarray(self):
        """Convert observed_data to xarray."""
        return self.data_to_xarray(self.observed_data, group="observed_data", dims=self.dims)

    @requires("constant_data")
    def constant_data_to_xarray(self):
        """Convert constant_data to xarray."""
        return self.data_to_xarray(self.constant_data, group="constant_data")

    @requires("predictions_constant_data")
    def predictions_constant_data_to_xarray(self):
        """Convert predictions_constant_data to xarray."""
        return self.data_to_xarray(
            self.predictions_constant_data, group="predictions_constant_data", dims=self.pred_dims
        )

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

        Note that if groups can not be created, then the InferenceData
        will not have those groups.
        """
        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(),
                "observed_data": self.observed_data_to_xarray(),
                "constant_data": self.constant_data_to_xarray(),
                "predictions_constant_data": self.predictions_constant_data_to_xarray(),
                "save_warmup": self.save_warmup,
                "attrs": self.attrs,
            }
        )


# pylint: disable=too-many-instance-attributes
def from_dict(
    posterior=None,
    *,
    posterior_predictive=None,
    predictions=None,
    sample_stats=None,
    log_likelihood=None,
    prior=None,
    prior_predictive=None,
    sample_stats_prior=None,
    observed_data=None,
    constant_data=None,
    predictions_constant_data=None,
    warmup_posterior=None,
    warmup_posterior_predictive=None,
    warmup_predictions=None,
    warmup_log_likelihood=None,
    warmup_sample_stats=None,
    save_warmup=None,
    index_origin: Optional[int] = None,
    coords=None,
    dims=None,
    pred_dims=None,
    pred_coords=None,
    attrs=None,
    **kwargs,
):
    """Convert Dictionary data into an InferenceData object.

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

    Parameters
    ----------
    posterior : dict, optional
    posterior_predictive : dict, optional
    predictions: dict, optional
    sample_stats : dict, optional
    log_likelihood : dict, optional
        For stats functions, log likelihood data should be stored here.
    prior : dict, optional
    prior_predictive : dict, optional
    observed_data : dict, optional
    constant_data : dict, optional
    predictions_constant_data: dict, optional
    warmup_posterior : dict, optional
    warmup_posterior_predictive : dict, optional
    warmup_predictions : dict, optional
    warmup_log_likelihood : dict, optional
    warmup_sample_stats : dict, optional
    save_warmup : bool, optional
        Save warmup iterations InferenceData object. If not defined, use default
        defined by the rcParams.
    index_origin : int, optional
    coords : dict of {str : list}, optional
        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 of {str : list of str}, optional
        A mapping from variables to a list of coordinate names for the variable.
    pred_dims : dict of {str : list of str}, optional
        A mapping from variables to a list of coordinate names for predictions.
    pred_coords : dict of {str : list}, optional
        A mapping from variables to a list of coordinate values for predictions.
    attrs : dict, optional
        A dictionary containing attributes for different groups.
    kwargs : dict, optional
        A dictionary containing group attrs. Accepted kwargs are:

        - posterior_attrs, posterior_warmup_attrs : attrs for posterior group
        - sample_stats_attrs, sample_stats_warmup_attrs : attrs for sample_stats group
        - log_likelihood_attrs, log_likelihood_warmup_attrs : attrs for log_likelihood group
        - posterior_predictive_attrs, posterior_predictive_warmup_attrs : attrs for
          posterior_predictive group
        - predictions_attrs, predictions_warmup_attrs : attrs for predictions group
        - prior_attrs : attrs for prior group
        - sample_stats_prior_attrs : attrs for sample_stats_prior group
        - prior_predictive_attrs : attrs for prior_predictive group

    Returns
    -------
    InferenceData
    """
    return DictConverter(
        posterior=posterior,
        posterior_predictive=posterior_predictive,
        predictions=predictions,
        sample_stats=sample_stats,
        log_likelihood=log_likelihood,
        prior=prior,
        prior_predictive=prior_predictive,
        sample_stats_prior=sample_stats_prior,
        observed_data=observed_data,
        constant_data=constant_data,
        predictions_constant_data=predictions_constant_data,
        warmup_posterior=warmup_posterior,
        warmup_posterior_predictive=warmup_posterior_predictive,
        warmup_predictions=warmup_predictions,
        warmup_log_likelihood=warmup_log_likelihood,
        warmup_sample_stats=warmup_sample_stats,
        save_warmup=save_warmup,
        index_origin=index_origin,
        coords=coords,
        dims=dims,
        pred_dims=pred_dims,
        pred_coords=pred_coords,
        attrs=attrs,
        **kwargs,
    ).to_inference_data()


from_pytree = from_dict
