import platform

from PyQt5 import QtCore, QtWidgets

from ost import geom, gfx, gui, seq
from ost import settings
from ost import LogError, mol
from ost.bindings import WrappedTMAlign
from ost.bindings import msms
from ost.seq import alg
import ost
from ost.gui.scene.query_editor import QueryEditorWidget,QueryDialog

class SelectRefDialog(QtWidgets.QDialog):
  def __init__(self, ent_list, parent=None):
    QtWidgets.QDialog.__init__(self, parent)
    self.ent_list_ = ent_list
    vb = QtWidgets.QVBoxLayout()
    self.setLayout(vb)
    self.setWindowTitle("Select Reference Object")
    self.label = QtWidgets.QLabel("Please Select the Reference Object")
    self.list = QtWidgets.QTableWidget(self)
    self.list.horizontalHeader().setStretchLastSection(True)
    self.list.setColumnCount(2)
    self.list.verticalHeader().setVisible(False)
    self.list.horizontalHeader().setVisible(False)
    self.list.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)
    self.list.setSelectionMode(QtWidgets.QAbstractItemView.SingleSelection)
    vb.addWidget(self.label)
    vb.addWidget(self.list)
    self.show_scores = QtWidgets.QCheckBox(self)
    self.show_scores.setText("Show Scores")
    self.show_scores.setChecked(True)
    vb.addWidget(self.show_scores)
    self.show_alignment = QtWidgets.QCheckBox(self)
    self.show_alignment.setText("Display Alignment")
    self.show_alignment.setChecked(False)
    vb.addWidget(self.show_alignment)
    hb = QtWidgets.QHBoxLayout()
    hb.setDirection(QtWidgets.QBoxLayout.LeftToRight)
    cancel_btn = QtWidgets.QPushButton("Cancel", self)
    load_btn = QtWidgets.QPushButton("Select", self)
    hb.addStretch(1)
    hb.addWidget(cancel_btn)
    hb.addWidget(load_btn)
    vb.addItem(hb)
    load_btn.setDefault(True)
    load_btn.clicked.connect(self.Select)
    cancel_btn.clicked.connect(self.reject)

    row = 0
    for ent in self.ent_list_:
      variant = QtCore.QVariant(ent)
      self.list.insertRow(row)
      new_item = QtWidgets.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 = QtWidgets.QTableWidgetItem(ent.GetName())
      new_item.setFlags(QtCore.Qt.ItemIsSelectable|QtCore.Qt.ItemIsEnabled)
      self.list.setItem(row, 1, new_item)
      row += 1  

    self.list.resizeColumnsToContents()
    
  def Select(self):
    items = self.list.selectedItems()
    for item in items:
      if(item.column()==0):
        ent = item.data(QtCore.Qt.UserRole)
        self.ent_list_.remove(ent)
        self.ent_list_.insert(0,ent)
    self.accept()
     
  def GetShowScores(self):
    return self.show_scores.isChecked()
  
  def GetDisplayAlignment(self):
    return self.show_alignment.isChecked()
  
  def GetEntities(self):
    return self.ent_list_

class ShowResultDialog(QtWidgets.QDialog):
  def __init__(self, ent_list, res_list, parent=None):
    QtWidgets.QDialog.__init__(self, parent)
    self.ent_list_ = ent_list
    vb = QtWidgets.QVBoxLayout()
    self.setLayout(vb)
    self.setWindowTitle("Alignment result")
    self.label = QtWidgets.QLabel("Alignment results with %s as reference"%ent_list[0].GetName())
    self.list = QtWidgets.QTableWidget(self)
    self.list.horizontalHeader().setStretchLastSection(True)
    self.list.setColumnCount(4)
    self.list.verticalHeader().setVisible(False)
    self.list.setHorizontalHeaderLabels (["Name","RMSD","TMScore",""])
    self.list.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)
    vb.addWidget(self.label)
    vb.addWidget(self.list)
    hb = QtWidgets.QHBoxLayout()
    hb.setDirection(QtWidgets.QBoxLayout.LeftToRight)
    ok_btn = QtWidgets.QPushButton("OK", self)
    hb.addStretch(1)
    hb.addWidget(ok_btn)
    vb.addItem(hb)
    ok_btn.setDefault(True)
    ok_btn.clicked.connect(self.accept)

    for i in range(0, len(res_list)):
      self.list.insertRow(i)
      new_item = QtWidgets.QTableWidgetItem(ent_list[i+1].GetName())
      new_item.setFlags(QtCore.Qt.ItemIsEnabled)
      self.list.setItem(i, 0, new_item)
      new_item = QtWidgets.QTableWidgetItem("%.2f"%res_list[i].rmsd)
      new_item.setFlags(QtCore.Qt.ItemIsEnabled)
      self.list.setItem(i, 1, new_item)
      new_item = QtWidgets.QTableWidgetItem("%.2f"%res_list[i].tm_score)
      new_item.setFlags(QtCore.Qt.ItemIsEnabled)
      self.list.setItem(i, 2, new_item)
      new_item = QtWidgets.QTableWidgetItem()
      new_item.setFlags(QtCore.Qt.ItemIsEnabled)
      self.list.setItem(i, 3, new_item)
          
    self.list.resizeColumnsToContents()

