reflow/reflowctl/reflowctl_gui.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

import sys, os, random, copy

from operator import attrgetter

import xml.etree.ElementTree as etree

import pylab
from PyQt4 import QtGui, QtCore

from numpy import arange, sin, pi, array, linspace, arange
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from matplotlib.lines import Line2D
from matplotlib.path import Path
import matplotlib.patches as patches
#from mpltools import annotation


progname = os.path.basename(sys.argv[0])
progversion = "0.1"

PSTEP_COLORS = ("lightgreen", "yellow", "orange", "red")


class State(object):
    def __init__(self, name, temp):
        self.name = name
        self.temp = temp


class Solder(object):
    def __init__(self):
        self.psteps = []
        self.durations = list()
        self.rates = list()
        self.name = None

        #start = self.add_state("start", 25)
        #ps = self.add_state("preheat start", 150)
        #pe = self.add_state("preheat end", 185)
        #tal = self.add_state("tal", 220)
        #peak = self.add_state("peak", 250)
        #end = self.add_state("end", 25)

        #self.add_duration((ps, pe), 100)
        #self.add_duration((tal, peak, tal), 100)

        #self.add_rate((start, ps), 1)
        #self.add_rate((ps, pe), 1)
        #self.add_rate((pe, tal), 1)
        #self.add_rate((tal, end), -2)

    def __unicode__(self):
        return unicode(self.name)

    def __str__(self):
        return self.name

    def color(self, index):
        return PSTEP_COLORS[index]

    def add_state(self, name, temp):
        s = State(name, temp)
        self.psteps.append(s)
        return s

    def add_rate(self, states, rate):
        self.rates.append((states, rate))

    def add_duration(self, states, duration):
        self.durations.append((states, duration))

    def get_state(self, name):
        for i in self.psteps:
            if i.name == name:
                return i
        return None

    def calc_profile(self):

        x = list()
        y = list()
        duration_points = dict()
        rate_points = dict()
        self.time = 0
        used_steps = set()
        for ix, pstep in enumerate(self.psteps):
            #print "-- ", repr(pstep.name), pstep.temp, pstep.used, self.time, x, y

            if pstep != self.psteps[0] and pstep not in used_steps:
                ix = self.psteps.index(pstep)
                raise Exception("step %r not connected to step %r or step %r" % (pstep.name, self.psteps[ix-1].name, self.psteps[ix+1].name))

            psteps = None
            duration = None
            for sts, dur in self.durations:
                if sts[0] == pstep:
                  duration = dur
                  psteps = sts
                  break

            if pstep not in used_steps:
                used_steps.add(pstep)
                x.append(self.time)
                y.append(pstep.temp)

            if duration is not None:
                if len(psteps) == 3:
                    used_steps.add(psteps[1])
                    used_steps.add(psteps[2])

                    self.time += duration / 2
                    x.append(self.time)
                    y.append(psteps[1].temp)

                    #print "3er duration", (self.time, psteps[1].temp)

                    self.time += duration / 2
                    x.append(self.time)
                    y.append(psteps[2].temp)

                    duration_points[ix] = (x[-3:], y[-3:])

                    #print "3er duration", (self.time, psteps[2].temp)
                else:
                    y.append(psteps[1].temp)
                    used_steps.add(psteps[1])
                    self.time += duration
                    x.append(self.time)
                    duration_points[ix] = (x[-2:], y[-2:])
                    #print "2er duration", (self.time, psteps[1].temp)
            else:
                for ex, (sts, rate) in enumerate(self.rates):
                    if sts[0] == pstep:
                        used_steps.add(sts[1])
                        duration = (sts[1].temp - pstep.temp) / rate
                        self.time += duration
                        x.append(self.time)
                        y.append(sts[1].temp)
                        #print "rate", (self.time, sts[1].temp)
                        rate_points[ex] = (x[-2:], y[-2:])


        return array(map(float, x)), array(map(float, y)), max(x), max(y), duration_points, rate_points


    @staticmethod
    def unpack(filename):
        xmltree = etree.parse(filename)
        root = xmltree.getroot()
        s = Solder()
        s.name = root[0].attrib["name"]
        for state in root[0].findall("state"):
            s.add_state(state.attrib["name"], int(state.attrib["temperature"]))
        for duration in root[0].findall("duration"):
            states = list()
            for state in duration:
                states.append(s.get_state(state.attrib["name"]))
            s.add_duration(states, int(duration.attrib["value"]))
        for rate in root[0].findall("rate"):
            #print rate
            states = list()
            for state in rate:
                states.append(s.get_state(state.attrib["name"]))
            s.add_rate(states, int(rate.attrib["value"]))

        return s

    def serialize(self, pstep_list):
        return ", ".join(map(attrgetter("name"), pstep_list))


