"""Plot annotations."""

from __future__ import annotations

__all__ = [
    "SI_PREFIXES",
    "SI_PREFIX_NAMES",
    "copy_mathtext",
    "fancy_labels",
    "integer_ticks",
    "label_subplot_properties",
    "label_subplots",
    "label_subplots_nature",
    "mark_points",
    "mark_points_outside",
    "plot_hv_text",
    "property_labels",
    "scale_units",
    "set_titles",
    "set_xlabels",
    "set_ylabels",
    "sizebar",
]

import io
import re
from collections.abc import Iterable
from typing import TYPE_CHECKING, Any, Literal, cast

import matplotlib
import matplotlib.figure
import matplotlib.font_manager
import matplotlib.mathtext
import matplotlib.pyplot as plt
import matplotlib.ticker
import matplotlib.transforms as mtransforms
import numpy as np

from erlab.accessors.utils import either_dict_or_kwargs
from erlab.plotting.colors import axes_textcolor

if TYPE_CHECKING:
    from collections.abc import Sequence

    import pyperclip
else:
    import lazy_loader as _lazy

    pyperclip = _lazy.load("pyperclip")

SI_PREFIXES: dict[int, str] = {
    24: "Y",
    21: "Z",
    18: "E",
    15: "P",
    12: "T",
    9: "G",
    6: "M",
    3: "k",
    2: "h",
    1: "da",
    0: "",
    -1: "d",
    -2: "c",
    -3: "m",
    -6: "μ",
    -9: "n",
    -12: "p",
    -15: "f",
    -18: "a",
    -21: "z",
    -24: "y",
}  #: Maps powers of 10 to valid SI prefix strings.

SI_PREFIX_NAMES: tuple[str, ...] = (
    "yotta",
    "zetta",
    "exa",
    "peta",
    "tera",
    "giga",
    "mega",
    "kilo",
    "hecto",
    "deca",
    "",
    "deci",
    "centi",
    "milli",
    "micro",
    "nano",
    "pico",
    "femto",
    "atto",
    "zepto",
    "yocto",
)  #: Names of the SI prefixes.

PRETTY_NAMES: dict[str, tuple[str, str]] = {
    "temperature": ("Temperature", "Temperature"),
    "T": (r"\ensuremath{T}", r"$T$"),
    "sample_temp": (r"\ensuremath{T}", r"$T$"),
    "t": (r"\ensuremath{t}", r"$t$"),
    "beta": (r"\ensuremath{\beta}", r"$\beta$"),
    "theta": (r"\ensuremath{\theta}", r"$\theta$"),
    "chi": (r"\ensuremath{\chi}", r"$\chi$"),
    "alpha": (r"\ensuremath{\alpha}", r"$\alpha$"),
    "psi": (r"\ensuremath{\psi}", r"$\psi$"),
    "phi": (r"\ensuremath{\phi}", r"$\phi$"),
    "xi": (r"\ensuremath{\xi}", r"$\xi$"),
    "Eb": (r"\ensuremath{E}", r"$E$"),
    "Ek": (r"\ensuremath{E_{\text{kin}}}", r"$E_{\text{kin}}$"),
    "eV": (r"\ensuremath{E-E_F}", r"$E-E_F$"),
    "kx": (r"\ensuremath{k_{x}}", r"$k_x$"),
    "ky": (r"\ensuremath{k_{y}}", r"$k_y$"),
    "kz": (r"\ensuremath{k_{z}}", r"$k_z$"),
    "kp": (r"\ensuremath{k_{\parallel}}", r"$k_\parallel$"),
    "hv": (r"\ensuremath{h\nu}", r"$h\nu$"),
}
"""Pretty names for labeling plots.

The first element is for LaTeX, and the second is for plain text.

"""

PRETTY_UNITS: dict[str, tuple[str, str]] = {
    "temperature": (r"K", r"K"),
    "T": (r"K", r"K"),
    "sample_temp": (r"K", r"K"),
    "t": (r"s", r"s"),
    "theta": (r"deg", r"deg"),
    "beta": (r"deg", r"deg"),
    "psi": (r"deg", r"deg"),
    "chi": (r"deg", r"deg"),
    "alpha": (r"deg", r"deg"),
    "phi": (r"deg", r"deg"),
    "xi": (r"deg", r"deg"),
    "Eb": (r"eV", r"eV"),
    "Ek": (r"eV", r"eV"),
    "eV": (r"eV", r"eV"),
    "hv": (r"eV", r"eV"),
    "kx": (r"Å\ensuremath{{}^{-1}}", r"Å${}^{-1}$"),
    "ky": (r"Å\ensuremath{{}^{-1}}", r"Å${}^{-1}$"),
    "kz": (r"Å\ensuremath{{}^{-1}}", r"Å${}^{-1}$"),
    "kp": (r"Å\ensuremath{{}^{-1}}", r"Å${}^{-1}$"),
}
"""Pretty units for labeling plots.

The first element is for LaTeX, and the second is for plain text.

"""


