"""Energy plot code."""
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 convert_to_dataset, rcParams
from arviz_base.validate import validate_dict_argument, 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 filter_aes, get_visual_kwargs, set_grid_layout
from arviz_plots.visuals import labelled_title, labelled_y, scatter_xy, vline


def plot_energy(
    dt,
    *,
    sample_dims=None,
    kind=None,
    show_bfmi=True,
    threshold=0.3,
    plot_collection=None,
    backend=None,
    labeller=None,
    aes_by_visuals: Mapping[
        Literal[
            "dist",
            "face",
            "title",
            "bfmi_points",
        ],
        Sequence[str],
    ] = None,
    visuals: Mapping[
        Literal[
            "dist",
            "face",
            "title",
            "legend",
            "remove_axis",
            "bfmi_points",
            "ref_line",
            "ylabel",
        ],
        Mapping[str, Any] | bool,
    ] = None,
    stats: Mapping[Literal["dist"], Mapping[str, Any] | xr.Dataset] = None,
    **pc_kwargs,
):
    r"""Plot energy distributions and bfmi from gradient-based algorithms.

    Generate a figure with the marginal energy distribution and the energy transition
    distribution. Optionally, include a BFMI panel to inspect chain-wise Bayesian Fraction
    of Missing Information values. Values below the threshold indicate poor exploration
    of the energy distribution.

    For details on BFMI and energy diagnostics see [1]_ for a more practical overview check
    the EABM chapter on MCMC diagnostic `of gradient-based algorithms <https://arviz-devs.github.io/EABM/Chapters/MCMC_diagnostics.html#diagnosis-of-gradient-based-algorithms>`_.


    Parameters
    ----------
    dt : DataTree
        ``sample_stats`` group with an ``energy`` variable is mandatory.
    sample_dims : sequence of str, optional
        Dimensions to consider as sample dimensions when computing BFMI.
        Defaults to ``rcParams["data.sample_dims"]``
    kind : {"kde", "hist", "dot", "ecdf"}, optional
        How to represent the marginal density.
        Defaults to ``rcParams["plot.density_kind"]``
    show_bfmi : bool, default True
        Whether to include the BFMI scatter plot. If ``False``, only the energy plot will be shown.
    threshold : float, default 0.3
        Reference threshold for BFMI values, values below this indicate poor exploration of the
        energy distribution.
    plot_collection : PlotCollection, optional
    backend : {"matplotlib", "bokeh", "plotly"}, 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`.

    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`
        * remove_axis -> not passed anywhere, can only be ``False`` to skip calling this function
        * title -> passed to :func:`~arviz_plots.visuals.labelled_title`
        * bfmi_points -> passed to :func:`~arviz_plots.visuals.scatter_xy` for BFMI scatter plot
        * ylabel -> passed to :func:`~arviz_plots.visuals.labelled_y` for BFMI column y-axis label
        * face -> :term:`visual` that fills the area under the energy distributions.

          Defaults to True. Depending on the value of `kind` it is passed to:

          * "kde" or "ecdf" -> passed to :func:`~arviz_plots.visuals.fill_between_y`
          * "hist" -> passed to :func:`~arviz_plots.visuals.hist`
          * dot -> ignored

    stats : mapping, optional
        Valid keys are:

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

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

    Returns
    -------
    PlotCollection

    Examples
    --------
    Plot an energy plot using ecdf for the energy distributions.

    .. plot::
        :context: close-figs

        >>> from arviz_plots import plot_energy, style
        >>> style.use("arviz-variat")
        >>> from arviz_base import load_arviz_data
        >>> data = load_arviz_data('non_centered_eight')
        >>> plot_energy(data, kind="ecdf")


    .. minigallery:: plot_energy

    References
    ----------
    .. [1] Betancourt. Diagnosing Suboptimal Cotangent Disintegrations in
        Hamiltonian Monte Carlo. (2016) https://arxiv.org/abs/1604.00695
    """  # pylint: disable=line-too-long
    visuals = validate_dict_argument(visuals, (plot_energy, "visuals"))
    aes_by_visuals = validate_dict_argument(aes_by_visuals, (plot_energy, "aes_by_visuals"))

    sample_stats = dt["sample_stats"]
    sample_dims = validate_sample_dims(sample_dims, sample_stats)
    energy_ds, bfmi_ds = _get_energy_ds(sample_stats, sample_dims=sample_dims)

    if backend is None:
        backend = rcParams["plot.backend"]
    plot_bknd = import_module(f".backend.{backend}", package="arviz_plots")

    if plot_collection is None:
        pc_kwargs["figure_kwargs"] = pc_kwargs.get("figure_kwargs", {}).copy()
        if show_bfmi:
            new_ds = energy_ds.expand_dims(column=2).assign_coords(column=["bfmi", "energy"])
            pc_kwargs["figure_kwargs"].setdefault("width_ratios", [1, 3])
            num_cols = 2
        else:
            new_ds = energy_ds.expand_dims(column=1).assign_coords(column=["energy"])
            num_cols = 1

        pc_kwargs["aes"] = pc_kwargs.get("aes", {}).copy()
        pc_kwargs.setdefault("cols", ["column"])
        pc_kwargs["aes"].setdefault("color", ["energy"])
        pc_kwargs = set_grid_layout(pc_kwargs, plot_bknd, new_ds, num_cols=num_cols, num_rows=1)

        plot_collection = PlotCollection.grid(
            new_ds,
            backend=backend,
            **pc_kwargs,
        )

    visuals.setdefault("credible_interval", False)
    visuals.setdefault("point_estimate", False)
    visuals.setdefault("point_estimate_text", False)
    visuals.setdefault("face", True)
    keys_to_remove = ("legend", "bfmi_points", "ref_line", "ylabel")

    # Energy distributions plot
    plot_collection.coords = {"column": "energy"}
    plot_collection = plot_dist(
        energy_ds,
        var_names=None,
        filter_vars=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={k: v for k, v in aes_by_visuals.items() if k not in keys_to_remove},
        visuals={k: v for k, v in visuals.items() if k not in keys_to_remove},
        stats=stats,
    )
    plot_collection.coords = None

    # legend for energy distributions
    legend_kwargs = get_visual_kwargs(visuals, "legend")
    if legend_kwargs is not False:
        legend_kwargs.setdefault("dim", ["energy"])
        legend_kwargs.setdefault("title", "")
        plot_collection.add_legend(**legend_kwargs)

    if show_bfmi:
        # Scatter plot of BFMI values
        bfmi_ms_kwargs = get_visual_kwargs(visuals, "bfmi_points")
        if bfmi_ms_kwargs is not False:
            _, _, bfmi_ignore = filter_aes(plot_collection, aes_by_visuals, "bfmi_points", [])

            bfmi_ms_kwargs.setdefault("color", "B1")

            plot_collection.coords = {"column": "bfmi"}
            plot_collection.map(
                scatter_xy,
                "bfmi_points",
                data=bfmi_ds,
                ignore_aes=bfmi_ignore,
                **bfmi_ms_kwargs,
            )
            plot_collection.coords = None

        # Reference line for BFMI threshold
        ref_line_kwargs = get_visual_kwargs(visuals, "ref_line")
        if ref_line_kwargs is not False:
            _, ref_aes, ref_ignore = filter_aes(
                plot_collection, aes_by_visuals, "ref_line", sample_dims
            )
            if "color" not in ref_aes:
                ref_line_kwargs.setdefault("color", "B2")
            if "linestyle" not in ref_aes:
                ref_line_kwargs.setdefault("linestyle", "C1")

            # Wrap threshold into an xr.Dataset for PlotCollection.map
            ref_ds = xr.Dataset({"ref_line": xr.DataArray(threshold)})
            plot_collection.coords = {"column": "bfmi"}
            plot_collection.map(
                vline, "ref_line", data=ref_ds, ignore_aes=ref_ignore, **ref_line_kwargs
            )
            plot_collection.coords = None

        # Add title for BFMI plot
        title_kwargs = get_visual_kwargs(visuals, "title")
        if title_kwargs is not False:
            _, title_aes, title_ignore = filter_aes(
                plot_collection, aes_by_visuals, "title", sample_dims
            )
            if "color" not in title_aes:
                title_kwargs.setdefault("color", "B1")
            plot_collection.coords = {"column": "bfmi"}
            plot_collection.map(
                labelled_title,
                "title",
                text="BFMI",
                ignore_aes=title_ignore,
                subset_info=True,
                labeller=labeller,
                **title_kwargs,
            )
            plot_collection.coords = None

        # Add ylabel for BFMI plot
        ylabel_kwargs = get_visual_kwargs(visuals, "ylabel")
        if ylabel_kwargs is not False:
            ylabel_kwargs.setdefault("text", "Chain")
            _, _, ylabel_ignore = filter_aes(plot_collection, {}, "ylabel", [])
            plot_collection.coords = {"column": "bfmi"}
            plot_collection.map(
                labelled_y,
                "ylabel",
                ignore_aes=ylabel_ignore,
                **ylabel_kwargs,
            )
            plot_collection.coords = None

    return plot_collection


def _get_energy_ds(sample_stats, sample_dims):
    """Extract energy and BFMI data from DataTree.

    Returns
    -------
    energy_ds : Dataset
        Dataset with Energy variable containing marginal and transition energy
    bfmi_ds : Dataset
        Dataset with bfmi variable containing BFMI values and chain indices
    """
    energy = sample_stats.energy.values
    bfmi_vals = sample_stats["energy"].azstats.bfmi(sample_dims=sample_dims)
    n_chains = len(bfmi_vals)
    chain_indices = np.arange(n_chains)

    bfmi_ds = xr.Dataset(
        {
            "bfmi": xr.DataArray(
                np.column_stack([bfmi_vals.values, chain_indices]),
                dims=["chain", "plot_axis"],
                coords={"chain": bfmi_vals.chain, "plot_axis": ["x", "y"]},
            )
        }
    )

    energy_ds = convert_to_dataset(
        {"Energy": np.dstack([energy - energy.mean(), np.diff(energy, append=np.nan)])},
        coords={"Energy_dim_0": ["marginal", "transition"]},
    ).rename({"Energy_dim_0": "energy"})

    return energy_ds, bfmi_ds
