"""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 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)


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)

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

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


@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()

    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)

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