"""Tools for automatic plotting with ``pyqtgraph``.

Contains a ``pyqtgraph``-specific implementation of :class:`.plot.base.FigureMaker`
(:class:`.FigureMaker`) and :class:`.plot.base.PlotWidget` (:class:`.AutoPlot`).
``FigureMaker`` can be used to quickly create a figure in a mostly automatic fashion
from known data.
For embedding in GUIs (such as plottr applications), ``AutoPlot`` is the main
object for plotting data automatically using ``pyqtgraph``.
"""

import logging
from pathlib import Path
import time
from dataclasses import dataclass
from typing import List, Optional, Any

import numpy as np
from pyqtgraph import mkPen

from plottr import QtWidgets, QtCore, Signal, Slot, \
    config_entry as getcfg
from plottr.data.datadict import DataDictBase
from .plots import Plot, PlotWithColorbar, PlotBase
from ..base import AutoFigureMaker as BaseFM, PlotDataType, \
    PlotItem, ComplexRepresentation, determinePlotDataType, \
    PlotWidgetContainer, PlotWidget

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
TIMESTRFORMAT = "%Y-%m-%dT%H%M%S"


class FigureWidget(QtWidgets.QWidget):
    """Widget that contains all plots generated by :class:`.FigureMaker`.

    Widget has a vertical layout, and plots can be added in a single column.
    """

    def __init__(self, parent: Optional[QtWidgets.QWidget] = None):
        """Constructor for :class:`.FigureMakerWidget`.

        :param parent: parent widget.
        """
        super().__init__(parent=parent)

        self.subPlots: List[PlotBase] = []

        self.title = QtWidgets.QLabel(parent=self)
        self.title.setAlignment(QtCore.Qt.AlignHCenter)

        self.split = QtWidgets.QSplitter(parent=self)
        self.split.setOrientation(QtCore.Qt.Vertical)

        layout = QtWidgets.QVBoxLayout()
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(2)
        layout.addWidget(self.title)
        layout.addWidget(self.split)
        self.setLayout(layout)

        self.setTitle('')

    def addPlot(self, plot: PlotBase) -> None:
        """Add a :class:`.PlotBase` widget.

        :param plot: plot widget
        :param title: title of the plot
        """
        self.split.addWidget(plot)
        self.subPlots.append(plot)

    def clearAllPlots(self) -> None:
        """Clear all plot contents."""
        for p in self.subPlots:
            p.clearPlot()

    def deleteAllPlots(self) -> None:
        """Delete all subplot widgets."""
        for p in self.subPlots:
            p.deleteLater()
        self.subPlots = []

    def setTitle(self, title: str = '') -> None:
        self.title.setText(title)
        if len(title.strip()) == 0:
            self.title.setVisible(False)
        else:
            self.title.setVisible(True)


