# -*- coding: utf-8 -*-

# ======================================================================================================================
#  Copyright (©) 2015-2021 LCS - Laboratoire Catalyse et Spectrochimie, Caen, France.
#  CeCILL-B FREE SOFTWARE LICENSE AGREEMENT - See full LICENSE agreement in the root directory
# ======================================================================================================================
from functools import partial
import re

import numpy as np
from matplotlib import pyplot as plt

import spectrochempy as scp

from spectrochempy_gui import pyqtgraph as pg
from spectrochempy_gui.pyqtgraph.Qt import QtCore, QtGui
from spectrochempy_gui.pyqtgraph.functions import mkPen, mkColor
from spectrochempy_gui.pyqtgraph import GraphicsLayoutWidget, ViewBox

# ----------------------------------------------------------------------------------------------------------------------
class CustomViewBox(ViewBox):
    """
    Subclass of ViewBox
    """
    sigPlotModeChanged = QtCore.Signal(object)
    sigColorMapChanged = QtCore.Signal(object)
    sigLineWidthChanged = QtCore.Signal(object)

    def __init__(self, parent=None, ndim=1, prefs=None):
        """
        Constructor of the CustomViewBox
        """
        super().__init__(parent)

        self.ndim = ndim
        self.prefs = prefs

        self.setRectMode()          # Set mouse mode to rect for convenient zooming
        self.menu = None            # Override pyqtgraph ViewBoxMenu
        self.menu = self.getMenu()  # Create the menu

    def raiseContextMenu(self, ev):
        """
        Raise the context menu
        """
        if not self.menuEnabled():
            return
        menu = self.menu # getMenu()
        pos = ev.screenPos()
        menu.popup(QtCore.QPoint(pos.x(), pos.y()))

    def getMenu(self):
        """
        Create the menu
        """
        if self.menu is None:
            self.menu = QtGui.QMenu()

        self.PlotModeMenu = QtGui.QMenu("Plot mode")
        self.plotModeCombo = QtGui.QComboBox()
        self.modeItems = ['pen', 'scatter', 'scatter+pen', 'bar'] if self.ndim < 2 else ['stack', 'map', 'image',
                                                                                      'surface']
        self.plotModeCombo.insertItems(1,self.modeItems)
        self.plotModeCombo.activated.connect(self.emitPlotModeChanged)
        self.plotModeAction = QtGui.QWidgetAction(None)
        self.plotModeAction.setDefaultWidget(self.plotModeCombo)
        self.PlotModeMenu.addAction(self.plotModeAction)
        self.menu.addMenu(self.PlotModeMenu)

        if self.ndim == 2:
            self.ColorMapMenu = QtGui.QMenu("Colormap")
            self.colorMapCombo = QtGui.QComboBox()
            self.colorMapItems = plt.colormaps()
            self.colorMapCombo.insertItems(1,self.colorMapItems)
            self.colorMapCombo.setCurrentIndex(self.colorMapItems.index(self.prefs.colormap))
            self.colorMapCombo.activated.connect(self.emitColorMapChanged)
            self.colorMapAction = QtGui.QWidgetAction(None)
            self.colorMapAction.setDefaultWidget(self.colorMapCombo)
            self.ColorMapMenu.addAction(self.colorMapAction)
            self.menu.addMenu(self.ColorMapMenu)

        self.LineMenu = QtGui.QMenu("Linewidth")
        self.lineSpinBox = pg.SpinBox(value=1, step=0.1, decimals=3)
        self.lineSpinBox.sigValueChanged.connect(self.emitLineWidthChanged)
        self.lineAction = QtGui.QWidgetAction(None)
        self.lineAction.setDefaultWidget(self.lineSpinBox)
        self.LineMenu.addAction(self.lineAction)
        self.menu.addMenu(self.LineMenu)

        self.viewAll = QtGui.QAction("Zoom reset", self.menu)
        self.viewAll.triggered.connect(self.autoRange)
        self.menu.addAction(self.viewAll)

        self.leftMenu = QtGui.QMenu("Left mouse click mode")
        group = QtGui.QActionGroup(self)
        pan = QtGui.QAction(u'Pan', self.leftMenu)
        zoom = QtGui.QAction(u'Zoom', self.leftMenu)
        self.leftMenu.addAction(pan)
        self.leftMenu.addAction(zoom)
        pan.triggered.connect(self.setPanMode)
        zoom.triggered.connect(self.setRectMode)
        pan.setCheckable(True)
        zoom.setCheckable(True)
        pan.setActionGroup(group)
        zoom.setActionGroup(group)
        self.menu.addMenu(self.leftMenu)

        self.menu.addSeparator()

        return self.menu

    def emitLineWidthChanged(self, val):
        self.sigLineWidthChanged.emit(float(val.value()))

    def emitPlotModeChanged(self, index):
        mode = self.modeItems[index]
        self.sigPlotModeChanged.emit(mode)

    def emitColorMapChanged(self, index):
        color = self.colorMapItems[index]
        self.sigColorMapChanged.emit(color)

    def setRectMode(self):
        """
        Set mouse mode to rect
        """
        self.setMouseMode(self.RectMode)

    def setPanMode(self):
        """
        Set mouse mode to pan
        """
        self.setMouseMode(self.PanMode)

