
from pymol.wizard import Wizard
from pymol import cmd
import pymol

sele_prefix = "_pf_s_"
sele_prefix_len = len(sele_prefix)

dist_prefix = "_pf_d_"


class Pair_fit(Wizard):

   def __init__(self):

      Wizard.__init__(self)
      
      self.memory = 0
      self.n_pair = 0
      self.status = 0 # 0 no atoms selections, 1 atom selected
      self.message = None
      
   def get_panel(self):
      return [
         [ 1, 'Pair Fitting',''],
         [ 2, 'Fit %d Pairs'%self.n_pair,'cmd.get_wizard().fit()'],
         [ 2, 'Delete Last Pair','cmd.get_wizard().remove_last()'],
         [ 2, 'Redraw','cmd.get_wizard().update_dashes()'],         
         [ 2, 'Clear','cmd.get_wizard().clear()'],
         [ 2, 'Done','cmd.set_wizard()'],
         ]

   def cleanup(self):
      self.clear()
      
   def clear(self):
      cmd.delete(sele_prefix+"*")
      cmd.delete(dist_prefix+"*")
      lst = cmd.get_names('selections')
      self.n_pair = 0
      self.status = 0
      self.message = None
      cmd.unpick()
      cmd.refresh_wizard()
      
   def get_prompt(self):
      self.prompt = None
      if self.status==0:
         self.prompt = [ 'Pick the mobile atom...']
      elif self.status==1:
         self.prompt = [ 'Pick the target atom...' ]
      if self.message!=None:
         self.prompt.append(self.message)
      return self.prompt
   
   def set_status(self,status):
      self.status = status
      cmd.refresh_wizard()

   def get_sele_list(self,mode='all'):
      lst = cmd.get_names('selections')
      lst = filter(lambda x:x[0:sele_prefix_len]==sele_prefix,lst)
      lst.sort()
      if mode == 'mobile': # mobile
         lst=filter(lambda x:x[-1:]=='b',lst) 
      elif mode == 'target': # target
         lst=filter(lambda x:x[-1:]=='a',lst)
      return lst
   
   def fit(self):
      # build up the pair-wise list of selections
      cmd.delete(dist_prefix+"*")
      lst = self.get_sele_list()
      c = 0
      args = []
      while 1:
         if not len(lst): break
         a = lst.pop()
         if not len(lst): break
         b = lst.pop()
         args.append(a)
         args.append(b)
      # do the fit
      if len(args):
         cmd.push_undo(args[0])
         dist = apply(cmd.pair_fit,args)
         self.message = "RMS over %d pairs = %5.3f"%(self.n_pair,dist)
         cmd.refresh_wizard()
      self.update_dashes()
                           
   def remove_last(self):
      # build up the pair-wise list of selections
      cmd.delete(dist_prefix+"*")
      lst = self.get_sele_list()
      if len(lst):
         cmd.delete(lst.pop())
         if len(lst):
            cmd.delete(lst.pop())         
         self.n_pair = self.n_pair - 1
      self.update_dashes()
      self.status=0
      cmd.refresh_wizard()
      
   def update_dashes(self):
      cmd.delete(dist_prefix+"*")
      lst = self.get_sele_list()
      c = 0
      while 1:
         if not len(lst): break
         a = lst.pop()
         if not len(lst): break
         b = lst.pop()
         name = dist_prefix+str(c)
         cmd.dist(name,a,b,width=7,length=0.05,gap=0.05)
         cmd.hide('label',name)
         cmd.enable(name)
         c = c + 1

   def check_same_object(self,lst,sele):
      if not len(lst):
         return 1
      else:
         if cmd.count_atoms("((byobj %s) and %s)"%(lst[0],sele),quiet=1):
            return 1
      return 0

   def check_different_object(self,lst,sele):
      if not len(lst):
         return 1
      else:
         if not cmd.count_atoms("((byobj %s) and %s)"%(lst[0],sele),quiet=1):
            return 1
      return 0

   def do_pick(self,bondFlag):
      if bondFlag:
         self.message = "Error: please select an atom, not a bond."
         print self.message
      else:
         if self.status==0:
            lst = self.get_sele_list(mode='mobile')
            if not self.check_same_object(lst,"(pk1)"):
               self.message = "Error: must select an atom in the same object as before."
               print self.message
            else:
               name = sele_prefix + "%02db"%self.n_pair # mobile end in 'b'
               cmd.select(name,"(pk1)")
               cmd.unpick()
               cmd.enable(name)
               self.status = 1
               self.message = None
         elif self.status==1:
            lst = self.get_sele_list(mode='target')
            if not self.check_same_object(lst,"(pk1)"):
               self.message = "Error: must select an atom in the same object as before."
               print self.message
            else:
               lst = self.get_sele_list(mode='mobile')
               if not self.check_different_object(lst,"(pk1)"):
                  self.message = "Error: target atom must be in a distinct object."
                  print self.message
               else:
                  name = sele_prefix + "%02da"%self.n_pair # target end in 'a'
                  cmd.select(name,"(pk1)")
                  cmd.unpick()
                  cmd.enable(name)
                  self.n_pair = self.n_pair + 1
                  self.status = 0
                  self.update_dashes()
            
      cmd.refresh_wizard()
