import typing
import logging
import time

import numpy as np
from matplotlib import colors, cm

from libertem.udf.base import UDFRunner

# Import here for backwards compatibility, refs #1031
from libertem.common.viz import encode_image  # NOQA: F401

if typing.TYPE_CHECKING:
    from libertem.common.buffers import BufferWrapper

logger = logging.getLogger(__name__)


def _get_norm(result, norm_cls=colors.Normalize, vmin=None, vmax=None, damage=None):
    if (vmin is not None) and (vmax is not None):
        return norm_cls(vmin=vmin, vmax=vmax)

    result = result.astype(np.float32)

    if damage is None:
        damage = (result != 0)

    damage = damage & np.isfinite(result)

    if damage.sum() == 0:
        return norm_cls(vmin=1, vmax=1)  # all-NaN or all-zero

    if vmin is None:
        vmin = np.min(result[damage])
    if vmax is None:
        vmax = np.max(result[damage])

    return norm_cls(vmin=vmin, vmax=vmax)


def visualize_simple(result, colormap=None, logarithmic=False, vmin=None, vmax=None, damage=None):
    """
    Normalize and visualize ``result`` with ``colormap`` and return the
    resulting RGBA data as an array.

    Parameters
    ----------
    result : numpy.ndarray
        2d array of intensity values

    colormap : matplotlib colormap or None
        colormap used for visualizing intensity values, defaults to ColormapCubehelix()

    Returns
    -------

    np.array
        A numpy array of shape (Y, X, 4) containing RGBA data, suitable for
        passing to `Image.fromarray` in PIL.
    """
    if logarithmic:
        # Convert to the smallest dtype that supports subtractions
        dtype = np.result_type(result, np.int8)
        result = result.astype(dtype)
        cnorm = colors.LogNorm
        result = result - np.min(result) + 1
    else:
        cnorm = colors.Normalize
    if colormap is None:
        colormap = cm.gist_earth
    norm = _get_norm(result, norm_cls=cnorm, vmin=vmin, vmax=vmax, damage=damage)
    shape = result.shape
    normalized = norm(result.reshape((-1,))).reshape(shape)
    colored = colormap(normalized, bytes=True)
    return colored


def get_plottable_2D_channels(buffers):
    return [
        k
        for k, buf in buffers.items()
        # 2D data, removing axes of size 1
        if len(buf.data.squeeze().shape) == 2
    ]


