#!/usr/bin/python

from PyQt4.QtGui import *
from PyQt4.QtCore import *
import sys, xml
import xml.dom.minidom as md

"""
 Author: Windel Bouwman
 Year: 2012
 Description: This script implements a diagram editor.
 run with python 3.x as:
  $ python [thisfile.py]
"""

def buildPath(pts):
   path = QPainterPath(pts[0])
   for pt in pts[1:]: path.lineTo(pt)
   return path

def equalSpace(n, l, offset=15):
   if n == 1:
      return [l / 2]
   elif n > 1:
      return [offset + (l - offset*2)/(n - 1)*i for i in range(n)]
   return []

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

class Connection(QGraphicsPathItem):
   """ Implementation of a connection between blocks """
   def __init__(self, fromPort, toPort):
      super(Connection, self).__init__()
      self.pos2 = self.fromPort = self.toPort = None
      self.setFlags(self.ItemIsSelectable | self.ItemClipsToShape)
      pen = QPen(Qt.blue, 2, cap=Qt.RoundCap)
      self.setPen(pen)
      self.arrowhead = QGraphicsPathItem(self)
      self.arrowhead.setPath(buildPath([QPointF(0.0, 0.0), QPointF(-6.0, 10.0), QPointF(6.0, 10.0), QPointF(0.0, 0.0)]))
      self.arrowhead.setPen(pen)
      self.arrowhead.setBrush(QBrush(pen.color()))
      self.vias = []
      self.setFromPort(fromPort)
      self.setToPort(toPort)
   def mouseDoubleClickEvent(self, event):
      pos = event.scenePos()
      pts = [self.getPos1()] + [v.pos() for v in self.vias] + [self.pos2]
      idx = 0
      tidx = 0
      for p1, p2 in zip(pts[0:-1], pts[1:]):
         l1 = QLineF(p1, p2)
         l2 = QLineF(p1, pos)
         l3 = QLineF(pos, p2)
         d = l2.length() + l3.length() - l1.length()
         if d < 5:
            tidx = idx
         idx += 1
      self.addHandle(pos, tidx)
   def addHandle(self, pos, idx=None):
      hi = HandleItem(self)
      if idx:
         self.vias.insert(idx, hi)
      else:
         self.vias.append(hi)
      def callback(p):
         self.updateLineStukken()
         return p
      hi.posChangeCallbacks.append(callback)
      hi.setPos(pos)
      self.updateLineStukken()
   def myDelete(self):
      scene = self.scene()
      if scene:
         self.setFromPort(None)
         self.setToPort(None)
         scene.removeItem(self)
   def setFromPort(self, fromPort):
      if self.fromPort:
         self.fromPort.posCallbacks.remove(self.setBeginPos)
         self.fromPort.connection = None
      self.fromPort = fromPort
      if self.fromPort:
         self.fromPort.connection = self
         self.updateLineStukken()
         self.fromPort.posCallbacks.append(self.setBeginPos)
   def setToPort(self, toPort):
      if self.toPort:
         self.toPort.posCallbacks.remove(self.setEndPos)
         self.toPort.connection = None
      self.toPort = toPort
      if self.toPort:
         self.setEndPos(toPort.scenePos())
         self.toPort.connection = self
         self.toPort.posCallbacks.append(self.setEndPos)
   def getPos1(self):
      if self.fromPort:
         return self.fromPort.scenePos()
   def setBeginPos(self, pos1): self.updateLineStukken()
   def setEndPos(self, endpos):
      self.pos2 = endpos
      self.updateLineStukken()
   def itemChange(self, change, value):
      if change == self.ItemSelectedHasChanged:
         for via in self.vias:
            via.setVisible(value)
      return super(Connection, self).itemChange(change, value)
   def shape(self): return self.myshape
   def updateLineStukken(self):
      """
         This algorithm determines the optimal routing of all signals.
         TODO: implement nice automatic line router
      """
      pos1 = self.getPos1()
      if pos1 is None or self.pos2 is None:
         return
      pts = [pos1] + [v.pos() for v in self.vias] + [self.pos2]
      self.arrowhead.setPos(self.pos2)
      if pts[-1].x() < pts[-2].x():
         self.arrowhead.setRotation(-90)
      else:
         self.arrowhead.setRotation(90)
      self.setPath(buildPath(pts))
      """ Create a shape outline using the path stroker """
      s = super(Connection, self).shape()
      pps = QPainterPathStroker()
      pps.setWidth(10)
      self.myshape = pps.createStroke(s).simplified()

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 PortItem(QGraphicsPathItem):
   """ Represents a port to a subsystem """
   def __init__(self, name, block, direction):
      super(PortItem, self).__init__(block)
      self.connection = None
      d = 10.0
      if direction == 'input':
         self.setPath(buildPath([QPointF(-d, -d), QPointF(0, 0), QPointF(-d, d)]))
      else:
         self.setPath(buildPath([QPointF(0.0, -d), QPointF(d, 0), QPointF(0.0, d)]))
      self.direction = direction
      self.block = block
      self.setCursor(QCursor(Qt.CrossCursor))
      self.setPen(QPen(Qt.blue, 2, cap=Qt.RoundCap))
      self.textItem = QGraphicsTextItem(self)
      self.name = name
      self.posCallbacks = []
      self.setFlag(self.ItemSendsScenePositionChanges, True)
   def getName(self): return self.textItem.toPlainText()
   def setName(self, name):
      self.textItem.setPlainText(name)
      rect = self.textItem.boundingRect()
      lw, lh = rect.width(), rect.height()
      lx = 3 if self.direction == 'input' else -3 - lw
      self.textItem.setPos(lx, -lh / 2)
   name = property(getName, setName)
   def itemChange(self, change, value):
      if change == self.ItemScenePositionHasChanged:
         for cb in self.posCallbacks:
            cb(value)
         return value
      return super(PortItem, self).itemChange(change, value)
   def mousePressEvent(self, event):
      if self.direction == 'output':
         self.scene().startConnection(self)

