# -*- coding: utf-8 -*-
#
# Licensed under the terms of the MIT License
# Copyright (c) 2021 Pierre Raybaut

"""
Simple example illustrating Qt Charts capabilities to plot curves with
a high number of points, using OpenGL accelerated series and comparing it
to PythonQwt and Matplotlib performance, with either PyQt5 or PySide2
"""

import os

# Uncomment one of the following lines to switch from PySide2 and PyQt5:
os.environ["QT_API"] = "pyside2"
# os.environ["QT_API"] = "pyqt5"

try:
    import matplotlib

    matplotlib.use("Qt5Agg")

    from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg
    from matplotlib.figure import Figure
except ImportError:
    FigureCanvasQTAgg = Figure = None

from qtpy.QtCharts import QtCharts as qtc
from qtpy import QtGui as QG
from qtpy import QtWidgets as QW
from qtpy import QtCore as QC
from qtpy import QT_API
from qtpy import PYSIDE2

Qt = QC.Qt

if PYSIDE2:
    import shiboken2
    import ctypes
    import PySide2

    PYTHON_QT_API = "PySide2 v" + PySide2.__version__
else:
    from PyQt5.QtCore import PYQT_VERSION_STR

    PYTHON_QT_API = "PyQt5 v" + PYQT_VERSION_STR

try:
    import qwt
except ImportError:
    qwt = None

import numpy as np
import time


def array2d_to_qpolygonf(xdata, ydata):
    """
    Utility function to convert two 1D-NumPy arrays representing curve data 
    (X-axis, Y-axis data) into a single polyline (QtGui.PolygonF object). 
    This feature is compatible with PyQt4, PyQt5 and PySide2 (requires QtPy).
    
    License/copyright: MIT License © Pierre Raybaut 2020.
    
    :param numpy.ndarray xdata: 1D-NumPy array (numpy.float64)
    :param numpy.ndarray ydata: 1D-NumPy array (numpy.float64)
    :return: Polyline
    :rtype: QtGui.QPolygonF
    """
    dtype = np.float
    if not (
        xdata.size == ydata.size == xdata.shape[0] == ydata.shape[0]
        and xdata.dtype == ydata.dtype == dtype
    ):
        raise ValueError("Arguments must be 1D, float64 NumPy arrays with same size")
    size = xdata.size
    polyline = QG.QPolygonF(size)
    if PYSIDE2:  # PySide2 (obviously...)
        address = shiboken2.getCppPointer(polyline.data())[0]
        buffer = (ctypes.c_double * 2 * size).from_address(address)
    else:  # PyQt4, PyQt5
        buffer = polyline.data()
        buffer.setsize(2 * size * np.finfo(dtype).dtype.itemsize)
    memory = np.frombuffer(buffer, dtype)
    memory[: (size - 1) * 2 + 1 : 2] = xdata
    memory[1 : (size - 1) * 2 + 2 : 2] = ydata
    return polyline


