"""PyVista Scalar bar module."""

from __future__ import annotations

import contextlib
import weakref

import numpy as np

import pyvista
from pyvista import MAX_N_COLOR_BARS
from pyvista._deprecate_positional_args import _deprecate_positional_args
from pyvista.core.utilities.misc import _NoNewAttrMixin

from . import _vtk
from .colors import Color
from .tools import parse_font_family


class ScalarBars(_NoNewAttrMixin):
    """Plotter Scalar Bars.

    Parameters
    ----------
    plotter : pyvista.Plotter
        Plotter that the scalar bars are associated with.

    """

    def __init__(self, plotter):
        """Initialize ScalarBars."""
        self._plotter = weakref.proxy(plotter)
        self._scalar_bar_ranges = {}
        self._scalar_bar_mappers = {}
        self._scalar_bar_actors = {}
        self._scalar_bar_widgets = {}

    def clear(self):
        """Remove all scalar bars and resets all scalar bar properties."""
        self._scalar_bar_ranges = {}
        self._scalar_bar_mappers = {}
        self._scalar_bar_actors = {}
        self._scalar_bar_widgets = {}

    def __repr__(self):
        """Nice representation of this class."""
        lines = []
        lines.append('Scalar Bar Title     Interactive')
        for title in self._scalar_bar_actors:
            interactive = title in self._scalar_bar_widgets
            title_quotes = f'"{title}"'
            lines.append(f'{title_quotes:20} {interactive!s:5}')
        return '\n'.join(lines)

    @_deprecate_positional_args(allowed=['actor'])
    def _remove_mapper_from_plotter(
        self,
        actor,
        reset_camera: bool = False,  # noqa: FBT001, FBT002
        render: bool = False,  # noqa: FBT001, FBT002
    ):  # numpydoc ignore=PR01,RT01
        """Remove an actor's mapper from the given plotter's _scalar_bar_mappers.

        This ensures that when actors are removed, their corresponding
        scalar bars are removed.

        """
        try:
            mapper = actor.GetMapper()
        except AttributeError:
            return

        # NOTE: keys to list to prevent iterator changing during loop
        for name in list(self._scalar_bar_mappers):
            with contextlib.suppress(ValueError):
                self._scalar_bar_mappers[name].remove(mapper)

            if not self._scalar_bar_mappers[name]:
                slot = self._plotter._scalar_bar_slot_lookup.pop(name, None)
                if slot is not None:
                    self._scalar_bar_mappers.pop(name)
                    self._scalar_bar_ranges.pop(name)
                    self._plotter.remove_actor(
                        self._scalar_bar_actors.pop(name),
                        reset_camera=reset_camera,
                        render=render,
                    )
                    self._plotter._scalar_bar_slots.add(slot)
            return

    @_deprecate_positional_args(allowed=['title'])
    def remove_scalar_bar(self, title=None, render: bool = True):  # noqa: FBT001, FBT002
        """Remove a scalar bar.

        Parameters
        ----------
        title : str, optional
            Title of the scalar bar to remove.  Required if there is
            more than one scalar bar.

        render : bool, default: True
            Render upon scalar bar removal.  Set this to ``False`` to
            stop the render window from rendering when a scalar bar
            is removed.

        Examples
        --------
        Remove a scalar bar from a plotter.

        >>> import pyvista as pv
        >>> mesh = pv.Sphere()
        >>> mesh['data'] = mesh.points[:, 2]
        >>> pl = pv.Plotter()
        >>> _ = pl.add_mesh(mesh, cmap='coolwarm')
        >>> pl.remove_scalar_bar()
        >>> pl.show()

        """
        if title is None:
            if len(self) > 1:
                titles = ', '.join(f'"{key}"' for key in self._scalar_bar_actors)
                msg = (
                    'Multiple scalar bars found.  Pick title of the'
                    f'scalar bar from one of the following:\n{titles}'
                )
                raise ValueError(msg)
            else:
                title = next(iter(self._scalar_bar_actors.keys()))

        actor = self._scalar_bar_actors.pop(title)
        self._plotter.remove_actor(actor, render=render)
        self._scalar_bar_ranges.pop(title)
        self._scalar_bar_mappers.pop(title)

        # add back in the scalar bar slot
        slot = self._plotter._scalar_bar_slot_lookup.pop(title, None)
        if slot is not None:
            self._plotter._scalar_bar_slots.add(slot)

        widget = self._scalar_bar_widgets.pop(title, None)
        if widget is not None:
            widget.SetEnabled(0)

    def __len__(self):
        """Return the number of scalar bar actors."""
        return len(self._scalar_bar_actors)

    def __getitem__(self, index):
        """Return a scalar bar actor."""
        return self._scalar_bar_actors[index]

    def keys(self):  # numpydoc ignore=RT01
        """Scalar bar keys."""
        return self._scalar_bar_actors.keys()

    def values(self):  # numpydoc ignore=RT01
        """Scalar bar values."""
        return self._scalar_bar_actors.values()

    def items(self):  # numpydoc ignore=RT01
        """Scalar bar items."""
        return self._scalar_bar_actors.items()

    def __contains__(self, key) -> bool:
        """Check if a title is a valid actors."""
        return key in self._scalar_bar_actors

    @_deprecate_positional_args(allowed=['title'])
    def add_scalar_bar(  # noqa: PLR0917
        self,
        title='',
        mapper=None,
        n_labels=5,
        italic: bool = False,  # noqa: FBT001, FBT002
        bold: bool = False,  # noqa: FBT001, FBT002
        title_font_size=None,
        label_font_size=None,
        color=None,
        font_family=None,
        shadow: bool = False,  # noqa: FBT001, FBT002
        width=None,
        height=None,
        position_x=None,
        position_y=None,
        vertical=None,
        interactive=None,
        fmt=None,
        use_opacity: bool = True,  # noqa: FBT001, FBT002
        outline: bool = False,  # noqa: FBT001, FBT002
        nan_annotation: bool = False,  # noqa: FBT001, FBT002
        below_label=None,
        above_label=None,
        background_color=None,
        n_colors=None,
        fill: bool = False,  # noqa: FBT001, FBT002
        render: bool = False,  # noqa: FBT001, FBT002
        theme=None,
        unconstrained_font_size: bool = False,  # noqa: FBT001, FBT002
    ):
        """Create scalar bar using the ranges as set by the last input mesh.

        Parameters
        ----------
        title : str, default: ""
            Title of the scalar bar.  Default is rendered as an empty title.

        mapper : :vtk:`vtkMapper`, optional
            Mapper used for the scalar bar.  Defaults to the last
            mapper created by the plotter.

        n_labels : int, default: 5
            Number of labels to use for the scalar bar.

        italic : bool, default: False
            Italicises title and bar labels.

        bold : bool, default: False
            Bolds title and bar labels.

        title_font_size : float, optional
            Sets the size of the title font.  Defaults to ``None`` and is sized
            according to :attr:`pyvista.plotting.themes.Theme.font`.

        label_font_size : float, optional
            Sets the size of the title font.  Defaults to ``None`` and is sized
            according to :attr:`pyvista.plotting.themes.Theme.font`.

        color : ColorLike, optional
            Either a string, rgb list, or hex color string.  Default
            set by :attr:`pyvista.plotting.themes.Theme.font`.  Can be
            in one of the following formats:

            * ``color='white'``
            * ``color='w'``
            * ``color=[1.0, 1.0, 1.0]``
            * ``color='#FFFFFF'``

        font_family : {'courier', 'times', 'arial'}
            Font family.  Default is set by
            :attr:`pyvista.plotting.themes.Theme.font`.

        shadow : bool, default: False
            Adds a black shadow to the text.

        width : float, optional
            The percentage (0 to 1) width of the window for the colorbar.
            Default set by
            :attr:`pyvista.plotting.themes.Theme.colorbar_vertical` or
            :attr:`pyvista.plotting.themes.Theme.colorbar_horizontal`
            depending on the value of ``vertical``.

        height : float, optional
            The percentage (0 to 1) height of the window for the
            colorbar.  Default set by
            :attr:`pyvista.plotting.themes.Theme.colorbar_vertical` or
            :attr:`pyvista.plotting.themes.Theme.colorbar_horizontal`
            depending on the value of ``vertical``.

        position_x : float, optional
            The percentage (0 to 1) along the windows's horizontal
            direction to place the bottom left corner of the colorbar.
            Default set by
            :attr:`pyvista.plotting.themes.Theme.colorbar_vertical` or
            :attr:`pyvista.plotting.themes.Theme.colorbar_horizontal`
            depending on the value of ``vertical``.

        position_y : float, optional
            The percentage (0 to 1) along the windows's vertical
            direction to place the bottom left corner of the colorbar.
            Default set by
            :attr:`pyvista.plotting.themes.Theme.colorbar_vertical` or
            :attr:`pyvista.plotting.themes.Theme.colorbar_horizontal`
            depending on the value of ``vertical``.

        vertical : bool, optional
            Use vertical or horizontal scalar bar.  Default set by
            :attr:`pyvista.plotting.themes.Theme.colorbar_orientation`.

        interactive : bool, optional
            Use a widget to control the size and location of the scalar bar.
            Default set by :attr:`pyvista.plotting.themes.Theme.interactive`.

        fmt : str, optional
            ``printf`` format for labels.
            Default set by :attr:`pyvista.plotting.themes.Theme.font`.

        use_opacity : bool, default: True
            Optionally display the opacity mapping on the scalar bar.

        outline : bool, default: False
            Optionally outline the scalar bar to make opacity mappings more
            obvious.

        nan_annotation : bool, default: False
            Annotate the NaN color.

        below_label : str, optional
            String annotation for values below the scalars range.

        above_label : str, optional
            String annotation for values above the scalars range.

        background_color : ColorLike, optional
            The color used for the background in RGB format.

        n_colors : int, optional
            The maximum number of color displayed in the scalar bar.

        fill : bool, default: False
            Draw a filled box behind the scalar bar with the
            ``background_color``.

        render : bool, default: False
            Force a render when True.

        theme : pyvista.plotting.themes.Theme, optional
            Plot-specific theme.  By default, calling from the
            ``Plotter``, will use the plotter theme.  Setting to
            ``None`` will use the global theme.

        unconstrained_font_size : bool, default: False
            Whether the font size of title and labels is unconstrained.
            When it is constrained, the size of the scalar bar will constrain the font size.
            When it is not, the size of the font will always be respected.
            Using custom labels will force this to be ``True``.

            .. versionadded:: 0.44.0

        Returns
        -------
        :vtk:`vtkScalarBarActor`
            Scalar bar actor.

        Notes
        -----
        Setting ``title_font_size``, or ``label_font_size`` disables
        automatic font sizing for both the title and label.

        See Also
        --------
        :ref:`scalar_bar_example`

        Examples
        --------
        Add a custom interactive scalar bar that is horizontal, has an
        outline, and has a custom formatting.

        >>> import pyvista as pv
        >>> sphere = pv.Sphere()
        >>> sphere['Data'] = sphere.points[:, 2]
        >>> plotter = pv.Plotter()
        >>> _ = plotter.add_mesh(sphere, show_scalar_bar=False)
        >>> _ = plotter.add_scalar_bar(
        ...     'Data',
        ...     interactive=True,
        ...     vertical=False,
        ...     title_font_size=35,
        ...     label_font_size=30,
        ...     outline=True,
        ...     fmt='%10.5f',
        ... )
        >>> plotter.show()

        """
        if mapper is None:
            msg = 'Mapper cannot be ``None`` when creating a scalar bar'
            raise ValueError(msg)

        if theme is None:
            theme = pyvista.global_theme

        if interactive is None:
            interactive = theme.interactive
        if font_family is None:
            font_family = theme.font.family
        if label_font_size is None:
            label_font_size = theme.font.label_size
        if title_font_size is None:
            title_font_size = theme.font.title_size
        if fmt is None:
            fmt = theme.font.fmt
        if vertical is None and theme.colorbar_orientation.lower() == 'vertical':
            vertical = True

        # Automatically choose size if not specified
        if width is None:
            width = theme.colorbar_vertical.width if vertical else theme.colorbar_horizontal.width
        if height is None:
            if vertical:
                height = theme.colorbar_vertical.height
            else:
                height = theme.colorbar_horizontal.height

        # Check that this data hasn't already been plotted
        if title in list(self._scalar_bar_ranges.keys()):
            clim = list(self._scalar_bar_ranges[title])
            newrng = mapper.scalar_range
            oldmappers = self._scalar_bar_mappers[title]
            # get max for range and reset everything
            clim[0] = min(newrng[0], clim[0])
            clim[1] = max(newrng[1], clim[1])
            for mh in oldmappers:
                mh.scalar_range = clim[0], clim[1]
            mapper.scalar_range = clim[0], clim[1]
            self._scalar_bar_mappers[title].append(mapper)
            self._scalar_bar_ranges[title] = clim
            self._scalar_bar_actors[title].SetLookupTable(mapper.lookup_table)
            # Color bar already present and ready to be used so returning
            return None

        # Automatically choose location if not specified
        if position_x is None or position_y is None:
            if not self._plotter._scalar_bar_slots:
                msg = f'Maximum number of color bars ({MAX_N_COLOR_BARS}) reached.'
                raise RuntimeError(msg)

            slot = min(self._plotter._scalar_bar_slots)
            self._plotter._scalar_bar_slots.remove(slot)
            self._plotter._scalar_bar_slot_lookup[title] = slot

            if position_x is None:
                if vertical:
                    position_x = theme.colorbar_vertical.position_x
                    position_x -= slot * (width + 0.2 * width)
                else:
                    position_x = theme.colorbar_horizontal.position_x

            if position_y is None:
                if vertical:
                    position_y = theme.colorbar_vertical.position_y
                else:
                    position_y = theme.colorbar_horizontal.position_y
                    position_y += slot * height

        # parse color
        color = Color(color, default_color=theme.font.color)

        # Create scalar bar
        scalar_bar = _vtk.vtkScalarBarActor()
        # self._scalar_bars.append(scalar_bar)

        if background_color is not None:
            background_color = np.array(Color(background_color).int_rgba)
            scalar_bar.GetBackgroundProperty().SetColor(background_color[0:3])

            if fill:
                scalar_bar.DrawBackgroundOn()

            lut = pyvista.LookupTable()
            lut.DeepCopy(mapper.lookup_table)
            ctable = _vtk.vtk_to_numpy(lut.GetTable())
            alphas = ctable[:, -1][:, np.newaxis] / 255.0
            use_table = ctable.copy()
            use_table[:, -1] = 255.0
            ctable = (use_table * alphas) + background_color * (1 - alphas)
            lut.SetTable(_vtk.numpy_to_vtk(ctable, array_type=_vtk.VTK_UNSIGNED_CHAR))
        else:
            lut = mapper.lookup_table

        scalar_bar.SetLookupTable(lut)
        if n_colors is None:
            # ensure the number of colors in the scalarbar's lookup table is at
            # least the number in the mapper
            n_colors = mapper.lookup_table.n_values

        scalar_bar.SetMaximumNumberOfColors(n_colors)

        if n_labels < 1:
            scalar_bar.SetDrawTickLabels(False)
        else:
            scalar_bar.SetDrawTickLabels(True)
            scalar_bar.SetNumberOfLabels(n_labels)

        if nan_annotation:
            scalar_bar.DrawNanAnnotationOn()

        if above_label is not None:
            scalar_bar.DrawAboveRangeSwatchOn()
            scalar_bar.SetAboveRangeAnnotation(above_label)
        elif lut.above_range_color:
            scalar_bar.DrawAboveRangeSwatchOn()
            scalar_bar.SetAboveRangeAnnotation('above')
        if below_label is not None:
            scalar_bar.DrawBelowRangeSwatchOn()
            scalar_bar.SetBelowRangeAnnotation(below_label)
        elif lut.below_range_color:
            scalar_bar.DrawBelowRangeSwatchOn()
            scalar_bar.SetBelowRangeAnnotation('below')

        # edit the size of the colorbar
        scalar_bar.SetHeight(height)
        scalar_bar.SetWidth(width)
        scalar_bar.SetPosition(position_x, position_y)

        if fmt is not None:
            scalar_bar.SetLabelFormat(fmt)

        if vertical:
            scalar_bar.SetOrientationToVertical()
        else:
            scalar_bar.SetOrientationToHorizontal()

        if label_font_size is not None or title_font_size is not None:
            scalar_bar.SetUnconstrainedFontSize(True)
            scalar_bar.SetAnnotationTextScaling(False)
        else:
            scalar_bar.SetAnnotationTextScaling(True)

        label_text = scalar_bar.GetLabelTextProperty()
        anno_text = scalar_bar.GetAnnotationTextProperty()
        label_text.SetColor(color.float_rgb)
        anno_text.SetColor(color.float_rgb)
        label_text.SetShadow(shadow)
        anno_text.SetShadow(shadow)

        # Set font
        label_text.SetFontFamily(parse_font_family(font_family))
        anno_text.SetFontFamily(parse_font_family(font_family))
        label_text.SetItalic(italic)
        anno_text.SetItalic(italic)
        label_text.SetBold(bold)
        anno_text.SetBold(bold)
        if label_font_size:
            label_text.SetFontSize(label_font_size)
            anno_text.SetFontSize(label_font_size)

        # Set properties
        self._scalar_bar_ranges[title] = mapper.scalar_range
        self._scalar_bar_mappers[title] = [mapper]

        scalar_bar.SetTitle(title)
        title_text = scalar_bar.GetTitleTextProperty()

        title_text.SetJustificationToCentered()

        title_text.SetItalic(italic)
        title_text.SetBold(bold)
        title_text.SetShadow(shadow)
        if title_font_size:
            title_text.SetFontSize(title_font_size)

        # Set font
        title_text.SetFontFamily(parse_font_family(font_family))

        # set color
        title_text.SetColor(color.float_rgb)

        self._scalar_bar_actors[title] = scalar_bar
        if interactive:
            scalar_widget = _vtk.vtkScalarBarWidget()
            scalar_widget.SetScalarBarActor(scalar_bar)
            scalar_widget.SetInteractor(self._plotter.iren.interactor)
            scalar_widget.SetEnabled(1)
            rep = scalar_widget.GetRepresentation()

            scalar_widget.On()
            if vertical is True or vertical is None:
                rep.SetOrientation(1)  # type: ignore[attr-defined] # 0 = Horizontal, 1 = Vertical
            else:
                # y position determined empirically
                y = -position_y / 2 - height - scalar_bar.GetPosition()[1]
                rep.GetPositionCoordinate().SetValue(width, y)  # type: ignore[attr-defined]
                rep.GetPosition2Coordinate().SetValue(height, width)  # type: ignore[attr-defined]
                rep.SetOrientation(0)  # type: ignore[attr-defined] # 0 = Horizontal, 1 = Vertical
            self._scalar_bar_widgets[title] = scalar_widget

        if use_opacity:
            scalar_bar.SetUseOpacity(True)

        if outline:
            scalar_bar.SetDrawFrame(True)
            frame_prop = scalar_bar.GetFrameProperty()
            frame_prop.SetColor(color.float_rgb)
        else:
            scalar_bar.SetDrawFrame(False)

        if unconstrained_font_size:
            scalar_bar.SetUnconstrainedFontSize(True)

        # finally, add to the actor and return the scalar bar
        self._plotter.add_actor(scalar_bar, reset_camera=False, pickable=False, render=render)

        return scalar_bar