from __future__ import annotations
import re
from copy import copy
from collections.abc import Sequence
from dataclasses import dataclass
from functools import partial
from typing import Any, Callable, Tuple, Optional, ClassVar

import numpy as np
import matplotlib as mpl
from matplotlib.ticker import (
    Locator,
    Formatter,
    AutoLocator,
    AutoMinorLocator,
    FixedLocator,
    LinearLocator,
    LogLocator,
    SymmetricalLogLocator,
    MaxNLocator,
    MultipleLocator,
    EngFormatter,
    FuncFormatter,
    LogFormatterSciNotation,
    ScalarFormatter,
    StrMethodFormatter,
)
from matplotlib.dates import (
    AutoDateLocator,
    AutoDateFormatter,
    ConciseDateFormatter,
)
from matplotlib.axis import Axis
from matplotlib.scale import ScaleBase
from pandas import Series

from seaborn._core.rules import categorical_order
from seaborn._core.typing import Default, default

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from seaborn._core.plot import Plot
    from seaborn._core.properties import Property
    from numpy.typing import ArrayLike, NDArray

    TransFuncs = Tuple[
        Callable[[ArrayLike], ArrayLike], Callable[[ArrayLike], ArrayLike]
    ]

    # TODO Reverting typing to Any as it was proving too complicated to
    # work out the right way to communicate the types to mypy. Revisit!
    Pipeline = Sequence[Optional[Callable[[Any], Any]]]


class Scale:
    """Base class for objects that map data values to visual properties."""

    values: tuple | str | list | dict | None

    _priority: ClassVar[int]
    _pipeline: Pipeline
    _matplotlib_scale: ScaleBase
    _spacer: staticmethod
    _legend: tuple[list[Any], list[str]] | None

    def __post_init__(self):

        self._tick_params = None
        self._label_params = None
        self._legend = None

    def tick(self):
        raise NotImplementedError()

    def label(self):
        raise NotImplementedError()

    def _get_locators(self):
        raise NotImplementedError()

    def _get_formatter(self, locator: Locator | None = None):
        raise NotImplementedError()

    def _get_scale(self, name: str, forward: Callable, inverse: Callable):

        major_locator, minor_locator = self._get_locators(**self._tick_params)
        major_formatter = self._get_formatter(major_locator, **self._label_params)

        class InternalScale(mpl.scale.FuncScale):
            def set_default_locators_and_formatters(self, axis):
                axis.set_major_locator(major_locator)
                if minor_locator is not None:
                    axis.set_minor_locator(minor_locator)
                axis.set_major_formatter(major_formatter)

        return InternalScale(name, (forward, inverse))

    def _spacing(self, x: Series) -> float:
        space = self._spacer(x)
        if np.isnan(space):
            # This happens when there is no variance in the orient coordinate data
            # Not exactly clear what the right default is, but 1 seems reasonable?
            return 1
        return space

    def _setup(
        self, data: Series, prop: Property, axis: Axis | None = None,
    ) -> Scale:
        raise NotImplementedError()

    def _finalize(self, p: Plot, axis: Axis) -> None:
        """Perform scale-specific axis tweaks after adding artists."""
        pass

    def __call__(self, data: Series) -> ArrayLike:

        trans_data: Series | NDArray | list

        # TODO sometimes we need to handle scalars (e.g. for Line)
        # but what is the best way to do that?
        scalar_data = np.isscalar(data)
        if scalar_data:
            trans_data = np.array([data])
        else:
            trans_data = data

        for func in self._pipeline:
            if func is not None:
                trans_data = func(trans_data)

        if scalar_data:
            return trans_data[0]
        else:
            return trans_data

    @staticmethod
    def _identity():

        class Identity(Scale):
            _pipeline = []
            _spacer = None
            _legend = None
            _matplotlib_scale = None

        return Identity()


