"""Predictive check using densities."""

import warnings
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 rcParams
from arviz_base.labels import BaseLabeller
from arviz_base.validate import (
    validate_dict_argument,
    validate_or_use_rcparam,
    validate_sample_dims,
)

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,
)
from arviz_plots.plots.utils_plot_types import warn_if_binary, warn_if_discrete


def plot_ppc_dist(
    dt,
    *,
    var_names=None,
    filter_vars=None,
    group="posterior_predictive",
    coords=None,
    sample_dims=None,
    kind=None,
    num_samples=50,
    plot_collection=None,
    backend=None,
    labeller=None,
    aes_by_visuals: Mapping[
        Literal["predictive_dist", "observed_dist", "title"], Sequence[str]
    ] = None,
    visuals: Mapping[
        Literal["predictive_dist", "observed_dist", "title", "remove_axis"],
        Mapping[str, Any] | bool,
    ] = None,
    stats: Mapping[
        Literal["predictive_dist", "observed_dist"], Mapping[str, Any] | xr.Dataset
    ] = None,
    **pc_kwargs,
):
    """
    Plot 1D marginals for the predictive distribution and the observed data.

    Parameters
    ----------
    dt : DataTree
        If group is "posterior_predictive", it should contain the ``posterior_predictive`` and
        ``observed_data`` groups. If group is "prior_predictive", it should contain the
        ``prior_predictive`` group.
    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 : str,
        Group to be plotted. Defaults to "posterior_predictive".
        It could also be "prior_predictive".
    coords : dict, optional
    sample_dims : str or sequence of hashable, optional
        Sampled dimensions used to overlay `num_samples` lines.
        Defaults to ``rcParams["data.sample_dims"]``
    kind : {"auto", "kde", "hist", "ecdf", "dot"}, optional
        How to represent the marginal density. Defaults to ``rcParams["plot.density_kind"]``
        If "dot" is selected, only the top points of the predictive draws are shown.
    num_samples : int, optional
        Number of samples to plot. Defaults to 100.
    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. Valid keys are the same as for `visuals` except "remove_axis".

        With a single model, no aesthetic mappings are generated by default,
        each variable+coord combination gets a :term:`plot` but they all look the same,
        unless there are user provided aesthetic mappings.
        With multiple models, ``plot_dist`` maps "color" and "y" to the "model" dimension.

        By default, all aesthetics but "y" are mapped to the distribution representation,
        and if multiple models are present, "color" and "y" are mapped to the
        credible interval and the point estimate.

        When "point_estimate" key is provided but "point_estimate_text" isn't,
        the values assigned to the first are also used for the second.
    visuals : mapping of {str : mapping or bool}, optional
        Valid keys are:

        * predictive_dist, observed_dist -> passed to a function that depends on
          the `kind` argument.

          * "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`
        * remove_axis -> not passed anywhere, can only be ``False`` to skip calling this function

        observed_dist defaults to False, no observed data is plotted, if group is
        "prior_predictive".

    stats : mapping, optional
        Valid keys are:

        * predictive_dist, observed_dist -> passed to kde, ecdf, ...

    **pc_kwargs
        Passed to :meth:`~arviz_plots.PlotCollection.wrap`

    Returns
    -------
    PlotCollection

    See Also
    --------
    :ref:`plots_intro` :
        General introduction to batteries-included plotting functions, common use and logic overview

    Examples
    --------
    Make a plot of the posterior predictive distribution vs the observed data.
    We used an ECDF representation customized the colors.

    .. plot::
        :context: close-figs

        >>> from arviz_plots import plot_ppc_dist, style
        >>> style.use("arviz-variat")
        >>> from arviz_base import load_arviz_data
        >>> radon = load_arviz_data('radon')
        >>> pc = plot_ppc_dist(
        >>>     radon,
        >>>     kind="ecdf",
        >>>     visuals={
        >>>         "predictive_dist": {"color":"C1"},
        >>>         "observed_dist": {"color":"C3"}
        >>>     },
        >>> )


    Faceting and aesthetics mappings happen on unique coordinate values. If there are repeated
    coordinate values they will be grouped and reduced along with `sample_dims`.
    This example updates the coordinate values to have repeated values and requests
    faceting along the "obs_id" dimension. It also keeps 90 out of the 919 observations
    in the original dataset; otherwise we'd end up with 85 :term:`plots` in the :term:`figure`

    .. plot::
        :context: close-figs

        >>> county = radon.constant_data["County"][radon.constant_data["county_idx"]]
        >>> reindexed_dt = radon.filter(
        >>>     lambda node: node.name in ("observed_data", "posterior_predictive")
        >>> ).map_over_datasets(
        >>>     lambda node: node.assign_coords(obs_id=county).isel(obs_id=slice(None, 90))
        >>> )
        >>> pc = azp.plot_ppc_dist(reindexed_dt, cols=["obs_id"], kind="auto")

    Note how counties with a lot of observations have a smoother ECDF whereas counties
    with only 2-3 observations have only 2-3 steps in their ECDF.

    .. minigallery:: plot_ppc_dist
    """
    kind = validate_or_use_rcparam(kind, "plot.density_kind")
    stats = validate_dict_argument(stats, (plot_ppc_dist, "stats"))
    visuals = validate_dict_argument(visuals, (plot_ppc_dist, "visuals"))

    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")

    rng = np.random.default_rng(4214)

    predictive_dist = process_group_variables_coords(
        dt, group=group, var_names=var_names, filter_vars=filter_vars, coords=coords
    )
    sample_dims = validate_sample_dims(sample_dims, data=predictive_dist)
    pp_dims = [dims for dims in dt[group].dims if dims not in sample_dims]

    if "observed_data" in dt:
        observed_dist = process_group_variables_coords(
            dt,
            group="observed_data",
            var_names=var_names,
            filter_vars=filter_vars,
            coords=coords,
        )
    else:
        observed_dist = None

    warn_if_binary(observed_dist, predictive_dist)
    warn_if_discrete(observed_dist, predictive_dist, kind)

    # Select a random subset of samples
    n_pp_samples = np.prod(
        [predictive_dist.sizes[dim] for dim in sample_dims if dim in predictive_dist.dims]
    )
    if num_samples > n_pp_samples:
        num_samples = n_pp_samples
        warnings.warn("num_samples is larger than the number of predictive samples.")

    if kind == "dot":
        stats.setdefault("predictive_dist", {"top_only": True})

    pp_sample_ix = rng.choice(n_pp_samples, size=num_samples, replace=False)
    predictive_dist = predictive_dist.stack(sample=sample_dims).isel(sample=pp_sample_ix)

    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("overlay_ppc", ["sample"])
        pc_kwargs.setdefault("cols", "__variable__")

        pc_kwargs = set_wrap_layout(pc_kwargs, plot_bknd, predictive_dist)

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

    aes_by_visuals = validate_dict_argument(aes_by_visuals, (plot_ppc_dist, "aes_by_visuals"))
    if labeller is None:
        labeller = BaseLabeller()

    # Plot the predictive density
    pred_density_kwargs = get_visual_kwargs(visuals, "predictive_dist")
    if pred_density_kwargs is not False:
        # We don't want credible_interval or point_estimate for the predictive distribution
        pred_density_kwargs.setdefault("alpha", 0.3)
        pred_visuals = {
            "dist": pred_density_kwargs,
            "credible_interval": False,
            "point_estimate": False,
            "point_estimate_text": False,
            "title": visuals.get("title", {}),
            "rug": False,
            "remove_axis": visuals.get("remove_axis", False),
        }
        pred_aes_by_visuals = {
            k.replace("predictive_", ""): v
            for k, v in aes_by_visuals.items()
            if k != "observed_dist"
        }

        plot_collection = plot_dist(
            predictive_dist,
            group=group,
            sample_dims=pp_dims,
            kind=kind,
            visuals=pred_visuals,
            aes_by_visuals=pred_aes_by_visuals,
            pc_kwargs=pc_kwargs,
            plot_collection=plot_collection,
            stats={"dist": stats.get("predictive_dist", {})},
        )
        plot_collection.rename_visuals(dist="predictive_dist")

    # Plot the observed density
    observed_density_kwargs = get_visual_kwargs(
        visuals, "observed_dist", False if group == "prior_predictive" else None
    )

    if observed_density_kwargs is not False:
        observed_density_kwargs.setdefault("color", "B1")
        observed_visuals = {
            "dist": observed_density_kwargs,
            "credible_interval": False,
            "point_estimate": False,
            "point_estimate_text": False,
            "title": False,
            "rug": False,
            "remove_axis": False,
        }
        obs_aes_by_visuals = (
            {"dist": aes_by_visuals["observed_dist"]} if "observed_dist" in aes_by_visuals else {}
        )
        plot_collection = plot_dist(
            observed_dist,
            group="observed_data",
            sample_dims=pp_dims,
            kind=kind,
            visuals=observed_visuals,
            aes_by_visuals=obs_aes_by_visuals,
            plot_collection=plot_collection,
            stats={"dist": stats.get("observed_dist", {})},
        )
        plot_collection.rename_visuals(dist="observed_dist")

    return plot_collection
