#!/usr/bin/python

from PyQt4.QtGui import *
from PyQt4.QtCore import *
import sys, json, base64

from diagramitems import Connection, ResizeSelectionHandle, Block, DiagramScene

"""
 Author: Windel Bouwman
 Year: 2012
 Description: This script implements a diagram editor.
 run with python 3.x as:
  $ python [thisfile.py]
"""
saveicon = b'iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAGGklEQVR42s2XeWxURRjAv7d3L9ay\nlJBCbAx/SIIpCWilhNLSFsrR0qUgyCmnnKVFa/xDjNbUKGIUSrlaLhEREWlZaQoC24VymBivoGJB\n/MuEtEpjoNt2z/H7Zt+8fdPdGk2MOu23MzvvzXy/+Y6ZWQX+46L8bwC2vLTlZte9e6MMRiMYjQYY\nPWYsZGXngM1iBkUxoERej1RYayJmYGrF1Dbj/31+P1xtuwQ3vvmaP6M/h8Nx69Xq6kclgDWrn2W3\nbt+G5KQkSE9Ph/E5eTB1xixIQID+yIraEMoVJYLGay4KGNS2ty8Aza4muHqpVZvClmCDHbW1igSw\neuUq1tPTA2EWRohkmJCbD9OLS+Fu14M4ZlP6M8lwSvSdEUNTwNXYCFc87iiAzQY7d9XJACuXr2Dd\n3d1gNpvBarXCxMmFUDzLCUk2S3zfKTrlqmIlYgqtTfWDXh80NZ6CNvdFnQUSYNfuXTLAimXLmV5B\nTn4hlDrLICUxHoAC+tHSTNz1jIcCw48HvX60AAK0RgGsaIE9e/fIAMuWPiMB5BZMAWfZHBiUaI1r\nAaZ+Mi32VKVqX5h/ZzwGzpxuRAtckFywt36fDLB08RIZoHAqlDhnQ0eXdwDlosFAP1APRGVoahK0\nYBBedp+XLNCwv0EGWLxwUT+AIphZ6ozNghhlEQCtrZqeqRC+QAjOftqEFogCkAX2HzwgAyxasFAC\nmFQwFaaVEIAJ5P2K6T5ln2uu0LnEHwzBZ2ea4Epr1AUU5AcPH5IBFsx/Wg5CBCgqdoLVbBxg5Xov\nDGyFQCgMF5tPI8B5CeDwkfdkgPlPzYsBmDLTCffu98b3f78OzS4s+g7Vg5Ot4G5xwdV+AEeOvi8D\nzJszV5p7IgLkT3equ9zAAPo4EMpFAIpnnrMuuKYDsCDA0WMfyABzy+bIAPlTIG9aKfzZecV0y5dc\nIdosMvzyORdc90RjwGKxwLHjH8oAZc7ZEoBj6DAYOeoxPJiMYKIDymTkbYPBoB5CkRpABByDcDjM\nJUQSDEIwFIIwyk83b0DXrx0SwPETH8kAs2eVxrhXKsrfPLnZwNOZEeDEyY9lgNLiEvYXpv1HCgGc\nPPWJDFAys/jfA8AD71RTowxQPH3GgAAK+t2IQv7X4oC+q5cSKiIORCyEyP9qLfr1AI2u0zLAjKJp\ncQHWblgPGRkZ0Q7G4uwFch8d6xXlm3jw0qEUCARgOF5yRDGZTOBqPiMDFBUUxsxLh8aa9evAbrfz\nVYuVRvVG2uKZ6COrvFj1Ao92fiL2eME+yK4pM6EFms+2yAAFeZNjAawWWF9eDqmpqVBT/ZqWZn1+\nH7y5dat2LxRmp1qY/pUtL/NgY9ju6e3lVz29Bc5dOC8D5OVMigNghfLKCg7wXOVm3kd53e31Qv3+\nBq6clAUx5/v7/I2a17mvOXBfHyTgLUgDQNdc8LTKADnZE2IAiHRz1fPgGDIEKjeW8z7aZPwBP9TW\n1WlK9QAC4p1tb2sAPp8PbLgYunFb8HgPBILgabssA0zIejIGgHy4uaoKARywacPGCABO7sOr9rs7\ntsdEOwnBkNRu34EAJnQBAz++TwB8V0WoIAZl2/VrMkDWuMdjAJISE2EjuoCCsBKjOuJrVISTbt32\nlqZMrJyiXQDt2VnHldGYYDCA7rSBGS0awvEmowmufH5dBhibOeYu+nSYHiABAdZhGg4igIoKbkqa\nnLKjpqZGWrneErTihn31XGEEIAgJOIZixqim5hdffRkFSElOUfDX0LgR6cNbMECGaBbAyF21dg2v\n+UBdkFGu6wH0z0jhgfp6vlIeuPiMrmGiUNr+0P5jcp/P5xUWoDoN/bT8kYczqjH4rCILFi5ZAvaH\n7JKPRZtMLtJPCE1Oz44cOsSh1X0KrBaLHiD47fffjcTmL/ojLhFlNCovtJjNDmRi2dnjk57IyhqM\nsWBBf1pxQjMJzmgke6rKmULJwVgQm36E8nd2dNxv9Xh+//nOHR9Ivxj4WjGRQu24+mb88psegNqE\nSQlrJH9lZmYacnNzjWlpaQpmBBdKLXXHU9QNiIkdEa0T9nq94c7OzpDb7Q63t7crTD6WFRXEj0J3\nveAfetNmUUsM6bsAAAAASUVORK5CYII=\n'

