"""Pair 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 rcParams, xarray_sel_iter
from arviz_base.labels import BaseLabeller
from arviz_base.validate import validate_dict_argument, validate_sample_dims

from arviz_plots.plot_matrix import PlotMatrix
from arviz_plots.plots.dist_plot import plot_dist
from arviz_plots.plots.utils import (
    filter_aes,
    get_group,
    get_visual_kwargs,
    process_group_variables_coords,
    set_grid_layout,
)
from arviz_plots.visuals import (
    label_plot,
    labelled_x,
    labelled_y,
    remove_matrix_axis,
    scatter_couple,
    set_ticklabel_visibility,
)


def plot_pair(
    dt,
    *,
    var_names=None,
    filter_vars=None,
    group="posterior",
    coords=None,
    sample_dims=None,
    marginal=True,
    marginal_kind=None,
    triangle="lower",
    plot_matrix=None,
    backend=None,
    labeller=None,
    aes_by_visuals: Mapping[
        Literal[
            "scatter",
            "divergence",
            "dist",
            "credible_interval",
            "point_estimate",
            "point_estimate_text",
            "label",
            "xlabel",
            "ylabel",
        ],
        Sequence[str],
    ] = None,
    visuals: Mapping[
        Literal[
            "scatter",
            "divergence",
            "dist",
            "credible_interval",
            "point_estimate",
            "point_estimate_text",
            "label",
            "xlabel",
            "ylabel",
            "remove_axis",
        ],
        Mapping[str, Any] | bool,
    ] = None,
    stats: Mapping[
        Literal[
            "dist",
            "credible_interval",
            "point_estimate",
        ],
        Mapping[str, Any] | xr.Dataset,
    ] = None,
    **pc_kwargs,
):
    """Plot all variables against each other in the dataset.

    Parameters
    ----------
    dt : DataTree
        Input data
    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 (default), 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, default "posterior"
        Group to use for plotting. Defaults to "posterior".
    coords : mapping, optional
        Coordinates to use for plotting.
    sample_dims : iterable, optional
        Dimensions to reduce unless mapped to an aesthetic.
        Defaults to ``rcParams["data.sample_dims"]``
    marginal : bool, default True
        Whether to plot marginal distributions on the diagonal.
    marginal_kind : {"kde", "hist", "ecdf", "dot"}, optional
        How to represent the marginal density.
        Defaults to ``rcParams["plot.density_kind"]``
    triangle : {"both", "upper", "lower"}, Defaults to "both"
        Which triangle of the pair plot to plot.
    plot_matrix : PlotMatrix, optional
    backend : {"matplotlib", "bokeh", "plotly", "none"}, optional
        Plotting backend to use. Defaults to ``rcParams["plot.backend"]``
    labeller : labeller, optional
    aes_by_visuals : mapping, optional
        Mapping of visuals to aesthetics that should use their mapping in `plot_matrix`
        when plotted. Valid keys are the same as for `visuals`.
        By default, there are no aesthetic mappings at all
    visuals : mapping of {str : mapping or bool}, optional
        Valid keys are:

        * scatter -> passed to :func:`~.visuals.scatter_couple`
        * divergence -> passed to :func:`~.visuals.scatter_couple`. Defaults to False.
        * dist -> depending on the value of `marginal_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`

        * credible_interval -> passed to :func:`~arviz_plots.visuals.line_x`
        * point_estimate -> passed to :func:`~arviz_plots.visuals.scatter_x`
        * point_estimate_text -> passed to :func:`~arviz_plots.visuals.point_estimate_text`
        * label -> Keyword arguments passed to :func:`~arviz_plots.visuals.label_plot`.

          Used to customize the variable name labels on the diagonal. Applied only
          if ``marginal=False``.

        * xlabel -> passed to :func:`~.visuals.labelled_x`.

          used to customize the xaxis labels on the bottom-most plots or diagonal plots depending
          upon the value of ``triangle``. If ``triangle`` is "lower" or "both" then it is used to
          map bottom-most row plots by using :meth:`arviz_plots.PlotMatrix.map_row` method and if
          ``triangle`` is "upper" then it is used to map diagonal plots by using
          :meth:`arviz_plots.PlotMatrix.map` method.It is applied only if ``marginal=True``, since
          in this case diagonal plots won't have labels to map variables to columns.

        * ylabel -> passed to :func:`~.visuals.labelled_y`.

          used to customize the yaxis labels on the left-most plots. It is applied, only if
          ``triangle`` is "lower" or "both" and ``marginal=True``, by using
          :meth:`arviz_plots.PlotMatrix.map_col` method. Not applied if ``triangle`` is
          "upper" or ``marginal=False``.

        * remove_axis -> not passed anywhere.

          It can only be set to ``False`` to disable the default removal of ``x`` and ``y`` axes
          from the plots of other half triangle. If ``triangle`` is "upper" then the lower triangle
          plot's axes will be removed and if ``triangle`` is "lower" then the upper triangle axes
          will be removed, in case if it is not set ``False`` manually.


    stats : mapping, optional
        Valid keys are:

        * dist -> passed to kde, ecdf, ...
        * credible_interval -> passed to eti or hdi
        * point_estimate -> passed to mean, median or mode

    **pc_kwargs
        Passed to :class:`arviz_plots.PlotMatrix`


    Returns
    -------
    PlotMatrix

    Examples
    --------
    plot_pair with ``triangle`` set to "upper" and ``marginal=True`` with ``marginal_kind`` set to
    "ecdf". In this case, since ``triangle`` is "upper", so the ``xlabels`` are mapped to the
    diagonal plots. ``marginals`` are plotted on the diagonal and the ``point_estimate`` and
    ``credible_interval`` are set to ``False`` by default. Also since ``marginal=True``, so
    ``sharex`` is set to "col", while ``sharey`` is not set to anything by default.

    .. plot::
        :context: close-figs

        >>> from arviz_plots import plot_pair, style
        >>> style.use("arviz-variat")
        >>> from arviz_base import load_arviz_data
        >>> dt = load_arviz_data('centered_eight')
        >>> plot_pair(
        >>>     dt,
        >>>     var_names=["mu", "tau"],
        >>>     visuals={"divergence": True},
        >>>     marginal=True,
        >>>     marginal_kind="ecdf",
        >>>     triangle="upper",
        >>> )

    plot_pair with `triangle` set to "both", so in this case the ``xlabels`` are mapped to the
    bottom-most plots and ``ylabels`` are mapped to the left-most plots. In this example we set
    ``color`` as "red" for ``credible_interval`` and ``point_estimate``, which enables
    ``credible_interval`` and ``point_estimate``. By default ``marginal`` is set to ``True`` and
    ``marginal_kind`` is set to ``rcParams["plot.density_kind"]``.

    .. plot::
        :context: close-figs

        >>> visuals = {"credible_interval":{"color":"red"},"point_estimate":{"color":"red"}}
        >>> plot_pair(
        >>>     dt,
        >>>     var_names=["mu", "tau"],
        >>>     visuals=visuals,
        >>>     triangle="both",
        >>> )

    plot_pair with ``marginal=False`` and ``triangle`` set to "upper". In this case, since
    ``marginal=False``, so ``xlabel`` and ``ylabel`` are disabled by default, and diagonal
    plots contain variable names as labels. ``xticks`` and ``yticks`` are also set on
    diagonal plots along with ``ticklabels``, to map ticks to rows and columns.
    Since ``marginal=False``, so ``sharex`` is set to "col" and ``sharey`` is set to "row"
    by default.

    .. plot::
        :context: close-figs

        >>> plot_pair(
        >>>     dt,
        >>>     coords = {"school":"Choate"},
        >>>     visuals={"divergence": True},
        >>>     marginal=False,
        >>>     triangle="upper",
        >>> )

    .. minigallery:: plot_pair

    """
    aes_by_visuals = validate_dict_argument(aes_by_visuals, (plot_pair, "aes_by_visuals"))
    visuals = validate_dict_argument(visuals, (plot_pair, "visuals"))
    if labeller is None:
        labeller = BaseLabeller()
    if backend is None:
        if plot_matrix is None:
            backend = rcParams["plot.backend"]
        else:
            backend = plot_matrix.backend

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

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

    if plot_matrix is None:
        pc_kwargs.setdefault(
            "facet_dims",
            ["__variable__"] + [dim for dim in distribution.dims if dim not in sample_dims],
        )
        pc_kwargs["figure_kwargs"] = pc_kwargs.get("figure_kwargs", {}).copy()
        pairs = tuple(
            xarray_sel_iter(
                distribution, skip_dims={dim for dim in distribution.dims if dim in sample_dims}
            )
        )
        n_pairs = len(pairs)
        pc_kwargs = set_grid_layout(
            pc_kwargs, plot_bknd, distribution, num_rows=n_pairs, num_cols=n_pairs
        )
        pc_kwargs["aes"] = pc_kwargs.get("aes", {}).copy()
        if "chain" in distribution:
            pc_kwargs["aes"].setdefault("overlay", ["chain"])
        pc_kwargs["figure_kwargs"].setdefault("sharex", "col")
        if not marginal:
            pc_kwargs["figure_kwargs"].setdefault("sharey", "row")

        plot_matrix = PlotMatrix(
            distribution,
            backend=backend,
            **pc_kwargs,
        )

    aes_by_visuals["scatter"] = {"overlay"}.union(
        aes_by_visuals.get("scatter", plot_matrix.aes_set)
    )
    aes_by_visuals["divergence"] = {"overlay"}.union(aes_by_visuals.get("divergence", {}))
    aes_by_visuals["dist"] = aes_by_visuals.get("dist", plot_matrix.aes_set.difference({"overlay"}))
    aes_by_visuals["credible_interval"] = aes_by_visuals.get("credible_interval", {})
    aes_by_visuals["point_estimate"] = aes_by_visuals.get("point_estimate", {})
    aes_by_visuals["point_estimate_text"] = aes_by_visuals.get("point_estimate_text", {})

    # scatter
    scatter_kwargs = get_visual_kwargs(visuals, "scatter")
    if scatter_kwargs is not False:
        _, scatter_aes, scatter_ignore = filter_aes(
            plot_matrix, aes_by_visuals, "scatter", sample_dims
        )

        if "color" not in scatter_aes:
            scatter_kwargs.setdefault("color", "C0")

        if "width" not in scatter_aes:
            scatter_kwargs.setdefault("width", 0)

        if "alpha" not in scatter_aes:
            scatter_kwargs.setdefault("alpha", 0.5)

        plot_matrix.map_triangle(
            scatter_couple,
            "scatter",
            triangle=triangle,
            data=distribution,
            ignore_aes=scatter_ignore,
            **scatter_kwargs,
        )

    # marginal
    if marginal is not False:
        if stats is None:
            stats = {}
        else:
            stats = stats.copy()

        dist_plot_visuals = {}
        dist_plot_aes_by_visuals = {}
        dist_plot_stats = {}
        marginal_dist_kwargs = get_visual_kwargs(visuals, "dist")
        marginal_ci_kwargs = get_visual_kwargs(visuals, "credible_interval", False)
        marginal_point_estimate_kwargs = get_visual_kwargs(visuals, "point_estimate", False)
        marginal_point_estimate_text_kwargs = get_visual_kwargs(
            visuals, "point_estimate_text", False
        )

        dist_plot_visuals["dist"] = marginal_dist_kwargs
        dist_plot_visuals["credible_interval"] = marginal_ci_kwargs
        dist_plot_visuals["point_estimate"] = marginal_point_estimate_kwargs
        dist_plot_visuals["point_estimate_text"] = marginal_point_estimate_text_kwargs
        dist_plot_visuals["title"] = False
        dist_plot_visuals["remove_axis"] = False
        dist_plot_visuals["rug"] = False

        dist_plot_aes_by_visuals["dist"] = aes_by_visuals.get(
            "dist", plot_matrix.aes_set.difference({"overlay"})
        )
        dist_plot_aes_by_visuals["credible_interval"] = aes_by_visuals.get("credible_interval", {})
        dist_plot_aes_by_visuals["point_estimate"] = aes_by_visuals.get("point_estimate", {})
        dist_plot_aes_by_visuals["point_estimate_text"] = aes_by_visuals.get(
            "point_estimate_text", {}
        )

        dist_plot_stats["dist"] = stats.get("dist", {})
        dist_plot_stats["credible_interval"] = stats.get("credible_interval", {})
        dist_plot_stats["point_estimate"] = stats.get("point_estimate", {})

        plot_matrix = plot_dist(
            distribution,
            sample_dims=sample_dims,
            kind=marginal_kind,
            plot_collection=plot_matrix,
            backend=backend,
            labeller=labeller,
            aes_by_visuals=dist_plot_aes_by_visuals,
            visuals=dist_plot_visuals,
            stats=dist_plot_stats,
        )

    # diagonal labels of rows and cols
    else:
        label_kwargs = get_visual_kwargs(visuals, "label")

        if label_kwargs is not False:
            lim_low = distribution.min(dim=sample_dims)
            lim_high = distribution.max(dim=sample_dims)
            text_center = (lim_high + lim_low) / 2
            _, _, label_ignore = filter_aes(plot_matrix, aes_by_visuals, "label", sample_dims)
            plot_matrix.map(
                label_plot,
                "label",
                subset_info=True,
                labeller=labeller,
                x=text_center,
                y=text_center,
                lim_low=lim_low,
                lim_high=lim_high,
                ignore_aes=label_ignore,
                **label_kwargs,
            )

    # divergence
    div_kwargs = get_visual_kwargs(visuals, "divergence", False)
    sample_stats = get_group(dt, "sample_stats", allow_missing=True)
    if (
        div_kwargs is not False
        and sample_stats is not None
        and "diverging" in sample_stats.data_vars
        and np.any(sample_stats.diverging)
    ):
        divergence_mask = dt.sample_stats.diverging
        _, div_aes, div_ignore = filter_aes(plot_matrix, aes_by_visuals, "divergence", sample_dims)
        if "color" not in div_aes:
            div_kwargs.setdefault("color", "C1")
        if "alpha" not in div_aes:
            div_kwargs.setdefault("alpha", 0.5)

        plot_matrix.map_triangle(
            scatter_couple,
            "divergence",
            triangle=triangle,
            data=distribution,
            ignore_aes=div_ignore,
            mask=divergence_mask,
            **div_kwargs,
        )

    # bottom plots xlabel and left plots ylabel
    if marginal and triangle in {"both", "lower"}:
        xlabel_kwargs = get_visual_kwargs(visuals, "xlabel")
        if xlabel_kwargs is not False:
            _, _, xlabel_ignore = filter_aes(plot_matrix, aes_by_visuals, "xlabel", sample_dims)
            plot_matrix.map_row(
                labelled_x,
                "xlabel",
                index=-1,
                data=distribution,
                ignore_aes=xlabel_ignore,
                labeller=labeller,
                subset_info=True,
                **xlabel_kwargs,
            )

        ylabel_kwargs = get_visual_kwargs(visuals, "ylabel")
        if ylabel_kwargs is not False:
            _, _, ylabel_ignore = filter_aes(plot_matrix, aes_by_visuals, "ylabel", sample_dims)
            plot_matrix.map_col(
                labelled_y,
                "ylabel",
                index=0,
                data=distribution,
                ignore_aes=ylabel_ignore,
                labeller=labeller,
                subset_info=True,
                **ylabel_kwargs,
            )
    elif marginal and triangle == "upper":
        xlabel_kwargs = get_visual_kwargs(visuals, "xlabel")

        if xlabel_kwargs is not False:
            _, _, xlabel_ignore = filter_aes(plot_matrix, aes_by_visuals, "xlabel", sample_dims)
            plot_matrix.map(
                labelled_x,
                "xlabel",
                subset_info=True,
                ignore_aes=xlabel_ignore,
                labeller=labeller,
                **xlabel_kwargs,
            )

    # set ticklabel visibility
    set_ticklabel_visibility_kwargs = {}
    if triangle == "upper":
        _, _, set_ticklabel_visibility_ignore = filter_aes(
            plot_matrix, aes_by_visuals, "set_ticklabel_visibility", sample_dims
        )
        plot_matrix.map(
            set_ticklabel_visibility,
            "set_ticklabel_visibility",
            axis="x",
            visible=True,
            ignore_aes=set_ticklabel_visibility_ignore,
            **set_ticklabel_visibility_kwargs,
        )

        if not marginal:
            plot_matrix.map(
                set_ticklabel_visibility,
                "set_ticklabel_visibility",
                axis="y",
                visible=True,
                ignore_aes=set_ticklabel_visibility_ignore,
                **set_ticklabel_visibility_kwargs,
            )

    #  removal of axis for better visualization
    remove_axis_bool = visuals.get("remove_axis", True)
    if remove_axis_bool:
        _, _, remove_axis_ignore = filter_aes(
            plot_matrix, aes_by_visuals, "remove_axis", sample_dims
        )
        # if triangle="upper" then remove the lower triangle axes
        if triangle == "upper":
            plot_matrix.map_triangle(
                remove_matrix_axis,
                "remove_axis",
                triangle="lower",
                axis="both",
                ignore_aes=remove_axis_ignore,
            )
        # if triangle="lower" then remove the upper triangle axes
        elif triangle == "lower":
            plot_matrix.map_triangle(
                remove_matrix_axis,
                "remove_axis",
                triangle="upper",
                axis="both",
                ignore_aes=remove_axis_ignore,
            )
    return plot_matrix