class SolderListModel(QtCore.QAbstractListModel):
    def __init__(self, parent=None, *args):
        """ datain: a list where each item is a row
        """
        super(SolderListModel, self).__init__(parent, *args)

        dirname = os.path.join(os.path.dirname(__file__), "solder_types")
        dirlisting = os.listdir(dirname)
        self.listdata = []
        for p in dirlisting:
            #try:
                self.listdata.append(Solder.unpack(os.path.join(dirname, p)))
            #except Exception:
                #pass

    def rowCount(self, parent=QtCore.QModelIndex()):
        return len(self.listdata)


    def headerData(self, col, orientation, role):
        if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
            return QtCore.QVariant("Solder Paste")
        return QtCore.QVariant()

    def data(self, index, role):
        if index.isValid() and role == QtCore.Qt.DisplayRole:
            return QtCore.QVariant(self.listdata[index.row()].name)
        else:
            return QtCore.QVariant()


class PStepModel(QtCore.QAbstractTableModel):
    def __init__(self, parent=None, *args):
        super(PStepModel, self).__init__(parent, *args)
        self.psteps = list()
        self.headerdata = [u"Name", u"Temperature (°C)"]

    def headerData(self, col, orientation, role):
        if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
            return QtCore.QVariant(self.headerdata[col])
        return QtCore.QVariant()

    def rowCount(self, parent):
        return len(self.psteps)

    def columnCount(self, parent):
        return 2

    def data(self, index, role):
        if not index.isValid():
            return QtCore.QVariant()

        if role == QtCore.Qt.DisplayRole:
          col = index.column()
          if col == 0:
              return QtCore.QVariant(self.psteps[index.row()].name)
          else:
              return QtCore.QVariant(self.psteps[index.row()].temp)

        if index.column() == 0 and role == QtCore.Qt.DecorationRole:
            p = QtGui.QPixmap(10,10)
            cr = row = index.row()
            color = index.row() in (0, len(self.psteps)-1) and QtGui.QColor("black") or QtGui.QColor(PSTEP_COLORS[cr-1])
            p.fill(color)
            return p

        return QtCore.QVariant()

    def remove_pstep(self, index):
        del self.psteps[index]
        self.reset()

    def add_pstep(self, index, pstep):
        self.psteps.insert(index, pstep)
        self.reset()

    def setSteps(self, steps):
        assert isinstance(steps, list)
        self.psteps = steps
        self.reset()



class MyDynamicMplCanvas(FigureCanvas):
    """A canvas that updates itself every second with a new plot."""
    def __init__(self, parent, myapp, width, height, dpi):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        super(MyDynamicMplCanvas, self).__init__(self.fig)
        self.axes = self.fig.add_subplot(111)

        ## We want the axes cleared every time plot() is called
        self.axes.set_axis_bgcolor('black')
        self.axes.set_title(u'reflow profile', size=12)
        self.axes.set_xlabel(u'time (seconds)', size=12)
        self.axes.set_ylabel(u'temperature (°C)', size=12)

        #pylab.setp(self.axes.get_xticklabels(), fontsize=8)
        #pylab.setp(self.axes.get_yticklabels(), fontsize=8)

        self.setParent(parent)
        self.myapp = myapp

        self.solder = None

        self.plot_data, = self.axes.plot([], linewidth=1.0, color=(0,0,1), zorder=10)
        self.selection_data, = self.axes.plot([], linewidth=1.0, color=(1,1,1), zorder=5)

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

        timer = QtCore.QTimer(self)

        self.counter = list()
        QtCore.QObject.connect(timer, QtCore.SIGNAL("timeout()"), self.update_figure)
        timer.start(1000)


    def update_figure(self):
        #Build a list of 4 random integers between 0 and 10 (both inclusive)
        x, y, xmax, ymax, duration_points, rate_points = self.solder.calc_profile()

        lines = list()
        legend = list()

        #for states, value in self.durations.iteritems():
            #annotation.slope_marker((states[0])

        self.fig.lines = lines

        self.axes.set_xbound(lower=0, upper=xmax + 20)
        self.axes.set_ybound(lower=0, upper=ymax + 20)
        self.axes.set_ymargin(0)
        self.axes.set_xmargin(0)

        self.axes.set_yticks([state.temp for state in self.solder.psteps])
        self.axes.set_xticks(x)

        self.plot_data.set_xdata(x)
        self.plot_data.set_ydata(y)
        self.plot_data.set_zorder(20)

        duration_widget = self.myapp.duration_widget

        #self.selection_data.set_xdata(array(da))
        #self.selection_data.set_ydata(array(db))

        for ix, i in enumerate(self.solder.psteps[1:-1]):
            line = Line2D([0, xmax + 20], [i.temp, i.temp],
                transform=self.axes.transData, figure=self.fig, color=self.solder.color(ix), label="name", zorder=1)
            lines.append(line)

        self.axes.legend(("Estimated profile",))
        self.draw()