loadicon = b'iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAACsElEQVR42u2Xf0hTURTHv++ZczlD\n6SeUREVRiRT9gFWM1rbAwr+iYAathPqjMosIQQqJQAhLLDI0CMwflRa1vyQst0oNLKxhCGWR9UfB\nwFCRaqC2vc6d7+1tr/7YnuwW5IEv95y79875vHvP7t4E/GUTpgGmAaL8vaRzpDlx3iuRXpEOkz5N\nFWDhrHR8bL+LtGXLKUqRJcoKyeUk1Q8Fges3gLIKeGlm+1QBHA4rPJ7WxG4eHqLlWoJRcvfJaPHY\nV1LPbwB5Dnja7if+BFUVQEdX7OqETdL48vD+C/DODze5e9iUCmAjgDt6FzJ+m/gG5GyD9OEz1lPY\nqwLQFrQ1JR+Amc0JPO0O941XBbAQQB0nABcBvNACbCGAWk4ABwngpRbATACXEksU7rsJGlnzUSZB\nkEdAbW9lXgnJtx0jAJ8WYCMBnE8MIESFx7+rRQRRLiaPESANmL2EAF6jgaLSCIA9Bx53sZokJpEY\nm4zFzP85BgTn0mUZUL9y7CPpz8CCvAJdA0DrE6D6NnxsLp9UaV2KVfW7Jy8QWcGUyUJhP0piFJA0\nmw6iAn19EAgAmdkICpRw5EEdsjabgXQDYo9gRZENVxVeIJO8Gjqs7y2wxoJ+IXUGxkcHkDrTqC+R\nXrt5D3AVoYU9yJs+L1bnruQLUFJO+34NpxlAc2MVCly7+ALkHQAedSKfAZwq3o/KK2X8ikvUT4ss\nkPyDWMwAtm5ai47uZn4Ag8PAAgvoxxzzGUCGMQ1DI89gMBr4ADymtwHHIXSSa1UOIt/zBqwz5/IB\nqG4Bjl9EDblFCkDN5ZM4csLJB+DoBaDWDTp3cVUBcDntaLx1hg/AzlKgvQc7yH2oAGSbjOg/WwjT\nvMzkFvdT65U34UdgDCtYGP1avoFUSMpK8gKw7q8n9bLgn/pj8n8C/ALihrxpNKi7hQAAAABJRU5E\nrkJggg==\n'

