"""Utilities for batteries included plots."""
import warnings
from copy import copy
from importlib import import_module

import numpy as np
import xarray as xr
from arviz_base import rcParams, references_to_dataset, xarray_sel_iter
from arviz_base.labels import BaseLabeller
from arviz_base.utils import _var_names
from arviz_stats import ecdf, histogram, kde, qds

from arviz_plots.plot_collection import concat_model_dict, process_facet_dims
from arviz_plots.visuals import annotate_xy, hline, hspan, vline, vspan


def get_group(data, group, allow_missing=False):
    """Get a group from a Datatree or Dataset if possible and return a Dataset.

    Also supports InferenceData or dictionaries of Datasets.

    Parameters
    ----------
    data : DataTree, Dataset, InferenceData or mapping of {str : Dataset}
        Object from which to extract `group`
    group : hashable
        Id to be extracted. It is checked against the ``name`` attribute
        and attempted to use as key to get the `group` item from `data`
    allow_missing : bool, default False
        Return ``None`` if `group` can't be extracted instead of raising an error.

    Returns
    -------
    Dataset

    Raises
    ------
    KeyError
        If unable to access `group` from `data` and ``allow_missing=False``.
    """
    if isinstance(data, xr.Dataset):
        return data
    if hasattr(data, "name") and data.name == group:
        return data.dataset
    try:
        data = data[group]
    except KeyError:
        if not allow_missing:
            raise
        return None
    if isinstance(data, xr.Dataset):
        return data
    return data.dataset


def get_visual_kwargs(visuals, name, default=None):
    """Get the kwargs for a visual from the visuals dict.

    Parameterss
    ----------
    visuals : dict
        Dictionary of visuals.
    name : str
        Name of the visual to get kwargs for.
    default : dict or False, optional
        Default kwargs to return if the visual is not found in `visuals`.

    Returns
    -------
    dict
        The kwargs for the visual if found, otherwise the default value.

    """
    if default is None:
        default = {}
    vis_kwargs = copy(visuals.get(name, default))
    if vis_kwargs is True:
        vis_kwargs = {}
    return vis_kwargs


def process_group_variables_coords(dt, group, var_names, filter_vars, coords, allow_dict=True):
    """Process main input arguments of batteries included plotting functions."""
    if coords is None:
        coords = {}
    if isinstance(dt, dict) and not allow_dict:
        raise ValueError("Input data as dictionary not supported")
    if isinstance(dt, dict):
        distribution = {}
        all_vars = []
        all_data_vars = []
        for key, value in dt.items():
            new_var_names = _var_names(
                var_names, get_group(value, group), filter_vars, check_if_present=False
            )
            group_ds = get_group(value, group)
            if new_var_names is not None:
                data_vars = group_ds.data_vars
                available_vars = [v for v in data_vars if v in new_var_names]
                distribution[key] = group_ds[available_vars].sel(coords)
                all_vars.extend(new_var_names)
                all_data_vars.extend(data_vars)
            else:
                distribution[key] = group_ds.sel(coords)

        if var_names is not None:
            missing_vars = set(all_vars).difference(set(all_data_vars))
            if missing_vars:
                plural = "" if len(missing_vars) == 1 else "s"
                raise KeyError(f"variable{plural} {missing_vars} not found in any dataset")

        distribution = concat_model_dict(distribution)
    else:
        distribution = get_group(dt, group)
        var_names = _var_names(var_names, distribution, filter_vars)
        if var_names is not None:
            distribution = distribution[var_names]
        distribution = distribution.sel(coords)
    return distribution


def filter_aes(pc, aes_by_visuals, visual, sample_dims):
    """Simplified version of `filter_aes_full` that doesn't return "active_dimensions"."""
    reduce_dims, _, artist_aes, ignore_aes = filter_aes_full(
        pc, aes_by_visuals, visual, sample_dims
    )
    return reduce_dims, artist_aes, ignore_aes


def filter_aes_full(pc, aes_by_visuals, visual, sample_dims):
    """Split aesthetics and get relevant dimensions.

    Returns
    -------
    reduce_dims : list
        Dimensions that should be reduced for this visual.
        That is, all dimensions in `sample_dims` that are not
        mapped to any aesthetic.
    active_dims : list
        Dimensions that have either faceting or aesthetic mappings
        active for that visual. Should not be reduced and should have
        a groupby performed on them if computing summaries.
    artist_aes : iterable
    ignore_aes : set
    """
    artist_aes = aes_by_visuals.get(visual, {})
    pc_aes = pc.aes_set
    ignore_aes = set(pc_aes).difference(artist_aes)
    _, all_loop_dims = pc.update_aes(ignore_aes=ignore_aes)
    reduce_dims = [dim for dim in sample_dims if dim not in all_loop_dims]
    return reduce_dims, all_loop_dims, artist_aes, ignore_aes


