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

from arviz_plots.plot_collection import PlotCollection
from arviz_plots.plots.utils import (
    filter_aes,
    get_group,
    get_visual_kwargs,
    process_group_variables_coords,
)
from arviz_plots.visuals import multiple_lines, set_xticks


def plot_parallel(
    dt,
    *,
    var_names=None,
    filter_vars=None,
    group="posterior",
    coords=None,
    sample_dims=None,
    norm_method=None,
    label_type="flat",
    plot_collection=None,
    backend=None,
    labeller=None,
    aes_by_visuals: Mapping[
        Literal[
            "line",
            "xticks",
        ],
        Sequence[str],
    ] = None,
    visuals: Mapping[
        Literal[
            "line",
            "xticks",
        ],
        Mapping[str, Any] | bool,
    ] = None,
    **pc_kwargs,
):
    """Plot parallel coordinates plot showing posterior points with and without divergences.

    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"]``
    norm_method : {None, "normal", "minmax", "rank"}, default None
        Transformation to apply to the samples before plotting.
    label_type : {"flat", "vert"}, default "flat"
        Indicator of which `labeller` method to use when generating the
        labels of the x axis.
    plot_collection : PlotCollection, 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_collection`
        when plotted. Valid keys are the same as for `visuals`.
        By default, there is a mapping from the value of `diverging` variable to
        color and alpha which is only active for the "line" visual.
    visuals : mapping of {str : mapping or bool}, optional
        Valid keys are:

        * line -> passed to :func:`~.visuals.multiple_lines`
        * xticks -> passed to :func:`~.visuals.set_xticks`. Defaults to False.

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


    Returns
    -------
    PlotCollection

    Examples
    --------
    Default plot_parallel without normalization and with default `label_type` as "flat".

    .. plot::
        :context: close-figs

        >>> from arviz_plots import plot_parallel, style
        >>> style.use("arviz-variat")
        >>> from arviz_base import load_arviz_data
        >>> dt = load_arviz_data('centered_eight')
        >>> plot_parallel(dt)

    parallel_plot with `norm_method` set to "normal" and `label_type` set to "vert"
    and rotation of xaxis labels set to 30 degrees.

    .. plot::
        :context: close-figs

        >>> plot_parallel(
        >>>     dt,
        >>>     norm_method="normal",
        >>>     label_type="vert",
        >>>     visuals={"xticks": {"rotation": 30}},
        >>> )

    .. minigallery:: plot_parallel

    """
    aes_by_visuals = validate_dict_argument(aes_by_visuals, (plot_parallel, "aes_by_visuals"))
    visuals = validate_dict_argument(visuals, (plot_parallel, "visuals"))

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

    if labeller is None:
        labeller = BaseLabeller()

    data = process_group_variables_coords(
        dt,
        group=group,
        var_names=var_names,
        filter_vars=filter_vars,
        coords=coords,
        allow_dict=False,
    )
    sample_dims = validate_sample_dims(sample_dims, data=data)
    if len(sample_dims) > 1:
        data = data.stack(sample=sample_dims)
        combined_dim = "sample"
    else:
        combined_dim = sample_dims[0]

    # create divergence mask
    sample_stats = get_group(dt, "sample_stats", allow_missing=True)
    if (
        sample_stats is not None
        and "diverging" in sample_stats.data_vars
        and np.any(sample_stats.diverging)
    ):
        div_mask = sample_stats.diverging
        if coords is not None:
            div_mask = div_mask.sel(
                {key: value for key, value in coords.items() if key in div_mask.dims}
            )
        if len(sample_dims) > 1:
            div_mask = div_mask.stack(sample=sample_dims)
    else:
        div_mask = xr.zeros_like(data.coords[combined_dim], dtype=bool)
    data = data.assign_coords(diverging=div_mask)
    data = dataset_to_dataarray(
        data, labeller=labeller, sample_dims=[combined_dim], label_type=label_type
    )
    if len(data.coords["label"].values) <= 1 or data.ndim != 2:
        raise ValueError(
            "Parallel plot requires at least two variables to plot. "
            "Please provide more than one variable."
        )

    # get labels and x-values
    x_labels = data.coords["label"].values
    x_values = np.arange(len(x_labels))

    # normalize data
    if norm_method is not None:
        if norm_method == "normal":
            mean = data.mean(dim=combined_dim)
            std_dev = data.std(dim=combined_dim)
            data = (data - mean) / std_dev
        elif norm_method == "minmax":
            min_val = data.min(dim=combined_dim)
            max_val = data.max(dim=combined_dim)
            data = (data - min_val) / (max_val - min_val)
        elif norm_method == "rank":
            data = data.azstats.compute_ranks(dim=combined_dim)
        else:
            raise ValueError(f"{norm_method} is not supported. Use normal, minmax or rank.")

    new_sample_dims = [combined_dim, "label"]
    plot_bknd = import_module(f".backend.{backend}", package="arviz_plots")
    if plot_collection is None:
        pc_kwargs.setdefault("cols", None)
        pc_kwargs["figure_kwargs"] = pc_kwargs.get("figure_kwargs", {}).copy()
        pc_kwargs["aes"] = pc_kwargs.get("aes", {}).copy()
        pc_kwargs["aes"].setdefault("color", ["diverging"])
        pc_kwargs["aes"].setdefault("alpha", ["diverging"])
        pc_kwargs["alpha"] = pc_kwargs.get("alpha", [0.03, 0.2])
        figsize = pc_kwargs.get("figure_kwargs", {}).get("figsize", None)
        figsize_units = pc_kwargs.get("figure_kwargs", {}).get("figsize_units", "inches")
        if figsize is None:
            figsize = plot_bknd.scale_fig_size(
                figsize,
                rows=2,
                cols=1 + 0.2 * len(x_labels),
                figsize_units=figsize_units,
            )
            label_max_len = np.vectorize(len)(x_labels).max()
            fontsize = 14
            figsize = (figsize[0], figsize[1] + fontsize * label_max_len)
            figsize_units = "dots"

        pc_kwargs["figure_kwargs"]["figsize"] = figsize
        pc_kwargs["figure_kwargs"]["figsize_units"] = figsize_units
        plot_collection = PlotCollection.wrap(
            data.to_dataset(),
            backend=backend,
            **pc_kwargs,
        )

    aes_by_visuals.setdefault("line", plot_collection.aes_set)

    # plot lines
    line_kwargs = get_visual_kwargs(visuals, "line")
    if line_kwargs is not False:
        _, _, line_ignore = filter_aes(plot_collection, aes_by_visuals, "line", new_sample_dims)
        plot_collection.map(
            multiple_lines,
            "line",
            data=data,
            x_dim="label",
            xvalues=x_values,
            ignore_aes=line_ignore,
            **line_kwargs,
        )

    # x-axis label
    xticks_kwargs = get_visual_kwargs(visuals, "xticks")
    if xticks_kwargs is not False:
        _, _, xticks_ignore = filter_aes(plot_collection, aes_by_visuals, "xticks", new_sample_dims)
        xticks_kwargs.setdefault("rotation", 90)
        plot_collection.map(
            set_xticks,
            "xticks",
            labels=x_labels,
            values=x_values,
            ignore_aes=xticks_ignore,
            artist_dims={"labels": len(x_labels)},
            **xticks_kwargs,
        )

    return plot_collection