def _alph_label(val, prefix, suffix, numeric, capital):
    """Generate labels from string or integer."""
    if isinstance(val, int | np.integer) or val.isdigit():
        if numeric:
            val = str(val)
        else:
            ref_char = "A" if capital else "a"
            val = chr(int(val) + ord(ref_char) - 1)
    elif not isinstance(val, str):
        raise TypeError("Input values must be integers or strings.")
    return prefix + val + suffix


def _unit_from_label(label: str) -> str | None:
    """Try to determine the unit from a given label.

    Returns None if it fails to determine the unit.

    """
    m = re.match(r".*\((.*)\)\s*?$", label)
    if m is None:
        return None
    try:
        return m.group(1)
    except IndexError:
        return None


def get_si_str(si: int) -> str:
    """Return the SI prefix string to be plotted by :mod:`matplotlib`.

    Parameters
    ----------
    si : int
        Exponent of 10 corresponding to a SI prefix.

    Returns
    -------
    str
        SI prefix corresponding to ``si``.

    """
    if plt.rcParams["text.usetex"] and si == -6:
        return "\\ensuremath{\\mu}"
    try:
        return SI_PREFIXES[si]
    except KeyError as e:
        raise ValueError("Invalid SI prefix.") from e


def name_for_dim(dim_name: str, escaped: bool = True) -> str:
    names: tuple[str, str] | None = PRETTY_NAMES.get(dim_name)

    if names is None:
        name = dim_name
    else:
        name = names[0] if plt.rcParams["text.usetex"] else names[1]

    if not escaped:
        name = name.replace("$", "")
    return name


def unit_for_dim(dim_name: str, deg2rad: bool = False) -> str:
    units: tuple[str, str] | None = PRETTY_UNITS.get(dim_name)

    if units is None:
        unit = ""
    else:
        unit = units[0] if plt.rcParams["text.usetex"] else units[1]

    if deg2rad:
        unit = unit.replace("deg", "rad")
    return unit


def label_for_dim(dim_name: str, deg2rad: bool = False, escaped: bool = True) -> str:
    name = name_for_dim(dim_name, escaped=escaped)
    unit = unit_for_dim(dim_name, deg2rad=deg2rad)
    if unit == "":
        return name
    return f"{name} ({unit})"


def parse_special_point(name: str) -> str:
    special_points = {"G": r"\Gamma", "D": r"\Delta"}

    if name in special_points:
        return special_points[name]

    return name


def parse_point_labels(name: str, roman: bool = True, bar: bool = False) -> str:
    name = parse_special_point(name)

    if name.endswith("*"):
        name = name[:-1]
        format_str = "\\mathdefault{{{}}}^*" if roman else "{}^*"
    elif name.endswith("'"):
        name = name[:-1]
        format_str = "\\mathdefault{{{}}}\\prime" if roman else "{}\\prime"
    elif roman:
        format_str = r"\mathdefault{{{}}}"
    else:
        format_str = r"{}"

    name = format_str.format(parse_special_point(name))

    return f"$\\overline{{{name}}}$" if bar else f"${name}$"