def set_wrap_layout(pc_kwargs, plot_bknd, ds):
    """Set the figure size and handle column wrapping.

    Parameters
    ----------
    pc_kwargs : dict
        Plot collection kwargs
    plot_bknd : str
        Backend for plotting
    ds : Dataset
        Dataset to be plotted
    """
    figsize = pc_kwargs["figure_kwargs"].get("figsize", None)
    figsize_units = pc_kwargs["figure_kwargs"].get("figsize_units", "inches")
    pc_kwargs.setdefault("col_wrap", 4)
    col_wrap = pc_kwargs["col_wrap"]
    if figsize is None:
        num_plots = process_facet_dims(ds, pc_kwargs["cols"])[0]
        if num_plots < col_wrap:
            cols = num_plots
            rows = 1
        else:
            div_mod = divmod(num_plots, col_wrap)
            rows = div_mod[0] + (div_mod[1] != 0)
            cols = col_wrap
        figsize = plot_bknd.scale_fig_size(
            figsize,
            rows=rows,
            cols=cols,
            figsize_units=figsize_units,
        )
        figsize_units = "dots"

    pc_kwargs["figure_kwargs"]["figsize"] = figsize
    pc_kwargs["figure_kwargs"]["figsize_units"] = figsize_units
    return pc_kwargs


def set_grid_layout(pc_kwargs, plot_bknd, ds, num_rows=None, num_cols=None):
    """Set the figure size for the given number of rows and columns.

    Parameters
    ----------
    pc_kwargs : dict
        Plot collection kwargs
    plot_bknd : str
        Backend for plotting
    ds : Dataset
        Dataset to be plotted
    num_rows, num_cols : int, optional
        Take the number of rows or columns as the provided one irrespective
        of pc_kwargs
    """
    figsize = pc_kwargs["figure_kwargs"].get("figsize", None)
    figsize_units = pc_kwargs["figure_kwargs"].get("figsize_units", "inches")
    if figsize is None:
        if num_cols is None:
            num_cols = process_facet_dims(ds, pc_kwargs["cols"])[0]
        if num_rows is None:
            num_rows = process_facet_dims(ds, pc_kwargs["rows"])[0]
        figsize = plot_bknd.scale_fig_size(
            figsize,
            rows=num_rows,
            cols=num_cols,
            figsize_units=figsize_units,
        )
        figsize_units = "dots"

    pc_kwargs["figure_kwargs"]["figsize"] = figsize
    pc_kwargs["figure_kwargs"]["figsize_units"] = figsize_units
    return pc_kwargs


def _compute_func_da(func, da, active_dims, reduce_dims, kwargs=None):
    if kwargs is None:
        kwargs = {}
    groupby_dims = [
        dim
        for dim in active_dims
        if dim in da.dims and (len(np.unique(da.coords[dim])) != da.sizes[dim])
    ]
    if groupby_dims and func.__name__ == "histogram":
        raise ValueError(
            "Histogram computation doesn't support groupby behaviour yet."
            "Make sure coordinate values are unique or that dimension is not used for "
            "faceting nor it has aesthetics mapped to it."
        )
    if groupby_dims:
        da = da.groupby(groupby_dims)
    reduce_dims = [dim for dim in reduce_dims if dim in da.dims]
    with warnings.catch_warnings():
        if "model" in da.dims:
            warnings.filterwarnings("ignore", message="Your data appears to have a single")
        out_da = func(da, dim=reduce_dims, **kwargs)
    func_to_name = {"qds": "dot", "histogram": "hist"}
    out_da.attrs["kind"] = func_to_name.get(func.__name__, func.__name__)
    return out_da


def _compute_func(func, data, active_dims, reduce_dims, var_names=None, kwargs=None):
    """Compute given function for taking into account active and dimensions to reduce."""
    if var_names is None:
        var_names = data.data_vars
    out_das = []
    for var_name in var_names:
        out_da = _compute_func_da(
            func, data[var_name], active_dims=active_dims, reduce_dims=reduce_dims, kwargs=kwargs
        )
        out_da.name = var_name
        out_das.append(out_da)
    if not out_das:
        return xr.Dataset()
    return xr.merge(out_das, compat="override", join="outer")


