ContextMenu, Improved alignment when more two structures are selected

git-svn-id: https://dng.biozentrum.unibas.ch/svn/openstructure/trunk@2047 5a81b35b-ba03-0410-adc8-b2c5c5119f08

modules/gui/pymod/init_context_menu.py

@@ -3,7 +3,98 @@ from PyQt4 import QtCore, QtGui
 from ost import geom, gfx, gui
 from ost import settings
 from ost.bindings import tmtools
-  
+
+class SelectRefDialog(QtGui.QDialog):
+  def __init__(self, ent_list, parent=None):
+    QtGui.QDialog.__init__(self, parent)
+    self.ent_list_ = ent_list
+    vb = QtGui.QVBoxLayout()
+    self.setLayout(vb)
+    self.setWindowTitle("Select Reference Object")
+    self.label = QtGui.QLabel("Please Select the Reference Object")
+    self.list = QtGui.QTableWidget(self)
+    self.list.setColumnCount(2)
+    self.list.setColumnWidth(0, 30)
+    self.list.setColumnWidth(1,220)
+    self.list.verticalHeader().setVisible(False)
+    self.list.horizontalHeader().setVisible(False)
+    self.list.setSelectionBehavior(QtGui.QAbstractItemView.SelectRows)
+    self.list.setSelectionMode(QtGui.QAbstractItemView.SingleSelection)
+    vb.addWidget(self.label)
+    vb.addWidget(self.list)
+    hb = QtGui.QHBoxLayout()
+    hb.setDirection(QtGui.QBoxLayout.LeftToRight)
+    cancel_btn = QtGui.QPushButton("Cancel", self)
+    load_btn = QtGui.QPushButton("Select", self)
+    hb.addStretch(1)
+    hb.addWidget(cancel_btn)
+    hb.addWidget(load_btn)
+    vb.addItem(hb)
+    load_btn.setDefault(True)
+    QtCore.QObject.connect(load_btn, QtCore.SIGNAL("clicked()"), self.Select)
+    QtCore.QObject.connect(cancel_btn, QtCore.SIGNAL("clicked()"), self.reject)
+
+    row = 0
+    for ent in self.ent_list_:
+      variant = QtCore.QVariant(ent)
+      self.list.insertRow(row)
+      new_item = QtGui.QTableWidgetItem("%i"%(row+1))
+      new_item.setFlags(QtCore.Qt.ItemIsSelectable|QtCore.Qt.ItemIsEnabled)
+      new_item.setData(QtCore.Qt.UserRole,variant)
+      self.list.setItem(row, 0, new_item)
+      new_item = QtGui.QTableWidgetItem(ent.GetName())
+      new_item.setFlags(QtCore.Qt.ItemIsSelectable|QtCore.Qt.ItemIsEnabled)
+      self.list.setItem(row, 1, new_item)
+      row += 1  
+
+  def Select(self):
+    items = self.list.selectedItems()
+    for item in items:
+      if(item.column()==0):
+        ent = item.data(QtCore.Qt.UserRole).toPyObject()
+        self.ent_list_.remove(ent)
+        self.ent_list_.insert(0,ent)
+    self.accept()
+      
+  def GetEntities(self):
+    return self.ent_list_
+
+class ShowResultDialog(QtGui.QDialog):
+  def __init__(self, ent_list, res_list, parent=None):
+    QtGui.QDialog.__init__(self, parent)
+    self.ent_list_ = ent_list
+    vb = QtGui.QVBoxLayout()
+    self.setLayout(vb)
+    self.setWindowTitle("Alignment result")
+    self.label = QtGui.QLabel("Alignment results with %s as reference"%ent_list[0].GetName())
+    self.list = QtGui.QTableWidget(self)
+    self.list.setColumnCount(3)
+    self.list.verticalHeader().setVisible(False)
+    self.list.setHorizontalHeaderLabels (["Name","RMSD","TMScore"])
+    self.list.setSelectionBehavior(QtGui.QAbstractItemView.SelectRows)
+    vb.addWidget(self.label)
+    vb.addWidget(self.list)
+    hb = QtGui.QHBoxLayout()
+    hb.setDirection(QtGui.QBoxLayout.LeftToRight)
+    ok_btn = QtGui.QPushButton("OK", self)
+    hb.addStretch(1)
+    hb.addWidget(ok_btn)
+    vb.addItem(hb)
+    ok_btn.setDefault(True)
+    QtCore.QObject.connect(ok_btn, QtCore.SIGNAL("clicked()"), self.accept)
+
+    for i in range(0, len(res_list)):
+      self.list.insertRow(i)
+      new_item = QtGui.QTableWidgetItem(ent_list[i+1].GetName())
+      new_item.setFlags(QtCore.Qt.ItemIsEnabled)
+      self.list.setItem(i, 0, new_item)
+      new_item = QtGui.QTableWidgetItem("%.2f"%res_list[i].rmsd)
+      new_item.setFlags(QtCore.Qt.ItemIsEnabled)
+      self.list.setItem(i, 1, new_item)
+      new_item = QtGui.QTableWidgetItem("%i"%res_list[i].tm_score)
+      new_item.setFlags(QtCore.Qt.ItemIsEnabled)
+      self.list.setItem(i, 2, new_item)
+
 class AlignmentContextMenu(QtCore.QObject):
   def __init__(self, context_menu):
     try:
@@ -12,27 +103,43 @@ class AlignmentContextMenu(QtCore.QObject):
 
       self.action = QtGui.QAction("Align", self)
       QtCore.QObject.connect(self.action,QtCore.SIGNAL("triggered()"), self.Align)
-      print gui.ContextActionType.ENTITY | gui.ContextActionType.MULTI
       context_menu.AddAction(self.action, gui.ContextActionType.ENTITY | gui.ContextActionType.MULTI)
     except settings.FileNotFound:
       return
     
   def Align(self):
     scene_selection = gui.SceneSelection.Instance()
-    if scene_selection.GetActiveNodeCount() >= 2:
-      node = scene_selection.GetActiveNode(0)
-      if isinstance(node, gfx.Entity):
-        ref = node.view.handle
-        for i in range(1,scene_selection.GetActiveNodeCount()):
-          node = scene_selection.GetActiveNode(i)
-          if isinstance(node, gfx.Entity):
-            print tmtools.TMAlign(node.view.handle, ref)
-            node.UpdatePositions()
-          
-      
-      
+    ent_list = list()
+    for i in range(0,scene_selection.GetActiveNodeCount()):
+      ent_list.append(scene_selection.GetActiveNode(i))
+    if len(ent_list) == 2:
+      self.__Align(ent_list)
+    elif len(ent_list) > 2:
+      sd = SelectRefDialog(ent_list)
+      if(sd.exec_()):
+        self.__Align(sd.GetEntities())
+        
+  def __Align(self, ent_list):
+    node = ent_list[0]
+    res_list = list()
+    if isinstance(node, gfx.Entity):
+      ref = node.view.handle
+      for i in range(1,len(ent_list)):
+        node = ent_list[i]
+        if isinstance(node, gfx.Entity):
+          res_list.append(tmtools.TMAlign(node.view.handle, ref))
+          node.UpdatePositions()
+    self.__ShowScore(ent_list, res_list)
+    
+  def __ShowScore(self, ent_list, res_list):
+    if(len(res_list)==1):
+      res = res_list[0]
+      string = "RMSD: %.2f, TMScore: %i"%(res.rmsd, res.tm_score)
+      gui.GostyApp.Instance().perspective.StatusMessage(string)
+    elif(len(res_list)>1):
+      ShowResultDialog(ent_list, res_list).exec_()
       
-
+    
 def _InitContextMenu():
   cm=gui.GostyApp.Instance().GetSceneWin().GetContextMenu()
   AlignmentContextMenu(cm)
\ No newline at end of file