class CalculateSurfaceSettingsDialog(QtWidgets.QDialog):
  def __init__(self, executable, parent=None):
    QtWidgets.QDialog.__init__(self, parent)
    vb = QtWidgets.QGridLayout()
    self.setLayout(vb)
    self.setWindowTitle("MSMS Surface Settings")
    msmsexe_label=QtWidgets.QLabel("executable")
    self.msmsexe_field=QtWidgets.QLineEdit()
    self.msmsexe_field.setText(executable)
    msmsexe_browsebutton=QtWidgets.QPushButton("Browse")
    vb.addWidget(msmsexe_label, 0, 0)
    vb.addWidget(self.msmsexe_field, 0, 1)
    vb.addWidget(msmsexe_browsebutton, 0, 2)
    surfname_label=QtWidgets.QLabel("surface name")
    self.surfname_field=QtWidgets.QLineEdit()
    self.surfname_field.setText("surface")
    vb.addWidget(surfname_label, 1, 0)
    vb.addWidget(self.surfname_field, 1, 1, 1, 2)
    density_label=QtWidgets.QLabel("density")
    self.density_spinbox=QtWidgets.QSpinBox()
    self.density_spinbox.setRange(1, 10)
    self.density_spinbox.setValue(4)
    vb.addWidget(density_label, 2, 0)
    vb.addWidget(self.density_spinbox, 2, 1, 1, 2)
    probe_label=QtWidgets.QLabel("probe radius")
    self.probe_spinbox=QtWidgets.QDoubleSpinBox()
    self.probe_spinbox.setDecimals(1)
    self.probe_spinbox.setSingleStep(0.1)
    self.probe_spinbox.setRange(0.3, 5.0)
    self.probe_spinbox.setValue(1.4)
    vb.addWidget(probe_label, 3, 0)
    vb.addWidget(self.probe_spinbox, 3, 1, 1, 2)
    selection_label=QtWidgets.QLabel("selection")
    self.selection_field=QtWidgets.QLineEdit()
    self.selection_field.setText("")
    vb.addWidget(selection_label, 4, 0)
    vb.addWidget(self.selection_field, 4, 1, 1, 2)
    self.noh_box=QtWidgets.QCheckBox("no hydrogens")
    vb.addWidget(self.noh_box, 5, 0)
    self.nohet_box=QtWidgets.QCheckBox("no hetatoms")
    vb.addWidget(self.nohet_box, 5, 1)
    self.nowat_box=QtWidgets.QCheckBox("no waters")
    vb.addWidget(self.nowat_box, 5, 2)
    
    cancel_btn = QtWidgets.QPushButton("Cancel", self)
    ok_btn = QtWidgets.QPushButton("OK", self)
    vb.addWidget(cancel_btn, 6, 1)
    vb.addWidget(ok_btn, 6, 2)
    
    msmsexe_browsebutton.clicked.connect(self.GetPath)
    ok_btn.clicked.connect(self.accept)
    cancel_btn.clicked.connect(self.reject)
    
  def GetPath(self):
    path, _ =QtWidgets.QFileDialog.getOpenFileName(self, "Choose MSMS Executable")
    if path!='':
      self.msmsexe_field.setText(path)