def compute_dist(data, reduce_dims, active_dims, kind=None, stats=None):
    """Compute marginal uncertainty visualization data for different kinds and facet/aes combos.

    Parameters
    ----------
    data : Dataset
    reduce_dims : sequence of hashable
        Dimensions of `data` that should be reduced. They will not be present in the output.
    active_dims : sequence of hashable
        Dimensions of `data` used for faceting or with an aesthetic mapped to it.
        These will be partially reduced through a groupby if needed so the output continues
        to have these dimensions but with unique coordinate values.
    kind : {"auto", "kde", "hist", "ecdf", "dot"}, optional
        Kind of 1d marginal uncertainty visualisation to use.
        If ``None``, the ``plot.density_kind`` rcParam will be used.
        If "auto" the type of visualization will be variable dependent through some heuristics.
    stats : mapping, optional
        Stats dictionary argument to plots like :func:`~arviz_plots.plot_dist`.
        If the values stored in the "dist" key are a :class:`~xarray.Dataset`
        they will be returned as is as they are understood to be pre-computed visualization data.
        This argument will be ignored if `kind` is "auto".
    """
    if stats is None:
        stats = {}
    # quick exit if pre-computed elements in `stats`
    stats_dist_value = stats.get("dist", {})
    if isinstance(stats_dist_value, xr.Dataset):
        return stats_dist_value
    if kind is None:
        kind = rcParams["plot.density_kind"]
    if set(reduce_dims).intersection(active_dims):
        raise ValueError("'reduce_dims' and 'active_dims' can't share elements")
    if kind == "auto":
        out_das = []
        for da in data.values():
            reduced_size = np.prod([da.sizes[dim] for dim in reduce_dims if dim in da.dims])
            groupby_dims = [dim for dim in active_dims if dim in da.dims]
            if groupby_dims:
                reduced_size *= np.prod(
                    [
                        np.min(np.unique(da.coords[dim], return_counts=True)[1])
                        for dim in groupby_dims
                    ]
                )
            if reduced_size < 100:
                out_das.append(
                    _compute_func_da(ecdf, da, active_dims=active_dims, reduce_dims=reduce_dims)
                )
            elif da.dtype.kind == "f":
                out_das.append(
                    _compute_func_da(kde, da, active_dims=active_dims, reduce_dims=reduce_dims)
                )
            else:
                out_das.append(
                    _compute_func_da(
                        histogram,
                        da,
                        active_dims=active_dims,
                        reduce_dims=reduce_dims,
                        kwargs={"density": True},
                    )
                )
        return xr.merge(
            out_das,
            compat="override",
            join="outer",
        )
    if kind == "dot":
        return _compute_func(
            qds, data, active_dims=active_dims, reduce_dims=reduce_dims, kwargs=stats_dist_value
        )
    if kind == "ecdf":
        return _compute_func(
            ecdf, data, active_dims=active_dims, reduce_dims=reduce_dims, kwargs=stats_dist_value
        )
    if kind == "hist":
        return _compute_func(
            histogram,
            data,
            active_dims=active_dims,
            reduce_dims=reduce_dims,
            kwargs=stats_dist_value,
        )
    if kind == "kde":
        return _compute_func(
            kde, data, active_dims=active_dims, reduce_dims=reduce_dims, kwargs=stats_dist_value
        )
    raise ValueError(f"provided 'kind' {kind} is not valid")