newicon = b'iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAADSUlEQVR42r2XO0xaURjHLzYmdXRi\nYXBwdDScJl3djQNLw0TSAR88FR9JU6P4QKMCaTqQdGgXsQQMQXFwcGXhCCG+4iNEo6mDMZL45kK/\n76THUO+9Qi/efskJcm4O3+/+v8f51Al/LBaLmTs6Or43gAkq7eHhoXh2dmZva2v7WusZHf9jZGTE\nPz4+blfjuFwuC6IoCsViUbi8vCwDhM1oNH75V4AAANjUAJRKJeYc183NjXB1dVU+Pj7+CIp+qxlg\neHg44PV6VQHg2z8+PjKAu7s7oampSQAVxL29PUtnZ+eP/waA6/r6WtDr9UyJ09NTcXd319LV1aUI\n8QQwNDQUmJiYqBsA5BcMBgPLC4Q4OTkRt7e3LSaTSRbiCWBwcDAwOTmpOgfQOVSBcHFxIbS0tLB9\nDgH5IGYyGYvZbJZAvAoAOsL4IwTEngGxH9fp2MJnhUJBTKfT5t7e3rAsgMfjYQB4oB4V+OJAuI+A\naGtra6nz8/P3o6OjJQnAwMBAYGpqSjUAh8Aw8JLE3OBqIMTGxgbNZrPE5/MVFQFUe6+A4M5xoWMe\nIg4ASksB+vv761aAG3eKMFwBBFhfX6ebm5sEyl0eAHOgjqvgRSBUA3KAghHouFIAt9vNFNACgIci\nmUxSKEcyNjYmBXC5XK8Wgudvz8OAAJgDigD1lGE155UAsiFwOp2vUgVKzisBZJMQAdR2wlqco62s\nrNBcLkfgzpECOByOuhXg5cc733NLJBIMQLYP2O12BqB0uJrjyk8li8fjdGtri4AfeQAMAUpXayJW\ne2MlANlWbLPZGAB2LPagCoQapZaXl+nOzg6ZmZmRB8CBhANoYTB5U5iQyOzsrDwAjmR4hWpl0WiU\nwpxI5ubmpAAwKLAQ4HWqlUUiEbq/v0/m5+flATAEWiqwtLREDw4OyMLCgjyA1iEIh8P08PCQ+P1+\nKUBPTw8D0DIEi4uL9OjoiASDQSlAd3c3A7i/v9cUAKsgFApJAaxWa2B6etp2e3v71yGlen++X60v\n4EwAAOlUKvUO+oEUoLW19UNfX9/nxsbGhsoOxy+Wl75X2+drdXX1J/zX9AkmI+lU3N7e/hYSxAAT\n0Rs+FdX69rXsA1y5ubn5Vz6fL1Q++w30VO4/0/9IewAAAABJRU5ErkJggg==\n'

def uniqify(name, names):
   newname, i = name, 1
   while newname in names: newname, i = name + str(i), i + 1
   return newname

def indent(lines):
   return ['   ' + line for line in lines]

class ParameterDialog(QDialog):
   def __init__(self, block, parent = None):
      super(ParameterDialog, self).__init__(parent)
      self.block = block
      self.button = QPushButton('Ok', self)
      self.nameEdit = QLineEdit(self.block.name)
      self.codeEdit = QTextEdit(self)
      self.codeEdit.setPlainText(self.block.code)
      l = QFormLayout(self)
      l.addRow('Name:', self.nameEdit)
      l.addRow('Code:', self.codeEdit)
      l.addWidget(self.button)
      self.button.clicked.connect(self.OK)
   def OK(self):
      self.block.name = self.nameEdit.text()
      self.block.code = self.codeEdit.toPlainText()
      self.close()