# contraint_list |  label - checkboxes
#                   label - value


class ConstraintWidget(QtGui.QWidget):
    def __init__(self, name):
        super(ConstraintWidget, self).__init__()
        self.name = name
        self.solder = None

        self.value = QtGui.QSpinBox(self)
        self.value.setRange(-300, 400)
        self.constraint_model = QtGui.QStringListModel(self) # constraint selection
        self.all_psteps = PStepModel(self) # pstep selection pool
        self.selected_psteps = PStepModel(self) # selected psteps

        #self.all_psteps.setSteps(self.solder.psteps)

        self.add_button = QtGui.QPushButton("Add", self)
        self.remove_button = QtGui.QPushButton("Remove", self)

        bg = QtGui.QWidget(self)
        gbl = QtGui.QVBoxLayout(bg)

        gbl.addWidget(self.add_button)
        gbl.addWidget(self.remove_button)
        gbl.addStretch(5)
        gbl.addWidget(self.value)


        self.constraint_list_view = QtGui.QListView(self)
        self.constraint_list_view.setModel(self.constraint_model)

        self.all_psteps_view = QtGui.QListView(self)
        self.all_psteps_view.setModel(self.all_psteps)

        self.selected_psteps_view = QtGui.QListView(self)
        self.selected_psteps_view.setModel(self.selected_psteps)

        #self.left = QtGui.QWidget(self)
        #gl = QtGui.QGridLayout(self.left)
        #gl.addWidget(QtGui.QLabel(u"Steps"), 0,0)
        #gl.addWidget(QtGui.QLabel(name), 1,0)
        #gl.addWidget(self.checkboxes_group, 0, 1)
        #gl.addWidget(self.value, 1, 1)
        #self.gl = gl


        h = QtGui.QHBoxLayout()
        h.addWidget(self.constraint_list_view)
        h.addWidget(self.all_psteps_view)
        h.addWidget(bg)
        h.addWidget(self.selected_psteps_view)
        self.setLayout(h)

        self.connect(
            self.constraint_list_view,
            QtCore.SIGNAL("clicked(QModelIndex)"),
            self.constraint_clicked)

        self.connect(self.add_button, QtCore.SIGNAL("clicked()"), self.add_constraint)
        self.connect(self.remove_button, QtCore.SIGNAL("clicked()"), self.remove_constraint)

    def add_constraint(self):
        raise NotImplementedError()

    def remove_constraint(self):
        raise NotImplementedError()

    def setData(self, solder):
        self.solder = solder
        #for k,v in self.checkboxes.iteritems():
            #self.cl.removeWidget(v)
        #self.checkboxes = dict()
        self.all_psteps.setSteps(self.solder.psteps)

        #for i in solder.psteps:
            #cb = QtGui.QCheckBox(self, checked=False)
            #label = QtGui.QLabel(i.name, self)
            #label.setBuddy(cb)
            #self.checkboxes[i] = cb
            ##self.cl.addWidget(label)
            ##self.cl.addWidget(cb)

        self.getConstraints()


    def getConstraints(self):
        raise NotImplementedError()

    def constraint_clicked(self, index):
        #for cb in self.checkboxes.itervalues():
            #cb.setChecked(False)

        self.handle_clicked(index)

    def handle_clicked(self, index):
        raise NotImplementedError()

class DurationConstraintWidget(ConstraintWidget):

    def getConstraints(self):
        tmp = QtCore.QStringList()
        for ix in xrange(len(self.solder.durations)):
            tmp << unicode(ix + 1)

        self.constraint_model.setStringList(tmp)
        k, t = self.solder.durations[0]

        self.value.setValue(t)
        self.constraint_list_view.setCurrentIndex(self.constraint_model.index(0, 0))

    def add_constraint(self):
        self.selected_psteps.psteps.append(self.all_psteps.psteps[self.all_psteps_view.currentIndex().row()])
        self.selected_psteps.reset()
        #self.selected_psteps_view.setCurrentIndex(QtCore.QModelIndex())
        self.selected_psteps_view.clearSelection()

    def remove_constraint(self):
        del self.selected_psteps.psteps[self.selected_psteps_view.currentIndex().row()]
        self.selected_psteps.reset()
        #self.selected_psteps_view.setCurrentIndex(QtCore.QModelIndex())
        self.selected_psteps_view.clearSelection()

    def handle_clicked(self, index):
        psteps, value = self.solder.durations[index.row()]
        self.selected_psteps.setSteps(psteps)
        self.value.setValue(value)