def add_lines(
    plot_collection,
    values,
    orientation="vertical",
    aes_by_visuals=None,
    visuals=None,
    sample_dims=None,
    ref_dim="ref_dim",
    **kwargs,
):
    """Add lines.

    This function adds lines to a plot collection based on the provided values.
    It supports both vertical and horizontal lines, depending on the specified orientation.

    Parameters
    ----------
    plot_collection : PlotCollection
        Plot collection to which the lines will be added.
    values : int, float, tuple, list or dict
        Positions for the lines.
    orientation : str, default "vertical"
        The orientation of the lines, either "vertical" or "horizontal".
    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`.

        The default is to use an "overlay_ref" aesthetic for all elements.

        It is possible to request aesthetics without mappings defined in the
        provided `plot_collection`. In those cases, a mapping of "ref_dim" to the requested
        aesthetic will be automatically added.
    visuals : mapping of {str : mapping or bool}, optional
        Valid keys are:

        * "ref_line" -> passed to :func:`~arviz_plots.visuals.vline` for vertical `orientation`
          and to :func:`~arviz_plots.visuals.hline` for horizontal `orientation`
        * "ref_text" -> TODO

    sample_dims : list, optional
        Dimensions that should not be added to the Dataset generated from
        `values` via :func:`arviz_base.references_to_dataset`.
        Defaults to all dimensions in ``plot_collection.data`` that are not ``facet_dims``
    ref_dim : str, optional
        Specifies the name of the dimension for the line values.
        Defaults to "ref_dim".
    **kwargs : mapping of {str : sequence}, optional
        Mapping of aesthetic keys to the values to be used in their mapping.
        See :func:`~arviz_plots.PlotCollection.generate_aes_dt` for more details.

    Returns
    -------
    plot_collection : PlotCollection
        Plot collection with the lines added.

    Examples
    --------
    Add lines at values 0 and 5 for all variables.

    .. plot::
        :context: close-figs

        >>> from arviz_plots import plot_dist, add_lines, style
        >>> style.use("arviz-variat")
        >>> from arviz_base import load_arviz_data
        >>> dt = load_arviz_data('centered_eight')
        >>> pc = plot_dist(
        >>>     dt,
        >>>     kind="ecdf",
        >>>     var_names=["mu"],
        >>> )
        >>> add_lines(pc, values=[0, 5])
    """
    if visuals is None:
        visuals = {}
    if aes_by_visuals is None:
        aes_by_visuals = {}
    else:
        aes_by_visuals = aes_by_visuals.copy()
    aes_by_visuals["ref_line"] = aes_by_visuals.get("ref_line", ["overlay_ref"])
    aes_by_visuals["ref_text"] = aes_by_visuals.get("ref_text", ["overlay_ref"])
    if sample_dims is None:
        sample_dims = list(set(plot_collection.data.dims).difference(plot_collection.facet_dims))
    if isinstance(sample_dims, str):
        sample_dims = [sample_dims]

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

    plot_func = vline if orientation == "vertical" else hline

    ref_ds = references_to_dataset(
        values, plot_collection.data, sample_dims=sample_dims, ref_dim=ref_dim
    )
    requested_aes = (
        set(aes_by_visuals["ref_line"])
        .union(aes_by_visuals["ref_text"])
        .difference(plot_collection.aes_set)
    )
    if ref_dim in ref_ds.dims:
        for aes_key in requested_aes:
            aes_values = np.array(plot_bknd.get_default_aes(aes_key, ref_ds.sizes[ref_dim], kwargs))
            plot_collection.update_aes_from_dataset(
                aes_key,
                xr.Dataset(
                    {
                        var_name: (ref_dim, aes_values)
                        for var_name in plot_collection.data.data_vars
                    },
                    coords={ref_dim: ref_ds[ref_dim]},
                ),
            )

    _, ref_aes, ref_ignore = filter_aes(plot_collection, aes_by_visuals, "ref_line", sample_dims)
    ref_kwargs = get_visual_kwargs(visuals, "ref_line")
    if ref_kwargs is not False:
        if "color" not in ref_aes:
            ref_kwargs.setdefault("color", "B2")
        if "linestyle" not in ref_aes:
            ref_kwargs.setdefault("linestyle", "C1")

        plot_collection.map(
            plot_func,
            "ref_line",
            data=ref_ds,
            ignore_aes=ref_ignore,
            **ref_kwargs,
        )
    return plot_collection


