import pyqtgraph as pg
from PyQt5.QtWidgets import *


class QuickPlotGenerator(pg.LayoutWidget):

    def __init__(self, ap, **kwargs):
        super().__init__(**kwargs)

        self.ap = ap

        self.title = QLabel('<h2> Quick Plot Generator <\h2>')

        self.addWidget(self.title, colspan=2)

        self.nextRow()

        self.newlayout = pg.LayoutWidget()

        self.newlayout.addWidget(QLabel('Make new plot'))
        self.newlayout.nextRow()
        self.newlayout.addWidget(QLabel('Title: '))

        self.title_le = QLineEdit('murks')
        self.newlayout.addWidget(self.title_le)

        self.newlayout.nextRow()
        self.bt1d = QPushButton('make 1D plot', self)
        self.newlayout.addWidget(self.bt1d)

        self.bt1d.clicked.connect(self.mk1dplot)

        self.btwaterfall = QPushButton('make waterfall plot', self)
        self.newlayout.addWidget(self.btwaterfall)

        self.btwaterfall.clicked.connect(self.mkwaterfallplot)

        self.bt2d = QPushButton('make 2d plot', self)
        self.newlayout.addWidget(self.bt2d)

        self.bt2d.clicked.connect(self.mk2dplot)

        self.addWidget(self.newlayout)

    def mk1dplot(self):
        title = self.title_le.text()
        self.ap.add_plot_dock(title, Quick1DPlot(self.ap), MultiDataExtractor(), closable=True)

    def mkwaterfallplot(self):
        title = self.title_le.text()
        self.ap.add_plot_dock(title, QuickWaterfallPlot(self.ap), MultiDataExtractor(), closable=True)

    def mk2dplot(self):
        title = self.title_le.text()
        self.ap.add_plot_dock(title, Quick2DPlot(self.ap), MultiDataExtractor(), closable=True)