@dataclass
class Boolean(Scale):
    """
    A scale with a discrete domain of True and False values.

    The behavior is similar to the :class:`Nominal` scale, but property
    mappings and legends will use a [True, False] ordering rather than
    a sort using numeric rules. Coordinate variables accomplish this by
    inverting axis limits so as to maintain underlying numeric positioning.
    Input data are cast to boolean values, respecting missing data.

    """
    values: tuple | list | dict | None = None

    _priority: ClassVar[int] = 3

    def _setup(
        self, data: Series, prop: Property, axis: Axis | None = None,
    ) -> Scale:

        new = copy(self)
        if new._tick_params is None:
            new = new.tick()
        if new._label_params is None:
            new = new.label()

        def na_safe_cast(x):
            # TODO this doesn't actually need to be a closure
            if np.isscalar(x):
                return float(bool(x))
            else:
                if hasattr(x, "notna"):
                    # Handle pd.NA; np<>pd interop with NA is tricky
                    use = x.notna().to_numpy()
                else:
                    use = np.isfinite(x)
                out = np.full(len(x), np.nan, dtype=float)
                out[use] = x[use].astype(bool).astype(float)
                return out

        new._pipeline = [na_safe_cast, prop.get_mapping(new, data)]
        new._spacer = _default_spacer
        if prop.legend:
            new._legend = [True, False], ["True", "False"]

        forward, inverse = _make_identity_transforms()
        mpl_scale = new._get_scale(str(data.name), forward, inverse)

        axis = PseudoAxis(mpl_scale) if axis is None else axis
        mpl_scale.set_default_locators_and_formatters(axis)
        new._matplotlib_scale = mpl_scale

        return new

    def _finalize(self, p: Plot, axis: Axis) -> None:

        # We want values to appear in a True, False order but also want
        # True/False to be drawn at 1/0 positions respectively to avoid nasty
        # surprises if additional artists are added through the matplotlib API.
        # We accomplish this using axis inversion akin to what we do in Nominal.

        ax = axis.axes
        name = axis.axis_name
        axis.grid(False, which="both")
        if name not in p._limits:
            nticks = len(axis.get_major_ticks())
            lo, hi = -.5, nticks - .5
            if name == "x":
                lo, hi = hi, lo
            set_lim = getattr(ax, f"set_{name}lim")
            set_lim(lo, hi, auto=None)

    def tick(self, locator: Locator | None = None):
        new = copy(self)
        new._tick_params = {"locator": locator}
        return new

    def label(self, formatter: Formatter | None = None):
        new = copy(self)
        new._label_params = {"formatter": formatter}
        return new

    def _get_locators(self, locator):
        if locator is not None:
            return locator
        return FixedLocator([0, 1]), None

    def _get_formatter(self, locator, formatter):
        if formatter is not None:
            return formatter
        return FuncFormatter(lambda x, _: str(bool(x)))


