#!/usr/bin/python

from PyQt4.QtGui import *
from PyQt4.QtCore import *
import sys
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]
"""

class Connection(QGraphicsPathItem):
   """ Implementation of a connection between blocks """
   def __init__(self, fromPort, toPort):
      super(Connection, self).__init__()
      self.pos1 = None
      self.pos2 = None
      self.fromPort = None
      self.toPort = None
      #self.setFlags(self.ItemIsSelectable | self.ItemIsMovable)
      self.setFlags(self.ItemIsSelectable)
      pen = QPen(Qt.blue, 2)
      pen.setCapStyle(Qt.RoundCap)
      self.setPen(pen)
      arrowPath = QPainterPath(QPointF(0.0, 0.0))
      arrowPath.lineTo(-6.0, 10.0)
      arrowPath.lineTo(6.0, 10.0)
      arrowPath.lineTo(0.0, 0.0)
      self.arrowhead = QGraphicsPathItem(arrowPath, self)
      self.arrowhead.setPen(pen)
      self.arrowhead.setBrush(QBrush(pen.color()))
      self.setFromPort(fromPort)
      self.setToPort(toPort)
   #def shape(self):
   #   return self.path()
   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.setBeginPos(fromPort.scenePos())
         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 releasePorts(self):
      self.setFromPort(None)
      self.setToPort(None)
   def setBeginPos(self, pos1):
      self.pos1 = pos1
      self.updateLineStukken()
   def setEndPos(self, endpos):
      self.pos2 = endpos
      self.updateLineStukken()
   def updateLineStukken(self):
      """
         This algorithm determines the optimal routing of all signals.
         TODO: implement nice automatic line router
      """
      if self.pos1 is None or self.pos2 is None:
         return
      # TODO: do not get the scene here?
      ds = editor.diagramScene
      self.arrowhead.setPos(self.pos2)

      # TODO: create pieces of lines.

      # Determine the begin and end positions:
      p1 = self.pos1
      p4 = self.pos2
      x1, y1 = self.pos1.x(), self.pos1.y()
      x4, y4 = self.pos2.x(), self.pos2.y()

      def stripHits(hits):
         """ Helper function that removes object hits """
         hits = [hit for hit in hits if type(hit) is BlockItem]
         if self in hits:
            hits.remove(self)
         if self.fromPort in hits:
            hits.remove(self.fromPort)
         if self.toPort in hits:
            hits.remove(self.toPort)
         return hits

      def viaPoints(pA, pB):
         # Construct intermediate points:

         pAB1 = QPointF(pA.x(), pB.y())
         pAB2 = QPointF(pB.x(), pA.y())
         path1 = QPainterPath(pA)
         path1.lineTo(pAB1)
         path2 = QPainterPath(pA)
         path2.lineTo(pAB2)
         if len(stripHits(ds.items(path1))) > len(stripHits(ds.items(path2))):
            pAB = pAB2
         else:
            pAB = pAB1
         return [pAB]

      # Determine left or right:
      dx = QPointF(20, 0)
      if len(stripHits(ds.items(p1 + dx))) > len(stripHits(ds.items(p1 - dx))):
         p2 = p1 - dx
      else:
         p2 = p1 + dx

      if len(stripHits(ds.items(p4 + dx))) > len(stripHits(ds.items(p4 - dx))):
         p3 = p4 - dx
         self.arrowhead.setRotation(90)
      else:
         p3 = p4 + dx
         self.arrowhead.setRotation(-90)

      path = QPainterPath(p1)
      path.lineTo(p2)
      # Now move from p2 to p3 without hitting blocks:
      pts = viaPoints(p2, p3)
      for pt in pts:
         path.lineTo(pt)
      path.lineTo(p3)
      path.lineTo(p4)

      hits = stripHits(ds.items(path))
      self.setPath(path)

class ParameterDialog(QDialog):
   def __init__(self, block, parent = None):
      super(ParameterDialog, self).__init__(parent)
      self.block = block
      self.button = QPushButton('Ok', self)
      l = QGridLayout(self)
      l.addWidget(QLabel('Name:', self), 0, 0)
      self.nameEdit = QLineEdit(self.block.name)
      l.addWidget(self.nameEdit, 0, 1)
      l.addWidget(self.button, 1, 0)
      self.button.clicked.connect(self.OK)
   def OK(self):
      self.block.setName(self.nameEdit.text())
      self.close()
# TODO: merge dialogs?

class AddPortDialog(QDialog):
   def __init__(self, block, parent = None):
      super(AddPortDialog, self).__init__(parent)
      self.block = block
      l = QGridLayout(self)
      l.addWidget(QLabel('Name:', self), 0, 0)
      self.nameEdit = QLineEdit('bla')
      l.addWidget(self.nameEdit, 0, 1)
      self.button = QPushButton('Ok', self)
      l.addWidget(self.button, 1, 0)
      self.button.clicked.connect(self.OK)

   def OK(self):
      name = self.nameEdit.text()
      self.block.addOutput(PortItem(name, self.block, 'output'))
      self.close()

class PortItem(QGraphicsPathItem):
   """ Represents a port to a subsystem """
   def __init__(self, name, block, direction):
      super(PortItem, self).__init__(block)
      #QRectF(-6,-6,12.0,12.0)
      self.connection = None
      path = QPainterPath()
      if direction == 'input':
         d = 5.0
         path.moveTo(-d, -d)
         path.lineTo(0.0, 0.0)
         path.lineTo(-d, d)
      else:
         d = 5.0
         path.moveTo(0.0, -d)
         path.lineTo(d, 0.0)
         path.lineTo(0.0, d)
      self.setPath(path)
      self.direction = direction
      self.block = block
      self.setCursor(QCursor(Qt.CrossCursor))
      #self.setBrush(QBrush(Qt.red))
      self.setPen(QPen(Qt.blue, 2))
      self.name = name
      self.textItem = QGraphicsTextItem(name, self)
      self.setName(name)
      self.posCallbacks = []
      self.setFlag(self.ItemSendsScenePositionChanges, True)
   def setName(self, name):
      self.name = name
      self.textItem.setPlainText(name)
      rect = self.textItem.boundingRect()
      lw, lh = rect.width(), rect.height()
      if self.direction == 'input':
         lx = 3
      else:
         lx = -3 - lw
      self.textItem.setPos(lx, -lh / 2) # TODO
   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':
         editor.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):
      super(HandleItem, self).__init__(QRectF(-4.0,-4.0,8.0,8.0), parent)
      self.posChangeCallbacks = []
      self.setBrush(QBrush(Qt.white))
      self.setFlag(self.ItemSendsScenePositionChanges, True)
      self.setFlag(self.ItemIsMovable, True)
      #self.setFlag(self.ItemIsSelectable, False)
      self.setCursor(QCursor(Qt.SizeFDiagCursor))

   def itemChange(self, change, value):
      if change == self.ItemPositionChange:
         x, y = value.x(), value.y()
         for cb in self.posChangeCallbacks:
            res = cb(x, y)
            if res:
               x, y = res
               value = QPointF(x, y)
         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)
      # Properties of the rectangle:
      self.setPen(QPen(Qt.blue, 2))
      self.setBrush(QBrush(Qt.lightGray))
      self.setFlags(self.ItemIsSelectable | self.ItemIsMovable)
      self.setBoundingRegionGranularity(1)
      self.setCursor(QCursor(Qt.PointingHandCursor))
      self.label = QGraphicsTextItem(name, self)
      self.name = name
      # Create corner for resize:
      self.sizer = HandleItem(self)
      self.sizer.posChangeCallbacks.append(self.changeSize) # Connect the callback
      #self.sizer.setVisible(False)

      # Inputs and outputs of the block:
      self.inputs = []
      self.outputs = []
      # Update size:
      self.sizer.setPos(60, 40) # This is a better resize function
      #self.changeSize(60, 40) # TODO: create a wrapper function
   def editParameters(self):
      pd = ParameterDialog(self, self.window())
      pd.exec_()
   def addPort(self):
      pd = AddPortDialog(self, self.window())
      pd.exec_()
   def setName(self, name):
      self.name = name
      self.label.setPlainText(name)
   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)
      pa = menu.addAction('Add port')
      pa.triggered.connect(self.addPort)
      menu.exec_(event.screenPos())

   def updateSize(self):
      rect = self.rect()
      h, w = rect.height(), rect.width()
      if len(self.inputs) == 1:
         self.inputs[0].setPos(0.0, h / 2)
      elif len(self.inputs) > 1:
         y = 15
         dy = (h - 30) / (len(self.inputs) - 1)
         for inp in self.inputs:
            inp.setPos(0.0, y)
            y += dy
      if len(self.outputs) == 1:
         self.outputs[0].setPos(w, h / 2)
      elif len(self.outputs) > 1:
         y = 15
         dy = (h - 30) / (len(self.outputs) - 1)
         for outp in self.outputs:
            outp.setPos(w, y)
            y += dy
      
   def changeSize(self, w, h):
      """ Resize block function """
      # Limit the block size:
      if h < 20:
         h = 20
      if w < 40:
         w = 40
      self.setRect(0.0, 0.0, w, h)
      # center label:
      rect = self.label.boundingRect()
      lw, lh = rect.width(), rect.height()
      lx = (w - lw) / 2
      ly = (h - lh) / 2
      self.label.setPos(lx, ly)
      # Update port positions:
      self.updateSize()
      return w, h

class EditorGraphicsView(QGraphicsView):
   def __init__(self, scene, parent=None):
      QGraphicsView.__init__(self, scene, parent)
      self.setDragMode(QGraphicsView.RubberBandDrag)
   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()
         b1 = BlockItem(name)
         b1.setPos(self.mapToScene(event.pos()))
         self.scene().addItem(b1)

class LibraryModel(QStandardItemModel):
   def __init__(self, parent=None):
      QStandardItemModel.__init__(self, parent)
   def mimeTypes(self):
      return ['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)
   def saveDiagram(self, filename):
      items = self.items()
      blocks = [item for item in items if type(item) is BlockItem]
      connections = [item for item in items if type(item) is Connection]

      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)))
         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)
         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
      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)
            # 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')
            fromPort = self.findPort(fromBlock, fromPort)
            toPort = self.findPort(toBlock, toPort)
            connection = Connection(fromPort, toPort)
            self.addItem(connection)
   def findPort(self, blockname, portname):
      items = self.items()
      blocks = [item for item in items if type(item) is BlockItem]
      for block in [b for b in blocks if b.name == blockname]:
         for port in block.inputs + block.outputs:
            if port.name == portname:
               return port
   def mouseMoveEvent(self, mouseEvent):
      editor.sceneMouseMoveEvent(mouseEvent)
      super(DiagramScene, self).mouseMoveEvent(mouseEvent)
   def mouseReleaseEvent(self, mouseEvent):
      editor.sceneMouseReleaseEvent(mouseEvent)
      super(DiagramScene, self).mouseReleaseEvent(mouseEvent)
   def deleteItem(self, item=None):
      if item:
         if type(item) is BlockItem:
            for p in item.inputs + item.outputs:
               if p.connection:
                  self.deleteItem(p.connection)
            self.removeItem(item)
         elif type(item) is Connection:
            item.releasePorts()
            self.removeItem(item)
      else:
         # No item was supplied, try to delete all currently selected items:
         items = self.selectedItems()
         connections = [item for item in items if type(item) is Connection]
         blocks = [item for item in items if type(item) is BlockItem]
         for item in connections + blocks:
            self.deleteItem(item)

class DiagramEditor(QWidget):
   def __init__(self, parent=None):
      QWidget.__init__(self, parent)
      self.setWindowTitle("Diagram editor")

      # Widget layout and child widgets:
      self.horizontalLayout = QHBoxLayout(self)
      self.libraryBrowserView = QListView(self)
      self.libraryBrowserView.setMaximumWidth(160)
      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.red)
      painter.drawEllipse(36, 2, 20, 20)
      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.libraryBrowserView.setModel(self.libraryModel)
      self.libraryBrowserView.setViewMode(self.libraryBrowserView.IconMode)
      self.libraryBrowserView.setDragDropMode(self.libraryBrowserView.DragOnly)

      self.diagramScene = DiagramScene(self)
      self.loadDiagram = self.diagramScene.loadDiagram
      self.diagramView = EditorGraphicsView(self.diagramScene, self)
      self.horizontalLayout.addWidget(self.libraryBrowserView)
      self.horizontalLayout.addWidget(self.diagramView)

      self.startedConnection = None
      fullScreenShortcut = QShortcut(QKeySequence("F11"), self)
      fullScreenShortcut.activated.connect(self.toggleFullScreen)
      saveShortcut = QShortcut(QKeySequence(QKeySequence.Save), self)
      saveShortcut.activated.connect(self.save)
      testShortcut = QShortcut(QKeySequence("F12"), self)
      testShortcut.activated.connect(self.test)
      zoomShortcut = QShortcut(QKeySequence("F8"), self)
      zoomShortcut.activated.connect(self.zoomAll)
      delShort = QShortcut(QKeySequence.Delete, self)
      delShort.activated.connect(self.diagramScene.deleteItem)

   def test(self):
      self.diagramView.rotate(30)
      self.zoomAll()
   def save(self):
      self.diagramScene.saveDiagram('diagram2.usd')
   def toggleFullScreen(self):
      self.setWindowState(self.windowState() ^ Qt.WindowFullScreen)
   def startConnection(self, port):
      self.startedConnection = Connection(port, None)
      self.diagramScene.addItem(self.startedConnection)
   def zoomAll(self):
      """ zoom to fit all items """
      rect = self.diagramScene.itemsBoundingRect()
      self.diagramView.fitInView(rect, Qt.KeepAspectRatio)
   def sceneMouseMoveEvent(self, event):
      if self.startedConnection:
         pos = event.scenePos()
         self.startedConnection.setEndPos(pos)
   def sceneMouseReleaseEvent(self, event):
      # Clear the actual connection:
      if self.startedConnection:
         items = self.diagramScene.items(event.scenePos())
         for item in items:
            if type(item) is PortItem:
               self.startedConnection.setToPort(item)
               self.startedConnection = None
               return
         self.diagramScene.deleteItem(self.startedConnection)
         del self.startedConnection
         self.startedConnection = None

if __name__ == '__main__':
   if sys.version_info.major != 3:
      print('Please use python 3.x')
      sys.exit(1)

   app = QApplication(sys.argv)
   global editor
   editor = DiagramEditor()
   editor.loadDiagram('diagram2.usd')
   editor.show()
   editor.resize(700, 500)
   editor.zoomAll()
   app.exec_()