class OutputPort(PortItem):
   # TODO: create a subclass OR make a member porttype
   pass

# Block part:
class HandleItem(QGraphicsEllipseItem):
   """ A handle that can be moved by the mouse """
   def __init__(self, parent=None):
      dx = 13.0
      super(HandleItem, self).__init__(QRectF(-0.5*dx,-0.5*dx,dx,dx), parent)
      self.posChangeCallbacks = []
      self.setBrush(QBrush(Qt.white))
      self.setFlags(self.ItemSendsScenePositionChanges | self.ItemIsMovable)
      self.setVisible(False)
      self.setCursor(QCursor(Qt.SizeFDiagCursor))
   def mouseMoveEvent(self, event):
      """ Move function without moving the other selected elements """
      p = self.mapToParent(event.pos())
      self.setPos(p)
   def itemChange(self, change, value):
      if change == self.ItemPositionChange:
         for cb in self.posChangeCallbacks:
            res = cb(value)
            if res:
               value = res
         return value
      # Call superclass method:
      return super(HandleItem, self).itemChange(change, value)

class BlockItem(QGraphicsRectItem):
   """ 
      Represents a block in the diagram
      Has an x and y and width and height
      width and height can only be adjusted with a tip in the lower right corner.

      - in and output ports
      - parameters
      - description
   """
   def __init__(self, name='Untitled', parent=None):
      super(BlockItem, self).__init__(parent)
      self.subModel = DiagramScene()
      self.subModel.containingBlock = self
      # Properties of the rectangle:
      self.setPen(QPen(Qt.blue, 2))
      self.setBrush(QBrush(Qt.lightGray))
      self.setFlags(self.ItemIsSelectable | self.ItemIsMovable | self.ItemSendsScenePositionChanges)
      self.setCursor(QCursor(Qt.PointingHandCursor))
      self.label = QGraphicsTextItem(name, self)
      self.name = name
      self.code = ''
      # Create corner for resize:
      self.sizer = HandleItem(self)
      self.sizer.posChangeCallbacks.append(self.changeSize) # Connect the callback
      button = QPushButton('+in')
      button.clicked.connect(self.newInputPort)
      self.buttonItemAddInput = QGraphicsProxyWidget(self)
      self.buttonItemAddInput.setWidget(button)
      self.buttonItemAddInput.setVisible(False)
      button = QPushButton('+out')
      button.clicked.connect(self.newOutputPort)
      self.buttonItemAddOutput = QGraphicsProxyWidget(self)
      self.buttonItemAddOutput.setWidget(button)
      self.buttonItemAddOutput.setVisible(False)
      # Inputs and outputs of the block:
      self.inputs = []
      self.outputs = []
      # Update size:
      self.sizer.setPos(QPointF(60, 40)) # This is a better resize function
   def editParameters(self):
      pd = ParameterDialog(self, self.window())
      pd.exec_()
   def mouseDoubleClickEvent(self, event):
      #self.editParameters()
      scene = self.scene()
      if scene:
         for view in scene.views():
            view.dScene = self.subModel
            view.zoomAll()
   def newInputPort(self):
      names = [i.name for i in self.inputs + self.outputs]
      self.addInput(PortItem(uniqify('in', names), self, 'input'))
   def newOutputPort(self):
      names = [i.name for i in self.inputs + self.outputs]
      self.addOutput(PortItem(uniqify('out', names), self, 'output'))
   def setName(self, name): self.label.setPlainText(name)
   def getName(self): return self.label.toPlainText()
   name = property(getName, setName)
   def addInput(self, i):
      self.inputs.append(i)
      self.updateSize()
   def addOutput(self, o):
      self.outputs.append(o)
      self.updateSize()
   def contextMenuEvent(self, event):
      menu = QMenu()
      pa = menu.addAction('Parameters')
      pa.triggered.connect(self.editParameters)
      menu.exec_(event.screenPos())
   def itemChange(self, change, value):
      if change == self.ItemSelectedHasChanged:
         for child in [self.sizer, self.buttonItemAddInput, self.buttonItemAddOutput]:
            child.setVisible(value)
      return super(BlockItem, self).itemChange(change, value)
   def myDelete(self):
         scene = self.scene()
         if scene:
            for p in self.inputs + self.outputs:
               if p.connection:
                  p.connection.myDelete()
            scene.removeItem(self)
   def updateSize(self):
      rect = self.rect()
      h, w = rect.height(), rect.width()
      self.buttonItemAddInput.setPos(0, h + 4)
      self.buttonItemAddOutput.setPos(w+10, h+4)
      for inp, y in zip(self.inputs, equalSpace(len(self.inputs), h)):
         inp.setPos(0.0, y)
      for outp, y in zip(self.outputs, equalSpace(len(self.outputs), h)):
         outp.setPos(w, y)
   def changeSize(self, p):
      h = 20 if p.y() < 20 else p.y()
      w = 40 if p.x() < 40 else p.x()
      self.setRect(0.0, 0.0, w, h)
      rect = self.label.boundingRect()
      self.label.setPos((w - rect.width()) / 2, (h - rect.height()) / 2)
      self.updateSize()
      return QPointF(w, h)