def add_bands(
    plot_collection,
    values,
    orientation="vertical",
    aes_by_visuals=None,
    visuals=None,
    sample_dims=None,
    ref_dim=None,
    **kwargs,
):
    """Add bands.

    This function adds bands (shared areas) to a plot collection based on the provided
    values. It supports both vertical and horizontal bands, depending on the
    specified orientation.

    Parameters
    ----------
    plot_collection : PlotCollection
        Plot collection to which the bands will be added.
    values : tuple, list or dict
        Start and end values for the bands to be plotted.
    orientation : str, default "vertical"
        The orientation of the bands, either "vertical" or "horizontal".
    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`.

        The default is to use an "overlay_band" aesthetic for all elements.

        It is possible to request aesthetics without mappings defined in the
        provided `plot_collection`. In those cases, a mapping of the dimensions in
        `ref_dim` minus its last element to the requested aesthetic will be
        automatically added.
    visuals : mapping of {str : mapping or bool}, optional
        Valid keys are:

        * "ref_band" -> passed to :func:`~arviz_plots.visuals.vspan` for vertical `orientation`
          and to :func:`~arviz_plots.visuals.hspan` for horizontal `orientation`

    sample_dims : list, optional
        Dimensions that should not be added to the Dataset generated from
        `values` via :func:`arviz_base.references_to_dataset`.
        Defaults to all dimensions in ``plot_collection.data`` that are not ``facet_dims``
    ref_dim : list, optional
        List of dimension names that define the axes along which the band values are stored.
        These dimensions are used to align or compare input data with band data.
        Defaults to ["ref_dim", "band_dim"].
    **kwargs : sequence, optional
        Mapping of aesthetic keys to the values to be used in their mapping.
        See :func:`~arviz_plots.PlotCollection.generate_aes_dt` for more details.

    Returns
    -------
    plot_collection : PlotCollection
        Plot collection with the bands added.

    Examples
    --------
    Add two bands for the theta variable, one from -2 to 2 and the other from -5 to 5.

    .. plot::
        :context: close-figs

        >>> from arviz_plots import plot_dist, add_bands, style
        >>> style.use("arviz-variat")
        >>> from arviz_base import load_arviz_data
        >>> dt = load_arviz_data('centered_eight')
        >>> pc = plot_dist(dt)
        >>> add_bands(pc, values=[[-2, 2], [-5, 5]])
    """
    if visuals is None:
        visuals = {}
    if aes_by_visuals is None:
        aes_by_visuals = {}
    else:
        aes_by_visuals = aes_by_visuals.copy()
    aes_by_visuals["ref_band"] = aes_by_visuals.get("ref_band", ["overlay_band"])
    if sample_dims is None:
        sample_dims = list(set(plot_collection.data.dims).difference(plot_collection.facet_dims))
    if isinstance(sample_dims, str):
        sample_dims = [sample_dims]
    if ref_dim is None:
        ref_dim = ["ref_dim", "band_dim"]

    plot_func = vspan if orientation == "vertical" else hspan

    ref_ds = references_to_dataset(
        values, plot_collection.data, sample_dims=sample_dims, ref_dim=ref_dim
    )

    requested_aes = set(aes_by_visuals["ref_band"]).difference(plot_collection.aes_set)
    *ref_dim, band_dim = ref_dim
    if ref_ds.sizes[band_dim] != 2:
        raise ValueError(
            f"Expected dimension '{band_dim}' in reference dataset to have size 2, "
            f"but found size {ref_ds.sizes[band_dim]}"
        )
    aes_dt = plot_collection.generate_aes_dt(
        {aes: ref_dim for aes in requested_aes}, ref_ds, **kwargs
    )
    for aes, child in aes_dt.children.items():
        plot_collection.update_aes_from_dataset(aes, child.dataset)

    _, ref_aes, ref_ignore = filter_aes(plot_collection, aes_by_visuals, "ref_band", sample_dims)
    ref_kwargs = get_visual_kwargs(visuals, "ref_band")
    if ref_kwargs is not False:
        if "color" not in ref_aes:
            ref_kwargs.setdefault("color", "B2")
        if "alpha" not in ref_aes:
            ref_kwargs.setdefault("alpha", 0.25)
        plot_collection.map(plot_func, "ref_band", data=ref_ds, ignore_aes=ref_ignore, **ref_kwargs)

    return plot_collection


