import math

from PyQt4 import QtCore, QtGui

from matplotlib.figure import Figure
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt4agg import NavigationToolbar2QT as NavigationToolbar


class MyFigure(Figure, QtCore.QThread):
    def __init__(self, parent, *args, **kwargs):
        QtCore.QThread.__init__(self, parent)
        Figure.__init__(self, *args, **kwargs)

        self.plot_data = list()

    def start_plotting_thread(self, plot_data, on_finish=None):
        """ Start plotting """
        self.plot_data = plot_data

        if on_finish is not None:
            self.finished.connect(on_finish)

        self.start()

    def run(self):
        """ Run as a thread """
        # Figure out rows and columns
        total_plots = len(self.plot_data)

        columns = int(math.sqrt(total_plots))
        if columns < 1:
            columns = 1

        rows = int(total_plots / columns)
        if (total_plots % columns) > 0:
            rows += 1
        if rows < 1:
            rows = 1

        # Plot Data
        for plot_index, _plot_data in enumerate(self.plot_data):
            plot_number = plot_index + 1
            args = (rows, columns, plot_number)
            kwargs = {
                'title': _plot_data['title'],
                'xlabel': _plot_data['xlabel'],
                'ylabel': _plot_data['ylabel']
            }

            figure = self.add_subplot(*args, **kwargs)

            figure.plot(_plot_data['data'])


class PlotDialog(QtGui.QDialog):
    def __init__(self, parent):
        super(PlotDialog, self).__init__(parent, QtCore.Qt.WindowMinMaxButtonsHint | QtCore.Qt.WindowCloseButtonHint)

        self.figure = MyFigure(self)
        self.canvas = FigureCanvas(self.figure)
        self.toolbar = NavigationToolbar(self.canvas, self)

        self.layout = QtGui.QGridLayout()
        self.setLayout(self.layout)

        layout = [
            [self.canvas],
            [self.toolbar],
        ]

        for row_index, columns in enumerate(layout):
            if type(columns) is list:
                for column_index, widget in enumerate(columns):
                    if widget is not None:
                        self.layout.addWidget(widget, row_index, column_index)

    def draw_plots(self, plot_data):
        """ Plot Plots """
        self.figure.start_plotting_thread(plot_data, on_finish=self.finish_drawing_plots)

    def finish_drawing_plots(self):
        """ Finish drawing plots """
        self.canvas.draw()
        self.show()