# -*- 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"] = 'pyqt5'

from PyQt5 import QtWidgets as QW
from PyQt5 import QtGui as QG
from PyQt5 import QtCore as QC
import numpy as np
import time
from PyQt5.QtCore import PYQT_VERSION_STR

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)
    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:
            print(self.polyline.size(), 'lines')
            qtlib = "PyQt5 v%s" % PYQT_VERSION_STR
            t0 = time.time()
            painter.drawPolyline(self.polyline)
            print("dt[%s] = %d ms" % (qtlib, ((time.time()-t0)*1e3)))


if __name__ == '__main__':
    app = QW.QApplication.instance()
    print(QC.qVersion())
    if app is None:
        app = QW.QApplication([])
    xdata = np.linspace(0, 800, 1000000)
    ydata = np.cos(xdata/100.)**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_()