#!/usr/bin/env python3

import numpy as np
import PyQt5  # make sure pyqtgraph imports Qt5
import pyqtgraph
from collections import deque # CircularBuffer

from artiq.applets.simple import TitleApplet


class RTPlot(pyqtgraph.PlotWidget):
    def __init__(self, args):
        pyqtgraph.PlotWidget.__init__(self)
        self.args = args
        self.xs = None
        self.ys = None
        self.y2s = None

    def data_changed(self, data, mods, title):
        try:
            y = data.get(self.args.y)[1]
        except KeyError:
            return

        y2 = data.get(self.args.y2, (False, None))[1]

        if self.xs is None:
            # Should be only first iteration
            num = int(self.args.num)
            self.xs = np.arange(0., num) # Fixed x-values
            self.ys = deque(np.full(num, 0.), num) # Y-vals are a CircBuff of len num
            if (y2 is not None): # make extra one if y2 is passed
                self.y2s = deque(np.full(num, 0.0), num)

            return

        self.ys.append(y[0])

        self.clear()
        self.plot(self.xs, self.ys, pen=None, symbol="x")
        if y2 is not None:
            self.y2s.append(y2[0])
            self.plot(self.xs, self.y2s, pen=None, symbol="+", symbolPen='r')

        self.setTitle(title)

def main():
    applet = TitleApplet(RTPlot)
    applet.add_dataset("num", "Number of measures")
    applet.add_dataset("y", "Y value")
    applet.add_dataset("y2", "Second Y value", required=False)
    applet.run()

if __name__ == "__main__":
    main()
