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

import sys, os, random

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 scipy.interpolate import *

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


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


class Solder(object):
    def __init__(self):
        self.psteps = []
        self.durations = dict()
        self.rates = dict()
        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 add_state(self, name, temp):
        s = State(name, temp)
        self.psteps.append(s)
        return s

    def add_rate(self, states, rate):
        self.rates[states] = rate

    def add_duration(self, states, duration):
        self.durations[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()
        self.time = 0
        used_steps = set()
        for pstep in 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.iteritems():
                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)

                    #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)
                    #print "2er duration", (self.time, psteps[1].temp)
            else:
                for sts, rate in self.rates.iteritems():
                    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)


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


    @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(tuple(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(tuple(states), int(rate.attrib["value"]))

        return s

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)
        self.listdata = [Solder.unpack(os.path.join("solder_types", p)) for p in os.listdir("solder_types")]

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

    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()

    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()
        elif role != QtCore.Qt.DisplayRole:
            return QtCore.QVariant()

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

    def removeRows(self, start, _count, _parent):
        print type(start), type(_count)
        self.beginRemoveRows(_parent, start, start + _count)
        self.psteps[:start].extend(self.psteps[start + _count:])
        self.endRemoveRows()


class MyDynamicMplCanvas(FigureCanvas):
    """A canvas that updates itself every second with a new plot."""
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        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)

        super(MyDynamicMplCanvas, self).__init__(self.fig)
        self.setParent(parent)

        self.solder = None

        self.compute_initial_figure()
        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 compute_initial_figure(self):

        """test foo bar"""

        #start_rate = 1
        #start_temp = 25
        #start_period = None

        ## warmup
        #warmup_rate = 1
        #warmup_temp = 155
        #preheat_period = None

        ## preheat start
        #preheat_start_rate = 1
        #preheat_start_temp = 155
        #preheat_start_period = 100

        ## preheat end
        #preheat_end_rate = 1
        #preheat_end_temp = 185
        #preheat_end_period = None

        #tal_rate = 1
        #tal_temp = 220
        #tal_duration = 60

        #peak_temp = 250
        #peak_rate = 1

        #x, y = self.solder.calc_profile()

        #print "x", repr(x)
        #print "y", repr(y)

        #dtp_min = 99999.
        #dtp_max = 0.
        #dtpi = -1

        #dtn_min = -99999.
        #dtn_max = 0.
        #dtni = -1
        #for i in xrange(1, len(y)):
          #tmp = (y[i] - y[i-1]) / (x[i] - x[i-1])
          #if tmp > 0:
            #if tmp < dtp_min:
              #dtp_min = tmp
              #dtpi = i
            #elif tmp > dtp_max:
              #dtp_max = tmp
              #dtpi = i
          #elif tmp < 0:
            #if tmp > dtn_min:
              #dtn_min = tmp
              #dtni = i
            #elif tmp < dtn_max:
              #dtn_max = tmp
              #dtni = i
        #print "max negative", dtn_min, dtn_max, dtni
        #print "max positive", dtp_min, dtp_max, dtpi

        self.plot_data, = self.axes.plot([], linewidth=1.0, color=(0,0,1))
        #self.axes.plot(p1x, p1y, 'r-o')

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

        lines = list()
        legend = list()

        cols = ("lightgreen", "yellow", "orange", "red")
        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=cols[ix], label="name")
            lines.append(line)

        self.fig.lines = lines

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

        #self.axes.grid(True, color='gray')

        pylab.setp(self.axes.get_xticklabels(), visible=True)

        self.plot_data.set_xdata(x)
        self.plot_data.set_ydata(y)
        self.axes.legend(("Estimated profile",))
        self.draw()


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.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.solder_view = QtGui.QListView()
        self.pstep_view = QtGui.QTableView()
        self.solder_view.setModel(self.solder_model)
        self.pstep_view.setModel(self.pstep_model)

        self.connect(self.solder_view, QtCore.SIGNAL("clicked(QModelIndex)"), self.solder_selected)
        pl.addWidget(self.solder_view)
        pl.addWidget(self.pstep_view)

        l = QtGui.QVBoxLayout(self.main_widget)
        self.dc = MyDynamicMplCanvas(self.main_widget, width=5, height=4, dpi=self.dpi)
        self.dc.solder = self.solder_model.listdata[0]
        l.addWidget(self.profile_widget, 1)
        l.addWidget(self.dc, 10)

        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.removeRows(0, self.pstep_model.rowCount(QtCore.QModelIndex()))
            self.pstep_model.insertRows()

    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
Copyright \N{COPYRIGHT SIGN} 2012 Stefan Kögl

reflowctl frontend"""
% {"prog": progname, "version": progversion})


qApp = QtGui.QApplication(sys.argv)

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