"""Contain functions for Bayes Factor plotting."""

from collections.abc import Mapping, Sequence
from importlib import import_module
from typing import Any, Literal

import numpy as np
import xarray as xr
from arviz_base import extract, rcParams
from arviz_base.validate import validate_dict_argument, validate_sample_dims
from xarray import concat

from arviz_plots.plot_collection import PlotCollection
from arviz_plots.plots.dist_plot import plot_dist
from arviz_plots.plots.utils import (
    get_visual_kwargs,
    process_group_variables_coords,
    set_wrap_layout,
)


def plot_prior_posterior(
    dt,
    *,
    var_names=None,
    filter_vars=None,
    group=None,  # pylint: disable=unused-argument
    coords=None,
    sample_dims=None,
    kind=None,
    plot_collection=None,
    backend=None,
    labeller=None,
    aes_by_visuals: Mapping[
        Literal[
            "dist",
            "credible_interval",
            "point_estimate",
            "point_estimate_text",
            "title",
            "rug",
        ],
        Sequence[str],
    ] = None,
    visuals: Mapping[
        Literal[
            "dist",
            "credible_interval",
            "point_estimate",
            "point_estimate_text",
            "title",
            "rug",
            "remove_axis",
        ],
        Mapping[str, Any] | bool,
    ] = None,
    stats: Mapping[
        Literal["dist", "credible_interval", "point_estimate"], Mapping[str, Any] | xr.Dataset
    ] = None,
    **pc_kwargs,
):
    r"""Plot 1D marginal densities for prior and posterior.

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

    Parameters
    ----------
    dt : DataTree or dict of {str : DataTree}
        Input data. In case of dictionary input, the keys are taken to be model names.
        In such cases, a dimension "model" is generated and can be used to map to aesthetics.
    var_names : str or list of str, optional
        One or more variables to be plotted.
        Prefix the variables by ~ when you want to exclude them from the plot.
    filter_vars : {None, "like", "regex"}, default=None
        If None, interpret var_names as the real variables names.
        If “like”, interpret var_names as substrings of the real variables names.
        If “regex”, interpret var_names as regular expressions on the real variables names.
    group : None
        This argument is ignored. Have it here for compatibility with other plotting functions.
    coords : dict, optional
    sample_dims : str or sequence of hashable, optional
        Dimensions to reduce unless mapped to an aesthetic.
        Defaults to ``rcParams["data.sample_dims"]``
    kind : {"kde", "hist", "dot", "ecdf"}, optional
        How to represent the marginal density.
        Defaults to ``rcParams["plot.density_kind"]``
    plot_collection : PlotCollection, optional
    backend : {"matplotlib", "bokeh"}, optional
    labeller : labeller, optional
    aes_by_visuals : mapping of {str : sequence of str}, optional
        Mapping of visuals to aesthetics that should use their mapping in `plot_collection`
        when plotted. The prior and posterior groups are combined creating a new
        dimension "group". By default, there is an aesthetic mapping from group to color.
        Valid keys are the same as for `visuals`.
    visuals : mapping of {str : mapping or bool}, optional
        Valid keys are:

        * dist -> depending on the value of `kind` passed to:

          * "kde" -> passed to :func:`~arviz_plots.visuals.line_xy`
          * "ecdf" -> passed to :func:`~arviz_plots.visuals.ecdf_line`
          * "hist" -> passed to :func: `~arviz_plots.visuals.step_hist`
          * "dot" -> passed to :func:`~arviz_plots.visuals.scatter_xy`

        * title -> passed to :func:`~arviz_plots.visuals.labelled_title`
        * legend -> passed to :class:`arviz_plots.PlotCollection.add_legend`

    stats : mapping, optional
        Valid keys are:

        * dist -> passed to kde, ecdf, ...

    **pc_kwargs
        Passed to :class:`arviz_plots.PlotCollection.wrap`

    Returns
    -------
    PlotCollection

    Examples
    --------
    Select two variables and plot them with an ecdf.

    .. plot::
        :context: close-figs

        >>> from arviz_plots import plot_prior_posterior, style
        >>> style.use("arviz-variat")
        >>> from arviz_base import load_arviz_data
        >>> dt = load_arviz_data('centered_eight')
        >>> plot_prior_posterior(dt, var_names=["mu", "tau"], kind="ecdf")


    .. minigallery:: plot_prior_posterior
    """
    aes_by_visuals = validate_dict_argument(aes_by_visuals, (plot_dist, "aes_by_visuals"))
    visuals = validate_dict_argument(visuals, (plot_dist, "visuals"))
    stats = validate_dict_argument(stats, (plot_dist, "stats"))

    if backend is None:
        if plot_collection is None:
            backend = rcParams["plot.backend"]
        else:
            backend = plot_collection.backend

    plot_bknd = import_module(f".backend.{backend}", package="arviz_plots")

    sample_dims_prior = validate_sample_dims(sample_dims, data=dt.prior)
    sample_dims = validate_sample_dims(sample_dims, data=dt.posterior)
    prior_size = np.prod([dt.prior.sizes[dim] for dim in sample_dims_prior])
    posterior_size = np.prod([dt.posterior.sizes[dim] for dim in sample_dims])
    num_samples = min(prior_size, posterior_size)

    ds_prior = extract(
        dt,
        group="prior",
        sample_dims=sample_dims_prior,
        combined=True,
        num_samples=num_samples,
        random_seed=0,
        keep_dataset=True,
    )
    prior_dims_drop = list(set(sample_dims_prior).union(ds_prior.attrs["sample_dims"]))
    sample_dims_prior = ds_prior.attrs["sample_dims"]
    ds_prior = ds_prior.drop_vars(prior_dims_drop).assign_coords(
        {sample_dims_prior[0]: np.arange(num_samples)}
    )
    ds_posterior = extract(
        dt,
        group="posterior",
        sample_dims=sample_dims,
        combined=True,
        num_samples=num_samples,
        random_seed=0,
        keep_dataset=True,
    )
    posterior_dims_drop = list(set(sample_dims).union(ds_posterior.attrs["sample_dims"]))
    sample_dims = ds_posterior.attrs["sample_dims"]
    ds_posterior = ds_posterior.drop_vars(posterior_dims_drop).assign_coords(
        {sample_dims[0]: np.arange(num_samples)}
    )

    distribution = concat([ds_prior, ds_posterior], dim="group").assign_coords(
        {"group": ["prior", "posterior"]}
    )

    distribution = process_group_variables_coords(
        distribution,
        group=None,
        var_names=var_names,
        filter_vars=filter_vars,
        coords=coords,
    )

    if plot_collection is None:
        pc_kwargs["figure_kwargs"] = pc_kwargs.get("figure_kwargs", {}).copy()

        pc_kwargs["aes"] = pc_kwargs.get("aes", {}).copy()
        pc_kwargs["aes"].setdefault("color", ["group"])
        pc_kwargs.setdefault("col_wrap", 4)
        pc_kwargs.setdefault(
            "cols",
            ["__variable__"]
            + [dim for dim in distribution.dims if dim not in sample_dims + ["group"]],
        )

        pc_kwargs = set_wrap_layout(pc_kwargs, plot_bknd, distribution)

        plot_collection = PlotCollection.wrap(
            distribution,
            backend=backend,
            **pc_kwargs,
        )

    visuals.setdefault("credible_interval", False)
    visuals.setdefault("point_estimate", False)
    visuals.setdefault("point_estimate_text", False)

    if kind == "hist":
        visuals.setdefault("dist", {})
        visuals.setdefault("remove_axis", True)

    plot_collection = plot_dist(
        distribution,
        var_names=None,
        group=None,
        coords=None,
        sample_dims=sample_dims,
        kind=kind,
        point_estimate=None,
        ci_kind=None,
        ci_prob=None,
        plot_collection=plot_collection,
        backend=backend,
        labeller=labeller,
        aes_by_visuals=aes_by_visuals,
        visuals=visuals,
        stats=stats,
        **pc_kwargs,
    )

    legend_kwargs = get_visual_kwargs(visuals, "legend")
    if legend_kwargs is not False:
        legend_kwargs.setdefault("dim", ["group"])
        plot_collection.add_legend(**legend_kwargs)

    return plot_collection