class SurfaceContextMenu(QtCore.QObject):
  def __init__(self, context_menu):
    try:
      settings_name="msms"
      self.executable=settings.Locate(settings_name)
    except settings.FileNotFound:
      self.executable=""
    QtCore.QObject.__init__(self, context_menu.qobject)
    self.action = QtWidgets.QAction("Calculate Surface", self)
    self.action.triggered.connect(self.CalculateSurface)
    context_menu.AddAction(self.action, gui.ContextActionType.ENTITY)
    
  def CalculateSurface(self):
    scene_selection = gui.SceneSelection.Instance()
    ent_list = list()
    for i in range(0,scene_selection.GetActiveNodeCount()):
      ent_list.append(scene_selection.GetActiveNode(i))
    cssd = CalculateSurfaceSettingsDialog(self.executable)
    if(cssd.exec_()):
      self.__CalculateSurface(ent_list,
                              str(cssd.surfname_field.text()),
                              str(cssd.msmsexe_field.text()),
                              cssd.density_spinbox.value(),
                              cssd.probe_spinbox.value(),
                              str(cssd.selection_field.text()),
                              cssd.noh_box.isChecked(),
                              cssd.nohet_box.isChecked(),
                              cssd.nowat_box.isChecked())
      
  def __CalculateSurface(self,ent_list,name,msms_exe,density,
                         radius,selection,noh,nohet,nowat):
    for entity in ent_list:
      if isinstance(entity, gfx.Entity):
        try:
          s=msms.CalculateSurface(entity.view,
                                  msms_exe=msms_exe,
                                  density=density,
                                  radius=radius,
                                  selection=selection,
                                  no_hydrogens=noh,
                                  no_hetatoms=nohet,
                                  no_waters=nowat)[0]
          gfx.Scene().Add(gfx.Surface("%s_%s"%(entity.GetName(),name),s))
        except (RuntimeError, msms.MsmsProcessError):
          LogError("WARNING: Surface could not be calculated")
          return
        except UserWarning:
          LogError("WARNING: Entry with the same name already present in scene")
          return


class SecStructContextMenu(QtCore.QObject):
  def __init__(self, context_menu):
    QtCore.QObject.__init__(self, context_menu.qobject)
    self.action = QtWidgets.QAction("Assign Secondary Structure", self)
    self.action.triggered.connect(self.AssignSecondaryStructure)
    context_menu.AddAction(self.action, gui.ContextActionType.ENTITY)
    
  def AssignSecondaryStructure(self):
    scene_selection = gui.SceneSelection.Instance()
    ent_list = list()
    for i in range(0,scene_selection.GetActiveNodeCount()):
      node = scene_selection.GetActiveNode(i)
      mol.alg.AssignSecStruct(node.view)
      node.UpdateView()

  
class AlignmentContextMenu(QtCore.QObject):

  def __init__(self, context_menu):
   
    QtCore.QObject.__init__(self, context_menu.qobject)
    self.action = QtWidgets.QAction("Align", self)
    self.action.triggered.connect(self.Align)
    context_menu.AddAction(self.action, gui.ContextActionType.ENTITY | gui.ContextActionType.MULTI)
    self.seq_viewer = None
    
  def Align(self):
    scene_selection = gui.SceneSelection.Instance()
    ent_list = list()
    for i in range(0,scene_selection.GetActiveNodeCount()):
      ent_list.append(scene_selection.GetActiveNode(i))
    sd = SelectRefDialog(ent_list)
    if(sd.exec_()):
      self.__Align(sd.GetEntities(),sd.GetShowScores(), sd.GetDisplayAlignment())
        
  def __Align(self, ent_list, show_scores=True, display_alignment=False):
    node = ent_list[0]
    res_list = list()
    if isinstance(node, gfx.Entity):
      ref = node.view.chains[0]
      for i in range(1,len(ent_list)):
        node = ent_list[i]
        if isinstance(node, gfx.Entity):
          tm_result = WrappedTMAlign(node.view.chains[0], ref)
          res_list.append(tm_result)
          node.view.handle.EditXCS().ApplyTransform(tm_result.transform)
          node.UpdatePositions()
    if show_scores:
      self.__ShowScore(ent_list, res_list)
    if display_alignment:
      self.__DisplayAlignment(ent_list, res_list)
    
  def __ShowScore(self, ent_list, res_list):
    if(len(res_list)==1):
      res = res_list[0]
      string = "RMSD: %.2f, TMScore: %.2f"%(res.rmsd, res.tm_score)
      gui.GostyApp.Instance().perspective.StatusMessage(string)
    elif(len(res_list)>1):
      ShowResultDialog(ent_list, res_list).exec_()
      
  def __DisplayAlignment(self, ent_list, res_list):
    if(len(res_list)>0):
      ref_seq = seq.CreateSequence("%s (ref)"%ent_list[0].GetName(),
                                   res_list[0].alignment.GetSequence(1).GetGaplessString())
      aln_list = seq.AlignmentList()
      if(ref_seq.IsValid()):
        for i in range(0, len(res_list)):
          # WrappedTMAlign returns an alignment with second sequence
          # being reference... let's swap...
          new_aln = seq.CreateAlignment()
          new_aln.AddSequence(res_list[i].alignment.GetSequence(1))
          new_aln.AddSequence(res_list[i].alignment.GetSequence(0))
          new_aln.SetSequenceName(1, ent_list[i+1].GetName())
          aln_list.append(new_aln)
        alignment = alg.MergePairwiseAlignments(aln_list, ref_seq)
        gosty = gui.GostyApp.Instance()
        main_area = gosty.perspective.GetMainArea()
        if self.seq_viewer:
          self.seq_viewer.qobject.close()
        self.seq_viewer = gui.SequenceViewer(True)
        self.seq_viewer.AddAlignment(alignment)
        self.seq_viewer.ChangeDisplayMode("Highlight conservation 1")
        self.seq_viewer.Show()