def copy_mathtext(
    s: str,
    fontsize: float
    | Literal["xx-small", "x-small", "small", "medium", "large", "x-large", "xx-large"]
    | None = None,
    fontproperties: matplotlib.font_manager.FontProperties | None = None,
    outline: bool = False,
    svg: bool = True,
    rcparams: dict | None = None,
    **mathtext_rc,
) -> str:
    if fontproperties is None:
        fontproperties = matplotlib.font_manager.FontProperties(size=fontsize)
    else:
        fontproperties.set_size(fontsize)
    if rcparams is None:
        rcparams = {}
    parser = matplotlib.mathtext.MathTextParser("path")
    width, height, depth, _, _ = parser.parse(s, dpi=72, prop=fontproperties)

    fig = matplotlib.figure.Figure(figsize=(width / 72, height / 72))
    fig.patch.set_facecolor("none")
    fig.text(0, depth / height, s, fontproperties=fontproperties)

    if svg:
        from matplotlib.backends.backend_svg import FigureCanvasSVG

        FigureCanvasSVG(fig)

    else:
        from matplotlib.backends.backend_pdf import FigureCanvasPdf

        FigureCanvasPdf(fig)

    for k, v in mathtext_rc.items():
        if k in ["bf", "cal", "it", "rm", "sf", "tt"] and isinstance(
            v, matplotlib.font_manager.FontProperties
        ):
            v = v.get_fontconfig_pattern()
        rcparams[f"mathtext.{k}"] = v

    with io.BytesIO() as buffer:
        if svg:
            rcparams.setdefault("svg.fonttype", "path" if outline else "none")
            rcparams.setdefault("svg.image_inline", True)
            with plt.rc_context(rcparams):
                fig.canvas.print_svg(buffer)  # type: ignore[attr-defined]
        else:
            rcparams.setdefault("pdf.fonttype", 3 if outline else 42)
            with plt.rc_context(rcparams):
                fig.canvas.print_pdf(buffer)  # type: ignore[attr-defined]

        buffer_str = buffer.getvalue().decode("utf-8")

    pyperclip.copy(buffer_str)
    return buffer_str


def fancy_labels(ax=None, deg2rad=False) -> None:
    if ax is None:
        ax = plt.gca()
    if np.iterable(ax):
        for axi in ax:
            fancy_labels(axi, deg2rad)
        return

    ax.set_xlabel(label_for_dim(dim_name=ax.get_xlabel(), deg2rad=deg2rad))
    ax.set_ylabel(label_for_dim(dim_name=ax.get_ylabel(), deg2rad=deg2rad))
    if hasattr(ax, "get_zlabel"):
        ax.set_zlabel(label_for_dim(dim_name=ax.get_zlabel(), deg2rad=deg2rad))


def property_labels(
    values: dict | None = None,
    decimals: int | None = None,
    si: int = 0,
    name: str | None = None,
    unit: str | None = None,
    order: Literal["C", "F", "A", "K"] = "C",
    **values_kwargs,
) -> list[str]:
    """Generate labels from a dictionary of values.

    Given a dictionary of values, this function generates a list of label strings  with
    the key as the dimension name and the value as the dimension value. Multiple
    key-value pairs will be separated by a newline character in each label. The name and
    unit will be automatically determined from the key.

    Parameters
    ----------
    values
        Key-value pair of annotations. The values may be a single value or an array-like
        of values. If a single value, it will be converted to a list of length 1. If an
        array of 2 or more dimensions is given, it will be flattened with the order
        given by `order`. All values must be of the same length when flattened.
    decimals
        Number of decimal places to round to. If decimals is None, no rounding is
        performed. If decimals is negative, it specifies the number of positions to the
        left of the decimal point.
    si
        Powers of 10 for automatic SI prefix setting.
    name
        When set, overrides automatic dimension name setting.
    unit
        When set, overrides automatic unit setting.
    order
        Order in which to flatten `ax`. 'C' means to flatten in row-major (C-style)
        order. 'F' means to flatten in column-major (Fortran-style) order. The default
        is 'C'.

    """
    values = either_dict_or_kwargs(values, values_kwargs, "property_labels")
    strlist: Any = []
    for k, v in values.items():
        if not isinstance(v, tuple | list | np.ndarray):
            v = [v]
        else:
            v = np.array(v).flatten(order=order)
        strlist.append(
            [
                property_label(
                    k, val, decimals=decimals, si=si, name=name, unit=unit
                ).strip()
                for val in v
            ]
        )
    strlist = list(zip(*strlist, strict=True))
    return ["\n".join(strlist[i]) for i in range(len(strlist))]


