# -*- coding: utf-8 -*-
"""
qtnumpy
-------

Module implementing some utilities between Qt and NumPy

License/copyright: MIT License © Pierre Raybaut 2020.
"""

import os

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

from qtpy import QtWidgets as QW
from qtpy import QtGui as QG
from qtpy import QtCore as QC
from qtpy import PYSIDE2
import numpy as np

if PYSIDE2:
    import shiboken2
    import ctypes


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):
    def __init__(self, polyline=None):
        super(TestWidget, self).__init__()
        self.polyline = polyline

    def paintEvent(self, event):
        painter = QG.QPainter(self)
        painter.setPen(QC.Qt.red)
        if self.polyline is not None:
            if PYSIDE2:
                from PySide2 import __version__

                qtlib = "PySide2 v%s" % __version__
            else:
                from PyQt5.QtCore import PYQT_VERSION_STR

                qtlib = "PyQt5 v%s" % PYQT_VERSION_STR
            import time

            t0 = time.time()
            painter.drawPoints(self.polyline)
            print("dt[%s] = %d ms" % (qtlib, ((time.time() - t0) * 1e3)))


if __name__ == "__main__":
    app = QW.QApplication.instance()
    if app is None:
        app = QW.QApplication([])
    xdata = np.linspace(0, 800, 1000000)
    ydata = np.cos(xdata / 100.0) ** 2 * 300
    polyline = array2d_to_qpolygonf(xdata, ydata)
    # for index, point in enumerate(polyline):
    #     assert point.x() == xdata[index] and point.y() == ydata[index]
    # print("%s <==> NumPy(%s, %s)" % (point, xdata[index], ydata[index]))
    widget = TestWidget(polyline)
    widget.show()
    app.exec_()