class FigureMaker(BaseFM):
    """``pyqtgraph`` implementation for :class:`.AutoFigureMaker`.

    Please see the base class for more general information.
    """

    # TODO: make scrollable when many figures (set min size)?
    # TODO: check for valid plot data

    def __init__(self, widget: Optional[FigureWidget] = None,
                 clearWidget: bool = True,
                 parentWidget: Optional[QtWidgets.QWidget] = None):
        """Constructor for :class:`.FigureMaker`.

        :param widget: a widget produced by an earlier use of this class
            may be supplied. This widget will then be used for generating the
            plot and avoids deleting and re-creating the widget if no structural
            changes to the plots is required (important case is live-plotting of
            data that keeps updating).
        :param clearWidget:
            if ``True`` force re-creation of all plot elements.
        :param parentWidget:
            parent of the main widget of FigureMaker.
        """
        super().__init__()

        self.widget: FigureWidget
        self.clearWidget = clearWidget

        if widget is None:
            self.widget = FigureWidget(parent=parentWidget)
        else:
            self.widget = widget

    # re-implementing to get correct type annotation.
    def __enter__(self) -> "FigureMaker":
        return self

    def subPlotFromId(self, subPlotId: int) -> PlotBase:
        """Get SubPlot from ID."""
        subPlots = self.subPlots[subPlotId].axes
        assert isinstance(subPlots, list) and len(subPlots) > 0 and \
               isinstance(subPlots[0], PlotBase)
        return subPlots[0]

    def makeSubPlots(self, nSubPlots: int) -> List[PlotBase]:
        """Create empty subplots in the widgets.

        If ``clearWidget`` was not set to ``True`` in the constructor,
        existing sub plot widgets are cleared, but not deleted and re-created.
        """
        plot: PlotBase
        if self.clearWidget:
            self.widget.deleteAllPlots()
            for i in range(nSubPlots):
                if max(self.dataDimensionsInSubPlot(i).values()) == 1:
                    plot = Plot(self.widget)
                    self.widget.addPlot(plot)
                elif max(self.dataDimensionsInSubPlot(i).values()) == 2:
                    plot = PlotWithColorbar(self.widget)
                    self.widget.addPlot(plot)
        else:
            self.widget.clearAllPlots()

        return self.widget.subPlots

    def formatSubPlot(self, subPlotId: int) -> None:
        """Set labels/legends for the given subplot."""
        if len(self.plotIdsInSubPlot(subPlotId)) == 0:
            return

        labels = self.subPlotLabels(subPlotId)
        subPlot = self.subPlotFromId(subPlotId)

        # label the x axis if there's only one x label
        if isinstance(subPlot, Plot):
            if len(set(labels[0])) == 1:
                subPlot.plot.setLabel("bottom", labels[0][0])

        if isinstance(subPlot, PlotWithColorbar):
            if len(set(labels[0])) == 1:
                subPlot.plot.setLabel("bottom", labels[0][0])

            if len(set(labels[1])) == 1:
                subPlot.plot.setLabel('left', labels[1][0])

            if len(set(labels[2])) == 1:
                subPlot.colorbar.setLabel('left', labels[2][0])

    def plot(self, plotItem: PlotItem) -> None:
        """Plot the given item."""
        if plotItem.plotDataType is PlotDataType.unknown:
            if len(plotItem.data) == 2:
                plotItem.plotDataType = PlotDataType.scatter1d
            elif len(plotItem.data) == 3:
                plotItem.plotDataType = PlotDataType.scatter2d

        if plotItem.plotDataType in [PlotDataType.scatter1d, PlotDataType.line1d]:
            self._1dPlot(plotItem)
        elif plotItem.plotDataType == PlotDataType.grid2d:
            self._colorPlot(plotItem)
        elif plotItem.plotDataType == PlotDataType.scatter2d:
            self._scatterPlot2d(plotItem)
        else:
            raise NotImplementedError('Cannot plot this data.')

    def _1dPlot(self, plotItem: PlotItem) -> None:
        colors = getcfg('main', 'pyqtgraph', 'line_colors', default=['r', 'b', 'g'])
        symbols = getcfg('main', 'pyqtgraph', 'line_symbols', default=['o'])
        symbolSize = getcfg('main', 'pyqtgraph', 'line_symbol_size', default=5)

        subPlot = self.subPlotFromId(plotItem.subPlot)

        assert len(plotItem.data) == 2
        x, y = plotItem.data

        color = colors[self.findPlotIndexInSubPlot(plotItem.id) % len(colors)]
        symbol = symbols[self.findPlotIndexInSubPlot(plotItem.id) % len(symbols)]

        if plotItem.plotDataType == PlotDataType.line1d:
            name = plotItem.labels[-1] if isinstance(plotItem.labels, list) else ''
            return subPlot.plot.plot(x.flatten(), y.flatten(), name=name,
                                     pen=mkPen(color, width=1), symbol=symbol, symbolBrush=color,
                                     symbolPen=None, symbolSize=symbolSize)
        else:
            name = plotItem.labels[-1] if isinstance(plotItem.labels, list) else ''
            return subPlot.plot.plot(x.flatten(), y.flatten(), name=name,
                                     pen=None, symbol=symbol, symbolBrush=color,
                                     symbolPen=None, symbolSize=symbolSize)

    def _colorPlot(self, plotItem: PlotItem) -> None:
        subPlot = self.subPlotFromId(plotItem.subPlot)
        assert isinstance(subPlot, PlotWithColorbar) and len(plotItem.data) == 3
        subPlot.setImage(*plotItem.data)

    def _scatterPlot2d(self, plotItem: PlotItem) -> None:
        subPlot = self.subPlotFromId(plotItem.subPlot)
        assert isinstance(subPlot, PlotWithColorbar) and len(plotItem.data) == 3
        subPlot.setScatter2d(*plotItem.data)