def label_subplot_properties(
    axes: matplotlib.axes.Axes | Iterable[matplotlib.axes.Axes],
    values: dict,
    decimals: int | None = None,
    si: int = 0,
    name: str | None = None,
    unit: str | None = None,
    order: Literal["C", "F", "A", "K"] = "C",
    **kwargs,
) -> None:
    r"""Labels subplots with automatically generated labels.

    Parameters
    ----------
    axes
        `matplotlib.axes.Axes` to label. If an array is given, the order will be
        determined by the flattening method given by `order`.
    values
        Key-value pair of annotations. The values may be a single value or an array-like
        of values. If a single value, it will be converted to a list of length 1. If an
        array of 2 or more dimensions is given, it will be flattened with the order
        given by `order`. All values must be of the same length when flattened.
    decimals
        Number of decimal places to round to. If decimals is None, no rounding is
        performed. If decimals is negative, it specifies the number of positions to the
        left of the decimal point.
    si
        Powers of 10 for automatic SI prefix setting.
    name
        When set, overrides automatic dimension name setting.
    unit
        When set, overrides automatic unit setting.
    order
        Order in which to flatten `ax`. 'C' means to flatten in row-major (C-style)
        order. 'F' means to flatten in column-major (Fortran-style) order. The default
        is 'C'.
    **kwargs
        Extra arguments to `erlab.plotting.annotations.label_subplots`.

    """
    kwargs.setdefault("fontweight", plt.rcParams["font.weight"])
    kwargs.setdefault("prefix", "")
    kwargs.setdefault("suffix", "")
    kwargs.setdefault("loc", "upper right")

    label_subplots(
        axes,
        property_labels(
            values, decimals=decimals, si=si, name=name, unit=unit, order=order
        ),
        order=order,
        **kwargs,
    )


def label_subplots(
    axes: matplotlib.axes.Axes | Iterable[matplotlib.axes.Axes],
    values: Iterable[int | str] | None = None,
    startfrom: int = 1,
    order: Literal["C", "F", "A", "K"] = "C",
    loc: Literal[
        "upper left",
        "upper center",
        "upper right",
        "center left",
        "center",
        "center right",
        "lower left",
        "lower center",
        "lower right",
    ] = "upper left",
    offset: tuple[float, float] = (0.0, 0.0),
    prefix: str = "",
    suffix: str = "",
    numeric: bool = False,
    capital: bool = False,
    fontweight: Literal[
        "ultralight",
        "light",
        "normal",
        "regular",
        "book",
        "medium",
        "roman",
        "semibold",
        "demibold",
        "demi",
        "bold",
        "heavy",
        "extra bold",
        "black",
    ] = "normal",
    fontsize: (
        float
        | Literal[
            "xx-small", "x-small", "small", "medium", "large", "x-large", "xx-large"
        ]
        | None
    ) = None,
    **kwargs,
) -> None:
    r"""Labels subplots with automatically generated labels.

    Parameters
    ----------
    axes
        `matplotlib.axes.Axes` to label. If an array is given, the order will be
        determined by the flattening method given by `order`.
    values
        Integer or string labels corresponding to each Axes in `axes` for
        manual labels.
    startfrom
        Start from this number when creating automatic labels. Has no
        effect when `values` is not `None`.
    order
        Order in which to flatten `ax`. 'C' means to flatten in
        row-major (C-style) order. 'F' means to flatten in column-major
        (Fortran-style) order. 'A' means to flatten in column-major
        order if a is Fortran contiguous in memory, row-major order
        otherwise. 'K' means to flatten a in the order the elements
        occur in memory. The default is 'C'.
    loc
        The box location. The default is ``'upper left'``.
    offset
        Values that are used to position the legend in conjunction with
        `loc`, given in display units.
    prefix
        String to prepend to the alphabet label.
    suffix
        String to append to the alphabet label.
    numeric
        Use integer labels instead of alphabets.
    capital
        Capitalize automatically generated alphabetical labels.
    fontweight
        Set the font weight. The default is ``'normal'``.
    fontsize
        Set the font size. The default is ``'medium'`` for axes, and ``'large'`` for
        figures.
    **kwargs
        Extra arguments to `matplotlib.text.Text`: refer to the `matplotlib`
        documentation for a list of all possible arguments.

    """
    kwargs["fontweight"] = fontweight
    if plt.rcParams["text.usetex"] & (fontweight == "bold"):
        prefix = "\\textbf{" + prefix
        suffix = suffix + "}"
        kwargs.pop("fontweight")

    axlist = np.array(axes, dtype=object).flatten(order=order)
    if values is None:
        value_arr = np.array(
            [i + startfrom for i in range(len(axlist))], dtype=np.int64
        )
    else:
        value_arr = np.array(values).flatten(order=order)
        if not (axlist.size == value_arr.size):
            raise IndexError(
                "The number of given values must match the number of given axes."
            )

    for i, ax in enumerate(axlist):
        if fontsize is None:
            fontsize = "large" if isinstance(ax, matplotlib.figure.Figure) else "medium"

        label_str = _alph_label(value_arr[i], prefix, suffix, numeric, capital)
        with plt.rc_context({"text.color": axes_textcolor(ax)}):
            at = matplotlib.offsetbox.AnchoredText(
                label_str,
                loc=loc,
                frameon=False,
                pad=0,
                borderpad=0.5,
                prop=dict(fontsize=fontsize, **kwargs),
                bbox_to_anchor=ax.bbox,
                bbox_transform=matplotlib.transforms.ScaledTranslation(
                    offset[0] / 72,
                    offset[1] / 72,
                    ax.get_figure().dpi_scale_trans,
                ),
                clip_on=False,
            )
        ax.add_artist(at)