class SelectMenuPoints(QtCore.QObject):
  def __init__(self, context_menu):
    QtCore.QObject.__init__(self, context_menu.qobject)
    action=QtWidgets.QAction("Select...", self)
    action.triggered.connect(self._Select)
    context_menu.AddAction(action, gui.ENTITY)
    action=QtWidgets.QAction("Copy Selection...", self)
    action.triggered.connect(self._CopyViews)
    context_menu.AddAction(action, gui.ENTITY)
    action=QtWidgets.QAction('Select...', self)
    action.triggered.connect(self._SelectViewsSameEntity)
    context_menu.AddAction(action, gui.ENTITY_VIEW|gui.VIEWS_SAME_OBJECT)    
  def _Select(self):
    scene_selection=gui.SceneSelection.Instance()
    ent=scene_selection.GetActiveNode(0)
    dialog=QueryDialog('Select...')
    if dialog.exec_():
      q=mol.Query(dialog.query)
      if q.IsValid():
        ent.selection=ent.view.Select(dialog.query, dialog.query_flags)
      else:
        ost.LogError("invalid query: %s" % q.error)

  def _UniqueName(self, ent):
    """
    Returns a name based on ent that is unique within the scene
    """
    ent_name=ent.GetName()
    num=2
    while True:
      candidate_name='%s-%d' % (ent_name, num)
      if not gfx.Scene().HasNode(candidate_name):
        return candidate_name
      num+=1

  def _SelectViewsSameEntity(self):

    union=gui.SceneSelection.Instance().GetViewUnion()
    dialog=QueryDialog('Select...')
    if dialog.exec_():
      q=mol.Query(dialog.query)    
      if q.IsValid():
        ve=gui.SceneSelection.Instance().GetViewEntity()
        ve.selection=union.Select(q, dialog.query_flags)
      else:    
        ost.LogError("invalid query: %s" % q.error)    
  
  def _CopyViews(self):
    views_to_add=[]
    scene_selection=gui.SceneSelection.Instance()
    ent=scene_selection.GetActiveNode(0)
    dialog=QueryDialog('Select...')
    if dialog.exec_():
      q=mol.Query(dialog.query)
      if q.IsValid():
        for i in range(scene_selection.GetActiveNodeCount()):
          ent=scene_selection.GetActiveNode(i)
          selected=ent.view.Select(q, dialog.query_flags)
          gfx_ent=gfx.Entity(self._UniqueName(ent),selected)
          gfx.Scene().Add(gfx_ent)
      else:    
        ost.LogError("invalid query: %s" % q.error)

def _InitContextMenu(app):
  _InitContextMenu.cm=app.scene_win.GetContextMenu()
  _InitContextMenu.am=AlignmentContextMenu(_InitContextMenu.cm)
  _InitContextMenu.sec=SecStructContextMenu(_InitContextMenu.cm)
  _InitContextMenu.sc=SurfaceContextMenu(_InitContextMenu.cm)
  _InitContextMenu.sm=SelectMenuPoints(_InitContextMenu.cm)