def format_coords_as_labels(data, skip_dims=None, labeller=None):
    """Format 1D or multi-D dataarray coords as string labels.

    Parameters
    ----------
    data : xr.DataArray
        DataArray whose coordinates will be converted to labels.
    skip_dims : str or list_like, optional
        Dimensions whose values should not be included in the labels.
    labeller : BaseLabeller, optional
        Labeller instance to use for formatting. If None, defaults to BaseLabeller().
        Users can pass different labellers to control whether indices or coordinate
        values are shown (e.g., IdxLabeller for indices).

    Returns
    -------
    ndarray of str
        Array of coordinate labels with the same flattened shape as the input.
    """
    if labeller is None:
        labeller = BaseLabeller()

    if skip_dims is None:
        skip_dims = []
    elif isinstance(skip_dims, str):
        skip_dims = [skip_dims]

    dims_to_include = [dim for dim in data.dims if dim not in skip_dims]

    if not dims_to_include:
        return np.array([], dtype=object)

    shape = [data.sizes[dim] for dim in dims_to_include]
    n_points = int(np.prod(shape))

    index_arrays = [
        arr.ravel() for arr in np.meshgrid(*[np.arange(s) for s in shape], indexing="ij")
    ]
    coord_arrays = [data.coords[dim].values for dim in dims_to_include]

    labels = [
        labeller.sel_to_str(
            {dim: coord_arrays[j][index_arrays[j][i]] for j, dim in enumerate(dims_to_include)},
            {dim: int(index_arrays[j][i]) for j, dim in enumerate(dims_to_include)},
        )
        for i in range(n_points)
    ]
    return np.array(labels, dtype=object)


def annotate_bin_text(da, target, x, y, count_da, n_da, bin_format, **kwargs):
    """Format and annotate bin text with count and percentage.

    Parameters
    ----------
    da : xr.DataArray
        Data array to annotate
    target : Axes
        Target axes for annotation
    x : float
        X-position for annotation
    y : float
        Y-position for annotation
    count_da : int or xr.DataArray
        Count value for the bin
    n_da : int or xr.DataArray
        Total count value
    bin_format : str
        Format string for bin text (supports {count} and {pct})
    **kwargs
        Additional keyword arguments passed to annotate_xy

    Returns
    -------
    Artist
        The annotation artist
    """
    if hasattr(count_da, "values"):
        arr = count_da.values
        count_val = int(arr[()] if arr.ndim == 0 else arr)
    else:
        count_val = int(count_da)

    if hasattr(n_da, "values"):
        arr = n_da.values
        n_val = int(arr[()] if arr.ndim == 0 else arr)
    else:
        n_val = int(n_da)

    pct = (count_val / n_val * 100) if n_val > 0 else 0.0
    text_str = bin_format.format(count=count_val, pct=pct)
    return annotate_xy(da, target, x=x, y=y, text=text_str, **kwargs)


def enable_hover_labels(backend, plot_collection, hover_format, labels, colors, values):
    """Set up interactive hover annotations for scatter plots on matplotlib backends.

    Parameters
    ----------
    backend : str
        The plotting backend being used. Only "matplotlib" is supported.
    plot_collection : PlotCollection
        Plot collection containing the visualization elements.
    hover_format : str
        Format string template for hover annotation text.
    labels : xr.DataArray
        Labels corresponding to each data point as a DataArray. It is subset per
        facet/aesthetic combination automatically.
    colors : xr.DataArray or None
        Colors for each data point as a DataArray, or None to extract from scatter artist.
    values : xr.DataArray or None
        Values to display in hover text as a DataArray, or None to use y-coordinates.
    """
    if backend != "matplotlib":
        return

    try:
        fig = plot_collection.viz["figure"].item()
    except KeyError:
        return
    if fig is None:
        return

    if not hasattr(fig.canvas, "mpl_connect"):
        warnings.warn(
            "hover labels are only available with interactive backends. "
            "To switch to an interactive backend from IPython or Jupyter, use `%matplotlib`.",
            UserWarning,
        )
        return

    try:
        scatter_tree = plot_collection.viz["khat"]
        scatter_da = scatter_tree["pareto_k"]
    except KeyError:
        return

    if scatter_da.ndim == 0:
        scatter = scatter_da.item()
        axis = plot_collection.get_target("pareto_k", {})

        if scatter is not None and axis is not None:
            label_arr = labels.to_numpy().ravel()
            label_list = [str(label) for label in label_arr]
            value_arr = values.to_numpy().ravel() if values is not None else None
            color_arr = colors.to_numpy().ravel() if colors is not None else None

            hover_labels(fig, axis, scatter, label_list, hover_format, color_arr, value_arr)
    else:
        scatter_ds = scatter_da.to_dataset(name="artist")
        for _, sel, _ in xarray_sel_iter(scatter_ds, skip_dims=set()):
            scatter = scatter_ds["artist"].sel(sel).item()
            if scatter is None:
                continue
            axis = plot_collection.get_target("pareto_k", sel)
            if axis is None:
                continue

            sel_args = {dim: val for dim, val in sel.items() if dim in labels.dims}
            label_subset = labels.sel(sel_args) if sel_args else labels
            label_arr = label_subset.to_numpy().ravel()

            offsets = scatter.get_offsets()
            if offsets is None or not offsets.size:
                continue
            if len(label_arr) != len(offsets):
                continue

            label_list = [str(label) for label in label_arr]

            value_arr = None
            if values is not None:
                sel_args = {dim: val for dim, val in sel.items() if dim in values.dims}
                value_subset = values.sel(sel_args) if sel_args else values
                value_arr = value_subset.to_numpy().ravel()

            color_arr = None
            if colors is not None:
                sel_args = {dim: val for dim, val in sel.items() if dim in colors.dims}
                color_subset = colors.sel(sel_args) if sel_args else colors
                color_arr = color_subset.to_numpy().ravel()

            hover_labels(fig, axis, scatter, label_list, hover_format, color_arr, value_arr)