def label_subplots_nature(
    axes: matplotlib.axes.Axes | Sequence[matplotlib.axes.Axes],
    values: Sequence[int | str] | None = None,
    startfrom: int = 1,
    order: Literal["C", "F", "A", "K"] = "C",
    offset: tuple[float, float] = (-20.0, 7.0),
    prefix: str = "",
    suffix: str = "",
    numeric: bool = False,
    capital: bool = False,
    fontweight: Literal[
        "ultralight",
        "light",
        "normal",
        "regular",
        "book",
        "medium",
        "roman",
        "semibold",
        "demibold",
        "demi",
        "bold",
        "heavy",
        "extra bold",
        "black",
    ] = "black",
    fontsize: (
        float
        | Literal[
            "xx-small", "x-small", "small", "medium", "large", "x-large", "xx-large"
        ]
    ) = "medium",
    **kwargs,
) -> None:
    r"""Labels subplots with automatically generated labels.

    Parameters
    ----------
    axes
        `matplotlib.axes.Axes` to label. If an array is given, the order will be
        determined by the flattening method given by `order`.
    values
        Integer or string labels corresponding to each Axes in `axes` for
        manual labels.
    startfrom
        Start from this number when creating automatic labels. Has no
        effect when `values` is not `None`.
    order
        Order in which to flatten `ax`. 'C' means to flatten in
        row-major (C-style) order. 'F' means to flatten in column-major
        (Fortran-style) order. 'A' means to flatten in column-major
        order if a is Fortran contiguous in memory, row-major order
        otherwise. 'K' means to flatten a in the order the elements
        occur in memory. The default is 'C'.
    offset
        Values that are used to position the labels, given in points.
    prefix
        String to prepend to the alphabet label.
    suffix
        String to append to the alphabet label.
    numeric
        Use integer labels instead of alphabets.
    capital
        Capitalize automatically generated alphabetical labels.
    fontweight
        Set the font weight. The default is ``'normal'``.
    fontsize
        Set the font size. The default is ``'medium'`` for axes, and ``'large'`` for
        figures.
    **kwargs
        Extra arguments to `matplotlib.text.Text`: refer to the `matplotlib`
        documentation for a list of all possible arguments.

    """
    kwargs["fontweight"] = fontweight
    if plt.rcParams["text.usetex"] & (fontweight == "bold"):
        prefix = "\\textbf{" + prefix
        suffix = suffix + "}"
        kwargs.pop("fontweight")

    axlist = np.array(axes, dtype=object).flatten(order=order)
    if values is None:
        value_arr = np.array(
            [i + startfrom for i in range(len(axlist))], dtype=np.int64
        )
    else:
        value_arr = np.array(values).flatten(order=order)
        if not (axlist.size == value_arr.size):
            raise IndexError(
                "The number of given values must match the number of given axes."
            )

    for i in range(len(axlist)):
        label_str = _alph_label(value_arr[i], prefix, suffix, numeric, capital)
        trans = matplotlib.transforms.ScaledTranslation(
            offset[0] / 72, offset[1] / 72, axlist[i].get_figure().dpi_scale_trans
        )
        axlist[i].figure.text(
            # axlist[i].text(
            0.0,
            1.0,
            label_str,
            transform=axlist[i].transAxes + trans,
            fontsize=fontsize,
            va="baseline",
            clip_on=False,
            **kwargs,
        )