class RateConstraintWidget(ConstraintWidget):

    def getConstraints(self):
        tmp = QtCore.QStringList()
        for ix in xrange(len(self.solder.durations)):
            tmp << unicode(ix + 1)

        self.constraint_model.setStringList(tmp)
        k, t = self.solder.rates[0]

        self.value.setValue(t)
        self.constraint_list_view.setCurrentIndex(self.constraint_model.index(0, 0))

    def handle_clicked(self, index):
        psteps, value = self.solder.durations[index.row()]
        self.selected_psteps.setSteps(psteps)
        self.value.setValue(value)

class ApplicationWindow(QtGui.QMainWindow):
    def __init__(self):
        QtGui.QMainWindow.__init__(self)
        self.setAttribute(QtCore.Qt.WA_DeleteOnClose)
        self.setWindowTitle("application main window")

        self.file_menu = QtGui.QMenu('&File', self)
        self.file_menu.addAction('&Quit', self.fileQuit,
                                 QtCore.Qt.CTRL + QtCore.Qt.Key_Q)
        self.file_menu.addAction('&Save plot', self.save_plot,
                                 QtCore.Qt.CTRL + QtCore.Qt.Key_S)
        self.menuBar().addMenu(self.file_menu)

        self.help_menu = QtGui.QMenu('&Help', self)
        self.menuBar().addSeparator()
        self.menuBar().addMenu(self.help_menu)

        self.help_menu.addAction('&About', self.about)

        self.main_widget = QtGui.QWidget(self)
        self.profile_widget = QtGui.QWidget(self)
        self.steps_box = QtGui.QGroupBox(self)

        self.tab_widget = QtGui.QTabWidget(self)
        self.duration_widget = DurationConstraintWidget(u"Duration (s)")
        self.rate_widget = RateConstraintWidget(u"Rate (°C/s)")
        self.tab_widget.addTab(self.duration_widget, u"Duration (s)")
        self.tab_widget.addTab(self.rate_widget, u"Rate (°C/s)")

        self.dpi = 100

        pl = QtGui.QHBoxLayout(self.profile_widget)
        sl = QtGui.QVBoxLayout(self.steps_box)

        self.solder_model = SolderListModel(self)
        self.pstep_model = PStepModel(self)

        self.pstep_view = QtGui.QTableView()
        self.pstep_view.setModel(self.pstep_model)
        self.pstep_view.verticalHeader().setVisible(False)
        self.pstep_view.resizeColumnsToContents()

        self.solder_view = QtGui.QListView()
        self.solder_view.setModel(self.solder_model)
        self.connect(
            self.solder_view,
            QtCore.SIGNAL("clicked(QModelIndex)"),
            self.solder_selected)

        pl.addWidget(self.solder_view, 1)
        pl.addWidget(self.pstep_view, 2)
        pl.addWidget(self.tab_widget, 7)
        #pl.addWidget(self.duration_widget)
        #pl.addWidget(self.rate_widget)

        l = QtGui.QVBoxLayout(self.main_widget)
        self.dc = MyDynamicMplCanvas(self, self, width=5, height=4, dpi=self.dpi)

        self.solder_view.setCurrentIndex(self.solder_model.index(0,0))
        self.solder_selected(self.solder_model.index(0,0))

        l.addWidget(self.profile_widget, 2)
        l.addWidget(self.dc, 8)

        self.main_widget.setFocus()
        self.setCentralWidget(self.main_widget)

        self.statusBar().showMessage("I'm in reflow heaven", 2000)

    def solder_selected(self, index):
        if index.isValid():
            self.dc.solder = self.solder_model.listdata[index.row()]
            self.pstep_model.setSteps(self.dc.solder.psteps)
            self.pstep_view.resizeColumnsToContents()
            self.duration_widget.setData(self.dc.solder)
            self.rate_widget.setData(self.dc.solder)

    def save_plot(self):
        file_choices = "PNG (*.png)|*.png"

        filename = QtGui.QFileDialog.getSaveFileName(self, 'Save File', 'qtplot.png')
        print type(filename), dir(filename)
        self.dc.print_figure(str(filename), dpi=self.dpi)

    def fileQuit(self):
        self.close()

    def closeEvent(self, ce):
        self.fileQuit()

    def about(self):
        QtGui.QMessageBox.about(self, "About %s" % progname,
        u"%(prog)s version %(version)s\n" \
        u"Copyright \N{COPYRIGHT SIGN} 2012 Stefan Kögl\n\n" \
        u"reflowctl frontend" % {"prog": progname, "version": progversion})


def main():

    qApp = QtGui.QApplication(sys.argv)

    aw = ApplicationWindow()
    aw.setWindowTitle("%s" % progname)
    aw.show()
    sys.exit(qApp.exec_())


if __name__ == '__main__':
    main()