class EditorGraphicsView(QGraphicsView):
   def __init__(self, parent=None):
      QGraphicsView.__init__(self, parent)
      self.dScene = DiagramScene(self)
      self.setDragMode(QGraphicsView.RubberBandDrag)
      delShort = QShortcut(QKeySequence.Delete, self)
      delShort.activated.connect(self.dScene.deleteItems)
   dScene = property(lambda s: s.scene(), lambda s, v: s.setScene(v))
   def save(self):
      self.diagramScene.saveDiagram('diagram2.usd')
   def load(self):
      filename = QFileDialog.getOpenFileName(self)
      self.diagramScene.loadDiagram(filename)
   def goDown(self):
      print('Down!')
      block = self.dScene.selected
   def goUp(self):
      print('Up!')
      if hasattr(self.dScene, 'containingBlock'):
         self.dScene = self.dScene.containingBlock.scene()
         self.zoomAll()
   def zoomAll(self):
      """ zoom to fit all items """
      rect = self.dScene.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)

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

class DiagramScene(QGraphicsScene):
   """ Save and load and deletion of item"""
   def __init__(self, parent=None):
      super(DiagramScene, self).__init__(parent)
      self.startedConnection = None
   blocks = property(lambda sel: [i for i in sel.items() if type(i) is BlockItem])
   connections = property(lambda sel: [i for i in sel.items() if type(i) is Connection])
   def saveDiagram(self, filename):
      blocks = self.blocks
      connections = self.connections
      doc = md.Document()
      modelElement = doc.createElement('system')
      doc.appendChild(modelElement)
      for block in blocks:
         blockElement = doc.createElement("block")
         x, y = block.scenePos().x(), block.scenePos().y()
         rect = block.rect()
         w, h = rect.width(), rect.height()
         blockElement.setAttribute("name", block.name)
         blockElement.setAttribute("x", str(int(x)))
         blockElement.setAttribute("y", str(int(y)))
         blockElement.setAttribute("width", str(int(w)))
         blockElement.setAttribute("height", str(int(h)))
         codeNode = doc.createCDATASection(block.code)
         codeElement = doc.createElement('code')
         codeElement.appendChild(codeNode)
         blockElement.appendChild(codeElement)
         for inp in block.inputs:
            portElement = doc.createElement("input")
            portElement.setAttribute("name", inp.name)
            blockElement.appendChild(portElement)
         for outp in block.outputs:
            portElement = doc.createElement("output")
            portElement.setAttribute("name", outp.name)
            blockElement.appendChild(portElement)
         modelElement.appendChild(blockElement)
      for connection in connections:
         connectionElement = doc.createElement("connection")
         fromPort = connection.fromPort.name
         toPort = connection.toPort.name
         fromBlock = connection.fromPort.block.name
         toBlock = connection.toPort.block.name
         connectionElement.setAttribute("fromBlock", fromBlock)
         connectionElement.setAttribute("fromPort", fromPort)
         connectionElement.setAttribute("toBlock", toBlock)
         connectionElement.setAttribute("toPort", toPort)
         for via in connection.vias:
            viaElement = doc.createElement('via')
            viaElement.setAttribute('x', str(int(via.x())))
            viaElement.setAttribute('y', str(int(via.y())))
            connectionElement.appendChild(viaElement)
         modelElement.appendChild(connectionElement)
      with open(filename, 'w') as f:
         f.write(doc.toprettyxml())
   def loadDiagram(self, filename):
      try:
         doc = md.parse(filename)
      except IOError as e:
         print('{0} not found'.format(filename))
         return
      except xml.parsers.expat.ExpatError as e:
         print('{0}'.format(e))
         return
      sysElements = doc.getElementsByTagName('system')
      blockElements = doc.getElementsByTagName('block')
      for sysElement in sysElements:
         blockElements = sysElement.getElementsByTagName('block')
         for blockElement in blockElements:
            x = float(blockElement.getAttribute('x'))
            y = float(blockElement.getAttribute('y'))
            w = float(blockElement.getAttribute('width'))
            h = float(blockElement.getAttribute('height'))
            name = blockElement.getAttribute('name')
            block = BlockItem(name)
            self.addItem(block)
            block.setPos(x, y)
            block.sizer.setPos(w, h)
            codeElements = blockElement.getElementsByTagName('code')
            if codeElements:
               cn = codeElements[0].childNodes
               cdatas = [cd for cd in cn if type(cd) is md.CDATASection]
               if len(cdatas) > 0:
                  block.code = cdatas[0].data
            # Load ports:
            portElements = blockElement.getElementsByTagName('input')
            for portElement in portElements:
               name = portElement.getAttribute('name')
               inp = PortItem(name, block, 'input')
               block.addInput(inp)
            portElements = blockElement.getElementsByTagName('output')
            for portElement in portElements:
               name = portElement.getAttribute('name')
               outp = PortItem(name, block, 'output')
               block.addOutput(outp)
         connectionElements = sysElement.getElementsByTagName('connection')
         for connectionElement in connectionElements:
            fromBlock = connectionElement.getAttribute('fromBlock')
            fromPort = connectionElement.getAttribute('fromPort')
            toBlock = connectionElement.getAttribute('toBlock')
            toPort = connectionElement.getAttribute('toPort')
            viaElements = connectionElement.getElementsByTagName('via')
            fromPort = self.findPort(fromBlock, fromPort)
            toPort = self.findPort(toBlock, toPort)
            connection = Connection(fromPort, toPort)
            for viaElement in viaElements:
               x = int(viaElement.getAttribute('x'))
               y = int(viaElement.getAttribute('y'))
               connection.addHandle(QPointF(x, y))
            self.addItem(connection)
   def findPort(self, blockname, portname):
      block = self.findBlock(blockname)
      if block:
         for port in block.inputs + block.outputs:
            if port.name == portname:
               return port
   def findBlock(self, blockname):
      items = self.items()
      blocks = [item for item in items if type(item) is BlockItem]
      for block in blocks:
         if block.name == blockname:
            return block
   def addNewBlock(self, pos, name):
      blocknames = [item.name for item in self.items() if type(item) is BlockItem]
      b1 = BlockItem(uniqify(name, blocknames))
      b1.setPos(pos)
      self.addItem(b1)
   def mouseMoveEvent(self, event):
      if self.startedConnection:
         pos = event.scenePos()
         self.startedConnection.setEndPos(pos)
      super(DiagramScene, self).mouseMoveEvent(event)
   def mouseReleaseEvent(self, event):
      if self.startedConnection:
         for item in self.items(event.scenePos()):
            if type(item) is PortItem:
               if item.direction == 'input' and item.connection == None:
                  self.startedConnection.setToPort(item)
                  self.startedConnection = None
                  return
         self.startedConnection.myDelete()
         self.startedConnection = None
      super(DiagramScene, self).mouseReleaseEvent(event)
   def startConnection(self, port):
      self.startedConnection = Connection(port, None)
      pos = port.scenePos()
      self.startedConnection.setEndPos(pos)
      self.addItem(self.startedConnection)
   def deleteItems(self):
      for item in list(self.selectedItems()):
         item.myDelete()

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:
      self.libItems = []
      self.libItems.append( QStandardItem(QIcon(pixmap), 'Block') )
      self.libItems.append( QStandardItem(QIcon(pixmap), 'Uber Unit') )
      self.libItems.append( QStandardItem(QIcon(pixmap), 'Device') )
      for i in self.libItems:
         self.libraryModel.appendRow(i)
      self.setModel(self.libraryModel)
      self.setViewMode(self.IconMode)
      self.setDragDropMode(self.DragOnly)