# --------------------------------------------------------------------
# Plot widget
# --------------------------------------------------------------------

class PlotWidget(GraphicsLayoutWidget):

    sigZoomReset = QtCore.Signal()

    _dataset = None  # current dataset associated to the plotwidget
    _autorange = True

    # ..................................................................................................................
    def __init__(self, parent):

        # Use the Qt's GraphicsView framework offered by Pyqtgraph
        super().__init__(title='')

        self.parent = parent

        # Prepare additionnal traces
        self.selected = None  # traces selected in 2D spectra
        self.selected_pen = None # original pen of a selected curve
        self.sigZoomReset.connect(self.zoomReset)

    # ..................................................................................................................
    def _masked(self, data, mask):
        """
        Utility function which returns a masked array.
        """
        if self._dataset is None:
            return None

        if not np.any(mask):
            mask = np.zeros(data.shape).astype(bool)
        data = np.ma.masked_where(mask, data)  # np.ma.masked_array(data, mask)
        return data

    # ..................................................................................................................
    @property
    def dataset(self):
        """
        Returns the current dataset
        """
        return self._dataset

    # ..................................................................................................................
    @dataset.setter
    def dataset(self, val):
        """
        Set the current dataset.
        """
        self._dataset = val

    # ..................................................................................................................
    @property
    def data(self):
        """
        Returns the current dataset masked data.
        """
        # z intensity (by default we plot real component of the data)

        return self._masked(np.real(self._dataset.data), self._dataset.mask)

    # ..................................................................................................................
    @property
    def processeddata(self):
        """
        Returns processeddata with the same mask as the original dataset.
        """
        if self._dataset.processeddata is None:
            return None
        return self._masked(np.real(self._dataset.processeddata), self._dataset.processedmask)

    # ..................................................................................................................
    @property
    def baselinedata(self):
        """
        Returns baselinedata with the same mask as the original dataset.
        """
        if self._dataset.baselinedata is None:
            return None
        return self._masked(np.real(self._dataset.baselinedata), self._dataset.mask)

    # ..................................................................................................................
    @property
    def referencedata(self):
        """
        Returns referencedata with the same mask as the original dataset.
        """
        if self._dataset.referencedata is None:
            return None
        return self._masked(np.real(self._dataset.referencedata), self._dataset.mask)