class AutoPlot(PlotWidget):
    """Widget for automatic plotting with pyqtgraph.

    Uses :class:`.FigureMaker` to produce subplots.
    """

    def __init__(self, parent: Optional[PlotWidgetContainer]) -> None:
        """Constructor for the pyqtgraph auto plot widget.

        :param parent: plot widget container
        """
        super().__init__(parent=parent)

        self.fmWidget: Optional[FigureWidget] = None
        self.figConfig: Optional[FigureConfigToolBar] = None
        self.figOptions: FigureOptions = FigureOptions()

        layout = QtWidgets.QVBoxLayout()
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)

        self.setLayout(layout)
        self.setMinimumSize(*getcfg('main', 'pyqtgraph', 'minimum_plot_size',
                                    default=(400, 400)))

    def setData(self, data: Optional[DataDictBase]) -> None:
        """Uses :class:`.FigureMaker` to populate the plot(s).

        This method aims to re-use existing plotwidgets if possible -- i.e.,
        when the type of data (generally speaking) has not changed.
        If changes to the data structure are detected, then all subplot
        widgets are re-created from scratch, however.

        :param data: input data
        :return: ``None``
        """
        super().setData(data)
        if self.data is None:
            return

        fmKwargs = {}  # {'widget': self.fmWidget}
        dc = self.dataChanges
        if not dc['dataTypeChanged'] and not dc['dataStructureChanged']:
            fmKwargs['clearWidget'] = False
        else:
            fmKwargs['clearWidget'] = True
        self._plotData(**fmKwargs)

    def _plotData(self, **kwargs: Any) -> None:
        if self.data is None:
            return

        with FigureMaker(parentWidget=self, widget=self.fmWidget,
                         **kwargs) as fm:

            fm.complexRepresentation = self.figOptions.complexRepresentation
            fm.combineTraces = self.figOptions.combineLinePlots

            for dep in self.data.dependents():
                inds = self.data.axes(dep)
                dvals = self.data.data_vals(dep)
                pdt = determinePlotDataType(self.data.extract([dep]))
                plotId = fm.addData(
                    *[np.asanyarray(self.data.data_vals(n)) for n in inds] + [dvals],
                    labels=[str(self.data.label(n)) for n in inds] + [str(self.data.label(dep))],
                    plotDataType=pdt,
                )

        if self.fmWidget is None:
            self.fmWidget = fm.widget
            self.layout().addWidget(self.fmWidget)
            self.figConfig = FigureConfigToolBar(self.figOptions,
                                                 parent=self)
            self.layout().addWidget(self.figConfig)
            self.figConfig.optionsChanged.connect(self._refreshPlot)
            self.figConfig.figCopied.connect(self.onfigCopied)
            self.figConfig.figSaved.connect(self.onfigSaved)

        if self.data.has_meta('title'):
            self.fmWidget.setTitle(self.data.meta_val('title'))

    @Slot()
    def _refreshPlot(self) -> None:
        self._plotData()

    @Slot()
    def onfigCopied(self) -> None:
        """
        Gets triggered when figCopied signal is emitted from self.figConfig
        Copy the current figuremaker widget to the clipboard.
        """
        assert isinstance(self.fmWidget, FigureWidget)
        screenshot = self.fmWidget.grab(rectangle=QtCore.QRect(QtCore.QPoint(0, 0), QtCore.QSize(-1, -1)))
        clipboard = QtWidgets.QApplication.clipboard()
        clipboard.setImage(screenshot.toImage())

    @Slot()
    def onfigSaved(self) -> None:
        """
        Gets triggered when figSaved signal is emitted from self.figConfig
        Save the current figuremaker widget to the data directory as a png
        with a timestamp at the front.
        """
        assert isinstance(self.fmWidget, FigureWidget)
        assert isinstance(self.data, DataDictBase)
        screenshot = self.fmWidget.grab(rectangle=QtCore.QRect(QtCore.QPoint(0, 0), QtCore.QSize(-1, -1)))
        path = Path(self.data.meta_val('title'))
        # add a timestamp here
        t = time.localtime()
        time_str = time.strftime(TIMESTRFORMAT, t)
        filename = time_str+'_'+str(path.stem)+'.png'
        screenshot.save(str(path.parent)+'/'+filename, format='PNG')

    # TODO: Allow for the option to choose filetypes and the name/directory