@dataclass
class Nominal(Scale):
    """
    A categorical scale without relative importance / magnitude.
    """
    # Categorical (convert to strings), un-sortable

    values: tuple | str | list | dict | None = None
    order: list | None = None

    _priority: ClassVar[int] = 4

    def _setup(
        self, data: Series, prop: Property, axis: Axis | None = None,
    ) -> Scale:

        new = copy(self)
        if new._tick_params is None:
            new = new.tick()
        if new._label_params is None:
            new = new.label()

        # TODO flexibility over format() which isn't great for numbers / dates
        stringify = np.vectorize(format, otypes=["object"])

        units_seed = categorical_order(data, new.order)

        # TODO move to Nominal._get_scale?
        # TODO this needs some more complicated rethinking about how to pass
        # a unit dictionary down to these methods, along with how much we want
        # to invest in their API. What is it useful for tick() to do here?
        # (Ordinal may be different if we draw that contrast).
        # Any customization we do to allow, e.g., label wrapping will probably
        # require defining our own Formatter subclass.
        # We could also potentially implement auto-wrapping in an Axis subclass
        # (see Axis.draw ... it already is computing the bboxes).
        # major_locator, minor_locator = new._get_locators(**new._tick_params)
        # major_formatter = new._get_formatter(major_locator, **new._label_params)

        class CatScale(mpl.scale.LinearScale):
            def set_default_locators_and_formatters(self, axis):
                ...
                # axis.set_major_locator(major_locator)
                # if minor_locator is not None:
                #     axis.set_minor_locator(minor_locator)
                # axis.set_major_formatter(major_formatter)

        mpl_scale = CatScale(data.name)
        if axis is None:
            axis = PseudoAxis(mpl_scale)

            # TODO Currently just used in non-Coordinate contexts, but should
            # we use this to (A) set the padding we want for categorial plots
            # and (B) allow the values parameter for a Coordinate to set xlim/ylim
            axis.set_view_interval(0, len(units_seed) - 1)

        new._matplotlib_scale = mpl_scale

        # TODO array cast necessary to handle float/int mixture, which we need
        # to solve in a more systematic way probably
        # (i.e. if we have [1, 2.5], do we want [1.0, 2.5]? Unclear)
        axis.update_units(stringify(np.array(units_seed)))

        # TODO define this more centrally
        def convert_units(x):
            # TODO only do this with explicit order?
            # (But also category dtype?)
            # TODO isin fails when units_seed mixes numbers and strings (numpy error?)
            # but np.isin also does not seem any faster? (Maybe not broadcasting in C)
            # keep = x.isin(units_seed)
            keep = np.array([x_ in units_seed for x_ in x], bool)
            out = np.full(len(x), np.nan)
            out[keep] = axis.convert_units(stringify(x[keep]))
            return out

        new._pipeline = [convert_units, prop.get_mapping(new, data)]
        new._spacer = _default_spacer

        if prop.legend:
            new._legend = units_seed, list(stringify(units_seed))

        return new

    def _finalize(self, p: Plot, axis: Axis) -> None:

        ax = axis.axes
        name = axis.axis_name
        axis.grid(False, which="both")
        if name not in p._limits:
            nticks = len(axis.get_major_ticks())
            lo, hi = -.5, nticks - .5
            if name == "y":
                lo, hi = hi, lo
            set_lim = getattr(ax, f"set_{name}lim")
            set_lim(lo, hi, auto=None)

    def tick(self, locator: Locator | None = None) -> Nominal:
        """
        Configure the selection of ticks for the scale's axis or legend.

        .. note::
            This API is under construction and will be enhanced over time.
            At the moment, it is probably not very useful.

        Parameters
        ----------
        locator : :class:`matplotlib.ticker.Locator` subclass
            Pre-configured matplotlib locator; other parameters will not be used.

        Returns
        -------
        Copy of self with new tick configuration.

        """
        new = copy(self)
        new._tick_params = {"locator": locator}
        return new

    def label(self, formatter: Formatter | None = None) -> Nominal:
        """
        Configure the selection of labels for the scale's axis or legend.

        .. note::
            This API is under construction and will be enhanced over time.
            At the moment, it is probably not very useful.

        Parameters
        ----------
        formatter : :class:`matplotlib.ticker.Formatter` subclass
            Pre-configured matplotlib formatter; other parameters will not be used.

        Returns
        -------
        scale
            Copy of self with new tick configuration.

        """
        new = copy(self)
        new._label_params = {"formatter": formatter}
        return new

    def _get_locators(self, locator):

        if locator is not None:
            return locator, None

        locator = mpl.category.StrCategoryLocator({})

        return locator, None

    def _get_formatter(self, locator, formatter):

        if formatter is not None:
            return formatter

        formatter = mpl.category.StrCategoryFormatter({})

        return formatter


@dataclass
class Ordinal(Scale):
    # Categorical (convert to strings), sortable, can skip ticklabels
    ...


@dataclass
class Discrete(Scale):
    # Numeric, integral, can skip ticks/ticklabels
    ...