#..................................................................................................................
    def draw_regions(self, reg=None):

        procs = self.parent.controller.params.param('processing').children()
        for proc in procs:
            if proc.name().startswith('define regions'):
                kind = proc.param('kind').value()
                if kind == 'undefined' or not hasattr(proc, 'regions'):
                    return
                dim = self.dataset.dims[-1]
                if not proc.param('regiongroup').childs[0].value().startswith(dim):
                    continue

                for el, par in proc.regions.regionItems.values():
                    if el._name.startswith(f"{dim}_{kind}_"):
                        if not proc.opts['expanded']:
                            self.p.removeItem(el)
                        else:
                            self.p.addItem(el, ignoreBounds=True)

    # ..................................................................................................................
    def changeColorMap(self, map):
        self.dataset.meta['colormap'] = map
        self.draw(self.dataset)

    def changePlotMode(self, mode):
        self.dataset.meta['plotmode'] = mode
        self.draw(self.dataset)

    def changeLineWidth(self, lw):
        self.dataset.meta['linewidth'] = lw
        self.draw(self.dataset)

    def zoomReset(self):
        self._autorange=True

    def inverted(self, lim):
        if lim[0] > lim[1]:
            return True
        return False

    # ..................................................................................................................
    def draw(self, dataset, zoom_reset=False):
        """
        Draw 1D or 2D dataset corresponding to the current dataset.

        Parameters
        ----------
        zoom_reset: bool, optional
            True if the x and y range must be reset to the full range when redrawing.
        """

        self.dataset = dataset
        if self._autorange:
            zoom_reset = True

        # Create the main plotItem
        if not hasattr(self, 'p'):
            vb = CustomViewBox(ndim=dataset.ndim, prefs=dataset.preferences)
            self.p = self.addPlot(row=0, col=0, viewBox=vb)

        # Draw main data
        scp.debug_('>>>>>>>>>> Draw')
        self._draw(zoom_reset=zoom_reset)

        # Draw processed
        self._draw_processed(zoom_reset=zoom_reset)

        self._autorange = False

    # ..................................................................................................................
    def _draw_processed(self, **kwargs):
        # Draw in a second viewbox the processed data.

        if self.dataset.processeddata is None:
            if not hasattr(self, 'proc'):
                return
            else:
                # Try to remove the processing plotItem if it exists
                self.removeItem(self.proc)
                del self.proc
                self.p.setTitle('')
                return

        if not hasattr(self, 'proc'):
            # we have processed data but not yet the corresponding plotItem: create one.
            vb = CustomViewBox(ndim=self.dataset.ndim, prefs=self.dataset.preferences)
            self.proc = self.addPlot(row=1, col=0, title='Processed dataset', viewBox=vb)
            self.p.setTitle('Original dataset')

        self._draw(plotitem=self.proc, processed=True, **kwargs)

        # self.p.vb.register(self.p.titleLabel.text)
        # self.proc.vb.register(self.proc.titleLabel.text)

    def _draw(self, **kwargs):

        # Prepare the viewbox
        plot = kwargs.get('plotitem', self.p)
        vb = plot.vb
        if self.dataset.ndim > 1:
            vb.sigColorMapChanged.connect(self.changeColorMap)
            vb.sigPlotModeChanged.connect(self.changePlotMode)
            vb.sigLineWidthChanged.connect(self.changeLineWidth)
        plot.clear()

        # Copy the dataset
        new = self.dataset.copy()
        processed = kwargs.get('processed', False)
        if processed :
            zdata = self.processeddata
        else:
            zdata = self.data

        # Get some preferences
        prefs = new.preferences
        lw = new.meta.get('linewidth', prefs.lines_linewidth)

        # Set axis
        # ========
        # Set the abscissa axis (x)
        # -------------------------

        # The actual dimension name is the last in the new.dims list
        dimx = new.dims[-1]

        # reduce data to the ROI
        x = getattr(new, dimx)
        lx, ux = x.roi
        if new.ndim > 1:
            new = new[:, lx:ux]
        else:
            new = new[lx:ux]

        # read again the x coordinate in case of ROI change
        x = getattr(new, dimx)
        if x is not None and x.implements('CoordSet'):
            # if several coords, take the default ones:
            x = x.default
        xsize = new.shape[-1]
        show_x_points = False
        if x is not None and hasattr(x, 'show_datapoints'):
            show_x_points = x.show_datapoints
        if show_x_points:
            # remove data and units for display
            x = scp.LinearCoord.arange(xsize)

        discrete_data = False
        if x is not None and (not x.is_empty or x.is_labeled):
            xdata = x.data
            if not np.any(xdata):
                if x.is_labeled:
                    discrete_data = True
                    # take into account the fact that sometimes axis have just labels
                    xdata = range(1, len(x.labels) + 1)
        else:
            xdata = range(xsize)

        xlim = [xdata[0], xdata[-1]]
        xlim.sort()

        vb.invertX(x.reversed)
        print(plot, x.reversed, xlim, x.limits, plot.getAxis('bottom').range)

        zoom_reset = kwargs.get('zoom_reset', False)
        if not zoom_reset:
            if sorted(plot.getAxis('bottom').range) != [0, 1] and x.title in plot.getAxis('bottom').labelText:
                range = plot.getAxis('bottom').range
                range= sorted(range, reverse=True)
                vb.setXRange(*range, padding=0)
                print('1 - setXrange (range)', range)
            else:
                vb.setXRange(*xlim, padding=0)
                print('2 - setXrange (xlim)', xlim)
        else:
            vb.setXRange(*xlim, padding=0)
            print('3 - setXrange (xlim)', xlim)

        ndim = new._squeeze_ndim
        if ndim > 1:

            # Set the ordinates axis (y)
            # --------------------------

            # The actual dimension name is the second in the new.dims list
            dimy = new.dims[-2]

            # Reduce to ROI
            y = getattr(new, dimy)
            ly, uy = y.roi
            new = new[ly:uy]

            y = getattr(new, dimy)
            if y is not None and y.implements('CoordSet'):
                # if several coords, take the default ones:
                y = y.default
            ysize = new.shape[-2]

            show_y_points = False
            if ysize > 1:

                # 2D (else it will be displayed as 1D)
                # ------------------------------------
                if y is not None and hasattr(y, 'show_datapoints'):
                    show_y_points = y.show_datapoints
                if show_y_points:
                    # remove data and units for display
                    y = scp.LinearCoord.arange(ysize)

                if y is not None and (not y.is_empty or y.is_labeled):
                    ydata = y.data

                    if not np.any(ydata):
                        if y.is_labeled:
                            ydata = range(1, len(y.labels) + 1)
                else:
                    ydata = range(ysize)

                yl = [ydata[0], ydata[-1]]
                yl.sort()

                ylim = list(kwargs.get("ylim", yl))
                ylim.sort()
                ylim[-1] = min(ylim[-1], yl[-1])
                ylim[0] = max(ylim[0], yl[0])

        # Amplitude (z)
        # -------------

        zlim = kwargs.get('zlim', (np.ma.min(zdata), np.ma.max(zdata)))

        method = new.meta.get('mode', prefs.method_2D) if ndim > 1 else 'stack'

        if method in ['stack']:  # For 2D and 1D plot

            # The z axis info
            # ---------------
            amp = 0
            zl = (np.min(np.ma.min(zdata) - amp), np.max(np.ma.max(zdata)) + amp)
            zlim = list(kwargs.get('zlim', zl))
            zlim.sort()
            z_reverse = kwargs.get('z_reverse', False)
            vb.invertY(z_reverse)

            # Set the limits
            # ---------------

            # if yscale == "log" and min(zlim) <= 0:
            #    # set the limits wrt smallest and largest strictly positive values
            #    ax.set_ylim(10 ** (int(np.log10(np.amin(np.abs(zdata)))) - 1),
            #                10 ** (int(np.log10(np.amax(np.abs(zdata)))) + 1))

            vb.setYRange(*zlim, padding=0)

        else:

            #TODO
            pass # not implemented

            # # the y axis info
            # # ----------------
            # # if data_only:
            # #    ylim = ax.get_ylim()
            #
            # ylim = list(kwargs.get('ylim', ylim))
            # ylim.sort()
            # y_reverse = kwargs.get('y_reverse', y.reversed if y else False)
            # if y_reverse:
            #     ylim.reverse()
            #
            # # set the limits
            # # ----------------
            # ax .set_ylim(ylim)


        # Log scale

        # yscale = kwargs.get("yscale", "linear")
        # ax.set_yscale(yscale)
        # xscale = kwargs.get("xscale", "linear")
        # ax.set_xscale(xscale)  # , nonpositive='mask')

        # Plot the dataset
        # ================

        #  ax.grid(prefs.axes_grid)  # TODO

        cmap = new.meta.get('colormap', prefs.colormap)
        self.cmap = pg.colormap.get(cmap, source='matplotlib', skipCache=True)

        if method in ['stack']:

            # if data.ndim == 1:
            #    data = data.at_least2d()

            ncurves = zdata.shape[0]
            colors = self.cmap.color
            if ncurves > 1:
                icolor = np.linspace(0, (colors.shape[0]-1), ncurves).astype(int)
                colors = colors[icolor]
            else:
                colors = [colors[0]] # [prefs('color')]
            self.colors = colors

            # self.curves = []
            if hasattr(zdata, 'mask'):
                mask = zdata.mask
                zdata[mask] = np.nan

            # Downsampling
            step = 1
            if ncurves > 250:
                step = int(ncurves / 250)

            for i in np.arange(0, ncurves, step):
                zdat = zdata[i:i + step].max(axis=0) if step > 1 else zdata[i]
                if np.alltrue(mask):
                    continue
                c = pg.PlotCurveItem(
                          x=xdata,
                          y=zdat,
                          pen=mkPen(mkColor(colors[i]), width=lw),
                          clickable=True,
                          connect='finite')
                plot.addItem(c)
                c.sigClicked.connect(partial(self._curveSelected, plot))


        # Display a title
        title = kwargs.get('title', None)
        if title:
            plot.setTitle(title)
        elif kwargs.get('plottitle', False):
            plot.setTitle(new.name)

        # Labels
        # ======

        def make_label(ss, label):
            if ss.units is not None and str(ss.units) != 'dimensionless':
                units = r"{:~P}".format(ss.units)
            else:
                units = ''
            label = f"{label} / {units}"
            return label

        # --------------------------------------------------------------------------------------------------------------
        # x label
        # --------------------------------------------------------------------------------------------------------------

        xlabel = kwargs.get("xlabel", None)
        if show_x_points:
            xlabel = 'data points'
        if not xlabel:
            xlabel = make_label(x, x.title)

        plot.setLabel('bottom', text=xlabel)

        # uselabelx = kwargs.get('uselabel_x', False)
        # if x and x.is_labeled and (uselabelx or not np.any(x.data)) and len(x.labels) < number_x_labels + 1:
        #     # TODO refine this to use different orders of labels
        #     ax.set_xticks(xdata)
        #     ax.set_xticklabels(x.labels)
        #

        if ndim > 1:
            # y label
            # --------
            ylabel = kwargs.get("ylabel", None)
            if show_y_points:
                ylabel = 'data points'
            if not ylabel:
                if method in ['stack']:
                    ylabel = make_label(new, y.title)
                else:
                    ylabel = make_label(y, new.dims[-2])

            # uselabely = kwargs.get('uselabel_y', False)
            # if y and y.is_labeled and (uselabely or not np.any(y.data)) and len(y.labels) < number_y_labels:
            #             # TODO refine this to use different orders of labels
            #             ax.set_yticks(ydata)
            #             ax.set_yticklabels(y.labels)

        # --------------------------------------------------------------------------------------------------------------
        # z label
        # --------------------------------------------------------------------------------------------------------------

        zlabel = kwargs.get("zlabel", None)
        if not zlabel:
            if method in ['stack']:
                zlabel = make_label(new, new.title)
            elif method in ['surface']:
                zlabel = make_label(new, 'values')
            plot.setLabel('left', text=zlabel)
        else:
            zlabel = make_label(new, 'z')

        if method in ['stack']:
            # do we display the ordinate axis?
            if kwargs.get('show_y', True):
                plot.setLabel('left', text=zlabel)
            else:
                plot.hideAxis('left')

        if not hasattr(self, 'label'):
            self.label = pg.LabelItem(parent=self.p, justify='left')

        # --------------------------------------------------------------------------------------------------------------
        # Regions
        # --------------------------------------------------------------------------------------------------------------

        # Restore regions
        self.draw_regions()

        # --------------------------------------------------------------------------------------------------------------
        # Vertical line
        # --------------------------------------------------------------------------------------------------------------

        vLine = pg.InfiniteLine(angle=90, movable=False)
        # hLine = pg.InfiniteLine(angle=0, movable=False)
        plot.addItem(vLine, ignoreBounds=True)
        # plot.addItem(hLine, ignoreBounds=True)

        scene = plot.scene()
        def mouseMoved(evt):
            pos = evt
            scene.blockSignals(True)
            if plot.sceneBoundingRect().contains(pos):
                mouse_point = vb.mapSceneToView(pos)
                coord = mouse_point.x()
                ll, hl = vb.state['viewRange'][0]
                lld, hld = x.roi
                if max(ll, lld) <= coord <= min(hl, hld):
                    ds = new[:, float(coord)]
                    vLine.setVisible(True)
                    if self.selected:
                        try:
                            # corresponding x index
                            index = x.loc2index(coord)
                            z = self.selected.yData[index] * x.units
                            zstr = f'{z:~0.2fP} '

                        except Exception:
                            vLine.setVisible(False)
                            scene.blockSignals(False)
                            return   # out of limits
                    else:
                        z = ds.value
                        zstr = f'{z:~0.2fP} '
                    if z.size > 1:
                        # mode than one element (2D)
                        z = z.squeeze()
                        zstr = f'{z.min():~0.2fP} -- {z.max():~0.2fP} '

                    coord = coord * x.units
                    coordstr = f'{coord:~0.2fP}'
                    self.label.setText(
                            f"<span style='background-color:#FFF; font-size: 12pt'>"
                            f"<span style='color: blue'>{x.title} = {coordstr}</span>"
                            f"<br/>"
                            f"<span style='color: green'>{ds.title} = {zstr}</span>"
                            f"</span>")
                    vLine.setPos(mouse_point.x())
                    # hLine.setPos(mouse_point.y())
                else:
                    self.label.setText('')
                    vLine.setPos(0)
                    vLine.setVisible(False)
                    # hLine.setPos(0)
            scene.blockSignals(False)

        plot.scene().sigMouseMoved.connect(mouseMoved)

    # ..................................................................................................................
    def _findCurveIndex(self, plot, curve):

        for index, c in enumerate(plot.curves):
            if not isinstance(c, pg.PlotCurveItem):
                continue
            if c is curve:
                break

        return index

    # ..................................................................................................................
    def _curveSelected(self, plot, curve):
        # Action when a curve is selected

        lw =  self.dataset.meta.get('linewidth', self.dataset.preferences.lines_linewidth)
        if self.dataset._squeeze_ndim < 2:
            return

        if self.selected is not None:
            # reset previous
            #index = self._findCurveIndex(plot, self.selected)
            pen = self.selected_pen
            self.selected.setPen(pen)

        if curve != self.selected:
            # set new selected
            self.selected_pen = curve.opts['pen']
            curve.setPen('k', width=lw*3)
            self.selected = curve

        else:
            # index = self._findCurveIndex(plot, curve)
            curve.setPen(self.selected_pen)
            self.selected = None