def mark_points(
    points: Sequence[float],
    labels: Sequence[str],
    y: float | Sequence[float] = 0.0,
    pad: tuple[float, float] = (0, 1.75),
    literal: bool = False,
    roman: bool = True,
    bar: bool = False,
    ax: matplotlib.axes.Axes | Iterable[matplotlib.axes.Axes] | None = None,
    **kwargs,
) -> None:
    """Mark points above the horizontal axis.

    Useful when annotating high symmetry points along a cut.

    Parameters
    ----------
    points
        Floats indicating the position of each label.
    labels
        Sequence of label strings indicating a high symmetry point. Must be the same
        length as `points`.
    y
        Position of the label in data coordinates
    pad
        Offset of the text in points.
    literal
        If `True`, take the input string literally.
    roman
        If ``False``, *True*, itallic fonts are used.
    bar
        If ``True``, prints a bar over the label.
    ax
        `matplotlib.axes.Axes` to annotate.

    """
    if ax is None:
        ax = plt.gca()

    if np.iterable(ax):
        for a in np.asarray(ax, dtype=object).flatten():
            mark_points(points, labels, y, pad, literal, roman, bar, a, **kwargs)
    else:
        ax = cast(matplotlib.axes.Axes, ax)  # to appease mypy
        fig = ax.get_figure()

        if fig is None:
            raise ValueError("Given axes does not belong to a figure")

        for k, v in {"ha": "center", "va": "baseline", "fontsize": "small"}.items():
            kwargs.setdefault(k, v)

        if not np.iterable(y):
            y = [y] * len(points)  # type: ignore[list-item]

        with plt.rc_context({"font.family": "serif"}):
            for xi, yi, label in zip(points, y, labels, strict=True):
                ax.text(
                    xi,
                    yi,
                    label if literal else parse_point_labels(label, roman, bar),
                    transform=ax.transData
                    + mtransforms.ScaledTranslation(
                        pad[0] / 72, pad[1] / 72, fig.dpi_scale_trans
                    ),
                    **kwargs,
                )


def mark_points_outside(
    points: Sequence[float],
    labels: Sequence[str],
    axis: Literal["x", "y"] = "x",
    roman: bool = True,
    bar: bool = False,
    ax: matplotlib.axes.Axes | Iterable[matplotlib.axes.Axes] | None = None,
    **kwargs,
) -> None:
    """Mark points above the horizontal axis.

    Useful when annotating high symmetry points along a cut.

    Parameters
    ----------
    points
        Floats indicating the position of each label.
    labels
        Sequence of label strings indicating a high symmetry point. Must be the same
        length as `points`.
    axis
        If ``'x'``, marks points along the horizontal axis. If ``'y'``, marks points
        along the vertical axis.
    roman
        If ``False``, itallic fonts are used.
    bar
        If ``True``, prints a bar over the label.
    ax
        `matplotlib.axes.Axes` to annotate.
    **kwargs
        Extra arguments to `matplotlib.text.Text`: refer to the `matplotlib`
        documentation for a list of all possible arguments.

    """
    if ax is None:
        ax = plt.gca()
    if np.iterable(ax):
        for a in np.asarray(ax, dtype=object).flatten():
            mark_points_outside(points, labels, axis, roman, bar, a)
    else:
        ax = cast(matplotlib.axes.Axes, ax)  # to appease mypy

        if axis == "x":
            label_ax = ax.twiny()
            label_ax.set_xlim(ax.get_xlim())
            label_ax.set_xticks(
                points,
                labels=[parse_point_labels(lab, roman, bar) for lab in labels],
                **kwargs,
            )
        else:
            label_ax = ax.twinx()
            label_ax.set_ylim(ax.get_ylim())
            label_ax.set_yticks(
                points,
                labels=[parse_point_labels(lab, roman, bar) for lab in labels],
                **kwargs,
            )
        label_ax.set_frame_on(False)


def mark_points_y(pts, labels, roman=True, bar=False, ax=None) -> None:
    if ax is None:
        ax = plt.gca()
    if not isinstance(ax, tuple | list | np.ndarray):
        ax = [ax]
    for a in np.array(ax, dtype=object).flatten():
        label_ax = a.twinx()
        label_ax.set_ylim(a.get_ylim())
        label_ax.set_yticks(pts)
        # label_ax.set_xlabel('')
        label_ax.set_yticklabels(
            [parse_point_labels(lab, roman, bar) for lab in labels]
        )
        # label_ax.set_zorder(a.get_zorder())
        label_ax.set_frame_on(False)


def plot_hv_text(ax, val, x=0.025, y=0.975, **kwargs) -> None:
    name = name_for_dim("hv", escaped=False)
    unit = unit_for_dim("hv")
    s = f"${name}={val}$ {unit}"
    ax.text(
        x,
        y,
        s,
        family="serif",
        horizontalalignment="left",
        verticalalignment="top",
        transform=ax.transAxes,
        **kwargs,
    )


def plot_hv_text_right(ax, val, x=1 - 0.025, y=0.975, **kwargs) -> None:
    name = name_for_dim("hv", escaped=False)
    unit = unit_for_dim("hv")
    s = f"${name}={val}$ {unit}"
    ax.text(
        x,
        y,
        s,
        family="serif",
        horizontalalignment="right",
        verticalalignment="top",
        transform=ax.transAxes,
        **kwargs,
    )