@dataclass
class ContinuousBase(Scale):

    values: tuple | str | None = None
    norm: tuple | None = None

    def _setup(
        self, data: Series, prop: Property, axis: Axis | None = None,
    ) -> Scale:

        new = copy(self)
        if new._tick_params is None:
            new = new.tick()
        if new._label_params is None:
            new = new.label()

        forward, inverse = new._get_transform()

        mpl_scale = new._get_scale(str(data.name), forward, inverse)

        if axis is None:
            axis = PseudoAxis(mpl_scale)
            axis.update_units(data)

        mpl_scale.set_default_locators_and_formatters(axis)
        new._matplotlib_scale = mpl_scale

        normalize: Optional[Callable[[ArrayLike], ArrayLike]]
        if prop.normed:
            if new.norm is None:
                vmin, vmax = data.min(), data.max()
            else:
                vmin, vmax = new.norm
            vmin, vmax = map(float, axis.convert_units((vmin, vmax)))
            a = forward(vmin)
            b = forward(vmax) - forward(vmin)

            def normalize(x):
                return (x - a) / b

        else:
            normalize = vmin = vmax = None

        new._pipeline = [
            axis.convert_units,
            forward,
            normalize,
            prop.get_mapping(new, data)
        ]

        def spacer(x):
            x = x.dropna().unique()
            if len(x) < 2:
                return np.nan
            return np.min(np.diff(np.sort(x)))
        new._spacer = spacer

        # TODO How to allow disabling of legend for all uses of property?
        # Could add a Scale parameter, or perhaps Scale.suppress()?
        # Are there other useful parameters that would be in Scale.legend()
        # besides allowing Scale.legend(False)?
        if prop.legend:
            axis.set_view_interval(vmin, vmax)
            locs = axis.major.locator()
            locs = locs[(vmin <= locs) & (locs <= vmax)]
            # Avoid having an offset / scientific notation in a legend
            # as we don't represent that anywhere so it ends up incorrect.
            # This could become an option (e.g. Continuous.label(offset=True))
            # in which case we would need to figure out how to show it.
            if hasattr(axis.major.formatter, "set_useOffset"):
                axis.major.formatter.set_useOffset(False)
            if hasattr(axis.major.formatter, "set_scientific"):
                axis.major.formatter.set_scientific(False)
            labels = axis.major.formatter.format_ticks(locs)
            new._legend = list(locs), list(labels)

        return new

    def _get_transform(self):

        arg = self.trans

        def get_param(method, default):
            if arg == method:
                return default
            return float(arg[len(method):])

        if arg is None:
            return _make_identity_transforms()
        elif isinstance(arg, tuple):
            return arg
        elif isinstance(arg, str):
            if arg == "ln":
                return _make_log_transforms()
            elif arg == "logit":
                base = get_param("logit", 10)
                return _make_logit_transforms(base)
            elif arg.startswith("log"):
                base = get_param("log", 10)
                return _make_log_transforms(base)
            elif arg.startswith("symlog"):
                c = get_param("symlog", 1)
                return _make_symlog_transforms(c)
            elif arg.startswith("pow"):
                exp = get_param("pow", 2)
                return _make_power_transforms(exp)
            elif arg == "sqrt":
                return _make_sqrt_transforms()
            else:
                raise ValueError(f"Unknown value provided for trans: {arg!r}")