class TestWidget(QW.QWidget):
    ENGINES = QTCHARTS, PYTHONQWT, MATPLOTLIB = "QtCharts", "PythonQwt", "Matplotlib"

    def __init__(self, ncurves=5, npoints=100000):
        super(TestWidget, self).__init__()
        self.setWindowTitle("Simple performance example - " + PYTHON_QT_API)
        self.results = [""] * len(self.ENGINES)
        self.formlayout = QW.QFormLayout()
        self.vlayout = QW.QVBoxLayout()
        self.stack = QW.QStackedWidget()
        self.fill_layouts(ncurves=ncurves, npoints=npoints)
        self.create_views()
        layout = QW.QGridLayout(self)
        layout.addWidget(self.stack, 0, 0, 1, 2)
        layout.addLayout(self.formlayout, 1, 0, 1, 1)
        layout.addLayout(self.vlayout, 1, 1, 1, 1)
        self.setLayout(layout)

    def fill_layouts(self, ncurves, npoints):
        self.view_cb = QW.QComboBox()
        self.view_cb.currentIndexChanged.connect(self.stack.setCurrentIndex)
        self.formlayout.addRow("Engine", self.view_cb)
        self.opengl_cb = QW.QCheckBox("Open-GL (performance!)", self)
        self.opengl_cb.setChecked(True)
        self.view_cb.currentIndexChanged.connect(
            lambda index=None: self.opengl_cb.setEnabled(index == 0)
        )
        self.formlayout.addRow("Rendering", self.opengl_cb)
        self.ncurves_sb = QW.QSpinBox(self)
        self.ncurves_sb.setValue(ncurves)
        self.formlayout.addRow("Number of curves", self.ncurves_sb)
        self.npoints_sb = QW.QSpinBox(self)
        self.npoints_sb.setSingleStep(10000)
        self.npoints_sb.setMaximum(int(100e6))
        self.npoints_sb.setValue(npoints)
        self.formlayout.addRow("Number of points", self.npoints_sb)
        self.results_lb = QW.QLabel("")
        self.results_lb.setAlignment(Qt.AlignRight)
        self.results_lb.setFont(QG.QFont("Consolas"))
        self.view_cb.currentIndexChanged.connect(
            lambda index=None: self.results_lb.setText(self.results[index])
        )
        button = QW.QPushButton("Replot", self)
        button.clicked.connect(self.replot)
        self.vlayout.addWidget(self.results_lb)
        self.vlayout.addWidget(button)

    def create_views(self):
        # Creating QtCharts view
        self.qtchart_view = qtc.QChartView(qtc.QChart())
        self.qtchart_view.chart().legend().setVisible(False)
        self.stack.addWidget(self.qtchart_view)
        self.view_cb.addItem(self.QTCHARTS)
        if qwt is not None:
            # Creating PythonQwt view
            self.qwt_view = qwt.QwtPlot()
            self.qwt_view.setCanvasBackground(Qt.white)
            self.qwt_view.setAutoReplot(False)
            qwt.QwtPlotGrid.make(
                self.qwt_view, color=Qt.lightGray, width=0.0, style=Qt.DotLine
            )
            self.stack.addWidget(self.qwt_view)
            self.view_cb.addItem(self.PYTHONQWT)
        else:
            self.qwt_view = None
        if Figure is not None:
            # Creating Matplotlib view
            self.mpl_fig = Figure(figsize=(5, 4), dpi=100)
            self.mpl_axes = self.mpl_fig.add_subplot(111)
            self.mpl_view = FigureCanvasQTAgg(self.mpl_fig)
            self.stack.addWidget(self.mpl_view)
            self.view_cb.addItem(self.MATPLOTLIB)
        else:
            self.mpl_fig = self.mpl_axes = self.mpl_axes = None

    def replot_all(self):
        for index in range(self.view_cb.count() - 1, -1, -1):
            self.view_cb.setCurrentIndex(index)
            self.replot()

    @property
    def current_engine(self):
        return self.view_cb.currentText()

    def replot(self):
        t0 = time.time()
        self.opengl_cb.setEnabled(self.current_engine == self.QTCHARTS)
        if self.current_engine == self.QTCHARTS:
            # Qt Charts
            plot = self.qtchart_view.chart()
            plot.removeAllSeries()
            set_title = plot.setTitle
        elif self.current_engine == self.PYTHONQWT:
            # PythonQwt
            plot = self.qwt_view
            plot.detachItems(qwt.QwtPlotItem.Rtti_PlotCurve)
            set_title = plot.setTitle
        elif self.current_engine == self.MATPLOTLIB:
            # Matplotlib
            self.mpl_axes.cla()
            set_title = self.mpl_fig.suptitle
        colors = (Qt.red, Qt.blue, Qt.green, Qt.yellow, Qt.cyan, Qt.magenta, Qt.gray)
        delta = []
        ncurves = self.ncurves_sb.value()
        npoints = self.npoints_sb.value()
        t1 = time.time()
        for index in range(ncurves):
            xdata, ydata = generate_arbitrary_data(npoints)
            title = "Curve #%d" % (index + 1)
            color = colors[index % len(colors)]
            polyline = array2d_to_qpolygonf(xdata, ydata)
            t2 = time.time()
            if self.current_engine == self.QTCHARTS:
                # Qt Charts
                curve = qtc.QLineSeries()
                curve.setUseOpenGL(self.opengl_cb.isChecked())
                pen = curve.pen()
                if color is not None:
                    pen.setColor(color)
                pen.setWidthF(1)
                curve.setName(title)
                curve.setPen(pen)
                curve.replace(polyline)
                plot.addSeries(curve)
            elif self.current_engine == self.PYTHONQWT:
                # PythonQwt
                qwt.QwtPlotCurve.make(
                    xdata, ydata, title, plot, linecolor=color, antialiased=True
                )
            elif self.current_engine == self.MATPLOTLIB:
                # Matplotlib
                self.mpl_axes.plot(xdata, ydata, color=QG.QColor(color).name())
            delta.append(time.time() - t2)
        t3 = time.time()
        if self.current_engine == self.QTCHARTS:
            plot.createDefaultAxes()
        elif self.current_engine == self.PYTHONQWT:
            plot.replot()
        elif self.current_engine == self.MATPLOTLIB:
            self.mpl_view.draw()
        t4 = time.time()
        self.results[self.view_cb.currentIndex()] = results = (
            "Total time: <b>{:>03} ms</b><br>"
            "Cleaning-up plot: <b>{:>03} ms</b><br>"
            "Refreshing plot: <b>{:>03} ms</b><br>"
            "Plotting curves: <u>Sum</u>: <b>{:>03} ms</b><br>"
            "(for each curve) <u>Avg</u>: <b>{:>03} ms</b>"
            "".format(
                int((t4 - t0) * 1e3),
                int((t1 - t0) * 1e3),
                int((t4 - t3) * 1e3),
                int(np.array(delta).sum() * 1e3),
                int(np.array(delta).mean() * 1e3),
            )
        )
        self.results_lb.setText(results)
        title = "{0} curves of {1} points".format(ncurves, npoints)
        print(
            "Plotting {0} with {1} ({2}):".format(
                title, self.current_engine, PYTHON_QT_API
            )
        )
        print(
            "    "
            + results.replace("<br>", os.linesep + "    ")
            .replace("<b>", "")
            .replace("</b>", "")
            .replace("<u>", "")
            .replace("</u>", "")
        )
        set_title(title)


def generate_arbitrary_data(npoints):
    xdata = np.linspace(2, 10.0, npoints)
    r1, r2, r3, r4 = np.random.rand(4)
    amp = r1 * 0.25 + 0.75
    phi = r2 * np.pi * 0.5
    pow = r3 * 0.5 + 0.5
    r4 = 0
    noise = np.random.rand(npoints) * amp * 0.1 * r4
    ydata = noise + amp * np.sin(xdata + phi) / xdata ** pow
    return xdata, ydata


if __name__ == "__main__":
    from qtpy.QtWidgets import QApplication

    app = QApplication([])

    widget = TestWidget()
    widget.show()
    widget.resize(500, 500)
    widget.replot_all()

    app.exec_()