class EditorGraphicsView(QGraphicsView):
   def __init__(self, parent=None):
      QGraphicsView.__init__(self, parent)
      self.setDragMode(QGraphicsView.RubberBandDrag)
      self.delShort = QShortcut(QKeySequence.Delete, self)
      self._model = None
      self.treeView = QTreeView()
   def setDiagram(self, d):
      self.setScene(d)
      self.delShort.activated.connect(d.deleteItems)
   def getModel(self): return self._model
   def setModel(self, m):
      self._model = m
      if m:
         self.treeView.setModel(m)
         self.diagram = m.rootDiagram
         self.model.modelReset.connect(self.treeView.expandAll)
   model = property(getModel, setModel)
   diagram = property(lambda s: s.scene(), setDiagram)
   def save(self):
      if self.model:
         if not self.model.filename:
            self.model.filename = QFileDialog.getSaveFileName(self)
         if self.model.filename:
            with open(self.model.filename, 'w') as f:
               f.write(json.dumps(self.model.Dict))
   def load(self):
      filename = QFileDialog.getOpenFileName(self)
      if filename:
         self.model = loadModel(filename)
   def newModel(self):
      self.model = ModelHierarchyModel()
   def goUp(self):
      if hasattr(self.diagram, 'containingBlock'):
         self.diagram = self.diagram.containingBlock.scene()
         self.zoomAll()
   def showCode(self):
      if self.model:
         c = self.model.gencode()
         c = '\n'.join(c)
         d = QDialog()
         l = QFormLayout(d)
         codeview = QTextEdit()
         codeview.setPlainText(c)
         l.addRow('code', codeview)
         runButton = QPushButton('Run')
         outputview = QTextEdit()
         l.addRow('Output', outputview)
         l.addWidget(runButton)
         def print2(txt):
            txt2 = outputview.toPlainText()
            outputview.setPlainText(txt2 + '\n' + txt)
         def runIt():
            outputview.clear()
            globs = {'print': print2}
            exec(codeview.toPlainText(), globs)
         runButton.clicked.connect(runIt)
         d.exec_()
   def zoomAll(self):
      """ zoom to fit all items """
      rect = self.diagram.itemsBoundingRect()
      self.fitInView(rect, Qt.KeepAspectRatio)
   def wheelEvent(self, event):
      pos = event.pos()
      posbefore = self.mapToScene(pos)
      degrees = event.delta() / 8.0
      sx = (100.0 + degrees) / 100.0
      self.scale(sx, sx)
      event.accept()
   def dragEnterEvent(self, event):
      if event.mimeData().hasFormat('component/name'): event.accept()
   def dragMoveEvent(self, event):
      if event.mimeData().hasFormat('component/name'): event.accept()
   def dropEvent(self, event):
      if event.mimeData().hasFormat('component/name'):
         name = bytes(event.mimeData().data('component/name')).decode()
         pos = self.mapToScene(event.pos())
         self.scene().addNewBlock(pos, name)
         # TODO: do this in a cleaner way:
         self.model.modelReset.emit()

class LibraryModel(QStandardItemModel):
   mimeTypes = lambda self: ['component/name']
   def mimeData(self, idxs):
      mimedata = QMimeData()
      for idx in idxs:
         if idx.isValid():
            txt = self.data(idx, Qt.DisplayRole)
            mimedata.setData('component/name', txt)
      return mimedata

class ModelHierarchyModel(QAbstractItemModel):
   def __init__(self):
      super(ModelHierarchyModel, self).__init__()
      self.rootDiagram = DiagramScene()
      self.filename = None
   def setDict(self, d):
      self.rootDiagram.Dict = d
      self.modelReset.emit()
   Dict = property(lambda s: s.rootDiagram.Dict, setDict)
   def gencode(self):
      c = ['def topLevel():']
      c += indent(self.rootDiagram.gencode())
      c.append('print("Running model")')
      c.append('topLevel()')
      c.append('print("Done")')
      return c
   def index(self, row, column, parent=None):
      if parent.isValid():
         parent = parent.internalPointer().subModel
      else:
         parent = self.rootDiagram
      blocks = sorted(parent.blocks, key=lambda b:  b.name)
      block = blocks[row]
      # Store the index to retrieve it later in the parent function.
      # TODO: solve this in a better way.
      block.index = self.createIndex(row, column, block)
      return block.index
   def parent(self, index):
      if index.isValid():
         block = index.internalPointer()
         if block.scene() == self.rootDiagram:
            return QModelIndex()
         else:
            outerBlock = block.scene().containingBlock
            return outerBlock.index
      print('parent: No valid index')
   def data(self, index, role):
      if index.isValid() and role == Qt.DisplayRole:
         b = index.internalPointer()
         return b.name
   def headerData(self, section, orientation, role):
      if orientation == Qt.Horizontal and role == Qt.DisplayRole:
         return "Element"
   def rowCount(self, parent):
      if parent.column() > 0: return 0
      if parent.isValid():
         parent = parent.internalPointer().subModel
      else:
         parent = self.rootDiagram
      return len(parent.blocks)
   def columnCount(self, parent):
      return 1


