#!/usr/bin/python

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

"""
  This script implements a basic diagram editor.
"""

class Connection(QGraphicsPathItem):
   """
    - fromPort
    - list of line items in between
    - toPort
   """
   def __init__(self, fromPort, toPort):
      super(Connection, self).__init__()
      self.pos1 = None
      self.pos2 = None
      #self.setFlags(self.ItemIsSelectable | self.ItemIsMovable)
      self.setFlags(self.ItemIsSelectable)
      pen = QPen(Qt.blue)
      pen.setWidth(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):
      self.fromPort = fromPort
      if self.fromPort:
         self.setBeginPos(fromPort.scenePos())
         self.fromPort.posCallbacks.append(self.setBeginPos)
   def setToPort(self, toPort):
      self.toPort = toPort
      if self.toPort:
         self.setEndPos(toPort.scenePos())
         self.toPort.posCallbacks.append(self.setEndPos)
   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)

   def delete(self):
      editor.diagramScene.removeItem(self)
      # Remove position update callbacks:

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.addInput(PortItem(name, self.block))
      self.close()

class PortItem(QGraphicsEllipseItem):
   """ Represents a port to a subsystem """
   def __init__(self, name, block):
      QGraphicsEllipseItem.__init__(self, QRectF(-6,-6,12.0,12.0), block)
      self.block = block
      self.setCursor(QCursor(Qt.CrossCursor))
      self.setBrush(QBrush(Qt.red))
      self.name = name
      self.textItem = QGraphicsTextItem(name, self)
      self.textItem.setPos(10, 0) # TODO
      self.posCallbacks = []
      self.setFlag(self.ItemSendsScenePositionChanges, True)
   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):
      editor.startConnection(self)

# 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.setFlags(self.ItemIsMovable | self.ItemSendsScenePositionChanges)
      self.setCursor(QCursor(Qt.SizeFDiagCursor))

   def itemChange(self, change, value):
      if change == self.ItemPositionChange:
         #value = value.toPointF()
         x, y = value.x(), value.y()
         # TODO: make this a signal?
         # This cannot be a signal because this is not a QObject
         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.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)
      self.sizer.setFlag(self.sizer.ItemIsSelectable, True)

      # Inputs and outputs of the block:
      self.inputs = []
      self.outputs = []
      # Update size:
      self.changeSize(60, 40)
   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()
      da = menu.addAction('Delete')
      da.triggered.connect(self.delete)
      pa = menu.addAction('Parameters')
      pa.triggered.connect(self.editParameters)
      pa = menu.addAction('Add port')
      pa.triggered.connect(self.addPort)
      menu.exec_(event.screenPos())

   def delete(self):
      editor.diagramScene.removeItem(self)
      # TODO: remove connections

   def updateSize(self):
      rect = self.rect()
      h, w = rect.height(), rect.width()
      if len(self.inputs) == 1:
         self.inputs[0].setPos(-4, h / 2)
      elif len(self.inputs) > 1:
         y = 8
         dy = (h - 16) / (len(self.inputs) - 1)
         for inp in self.inputs:
            inp.setPos(-4, y)
            y += dy
      if len(self.outputs) == 1:
         self.outputs[0].setPos(w+4, h / 2)
      elif len(self.outputs) > 1:
         y = 8
         dy = (h - 16) / (len(self.outputs) + 0)
         for outp in self.outputs:
            outp.setPos(w+4, 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)
   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).toByteArray() # python2
            txt = self.data(idx, Qt.DisplayRole) # python 3
            mimedata.setData('component/name', txt)
      return mimedata

class DiagramScene(QGraphicsScene):
   def __init__(self, parent=None):
      super(DiagramScene, self).__init__(parent)
   def mouseMoveEvent(self, mouseEvent):
      editor.sceneMouseMoveEvent(mouseEvent)
      super(DiagramScene, self).mouseMoveEvent(mouseEvent)
   def mouseReleaseEvent(self, mouseEvent):
      editor.sceneMouseReleaseEvent(mouseEvent)
      super(DiagramScene, self).mouseReleaseEvent(mouseEvent)

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

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

   def test(self):
      self.diagramView.rotate(11)
   def save(self):
      self.saveDiagram('diagram2.usd')

   def saveDiagram(self, filename):
      items = self.diagramScene.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.diagramScene.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)
               block.addInput(inp)
            portElements = blockElement.getElementsByTagName('output')
            for portElement in portElements:
               name = portElement.getAttribute('name')
               outp = PortItem(name, block)
               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.diagramScene.addItem(connection)
            
      self.zoomAll()

   def findPort(self, blockname, portname):
      items = self.diagramScene.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 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.startedConnection.delete()
         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_()