"""Interoperability with `matplotlib`."""
# pylint: disable=import-error

# This module was adapted from both the `astropy` and `pint` package's
# implementations.

__all__ = ["setup_matplotlib_support_for_unxt", "UnxtConverter"]


from collections.abc import Iterable, Sized
from dataclasses import dataclass, field
from typing import Any

import matplotlib.units
from jaxtyping import Array
from matplotlib.axes import Axes
from zeroth import zeroth

from unxt._quantity.base import AbstractQuantity
from unxt._quantity.fast import UncheckedQuantity


@dataclass
class UnxtConverter(matplotlib.units.ConversionInterface):  # type: ignore[misc]
    """Support `unxt` in `matplotlib`'s unit conversion framework.

    This class is a subclass of `matplotlib.units.ConversionInterface`
    and is used to convert `unxt.Quantity` instances for use with
    `matplotlib`.

    """

    axisinfo_kw: dict[str, Any] = field(
        default_factory=lambda: {"format": "latex_inline"}
    )
    """Keyword arguments to use when making the :meth:`matplotlib.units.AxisInfo`."""

    def convert(self, obj: Any, unit: Any, axis: Axes) -> Array | list[Array]:
        """Convert *obj* using *unit* for the specified *axis*."""
        # Hot-path Quantity
        if isinstance(obj, AbstractQuantity):
            return obj.to_units_value(unit)
        # Need to recurse (singly) into iterables
        if isinstance(obj, Iterable):
            return [self._convert_value(v, unit, axis) for v in obj]

        return self._convert_value(obj, unit, axis)

    @staticmethod
    def _convert_value(obj: Any, unit: Any, axis: Axes) -> Array:
        """Handle converting using attached unit or falling back to axis units."""
        if isinstance(obj, AbstractQuantity):
            return obj.to_units_value(unit)

        return UncheckedQuantity.constructor(obj, axis.get_units()).to_units_value(unit)

    def axisinfo(self, unit: Any, _: Axes) -> matplotlib.units.AxisInfo:
        """Return axis information for this particular unit."""
        fmt = self.axisinfo_kw.get("format", "latex_inline")
        return matplotlib.units.AxisInfo(label=unit.to_string(fmt))

    @staticmethod
    def default_units(x: Any, _: Axes) -> Any:
        """Get the default unit to use for the given combination of unit and axis."""
        if hasattr(x, "unit"):
            return x.unit
        if isinstance(x, Iterable) and isinstance(x, Sized):
            x = zeroth(x)
        return getattr(x, "units", None)


def setup_matplotlib_support_for_unxt(*, enable: bool = True) -> None:
    """Set up matplotlib's unit support for `unxt`.

    Parameters
    ----------
    enable : bool, optional keyword-only
        Whether support should be enabled or disabled.

    """
    if enable:
        matplotlib.units.registry[AbstractQuantity] = UnxtConverter()
    else:
        matplotlib.units.registry.pop(AbstractQuantity, None)