class LibraryWidget(QListView):
   def __init__(self):
      super(LibraryWidget, self).__init__(None)
      self.libraryModel = LibraryModel(self)
      self.libraryModel.setColumnCount(1)
      # Create an icon with an icon:
      pixmap = QPixmap(60, 60)
      pixmap.fill()
      painter = QPainter(pixmap)
      painter.fillRect(10, 10, 40, 40, Qt.blue)
      painter.setBrush(Qt.yellow)
      painter.drawEllipse(20, 20, 20, 20)
      painter.end()
      # Fill library:
      for name in ['Block', 'Uber unit', 'Device']:
         self.libraryModel.appendRow(QStandardItem(QIcon(pixmap), name))
      self.setModel(self.libraryModel)
      self.setViewMode(self.IconMode)
      self.setDragDropMode(self.DragOnly)

def warning(txt):
   QMessageBox.warning(None, "Warning", txt)

def loadModel(filename):
   try:
      m = ModelHierarchyModel()
      with open(filename, 'r') as f: data = f.read()
      m.filename = filename
      m.Dict = json.loads(data)
      return m
   except KeyError:
      warning('Corrupt model: {0}'.format(filename))
   except ValueError:
      warning('Corrupt model: {0}'.format(filename))
   except FileNotFoundError:
      warning('File [{0}] not found'.format(filename))

class Main(QMainWindow):
   def __init__(self):
      super(Main, self).__init__(None)
      self.editor = EditorGraphicsView()
      self.setCentralWidget(self.editor)
      self.setWindowTitle("Diagram editor")
      def buildIcon(b64):
         icon = base64.decodestring(b64)
         pm = QPixmap()
         pm.loadFromData(icon)
         return QIcon(pm)
      toolbar = self.addToolBar('Tools')
      toolbar.setObjectName('Tools')
      def act(name, shortcut, callback, icon=None):
         a = QAction(icon, name, self) if icon else QAction(name, self)
         a.setShortcuts(shortcut)
         a.triggered.connect(callback)
         toolbar.addAction(a)
      act('New', QKeySequence.New, self.editor.newModel, buildIcon(newicon))
      act('Save', QKeySequence.Save, self.editor.save, buildIcon(saveicon))
      act('Load', QKeySequence.Open, self.editor.load, buildIcon(loadicon))
      act('Full screen', QKeySequence("F11"), self.toggleFullScreen)
      act('Fit in view', QKeySequence("F8"), self.editor.zoomAll)
      act('Go up', QKeySequence(Qt.Key_Up), self.editor.goUp)
      act('Model code', QKeySequence("F7"), self.editor.showCode)
      def addDock(name, widget):
         dock = QDockWidget(name, self)
         dock.setObjectName(name)
         dock.setWidget(widget)
         self.addDockWidget(Qt.LeftDockWidgetArea, dock)
      addDock('Library', LibraryWidget())
      addDock('Model tree', self.editor.treeView)
      self.settings = QSettings('windelsoft', 'diagrameditor')
      self.loadSettings()
   def toggleFullScreen(self):
      self.setWindowState(self.windowState() ^ Qt.WindowFullScreen)
      self.editor.zoomAll()
   def loadSettings(self):
      if self.settings.contains('mainwindowstate'):
         self.restoreState(self.settings.value('mainwindowstate'))
      if self.settings.contains('mainwindowgeometry'):
         self.restoreGeometry(self.settings.value('mainwindowgeometry'))
      if self.settings.contains('openedmodel'):
         modelfile = self.settings.value('openedmodel')
         self.editor.model = loadModel(modelfile)
   def closeEvent(self, ev):
      self.settings.setValue('mainwindowstate', self.saveState())
      self.settings.setValue('mainwindowgeometry', self.saveGeometry())
      if self.editor.model and self.editor.model.filename:
         self.settings.setValue('openedmodel', self.editor.model.filename)
         # TODO: ask for save of opened files
      else:
         self.settings.remove('openedmodel')
      ev.accept()

if __name__ == '__main__':
   if sys.version_info.major != 3:
      print('Please use python 3.x')
      sys.exit(1)
   app = QApplication(sys.argv)
   main = Main()
   main.show()
   app.exec_()