class Main(QMainWindow):
   def __init__(self):
      super(Main, self).__init__(None)
      self.editor = EditorGraphicsView()
      self.setCentralWidget(self.editor)
      self.setWindowTitle("Diagram editor")
      toolbar = self.addToolBar('Tools')
      def act(name, shortcut, callback):
         a = QAction(name, self)
         a.setShortcuts(shortcut)
         a.triggered.connect(callback)
         toolbar.addAction(a)
      act('Save', QKeySequence.Save, self.editor.save)
      act('Load', QKeySequence.Open, self.editor.load)
      act('Full screen', QKeySequence("F11"), self.toggleFullScreen)
      act('Fit in view', QKeySequence("F8"), self.editor.zoomAll)
      act('Go down', QKeySequence(Qt.Key_Down), self.editor.goDown)
      act('Go up', QKeySequence(Qt.Key_Up), self.editor.goUp)
      self.library = LibraryWidget() 
      libraryDock = QDockWidget('Library', self)
      libraryDock.setWidget(self.library)
      self.addDockWidget(Qt.LeftDockWidgetArea, libraryDock)
      self.editor.scene().loadDiagram('diagram2.usd')
   def toggleFullScreen(self):
      self.setWindowState(self.windowState() ^ Qt.WindowFullScreen)
      self.editor.zoomAll()

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()
   main.resize(700, 500)
   main.editor.zoomAll()
   app.exec_()