"""Predictive intervals plot."""
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.utils import (
    filter_aes,
    get_visual_kwargs,
    process_group_variables_coords,
    set_wrap_layout,
)
from arviz_plots.visuals import ci_bound_y, labelled_title, labelled_x, labelled_y, point_y


def plot_ppc_interval(
    dt,
    *,
    var_names=None,
    filter_vars=None,
    group="posterior_predictive",
    coords=None,
    sample_dims=None,
    point_estimate=None,
    ci_kind=None,
    ci_probs=None,
    plot_collection=None,
    backend=None,
    labeller=None,
    aes_by_visuals: Mapping[
        Literal[
            "trunk",
            "twig",
            "observed_markers",
            "prediction_markers",
            "xlabel",
            "ylabel",
            "title",
        ],
        Sequence[str] | bool,
    ] = None,
    visuals: Mapping[
        Literal[
            "trunk",
            "twig",
            "observed_markers",
            "prediction_markers",
            "xlabel",
            "ylabel",
            "title",
        ],
        Mapping[str, Any] | bool,
    ] = None,
    stats: Mapping[
        Literal["trunk", "twig", "point_estimate"], Mapping[str, Any] | xr.Dataset
    ] = None,
    **pc_kwargs,
):
    """Plot posterior predictive intervals with observed data overlaid.

    Displays observed data as a point and predicted data as a point estimate plus two
    credible intervals.

    Parameters
    ----------
    dt : DataTree
        Input data. It should contain the ``posterior_predictive`` and
        ``observed_data`` groups.
    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
        Coordinates of `var_names` to be plotted.
    sample_dims : str or sequence of hashable, optional
        Dimensions to reduce unless mapped to an aesthetic.
        Defaults to ``rcParams["data.sample_dims"]``
    point_estimate : {"mean", "median", "mode"}, optional
        Which point estimate to plot for the predictive distribution.
        Defaults to rcParam ``stats.point_estimate``.
    ci_kind : {"hdi", "eti"}, optional
        Which credible interval to use. Defaults to ``rcParams["stats.ci_kind"]``.
    ci_probs : (float, float), optional
        Indicates the probabilities for the inner (twig) and outer (trunk) credible intervals.
        Defaults to ``(0.5, rcParams["stats.ci_prob"])``. It's assumed that
        ``ci_probs[0] < ci_probs[1]``, otherwise they are sorted.
    plot_collection : PlotCollection, optional
    backend : {"matplotlib", "bokeh", "plotly", "none"}, optional
    labeller : labeller, optional
    aes_by_visuals : mapping of {str : sequence of str or False}, optional
        Mapping of visuals to aesthetics that should use their mapping in `plot_collection`
        when plotted.
    visuals : mapping of {str : mapping or bool}, optional
        Valid keys are:

        * trunk, twig -> passed to :func:`~arviz_plots.visuals.ci_bound_y`
        * observed_markers -> passed to :func:`~arviz_plots.visuals.point_y`
        * prediction_markers -> passed to :func:`~arviz_plots.visuals.point_y`
        * xlabel -> passed to :func:`~arviz_plots.visuals.labelled_x`
        * ylabel -> passed to :func:`~arviz_plots.visuals.labelled_y`
        * title -> passed to :func:`~arviz_plots.visuals.labelled_title` defaults to False

    stats : mapping, optional
        Valid keys are:

        * trunk, twig -> passed to eti or hdi
        * point_estimate -> passed to mean, median or mode

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

    Returns
    -------
    PlotCollection

    See Also
    --------
    plot_loo_interval : Plot LOO posterior predictive intervals and observed data.
    plot_ppc_dist : Plot 1D marginals for the posterior/prior predictive and observed data.
    plot_ppc_rootogram : Plot ppc rootogram for discrete (count) data
    plot_forest : Plot forest plot for posterior/prior groups.

    Examples
    --------
    Plot posterior predictive intervals for the radon dataset, with custom styling.

    .. plot::
        :context: close-figs

        >>> from arviz_base import load_arviz_data
        >>> import arviz_plots as azp
        >>> azp.style.use("arviz-variat")
        >>> data = load_arviz_data("rugby")
        >>> pc = azp.plot_ppc_interval(data)

    .. minigallery:: plot_ppc_interval
    """
    point_estimate = validate_or_use_rcparam(point_estimate, "stats.point_estimate")
    ci_kind = validate_or_use_rcparam(ci_kind, "stats.ci_kind")
    aes_by_visuals = validate_dict_argument(aes_by_visuals, (plot_ppc_interval, "aes_by_visuals"))
    visuals = validate_dict_argument(visuals, (plot_ppc_interval, "visuals"))
    stats = validate_dict_argument(stats, (plot_ppc_interval, "stats"))
    if ci_probs is None:
        rc_ci_prob = rcParams["stats.ci_prob"]
        ci_probs = (0.5, rc_ci_prob)
    ci_probs = np.array(ci_probs)
    if ci_probs.size != 2:
        raise ValueError("ci_probs must have two elements for twig and trunk intervals.")
    if np.any(ci_probs < 0) or np.any(ci_probs > 1):
        raise ValueError("ci_probs must be between 0 and 1.")
    ci_probs.sort()

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

    if labeller is None:
        labeller = BaseLabeller()

    visuals.setdefault("title", False)

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

    ds_predictive = 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=ds_predictive)

    # Extract observed data
    if "observed_data" in dt:
        observed_data = process_group_variables_coords(
            dt,
            group="observed_data",
            var_names=var_names,
            filter_vars=filter_vars,
            coords=coords,
        )
    else:
        observed_data = None

    if ci_kind == "eti":
        ci_fun = ds_predictive.azstats.eti
    elif ci_kind == "hdi":
        ci_fun = ds_predictive.azstats.hdi

    ci_trunk = ci_fun(prob=ci_probs[1], dim=sample_dims, **stats.get("trunk", {}))
    ci_twig = ci_fun(prob=ci_probs[0], dim=sample_dims, **stats.get("twig", {}))

    pe_stats = stats.get("point_estimate", {})
    pe_stats.setdefault("round_to", "none")
    if point_estimate == "median":
        point = ds_predictive.azstats.median(dim=sample_dims, **pe_stats)
    elif point_estimate == "mean":
        point = ds_predictive.azstats.mean(dim=sample_dims, **pe_stats)
    elif point_estimate == "mode":
        point = ds_predictive.azstats.mode(dim=sample_dims, **pe_stats)
    else:
        raise ValueError(
            f"point_estimate must be 'mean', 'median' or 'mode', but {point_estimate} was passed."
        )

    return _plot_interval(
        ds_predictive,
        ci_trunk,
        ci_twig,
        point,
        observed_data,
        sample_dims,
        group,
        plot_collection,
        backend,
        pc_kwargs,
        plot_bknd,
        visuals,
        aes_by_visuals,
        labeller,
    )


