#!/usr/bin/python

from PyQt4.QtGui import *
from PyQt4.QtCore import *
import sys
#from elementtree.SimpleXMLWriter import XMLWriter
import xml.dom.minidom as md

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

class Connection:
   """
    - fromPort
    - list of line items in between
    - toPort
   """
   def __init__(self, fromPort, toPort):
      self.pos1 = None
      self.pos2 = None
      self.pathItem = QGraphicsPathItem()
      pen = QPen()
      pen.setWidth(4)
      pen.setColor(Qt.blue)
      pen.setCapStyle(Qt.RoundCap)
      self.pathItem.setPen(pen)
      #self.pathItem.setFlags(self.pathItem.ItemIsSelectable)
      editor.diagramScene.addItem(self.pathItem)
      self.setFromPort(fromPort)
      self.setToPort(toPort)
   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

      ds = editor.diagramScene

      # 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.pathItem in hits:
            hits.remove(self.pathItem)
         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(10, 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
      else:
         p3 = p4 + dx

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

      #print('Items:', hits) 
      self.pathItem.setPath(path)

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

class ParameterDialog(QDialog):
   def __init__(self, parent = None):
      super(ParameterDialog, self).__init__(parent)
      self.button = QPushButton('Ok', self)
      l = QVBoxLayout(self)
      l.addWidget(self.button)
      self.button.clicked.connect(self.OK)
   def OK(self):
      self.close()

class PortItem(QGraphicsEllipseItem):
   """ Represents a port to a subsystem """
   def __init__(self, name, parent=None):
      QGraphicsEllipseItem.__init__(self, QRectF(-6,-6,12.0,12.0), parent)
      self.setCursor(QCursor(Qt.CrossCursor))
      # Properties:
      self.setBrush(QBrush(Qt.red))
      # Name:
      self.name = name
      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)
      w = 60.0
      h = 40.0
      # 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))
      # Label:
      self.label = QGraphicsTextItem(name, self)
      # Create corner for resize:
      self.sizer = HandleItem(self)
      self.sizer.setPos(w, h)
      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.inputs.append( PortItem('a', self) )
      self.inputs.append( PortItem('b', self) )
      self.inputs.append( PortItem('c', self) )
      self.outputs = []
      self.outputs.append( PortItem('y', self) )
      # Update size:
      self.changeSize(w, h)
   def editParameters(self):
      pd = ParameterDialog(self.window())
      pd.exec_()

   def contextMenuEvent(self, event):
      menu = QMenu()
      da = menu.addAction('Delete')
      da.triggered.connect(self.delete)
      pa = menu.addAction('Parameters')
      pa.triggered.connect(self.editParameters)
      menu.exec_(event.screenPos())

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

   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:
      if len(self.inputs) == 1:
         self.inputs[0].setPos(-4, h / 2)
      elif len(self.inputs) > 1:
         y = 5
         dy = (h - 10) / (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 = 5
         dy = (h - 10) / (len(self.outputs) + 0)
         for outp in self.outputs:
            outp.setPos(w+4, y)
            y += dy
      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 = str(event.mimeData().data('component/name'))
         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)
      self.loadDiagram('diagram2.usd')

   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]

      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.label.toPlainText())
         blockElement.setAttribute("x", str(x))
         blockElement.setAttribute("y", str(y))
         blockElement.setAttribute("width", str(w))
         blockElement.setAttribute("height", str(h))
         modelElement.appendChild(blockElement)
      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')
      print('systems:',sysElements)
      blockElements = doc.getElementsByTagName('block')
      print('blocks:', blockElements)
      for sysElement in sysElements:
         print('system:', sysElement)
         blockElements = sysElement.getElementsByTagName('block')
         print('blocks:', blockElements)
         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)

   def toggleFullScreen(self):
      self.setWindowState(self.windowState() ^ Qt.WindowFullScreen)
   def startConnection(self, port):
      self.startedConnection = Connection(port, None)
   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.show()
   editor.resize(700, 800)
   app.exec_()