import sys

from PyQt5.QtWidgets import QApplication, QMainWindow, QMenu, QVBoxLayout, QSizePolicy, QMessageBox, QWidget, \
    QPushButton
from PyQt5.QtGui import QIcon

from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import random

import rigol_ds1104b

osc = rigol_ds1104b.RigolDS1104B("/dev/usbtmc0")


class App(QMainWindow):
    def __init__(self):
        super(App, self).__init__()
        self.left = 10
        self.top = 10
        self.title = 'PyQt5 matplotlib example - pythonspot.com'
        self.width = 640
        self.height = 400
        self.initUI()

    def initUI(self):
        self.setWindowTitle(self.title)
        self.setGeometry(self.left, self.top, self.width, self.height)

        m = PlotCanvas(self, width=8, height=6)
        m.move(0, 0)

        button = QPushButton('PyQt5 button', self)
        button.setToolTip('This s an example button')
        button.move(500, 0)
        button.resize(140, 100)

        self.show()


class PlotCanvas(FigureCanvas):
    def __init__(self, parent=None, width=6, height=5, dpi=100):
        fig = Figure(figsize=(width, height), dpi=dpi)
        self.axes = fig.add_subplot(111)

        FigureCanvas.__init__(self, fig)
        self.setParent(parent)

        FigureCanvas.setSizePolicy(self,
                                   QSizePolicy.Expanding,
                                   QSizePolicy.Expanding)
        FigureCanvas.updateGeometry(self)
        self.plot()

    def plot(self):
        ax = self.figure.add_subplot(111)

        def animate(*args):
            try:
                x, y = osc.get_data()
                ax.clear()
                ax.plot(x, y, 'r-')
                ax.set_xlabel("time [s]")
                ax.set_ylabel("voltage [V]")
                ax.set_xlabel("time [s]", color="red")
                ax.set_ylabel("voltage [V]", color="red")
                ax.set_axis_bgcolor('black')
                plt.grid(True, color="red")
            except Exception as err:
                print(err)

        ani = animation.FuncAnimation(self.figure, animate, interval=10)
        plt.show()
        self.draw()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = App()
    sys.exit(app.exec_())