import numpy as np
import pyqtgraph as pg
from pyqtgraph.Qt import QtCore, QtGui

from time import perf_counter

def generate_pattern(nsegs, size):
    rng = np.random.default_rng()
    x = rng.random(nsegs) * size
    y = rng.random(nsegs) * size
    arr = np.empty((nsegs, 4), dtype=np.float64)
    arr[:, 0] = x
    arr[:, 1] = y
    arr[:, 2] = x + 2
    arr[:, 3] = y + 2
    return arr

nsegs = 10_000
size = 500
nframes = 100
pattern = generate_pattern(nsegs, size)

def run(use_array):
    # generate lines once
    segments = pg.Qt.internals.PrimitiveArray(QtCore.QLineF, 4, use_array=use_array)
    segments.resize(nsegs)
    memory = segments.ndarray()
    memory[:] = pattern

    # draw multiple frames using the same lines array
    t0 = perf_counter()
    for _ in range(nframes):
        qimg = QtGui.QImage(size, size, QtGui.QImage.Format.Format_RGB32)
        qimg.fill(QtCore.Qt.GlobalColor.transparent)
        painter = QtGui.QPainter(qimg)
        painter.setPen(QtCore.Qt.GlobalColor.cyan)
        painter.drawLines(*segments.drawargs())
        painter.end()
    t1 = perf_counter()
    return t1 - t0

for use_array in [False, None, True]:
    duration = run(use_array)
    fps = int(nframes / duration)
    print(f'{use_array=} {duration=:.3f} {fps=}')