@dataclass
class FigureOptions:
    """Dataclass that describes the configuration options for the figure."""

    #: whether to plot all 1D traces into a single panel
    combineLinePlots: bool = False

    #: how to represent complex data
    complexRepresentation: ComplexRepresentation = ComplexRepresentation.realAndImag


class FigureConfigToolBar(QtWidgets.QToolBar):
    """Simple toolbar to configure the figure."""

    # TODO: find better config system that generates GUI automatically and
    #   links updates easier.

    #: Signal() -- emitted when options have been changed in the GUI.
    optionsChanged = Signal()
    #: Signal() -- emitted when the copy figure button has been pressed
    figCopied = Signal()
    #: Signal() -- emitted when the save figure button has been pressed
    figSaved = Signal()

    def __init__(self, options: FigureOptions,
                 parent: Optional[QtWidgets.QWidget] = None) -> None:
        """Constructor.

        :param options: options object. GUI interaction will make changes
            in-place to this object.
        :param parent: parent Widget
        """
        super().__init__(parent)

        self.options = options

        combineLinePlots = self.addAction("Combine 1D")
        combineLinePlots.setCheckable(True)
        combineLinePlots.setChecked(self.options.combineLinePlots)
        combineLinePlots.triggered.connect(
            lambda: self._setOption('combineLinePlots',
                                    combineLinePlots.isChecked())
        )

        complexOptions = QtWidgets.QMenu(parent=self)
        complexGroup = QtWidgets.QActionGroup(complexOptions)
        complexGroup.setExclusive(True)
        for k in ComplexRepresentation:
            a = QtWidgets.QAction(k.label, complexOptions)
            a.setCheckable(True)
            complexGroup.addAction(a)
            complexOptions.addAction(a)
            a.setChecked(k == self.options.complexRepresentation)
        complexGroup.triggered.connect(
            lambda _a: self._setOption('complexRepresentation',
                                       ComplexRepresentation.fromLabel(_a.text()))
        )
        complexButton = QtWidgets.QToolButton()
        complexButton.setToolButtonStyle(QtCore.Qt.ToolButtonTextOnly)
        complexButton.setText('Complex')
        complexButton.setPopupMode(QtWidgets.QToolButton.InstantPopup)
        complexButton.setMenu(complexOptions)
        self.addWidget(complexButton)

        # Adding functionality to copy and save the graph
        self.copyFig = self.addAction('Copy Figure', self._copyFig)
        self.saveFig = self.addAction('Save Figure', self._saveFig)

    def _setOption(self, option: str, value: Any) -> None:
        setattr(self.options, option, value)
        self.optionsChanged.emit()

    def _copyFig(self) -> None:
        self.figCopied.emit()

    def _saveFig(self) -> None:
        self.figSaved.emit()