class Live2DPlot:
    """
    Base plotting class for interactive use. Please see the subclasses for
    concrete details.

    .. versionadded:: 0.7.0

    Parameters
    ----------
    dataset : DataSet
        The dataset on which the UDF will be run. This allows
        to determine the shape of the plots for initialization.
    udf : UDF
        The UDF instance this plot is associated to. This needs to be
        the same instance that is passed to
        :meth:`~libertem.api.Context.run_udf`.
    roi : numpy.ndarray or None
        Region of interest (ROI) that the UDF will
        be run on. This is necessary for UDFs where the `extra_shape`
        parameter of result buffers is a function of the ROI, such as
        :class:`~libertem.udf.raw.PickUDF`.
    channel : misc
        Indicate the channel to be plotted.

        - :code:`None`: The first plottable (2D) channel of the UDF is plotted.
        - :code:`str`: The UDF result buffer name that should be plotted.
        - :code:`tuple(str, function(ndarray) -> ndarray)`: The UDF result buffer name that
          should be plotted together with a function that extracts a plottable result
        - :code:`function(udf_result, damage) -> (ndarray, damage)`: Function that derives a
          plottable 2D ndarray and damage indicator from the full
          UDF results and the processed nav space. See :ref:`plotting` for more details!
    title : str
        The plot title. By default UDF class name and channel name.
    min_delta : float
        Minimum time span in seconds between updates to reduce overheads for slow plotting.
    udfresult : UDFResults, optional
        UDF result to initialize the plot data and determine plot shape. If None (default),
        this is determined using :meth:`~libertem.udf.base.UDFRunner.dry_run` on the dataset,
        UDF and ROI. This parameter allows re-using buffers to avoid unnecessary dry runs.
    """
    def __init__(
            self, dataset, udf, roi=None, channel=None, title=None, min_delta=0, udfresult=None
    ):
        if udfresult is None:
            udfresult = UDFRunner.dry_run([udf], dataset, roi)
        eligible_channels = get_plottable_2D_channels(udfresult.buffers[0])
        if channel is None:
            if not eligible_channels:
                raise ValueError(f"No plottable 2D channel found for {udf.__class__.__name__}")
            channel = eligible_channels[0]
            channel_title = channel

        if callable(channel):
            extract = channel
            channel_title = channel.__name__
            channel = None
        elif isinstance(channel, (tuple, list)):
            channel, func = channel
            if channel not in udfresult.buffers[0]:
                raise ValueError(
                    f"channel {channel} not found, have: {udfresult.buffers[0].keys()}"
                )
            kind = udfresult.buffers[0][channel].kind
            if kind == 'nav':
                def extract(udf_results, damage):
                    return (func(udf_results[channel].data), damage)
            else:
                def extract(udf_results, damage):
                    return (func(udf_results[channel].data), True)

            channel_title = f"{func.__name__}({channel})"
        else:
            extract = None
            if channel not in eligible_channels:
                raise ValueError("channel {} not found or not plottable, have: {!r}".format(
                    channel, eligible_channels
                ))
            channel_title = channel

        self._extract = extract
        self.channel = channel
        if title is None:
            title = f"{udf.__class__.__name__}: {channel_title}"
        self.title = title
        self.data, _ = self.extract(udfresult.buffers[0], udfresult.damage)
        self.udf = udf
        self.last_update = 0
        self.min_delta = min_delta

    def get_udf(self):
        """
        Returns the associated UDF instance
        """
        return self.udf

    def extract(
        self, udf_results: dict[str, "BufferWrapper"], damage: "BufferWrapper"
    ) -> tuple[np.ndarray, np.ndarray]:
        """
        Extract plotting data from UDF result.


        Parameters
        ----------

        udf_results : UDF result
            (Partial) UDF result

        damage : BufferWrapper
            :class:`~libertem.common.buffers.BufferWraper` with :code:`kind='nav'`
            and :code:`dtype=bool` that indicates the area of the nav dimension that
            has been processed by the UDF already.

        Returns
        -------
        (numpy.ndarray, damage)
            It returns the data and damage of a UDF result buffer indicated by :code:`channel`
            if that is a string, or a numpy.ndarray and damage derived from the UDF results by
            calling :code:`channel` with this method's arguments if it is callable.
        """
        if self._extract is None:
            buffer = udf_results[self.channel]
            squeezed = buffer.masked_data.squeeze()
            res_damage = buffer.valid_mask.squeeze()
            return (squeezed, res_damage)
        else:
            return self._extract(udf_results, damage)

    def new_data(self, udf_results: dict[str, "BufferWrapper"], damage, force: bool = False):
        """
        This method is called with the raw `udf_results` any time a new
        partition has finished processing.

        The :code:`damage` parameter is filtered to only cover finite
        values of :code:`self.data` and passed to :meth:`self.update`,
        which should then be implemented by a subclass.
        """
        t0 = time.time()
        delta = t0 - self.last_update
        if (not force) and delta < self.min_delta:
            return  # don't update plot if we recently updated
        (self.data, damage) = self.extract(udf_results, damage)
        finite = np.isfinite(self.data)
        # duck typing "array-like" since damage can also be bool
        if hasattr(damage, "shape"):
            start_axis = len(damage.shape)
            stop_axis = len(self.data.shape)
            # Reduce surplus axes, for example RGB plotting #1052
            if stop_axis > start_axis:
                finite = np.all(finite, axis=tuple(range(start_axis, stop_axis)))
        damage = damage & finite
        self.update(damage, force=force)
        self.last_update = time.time()
        logger.debug("%s updated in %.3f seconds", self.__class__.__name__, self.last_update - t0)

    def update(self, damage, force=False):
        """
        Update the plot based on `self.data`.

        Parameters
        ----------
        damage : numpy.ndarray Boolean array with the shape of
            :code:`self.data`. It is :code:`True` for all positions in
            :code:`self.data` that contain finite values and have potentially
            been touched by the UDF. This can be used to extract the correct
            plot range by ignoring invalid buffer portions.

        force : bool Force an update, disabling any throttling mechanisms
        """
        raise NotImplementedError()

    def display(self):
        """
        Show the plot, for example in the current Jupyter cell.
        """
        raise NotImplementedError()


class Dummy2DPlot(Live2DPlot):
    '''
    No-op plot. This is useful for test and example code to not
    attempt displaying :code:`matplotlib` plots in a headless environment
    or during batch operation.
    '''
    def update(self, damage, force=False):
        pass

    def display(self):
        pass