def _plot_interval(
    ds_predictive,
    ci_trunk,
    ci_twig,
    point,
    observed_data,
    sample_dims,
    group,
    plot_collection,
    backend,
    pc_kwargs,
    plot_bknd,
    visuals,
    aes_by_visuals,
    labeller,
):
    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.setdefault("cols", "__variable__")

        pc_kwargs = set_wrap_layout(pc_kwargs, plot_bknd, ds_predictive)

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

    ## trunk intervals
    ci_trunk_kwargs = get_visual_kwargs(visuals, "trunk")
    _, ci_trunk_aes, ci_trunk_ignore = filter_aes(
        plot_collection, aes_by_visuals, "trunk", sample_dims
    )

    if ci_trunk_kwargs is not False:
        if "color" not in ci_trunk_aes:
            ci_trunk_kwargs.setdefault("color", "C0")

        ci_trunk_kwargs.setdefault("alpha", 0.5)
        ci_trunk_kwargs.setdefault("width", 3)

        plot_collection.map(
            ci_bound_y,
            "trunk",
            data=ci_trunk,
            ignore_aes=ci_trunk_ignore,
            **ci_trunk_kwargs,
        )

    ## twig intervals
    ci_twig_kwargs = get_visual_kwargs(visuals, "twig")
    _, ci_twig_aes, ci_twig_ignore = filter_aes(
        plot_collection, aes_by_visuals, "twig", sample_dims
    )

    if ci_twig_kwargs is not False:
        if "color" not in ci_twig_aes:
            ci_twig_kwargs.setdefault("color", "C0")

        ci_twig_kwargs.setdefault("width", 3)

        plot_collection.map(
            ci_bound_y,
            "twig",
            data=ci_twig,
            ignore_aes=ci_twig_ignore,
            **ci_twig_kwargs,
        )

    ## prediction_markers
    prediction_ms_kwargs = get_visual_kwargs(visuals, "prediction_markers")

    if prediction_ms_kwargs is not False:
        _, _, prediction_ms_ignore = filter_aes(
            plot_collection, aes_by_visuals, "prediction_markers", sample_dims
        )
        prediction_ms_kwargs.setdefault("color", "C0")
        prediction_ms_kwargs.setdefault("marker", "C0")

        plot_collection.map(
            point_y,
            "prediction_markers",
            data=point,
            ignore_aes=prediction_ms_ignore,
            **prediction_ms_kwargs,
        )

    ## observed_markers
    observed_ms_kwargs = get_visual_kwargs(
        visuals,
        "observed_markers",
        False if group == "prior_predictive" or observed_data is None else None,
    )

    if observed_ms_kwargs is not False:
        _, _, observed_ms_ignore = filter_aes(
            plot_collection, aes_by_visuals, "observed_markers", sample_dims
        )
        observed_ms_kwargs.setdefault("color", "B1")
        observed_ms_kwargs.setdefault("marker", "C6")

        plot_collection.map(
            point_y,
            "observed_markers",
            data=observed_data,
            ignore_aes=observed_ms_ignore,
            **observed_ms_kwargs,
        )

    # set xlabel
    _, xlabels_aes, xlabels_ignore = filter_aes(
        plot_collection, aes_by_visuals, "xlabel", sample_dims
    )
    xlabel_kwargs = get_visual_kwargs(visuals, "xlabel")
    if xlabel_kwargs is not False:
        if "color" not in xlabels_aes:
            xlabel_kwargs.setdefault("color", "B1")

        xlabel_kwargs.setdefault("text", "data point (index)")

        plot_collection.map(
            labelled_x,
            "xlabel",
            ignore_aes=xlabels_ignore,
            subset_info=True,
            **xlabel_kwargs,
        )

    # set ylabel
    _, ylabels_aes, ylabels_ignore = filter_aes(
        plot_collection, aes_by_visuals, "ylabel", sample_dims
    )
    ylabel_kwargs = get_visual_kwargs(visuals, "ylabel")
    if ylabel_kwargs is not False:
        if "color" not in ylabels_aes:
            ylabel_kwargs.setdefault("color", "B1")

        plot_collection.map(
            labelled_y,
            "ylabel",
            ignore_aes=ylabels_ignore,
            subset_info=True,
            labeller=labeller,
            **ylabel_kwargs,
        )

    # set title
    _, title_aes, title_ignore = filter_aes(plot_collection, aes_by_visuals, "title", sample_dims)
    title_kwargs = get_visual_kwargs(visuals, "title")
    if title_kwargs is not False:
        if "color" not in title_aes:
            title_kwargs.setdefault("color", "B1")

        plot_collection.map(
            labelled_title,
            "title",
            ignore_aes=title_ignore,
            subset_info=True,
            labeller=labeller,
            **title_kwargs,
        )

    return plot_collection