@dataclass
class Continuous(ContinuousBase):
    """
    A numeric scale supporting norms and functional transforms.
    """
    values: tuple | str | None = None
    trans: str | TransFuncs | None = None

    # TODO Add this to deal with outliers?
    # outside: Literal["keep", "drop", "clip"] = "keep"

    _priority: ClassVar[int] = 1

    def tick(
        self,
        locator: Locator | None = None, *,
        at: Sequence[float] | None = None,
        upto: int | None = None,
        count: int | None = None,
        every: float | None = None,
        between: tuple[float, float] | None = None,
        minor: int | None = None,
    ) -> Continuous:
        """
        Configure the selection of ticks for the scale's axis or legend.

        Parameters
        ----------
        locator : :class:`matplotlib.ticker.Locator` subclass
            Pre-configured matplotlib locator; other parameters will not be used.
        at : sequence of floats
            Place ticks at these specific locations (in data units).
        upto : int
            Choose "nice" locations for ticks, but do not exceed this number.
        count : int
            Choose exactly this number of ticks, bounded by `between` or axis limits.
        every : float
            Choose locations at this interval of separation (in data units).
        between : pair of floats
            Bound upper / lower ticks when using `every` or `count`.
        minor : int
            Number of unlabeled ticks to draw between labeled "major" ticks.

        Returns
        -------
        scale
            Copy of self with new tick configuration.

        """
        # Input checks
        if locator is not None and not isinstance(locator, Locator):
            raise TypeError(
                f"Tick locator must be an instance of {Locator!r}, "
                f"not {type(locator)!r}."
            )
        log_base, symlog_thresh = self._parse_for_log_params(self.trans)
        if log_base or symlog_thresh:
            if count is not None and between is None:
                raise RuntimeError("`count` requires `between` with log transform.")
            if every is not None:
                raise RuntimeError("`every` not supported with log transform.")

        new = copy(self)
        new._tick_params = {
            "locator": locator,
            "at": at,
            "upto": upto,
            "count": count,
            "every": every,
            "between": between,
            "minor": minor,
        }
        return new

    def label(
        self,
        formatter: Formatter | None = None, *,
        like: str | Callable | None = None,
        base: int | None | Default = default,
        unit: str | None = None,
    ) -> Continuous:
        """
        Configure the appearance of tick labels for the scale's axis or legend.

        Parameters
        ----------
        formatter : :class:`matplotlib.ticker.Formatter` subclass
            Pre-configured formatter to use; other parameters will be ignored.
        like : str or callable
            Either a format pattern (e.g., `".2f"`), a format string with fields named
            `x` and/or `pos` (e.g., `"${x:.2f}"`), or a callable with a signature like
            `f(x: float, pos: int) -> str`. In the latter variants, `x` is passed as the
            tick value and `pos` is passed as the tick index.
        base : number
            Use log formatter (with scientific notation) having this value as the base.
            Set to `None` to override the default formatter with a log transform.
        unit : str or (str, str) tuple
            Use  SI prefixes with these units (e.g., with `unit="g"`, a tick value
            of 5000 will appear as `5 kg`). When a tuple, the first element gives the
            separator between the number and unit.

        Returns
        -------
        scale
            Copy of self with new label configuration.

        """
        # Input checks
        if formatter is not None and not isinstance(formatter, Formatter):
            raise TypeError(
                f"Label formatter must be an instance of {Formatter!r}, "
                f"not {type(formatter)!r}"
            )
        if like is not None and not (isinstance(like, str) or callable(like)):
            msg = f"`like` must be a string or callable, not {type(like).__name__}."
            raise TypeError(msg)

        new = copy(self)
        new._label_params = {
            "formatter": formatter,
            "like": like,
            "base": base,
            "unit": unit,
        }
        return new

    def _parse_for_log_params(
        self, trans: str | TransFuncs | None
    ) -> tuple[float | None, float | None]:

        log_base = symlog_thresh = None
        if isinstance(trans, str):
            m = re.match(r"^log(\d*)", trans)
            if m is not None:
                log_base = float(m[1] or 10)
            m = re.match(r"symlog(\d*)", trans)
            if m is not None:
                symlog_thresh = float(m[1] or 1)
        return log_base, symlog_thresh

    def _get_locators(self, locator, at, upto, count, every, between, minor):

        log_base, symlog_thresh = self._parse_for_log_params(self.trans)

        if locator is not None:
            major_locator = locator

        elif upto is not None:
            if log_base:
                major_locator = LogLocator(base=log_base, numticks=upto)
            else:
                major_locator = MaxNLocator(upto, steps=[1, 1.5, 2, 2.5, 3, 5, 10])

        elif count is not None:
            if between is None:
                # This is rarely useful (unless you are setting limits)
                major_locator = LinearLocator(count)
            else:
                if log_base or symlog_thresh:
                    forward, inverse = self._get_transform()
                    lo, hi = forward(between)
                    ticks = inverse(np.linspace(lo, hi, num=count))
                else:
                    ticks = np.linspace(*between, num=count)
                major_locator = FixedLocator(ticks)

        elif every is not None:
            if between is None:
                major_locator = MultipleLocator(every)
            else:
                lo, hi = between
                ticks = np.arange(lo, hi + every, every)
                major_locator = FixedLocator(ticks)

        elif at is not None:
            major_locator = FixedLocator(at)

        else:
            if log_base:
                major_locator = LogLocator(log_base)
            elif symlog_thresh:
                major_locator = SymmetricalLogLocator(linthresh=symlog_thresh, base=10)
            else:
                major_locator = AutoLocator()

        if minor is None:
            minor_locator = LogLocator(log_base, subs=None) if log_base else None
        else:
            if log_base:
                subs = np.linspace(0, log_base, minor + 2)[1:-1]
                minor_locator = LogLocator(log_base, subs=subs)
            else:
                minor_locator = AutoMinorLocator(minor + 1)

        return major_locator, minor_locator

    def _get_formatter(self, locator, formatter, like, base, unit):

        log_base, symlog_thresh = self._parse_for_log_params(self.trans)
        if base is default:
            if symlog_thresh:
                log_base = 10
            base = log_base

        if formatter is not None:
            return formatter

        if like is not None:
            if isinstance(like, str):
                if "{x" in like or "{pos" in like:
                    fmt = like
                else:
                    fmt = f"{{x:{like}}}"
                formatter = StrMethodFormatter(fmt)
            else:
                formatter = FuncFormatter(like)

        elif base is not None:
            # We could add other log options if necessary
            formatter = LogFormatterSciNotation(base)

        elif unit is not None:
            if isinstance(unit, tuple):
                sep, unit = unit
            elif not unit:
                sep = ""
            else:
                sep = " "
            formatter = EngFormatter(unit, sep=sep)

        else:
            formatter = ScalarFormatter()

        return formatter


