# $Id$
#
# Copyright (C) 2004-2012 Greg Landrum and Rational Discovery LLC
#
#   @@ All Rights Reserved @@
#  This file is part of the RDKit.
#  The contents are covered by the terms of the BSD license
#  which is included in the file license.txt, found at the root
#  of the RDKit source tree.
#
""" uses pymol to interact with molecules

"""
from rdkit import Chem
import xmlrpclib,os,tempfile


_server=None
class MolViewer(object):
  def __init__(self,host=None,port=9123,force=0,**kwargs):
    global _server
    if not force and _server is not None:
      self.server=_server
    else:
      if not host:
        host=os.environ.get('PYMOL_RPCHOST','localhost')
      _server=None
      serv = xmlrpclib.Server('http://%s:%d'%(host,port))
      serv.ping()
      _server = serv
      self.server=serv
    self.InitializePyMol()
    
  def InitializePyMol(self):
    """ does some initializations to set up PyMol according to our
    tastes

    """
    self.server.do('set valence,1')
    self.server.do('set stick_rad,0.15')
    self.server.do('set mouse_selection_mode,0')
    self.server.do('set line_width,2')
    self.server.do('set selection_width,10')
    self.server.do('set auto_zoom,0')
    

  def DeleteAll(self):
    " blows out everything in the viewer "
    self.server.deleteAll()

  def DeleteAllExcept(self,excludes):
    " deletes everything except the items in the provided list of arguments "
    allNames = self.server.getNames('*',False)
    for nm in allNames:
      if nm not in excludes:
        self.server.deleteObject(nm)

  def LoadFile(self,filename,name,showOnly=False):
    """ calls pymol's "load" command on the given filename; the loaded object
    is assigned the name "name"
    """
    if showOnly:
      self.DeleteAll()
    id = self.server.loadFile(filename,name)
    return id

  def ShowMol(self,mol,name='molecule',showOnly=True,highlightFeatures=[],
              molB="",confId=-1,zoom=True,forcePDB=False):
    """ special case for displaying a molecule or mol block """
  
    server = self.server
    if not zoom:
      self.server.do('view rdinterface,store')
    if showOnly:
      self.DeleteAll()

    if not forcePDB and mol.GetNumAtoms()<999 :
      if not molB:
        molB = Chem.MolToMolBlock(mol,confId=confId)
      mid = server.loadMolBlock(molB,name)
    else:
      if not molB:
        molB = Chem.MolToPDBBlock(mol,confId=confId)
      mid = server.loadPDB(molB,name)
      
    if highlightFeatures:
      nm = name+'-features'
      conf = mol.GetConformer(confId)
      for feat in highlightFeatures:
        pt = [0.0,0.0,0.0]
        for idx in feat:
          loc = conf.GetAtomPosition(idx)
          pt[0] += loc[0]/len(feat)
          pt[1] += loc[1]/len(feat)
          pt[2] += loc[2]/len(feat)
        server.sphere(pt,0.2,(1,1,1),nm)
    if zoom:
      server.zoom('visible')
    else:
      self.server.do('view rdinterface,recall')
    return mid

  def GetSelectedAtoms(self,whichSelection=None):
    " returns the selected atoms "
    if not whichSelection:
      sels = self.server.getNames('selections')
      if sels:
        whichSelection = sels[-1]
      else:
        whichSelection=None
    if whichSelection:   
      items = self.server.index(whichSelection)
    else:
      items = []
    return items


  def SelectAtoms(self,itemId,atomIndices,selName='selection'):
    " selects a set of atoms "
    ids = '(id '
    ids += ','.join(['%d'%(x+1) for x in atomIndices])
    ids += ')'
    cmd = 'select %s,%s and %s'%(selName,ids,itemId)
    self.server.do(cmd)
  
  def HighlightAtoms(self,indices,where,extraHighlight=False):
    " highlights a set of atoms " 
    if extraHighlight:
      idxText = ','.join(['%s and (id %d)'%(where,x) for x in indices])
      self.server.do('edit %s'%idxText)
    else:
      idxText = ' or '.join(['id %d'%x for x in indices])
      self.server.do('select selection, %s and (%s)'%(where,idxText))

  def SetDisplayStyle(self,obj,style=''):
    " change the display style of the specified object "
    self.server.do('hide everything,%s'%(obj,))
    if style:
      self.server.do('show %s,%s'%(style,obj))

  def SelectProteinNeighborhood(self,aroundObj,inObj,distance=5.0,
                                name='neighborhood',showSurface=False):
    """ selects the area of a protein around a specified object/selection name;
    optionally adds a surface to that """
    self.server.do('select %(name)s,byres (%(aroundObj)s around %(distance)f) and %(inObj)s'%locals())
    

    if showSurface:
      self.server.do('show surface,%s'%name)
      self.server.do('disable %s'%name)

  def AddPharmacophore(self,locs,colors,label,sphereRad=0.5):
    " adds a set of spheres "
    self.server.do('view rdinterface,store')
    self.server.resetCGO(label)
    for i,loc in enumerate(locs):
      self.server.sphere(loc,sphereRad,colors[i],label,1)
    self.server.do('enable %s'%label)
    self.server.do('view rdinterface,recall')
    

  def SetDisplayUpdate(self,val):
    if not val:
      self.server.do('set defer_update,1')
    else:
      self.server.do('set defer_update,0')

  def GetAtomCoords(self,sels):
    " returns the coordinates of the selected atoms " 
    res = {}
    for label,idx in sels:
      coords = self.server.getAtomCoords('(%s and id %d)'%(label,idx))
      res[(label,idx)] = coords
    return res

  def HideAll(self):
    self.server.do('disable all')
  def HideObject(self,objName):
    self.server.do('disable %s'%objName)
  def DisplayObject(self,objName):
    self.server.do('enable %s'%objName)

  def Redraw(self):
    self.server.do('refresh')
  def Zoom(self,objName):
    self.server.zoom(objName)

  def DisplayHBonds(self,objName,molName,proteinName,
                    molSelText='(%(molName)s)',
                    proteinSelText='(%(proteinName)s and not het)'):
    " toggles display of h bonds between the protein and a specified molecule "
    cmd = "delete %(objName)s;\n"
    cmd += "dist %(objName)s," + molSelText+","+proteinSelText+",mode=2;\n"
    cmd += "enable %(objName)s;"
    cmd = cmd%locals()

    self.server.do(cmd)

  def DisplayCollisions(self,objName,molName,proteinName,distCutoff=3.0,
                        color='red',
                        molSelText='(%(molName)s)',
                        proteinSelText='(%(proteinName)s and not het)'):
    " toggles display of collisions between the protein and a specified molecule "
    cmd = "delete %(objName)s;\n"
    cmd += "dist %(objName)s," + molSelText+","+proteinSelText+",%(distCutoff)f,mode=0;\n"
    cmd += """enable %(objName)s
    color %(color)s, %(objName)s"""
    cmd = cmd%locals()
    self.server.do(cmd)

  def GetPNG(self,h=None,w=None,preDelay=0):
    try:
      import Image
    except ImportError:
      from PIL import Image
    import time
    if preDelay>0:
      time.sleep(preDelay)
    fd = tempfile.NamedTemporaryFile(suffix='.png',delete=False)
    fd.close()
    self.server.do('png %s'%fd.name)
    time.sleep(0.2)  # <- wait a short period so that PyMol can finish
    for i in range(10):
      try:
        img = Image.open(fd.name)
        break
      except IOError:
        time.sleep(0.1)
    os.unlink(fd.name)
    fd=None
    if h is not None or w is not None:
      sz = img.size
      if h is None:
        h=sz[1]
      if w is None:
        w=sz[0]
      if h<sz[1]:
        frac = float(h)/sz[1]
        w *= frac
        w = int(w)
        img=img.resize((w,h),True)
      elif w<sz[0]:
        frac = float(w)/sz[0]
        h *= frac
        h = int(h)
        img=img.resize((w,h),True)        
    return img