def hover(
    event, annot, ax, scatter, fig, offsets, labels, values, hover_format, colors, offset_distance
):
    """Handle mouse hover events to show or hide data point annotations.

    Parameters
    ----------
    event : MouseEvent
        Matplotlib mouse event containing cursor position.
    annot : Annotation
        Matplotlib annotation object to show/hide and update.
    ax : Axes
        Matplotlib axes containing the scatter plot.
    scatter : PathCollection
        Scatter plot artist whose points are being tracked.
    fig : Figure
        Matplotlib figure containing the plot.
    offsets : ndarray
        Array of (x, y) coordinates for all scatter plot points.
    labels : array_like
        Array of string labels for each data point.
    values : array_like or None
        Array of numeric values to display. If None, y-coordinates are used.
    hover_format : str
        Format string template for the annotation text.
    colors : ndarray or None
        Array of RGBA colors for annotation backgrounds.
    offset_distance : float
        Distance in points to offset the annotation from the cursor position.
    """
    vis = annot.get_visible()
    if event.inaxes == ax:
        cont, ind = scatter.contains(event)
        if cont:
            idx = ind["ind"][0]
            pos = offsets[idx]
            label = labels[idx] if idx < len(labels) else str(idx)
            value = values[idx] if values is not None and idx < len(values) else pos[1]
            annot.xy = pos
            xmid = np.mean(ax.get_xlim())
            ymid = np.mean(ax.get_ylim())
            annot.set_position(
                (
                    -offset_distance if pos[0] > xmid else offset_distance,
                    -offset_distance if pos[1] > ymid else offset_distance,
                )
            )
            annot.set_text(_format_hover_text(hover_format, idx, label, value))

            if colors is not None and idx < len(colors):
                annot.get_bbox_patch().set_facecolor(colors[idx])

            annot.set_ha("right" if pos[0] > xmid else "left")
            annot.set_va("top" if pos[1] > ymid else "bottom")
            annot.set_visible(True)
            fig.canvas.draw_idle()
        elif vis:
            annot.set_visible(False)
            fig.canvas.draw_idle()


def hover_labels(fig, ax, scatter, labels, hover_format, colors, values):
    """Configure hover annotation display for scatter plot data points.

    Parameters
    ----------
    fig : Figure
        Matplotlib figure to attach the hover event handler.
    ax : Axes
        Matplotlib axes containing the scatter plot.
    scatter : PathCollection
        Scatter plot artist whose data points trigger hover annotations.
    labels : array_like
        Array of string labels for each data point.
    hover_format : str
        Format string template for annotation text.
    colors : ndarray or None
        Array of RGBA colors for annotation box backgrounds.
    values : array_like or None
        Array of numeric values to display. If None, y-coordinates are used.
    """
    offsets = scatter.get_offsets()
    if offsets is None or not offsets.size:
        return

    annot = ax.annotate(
        "",
        xy=(0, 0),
        xytext=(0, 0),
        textcoords="offset points",
        bbox={"boxstyle": "round", "fc": "w", "alpha": 0.4},
        arrowprops={"arrowstyle": "->"},
    )
    annot.set_visible(False)
    offset_distance = 10

    fig.canvas.mpl_connect(
        "motion_notify_event",
        lambda event: hover(
            event,
            annot,
            ax,
            scatter,
            fig,
            offsets,
            labels,
            values,
            hover_format,
            colors,
            offset_distance,
        ),
    )


def _format_hover_text(template, index, label, value):
    """Format hover annotation text using named placeholders."""
    if hasattr(value, "item"):
        value = value.item()
    return template.format(index=index, label=label, value=value)