@dataclass
class Temporal(ContinuousBase):
    """
    A scale for date/time data.
    """
    # TODO date: bool?
    # For when we only care about the time component, would affect
    # default formatter and norm conversion. Should also happen in
    # Property.default_scale. The alternative was having distinct
    # Calendric / Temporal scales, but that feels a bit fussy, and it
    # would get in the way of using first-letter shorthands because
    # Calendric and Continuous would collide. Still, we haven't implemented
    # those yet, and having a clear distinction betewen date(time) / time
    # may be more useful.

    trans = None

    _priority: ClassVar[int] = 2

    def tick(
        self, locator: Locator | None = None, *,
        upto: int | None = None,
    ) -> Temporal:
        """
        Configure the selection of ticks for the scale's axis or legend.

        .. note::
            This API is under construction and will be enhanced over time.

        Parameters
        ----------
        locator : :class:`matplotlib.ticker.Locator` subclass
            Pre-configured matplotlib locator; other parameters will not be used.
        upto : int
            Choose "nice" locations for ticks, but do not exceed this number.

        Returns
        -------
        scale
            Copy of self with new tick configuration.

        """
        if locator is not None and not isinstance(locator, Locator):
            err = (
                f"Tick locator must be an instance of {Locator!r}, "
                f"not {type(locator)!r}."
            )
            raise TypeError(err)

        new = copy(self)
        new._tick_params = {"locator": locator, "upto": upto}
        return new

    def label(
        self,
        formatter: Formatter | None = None, *,
        concise: bool = False,
    ) -> Temporal:
        """
        Configure the appearance of tick labels for the scale's axis or legend.

        .. note::
            This API is under construction and will be enhanced over time.

        Parameters
        ----------
        formatter : :class:`matplotlib.ticker.Formatter` subclass
            Pre-configured formatter to use; other parameters will be ignored.
        concise : bool
            If True, use :class:`matplotlib.dates.ConciseDateFormatter` to make
            the tick labels as compact as possible.

        Returns
        -------
        scale
            Copy of self with new label configuration.

        """
        new = copy(self)
        new._label_params = {"formatter": formatter, "concise": concise}
        return new

    def _get_locators(self, locator, upto):

        if locator is not None:
            major_locator = locator
        elif upto is not None:
            major_locator = AutoDateLocator(minticks=2, maxticks=upto)

        else:
            major_locator = AutoDateLocator(minticks=2, maxticks=6)
        minor_locator = None

        return major_locator, minor_locator

    def _get_formatter(self, locator, formatter, concise):

        if formatter is not None:
            return formatter

        if concise:
            # TODO ideally we would have concise coordinate ticks,
            # but full semantic ticks. Is that possible?
            formatter = ConciseDateFormatter(locator)
        else:
            formatter = AutoDateFormatter(locator)

        return formatter


# ----------------------------------------------------------------------------------- #


# TODO Have this separate from Temporal or have Temporal(date=True) or similar?
# class Calendric(Scale):

# TODO Needed? Or handle this at layer (in stat or as param, eg binning=)
# class Binned(Scale):

# TODO any need for color-specific scales?
# class Sequential(Continuous):
# class Diverging(Continuous):
# class Qualitative(Nominal):


# ----------------------------------------------------------------------------------- #