def property_label(key, value, decimals=None, si=0, name=None, unit=None) -> str:
    delim = "" if name == "" else " = "
    if name is None:
        name = name_for_dim(key, escaped=False)
        if name is None:
            name = ""
    if unit is None:
        unit = unit_for_dim(key)
        if unit is None:
            unit = ""

    unit = get_si_str(si) + unit
    value /= 10**si
    if decimals is not None:
        value = np.around(value, decimals=decimals)
    if int(value) == value:
        value = int(value)

    if key == "Eb":
        if value == 0:
            if delim == "":
                return "$E_F$"
            return "$E = E_F$"
        if delim == "":
            name = "E_F"
        else:
            delim += "E_F"
        if value > 0:
            delim += "+"

    base = "${}" + delim + "{}$ {}"
    return str(base.format(name, value, unit))


class _SIFormatter(matplotlib.ticker.ScalarFormatter):
    def __init__(self, si: int = 0, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._si_exponent = int(si)

    def __call__(self, x, pos=None):
        self.orderOfMagnitude += self._si_exponent

        match_format = re.match(r".*%1.(\d+)f", self.format)
        if match_format is None:
            # Match failed, may be due to changes in matplotlib
            raise RuntimeError("Failed to match format string. Please report this bug")

        sigfigs = int(match_format.group(1))
        self.format = self.format.replace(
            f"%1.{sigfigs}f", f"%1.{max(0, sigfigs + self._si_exponent)}f"
        )
        val = super().__call__(x, pos)

        self.orderOfMagnitude -= self._si_exponent
        return val


def scale_units(
    ax: matplotlib.axes.Axes | Iterable[matplotlib.axes.Axes],
    axis: Literal["x", "y", "z"],
    si: int = 0,
    *,
    prefix: bool = True,
    power: bool = False,
) -> None:
    """Rescales ticks and adds an SI prefix to the axis label.

    Useful when you want to rescale the ticks without actually rescaling the data. For
    example, when plotting a cut from a low pass energy scan, you might want to convert
    the energy units from eV to meV.

    Using this function on an axis where the major locator is not the default formatter
    `matplotlib.ticker.ScalarFormatter` will result in undefined behavior.

    Parameters
    ----------
    ax
        Axes or iterable of Axes to rescale.
    axis
        The axis you wish to rescale.
    si
        Exponent of 10 corresponding to a SI prefix.
    prefix
        If True, tries to detect the unit from the axis label and scales it accordingly.
        The scaling behaviour is controlled by the `power` argument. If no units are
        found in the axis label, it is silently ignored.
    power
        If False, prefixes the detected unit on the axis label with a SI prefix
        corresponding to `si`. If True, the unit is prefixed with a scientific notation
        instead.

    """
    if np.iterable(ax):
        for a in np.asarray(ax, dtype=object).flat:
            scale_units(a, axis, si, prefix=prefix, power=power)
        return

    getlabel = getattr(ax, f"get_{axis}label")
    setlabel = getattr(ax, f"set_{axis}label")

    label = getlabel()
    unit = _unit_from_label(label)

    getattr(ax, f"{axis}axis").set_major_formatter(_SIFormatter(si))

    if prefix and (unit is not None):
        if power:
            setlabel(label.replace(f"({unit})", f"($\\times{{{10}}}^{{{si}}}$ {unit})"))
        else:
            setlabel(label.replace(f"({unit})", f"({get_si_str(si)}{unit})"))


def integer_ticks(
    ax: matplotlib.axes.Axes | Iterable[matplotlib.axes.Axes],
) -> None:
    """Set the ticks on the x and y axes to only display integer values.

    Parameters
    ----------
    ax
        Axes or iterable of Axes.

    """
    if isinstance(ax, Iterable):
        for a in np.asarray(ax, dtype=object).flat:
            integer_ticks(a)
        return
    ax.set_xticks(
        [
            t
            for t in ax.get_xticks()
            if t.is_integer() and t >= ax.get_xlim()[0] and t <= ax.get_xlim()[1]
        ]
    )
    ax.set_yticks(
        [
            t
            for t in ax.get_yticks()
            if t.is_integer() and t >= ax.get_ylim()[0] and t <= ax.get_ylim()[1]
        ]
    )


def set_titles(axes, labels: Iterable[str] | str, order="C", **kwargs) -> None:
    """Set titles for multiple axes.

    Parameters
    ----------
    axes
        Axes to set titles for.
    labels
        Title strings to set. If a single string is given, it will be used for all axes.
    order
        Order in which to flatten `ax`. 'C' means to flatten in row-major (C-style)
        order. 'F' means to flatten in column-major (Fortran-style) order. The default
        is 'C'.
    **kwargs
        Extra arguments to `matplotlib.axes.Axes.set_title`: refer to the `matplotlib`
        documentation for a list of all possible arguments.

    """
    axlist = np.array(axes, dtype=object).flatten(order=order)
    if isinstance(labels, str):
        labels = [labels] * len(axlist)
    labels = np.asarray(labels)
    for ax, label in zip(axlist.flat, labels.flat, strict=True):
        ax.set_title(label, **kwargs)


def set_xlabels(axes, labels: Iterable[str] | str, order="C", **kwargs) -> None:
    """Set x-axis labels for multiple axes.

    Parameters
    ----------
    axes
        Axes to set x-axis labels for.
    labels
        Label strings to set. If a single string is given, it will be used for all axes.
    order
        Order in which to flatten `ax`. 'C' means to flatten in row-major (C-style)
        order. 'F' means to flatten in column-major (Fortran-style) order. The default
        is 'C'.
    **kwargs
        Extra arguments to `matplotlib.axes.Axes.set_xlabel`: refer to the `matplotlib`
        documentation for a list of all possible arguments.
    """
    axlist = np.array(axes, dtype=object).flatten(order=order)
    if isinstance(labels, str):
        labels = [labels] * len(axlist)
    labels = np.asarray(labels)
    for ax, label in zip(axlist.flat, labels.flat, strict=True):
        ax.set_xlabel(label, **kwargs)


def set_ylabels(axes, labels: Iterable[str] | str, order="C", **kwargs) -> None:
    """Set y-axis labels for multiple axes.

    Parameters
    ----------
    axes
        Axes to set y-axis labels for.
    labels
        Label strings to set. If a single string is given, it will be used for all axes.
    order
        Order in which to flatten `ax`. 'C' means to flatten in row-major (C-style)
        order. 'F' means to flatten in column-major (Fortran-style) order. The default
        is 'C'.
    **kwargs
        Extra arguments to `matplotlib.axes.Axes.set_ylabel`: refer to the `matplotlib`

    """
    axlist = np.array(axes, dtype=object).flatten(order=order)
    if isinstance(labels, str):
        labels = [labels] * len(axlist)
    labels = np.asarray(labels)
    for ax, label in zip(axlist.flat, labels.flat, strict=True):
        ax.set_ylabel(label, **kwargs)


def sizebar(
    ax: matplotlib.axes.Axes,
    value: float,
    unit: str,
    si: int = 0,
    resolution: float = 1.0,
    decimals: int = 0,
    label: str | None = None,
    loc: Literal[
        "upper left",
        "upper center",
        "upper right",
        "center left",
        "center",
        "center right",
        "lower left",
        "lower center",
        "lower right",
    ] = "lower right",
    pad: float = 0.1,
    borderpad: float = 0.5,
    sep: float = 3.0,
    frameon: bool = False,
    **kwargs,
):
    """Add a size bar to an axes.

    Parameters
    ----------
    ax
        The `matplotlib.axes.Axes` instance to place the size bar in.
    value
        Length of the size bar in terms of `unit`.
    unit
        An SI unit string without prefixes.
    si
        Exponents that have a corresponding SI prefix
    resolution
        Value to scale the data coordinates in terms of `unit`.
    decimals
        Number of decimals on the size bar label.
    label
        When provided, overrides the automatically generated label string.
    loc
        Location of the size bar.
    pad
        Padding around the label and size bar, in fraction of the font size.
    borderpad
        Border padding, in fraction of the font size.
    sep
        Separation between the label and the size bar, in points.
    frameon
        If True, draw a box around the horizontal bar and label.
    **kwargs
        Keyword arguments forwarded to
        `mpl_toolkits.axes_grid1.anchored_artists.AnchoredSizeBar`.

    """
    from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar

    size = value / resolution
    unit = get_si_str(si) + unit
    value = np.around(value / 10**si, decimals=decimals)
    if int(value) == value:
        value = int(value)
    if label is None:
        label = f"{value} {unit}"

    kwargs.setdefault("color", axes_textcolor(ax))

    asb = AnchoredSizeBar(
        ax.transData,
        size,
        label,
        loc=loc,
        pad=pad,
        borderpad=borderpad,
        sep=sep,
        frameon=frameon,
        **kwargs,
    )
    ax.add_artist(asb)
    return asb