class PseudoAxis:
    """
    Internal class implementing minimal interface equivalent to matplotlib Axis.

    Coordinate variables are typically scaled by attaching the Axis object from
    the figure where the plot will end up. Matplotlib has no similar concept of
    and axis for the other mappable variables (color, etc.), but to simplify the
    code, this object acts like an Axis and can be used to scale other variables.

    """
    axis_name = ""  # Matplotlib requirement but not actually used

    def __init__(self, scale):

        self.converter = None
        self.units = None
        self.scale = scale
        self.major = mpl.axis.Ticker()
        self.minor = mpl.axis.Ticker()

        # It appears that this needs to be initialized this way on matplotlib 3.1,
        # but not later versions. It is unclear whether there are any issues with it.
        self._data_interval = None, None

        scale.set_default_locators_and_formatters(self)
        # self.set_default_intervals()  Is this ever needed?

    def set_view_interval(self, vmin, vmax):
        self._view_interval = vmin, vmax

    def get_view_interval(self):
        return self._view_interval

    # TODO do we want to distinguish view/data intervals? e.g. for a legend
    # we probably want to represent the full range of the data values, but
    # still norm the colormap. If so, we'll need to track data range separately
    # from the norm, which we currently don't do.

    def set_data_interval(self, vmin, vmax):
        self._data_interval = vmin, vmax

    def get_data_interval(self):
        return self._data_interval

    def get_tick_space(self):
        # TODO how to do this in a configurable / auto way?
        # Would be cool to have legend density adapt to figure size, etc.
        return 5

    def set_major_locator(self, locator):
        self.major.locator = locator
        locator.set_axis(self)

    def set_major_formatter(self, formatter):
        self.major.formatter = formatter
        formatter.set_axis(self)

    def set_minor_locator(self, locator):
        self.minor.locator = locator
        locator.set_axis(self)

    def set_minor_formatter(self, formatter):
        self.minor.formatter = formatter
        formatter.set_axis(self)

    def set_units(self, units):
        self.units = units

    def update_units(self, x):
        """Pass units to the internal converter, potentially updating its mapping."""
        self.converter = mpl.units.registry.get_converter(x)
        if self.converter is not None:
            self.converter.default_units(x, self)

            info = self.converter.axisinfo(self.units, self)

            if info is None:
                return
            if info.majloc is not None:
                self.set_major_locator(info.majloc)
            if info.majfmt is not None:
                self.set_major_formatter(info.majfmt)

            # This is in matplotlib method; do we need this?
            # self.set_default_intervals()

    def convert_units(self, x):
        """Return a numeric representation of the input data."""
        if np.issubdtype(np.asarray(x).dtype, np.number):
            return x
        elif self.converter is None:
            return x
        return self.converter.convert(x, self.units, self)

    def get_scale(self):
        # Note that matplotlib actually returns a string here!
        # (e.g., with a log scale, axis.get_scale() returns "log")
        # Currently we just hit it with minor ticks where it checks for
        # scale == "log". I'm not sure how you'd actually use log-scale
        # minor "ticks" in a legend context, so this is fine....
        return self.scale

    def get_majorticklocs(self):
        return self.major.locator()


# ------------------------------------------------------------------------------------ #
# Transform function creation


def _make_identity_transforms() -> TransFuncs:

    def identity(x):
        return x

    return identity, identity


def _make_logit_transforms(base: float | None = None) -> TransFuncs:

    log, exp = _make_log_transforms(base)

    def logit(x):
        with np.errstate(invalid="ignore", divide="ignore"):
            return log(x) - log(1 - x)

    def expit(x):
        with np.errstate(invalid="ignore", divide="ignore"):
            return exp(x) / (1 + exp(x))

    return logit, expit


def _make_log_transforms(base: float | None = None) -> TransFuncs:

    fs: TransFuncs
    if base is None:
        fs = np.log, np.exp
    elif base == 2:
        fs = np.log2, partial(np.power, 2)
    elif base == 10:
        fs = np.log10, partial(np.power, 10)
    else:
        def forward(x):
            return np.log(x) / np.log(base)
        fs = forward, partial(np.power, base)

    def log(x: ArrayLike) -> ArrayLike:
        with np.errstate(invalid="ignore", divide="ignore"):
            return fs[0](x)

    def exp(x: ArrayLike) -> ArrayLike:
        with np.errstate(invalid="ignore", divide="ignore"):
            return fs[1](x)

    return log, exp


def _make_symlog_transforms(c: float = 1, base: float = 10) -> TransFuncs:

    # From https://iopscience.iop.org/article/10.1088/0957-0233/24/2/027001

    # Note: currently not using base because we only get
    # one parameter from the string, and are using c (this is consistent with d3)

    log, exp = _make_log_transforms(base)

    def symlog(x):
        with np.errstate(invalid="ignore", divide="ignore"):
            return np.sign(x) * log(1 + np.abs(np.divide(x, c)))

    def symexp(x):
        with np.errstate(invalid="ignore", divide="ignore"):
            return np.sign(x) * c * (exp(np.abs(x)) - 1)

    return symlog, symexp


def _make_sqrt_transforms() -> TransFuncs:

    def sqrt(x):
        return np.sign(x) * np.sqrt(np.abs(x))

    def square(x):
        return np.sign(x) * np.square(x)

    return sqrt, square


def _make_power_transforms(exp: float) -> TransFuncs:

    def forward(x):
        return np.sign(x) * np.power(np.abs(x), exp)

    def inverse(x):
        return np.sign(x) * np.power(np.abs(x), 1 / exp)

    return forward, inverse


def _default_spacer(x: Series) -> float:
    return 1