""" an XML-RPC server to allow remote control of PyMol
 
  Author: Greg Landrum (glandrum@users.sourceforge.net)
  Created:       January 2002
  $LastChangedDate$
  License:  PyMol
  Requires:
            - a python xmlrpclib distribution containing the SimpleXMLRPCServer
              module (1.0 or greater should be fine)
            - python with threading enabled  
 
  RD Version: $Rev$            

  Modified 2013-04-17 Thomas Holder, Schrodinger, Inc.
"""

import sys
import xmlrpc.server as SimpleXMLRPCServer
import threading,os,tempfile
from pymol import cmd,cgo

# initial port to try for the server
_xmlPort=9123
# number of alternate ports to try if the first fails
_nPortsToTry=5

def rpcPing():
  """ Used to establish whether or not the server is alive.
 
  This is a good thing to call after establishing a connection just to
  make sure that everything is ok.
 
    Returns 1
 
  """
  return 1

def rpcLabel(pos,labelText,id='lab1',color=(1,1,1)):
  """ create a text label
 
    Arguments:
      pos: a 3 tuple with the position of the label
      text: a string with the label
      color: a 3 tuple with the color of the label. (1,1,1) is white
      id: (OPTIONAL) the name of the object to be created
 
    NOTE:
      at the moment this is, how you say, a hack
 
  """
  cmd.pseudoatom(id, label=repr(labelText), elem='C', pos=pos)
  cmd.set_color("%s-color"%id,color)
  cmd.color("%s-color"%id,id)
  return 1

def rpcResetCGO(id):
  """ removes a CGO from the local dictionary
 
  """
  global cgoDict
  if id=="*":
    cgoDict={}
    res = 1
  elif id in cgoDict:
    del(cgoDict[id])
    res = 1
  else:
    res = 0
  return res

def rpcSphere(pos,rad,color,id='cgo',extend=1,
              transparent=0,transparency=0.5):
  """ create a sphere
 
    Arguments:
      pos: a 3 tuple with the position of the sphere
      rad: a float with the radius
      color: a 3 tuple with the color of the sphere. (1,1,1) is white
      id: (OPTIONAL) the name of the object to be created
      extend: (OPTIONAL) if this is nonzero, the object will be cleared
        before adding the new sphere.  Otherwise the sphere is appended
        to the ojbect
      transparent: (OPTIONAL) sets the object to be transparent
      transparency: (OPTIONAL) the percent transparency of the object
  """
  r,g,b = color
  x,y,z = pos
  if extend:
    obj = cgoDict.get(id,[])
  else:
    obj = []
  if not transparent:
    o = []
  else:
    o = [cgo.ALPHA,1-transparency]
  o.extend([cgo.COLOR,r,g,b,cgo.SPHERE,x,y,z,rad])
  obj.extend(o)
  cgoDict[id] = obj
  cmd.load_cgo(obj,id,1)
  return 1

def rpcRenderCGO(cgoV,id='cgo',extend=1):

  """ renders a CGO vector
 
    Arguments:
      cgoV: a vector of floats
      id: (OPTIONAL) the name of the object to be created
      extend: (OPTIONAL) if this is nonzero, the object will be cleared
        before adding the new sphere.  Otherwise the sphere is appended
        to the ojbect
  """
  if extend:
    obj = cgoDict.get(id,[])
  else:
    obj = []
  obj.extend(cgoV)
  cmd.load_cgo(obj,id,1)
  return 1


def rpcSpheres(sphereD,id='cgo',extend=1):
  """ create a sphere
 
    Arguments:
      sphereD: a series of (pos,rad,color,transparent,transparency) tuples
      id: (OPTIONAL) the name of the object to be created
      extend: (OPTIONAL) if this is nonzero, the object will be cleared
        before adding the new sphere.  Otherwise the sphere is appended
        to the ojbect
  """
  if extend:
    obj = cgoDict.get(id,[])
  else:
    obj = []
  for pos,rad,color,transparent,transparency in sphereD:
    r,g,b = color
    x,y,z = pos
    if not transparent:
      o = []
    else:
      o = [cgo.ALPHA,1-transparency]
    o.extend([cgo.COLOR,r,g,b,cgo.SPHERE,x,y,z,rad])
    obj.extend(o)
  cgoDict[id] = obj
  cmd.load_cgo(obj,id,1)
  return 1

def rpcCylinder(end1,end2,rad,color1,id='cgo',color2=None,extend=1,
                transparent=0,transparency=0.5):
  """ create a cylinder
 
    Arguments:
      end1: a 3 tuple with the position of end1 of the sphere
      end2: a 3 tuple with the position of end1 of the sphere
      rad: a float with the radius
      color1: a 3 tuple with the color of end1 of the sphere. (1,1,1) is white
      id: (OPTIONAL) the name of the object to be created
      color2: (OPTIONAL) a 3 tuple with the color of end2 of the sphere. (1,1,1) 
is white
      extend: (OPTIONAL) if this is nonzero, the object will be cleared
        before adding the new sphere.  Otherwise the sphere is appended
        to the ojbect
      transparent: (OPTIONAL) sets the object to be transparent
      transparency: (OPTIONAL) the percent transparency of the object
 
    NOTE: the reason that color2 follows id is that I think clients are
    going to be interested in setting the id more often than they are going
    to care about the second color.
 
  """
  global cgoDict

  if color2 is None: color2 = color1
  r1,g1,b1 = color1
  r2,g2,b2 = color2
  x1,y1,z1 = end1
  x2,y2,z2 = end2
  if extend:
    obj = cgoDict.get(id,[])
  else:
    obj = []
  if not transparent:
    o = []
  else:
    o = [cgo.ALPHA,1-transparency]
  o.extend([cgo.CYLINDER,x1,y1,z1,x2,y2,z2,rad,r1,g1,b1,r2,g2,b2,])
  obj.extend(o)
  cgoDict[id] = obj
  cmd.load_cgo(obj,id,1)
  return 1

def rpcDeleteObject(objName):
  """ deletes an object """
  try:
    cmd.delete(objName)
  except:
    res = 0
  else:
    res = 1
  return res

def rpcDeleteAll():
  """ deletes all objects """
  res = cmd.delete('all')
  if res is not None:
    return res
  else:
    return ''

def colorObj(objName,colorScheme):
  """ sets an molecule's color scheme
    Arguments:
      - objName: the object (molecule) to change
      - colorScheme: name of the color scheme to use
        for the object (should be either 'std' or one of the
        color schemes defined in pymol.utils)
  """
  if colorScheme:
    if colorScheme == 'std':
      # this is an adaptation of the cbag scheme from util.py, but
      # with a gray carbon.
      cmd.color("magenta","("+objName+")",quiet=1)
      cmd.color("oxygen","(elem O and "+objName+")",quiet=1)
      cmd.color("nitrogen","(elem N and "+objName+")",quiet=1)
      cmd.color("sulfur","(elem S and "+objName+")",quiet=1)
      cmd.color("hydrogen","(elem H and "+objName+")",quiet=1)
      cmd.color("gray","(elem C and "+objName+")",quiet=1)
    elif hasattr(utils,colorScheme):
      fn = getattr(utils,colorScheme)
      fn(objName,quiet=1)
    res = 1
  else:
    res = 0
  return res

def rpcLoadPDB(data,objName,colorScheme='',replace=1):
  """ loads a molecule from a pdb string
 
    Arguments:
      data: the mol block
      objName: name of the object to create
      colorScheme: (OPTIONAL) name of the color scheme to use
        for the molecule (should be either 'std' or one of the
        color schemes defined in pymol.utils)
      replace: (OPTIONAL) if an object with the same name already
        exists, delete it before adding this one
 
 
  """
  from pymol import util
  if replace:
    cmd.delete(objName)
  res = cmd.read_pdbstr(data,objName)
  colorObj(objName,colorScheme)
  if res is not None:
    return res
  else:
    return ''


def rpcLoadMolBlock(data,objName,colorScheme='',replace=1):
  """ loads a molecule from a mol block
 
    Arguments:
      data: the mol block
      objName: name of the object to create
      colorScheme: (OPTIONAL) name of the color scheme to use
        for the molecule (should be either 'std' or one of the
        color schemes defined in pymol.utils)
      replace: (OPTIONAL) if an object with the same name already
        exists, delete it before adding this one
 
  """
  from pymol import util
  if replace:
    cmd.delete(objName)
  res = cmd.read_molstr(data,objName)
  colorObj(objName,colorScheme)
  if res is not None:
    return res
  else:
    return ''


def rpcLoadFile(fileName,objName='',format='',colorScheme='',replace=1):
  """ loads an object from a file
 
    Arguments:
      fileName: the file to load
      objName: (OPTIONAL) name of the object to create
      format: (OPTIONAL) the format of the input file
      colorScheme: (OPTIONAL) name of the color scheme to use
        for the object (should be either 'std' or one of the
        color schemes defined in pymol.utils)
      replace: (OPTIONAL) if an object with the same name already
        exists, delete it before adding this one
  """
  if not objName:
    objName = fileName.split('.')[0]
  if replace:
    cmd.delete(objName)
  res = cmd.load(fileName,objName,format=format)
  colorObj(objName,colorScheme)
  if res is not None:
    return res
  else:
    return ''


def rpcLoadSurface(fileName,objName,format='',surfaceLevel=1.0):
  """ loads surface data from a file and adds an isosurface
 
    Arguments:
      fileName: the file to load
      objName: (OPTIONAL) name of the object to create
      format: (OPTIONAL) the format of the input file
      surfaceLevel: (OPTIONAL) the isosurface level
 
   """
  if not objName:
    objName = fileName.split('.')[0]
  gridName = 'grid-%s'%objName
  res = cmd.load(fileName,gridName,format='')
  cmd.isosurface(objName,gridName,level=surfaceLevel)
  if res is not None:
    return res
  else:
    return ''

def rpcLoadSurfaceData(data,objName='surface',format='',surfaceLevel=1.0):
  """ loads surface data from a string and adds an isosurface
 
    Arguments:
      data: the data to load
      objName: (OPTIONAL) name of the object to create
      format: (OPTIONAL) the format of the input file
      surfaceLevel: (OPTIONAL) the isosurface level
 
   """
  gridName = 'grid-%s'%objName
  # it would be nice if we didn't have to go by way of the temporary file,
  # but at the moment pymol will only read shapes from files
  tempnm = tempfile.mktemp('.grd')
  open(tempnm,'w+').write(data)
  res = rpcLoadSurface(tempnm,objName,format='',surfaceLevel=surfaceLevel)
  os.unlink(tempnm)
  if res is not None:
    return res
  else:
    return ''


def rpcRotate(vect,objName='',state=-1):
  """ rotates objects
 
    Arguments:
     - vect: a sequence with x y and z rotations
     - objName: (OPTIONAL) object to be rotated
     - state: (OPTIONAL) if zero only visible states are rotated,
       if -1 (the default), all states are rotated
 
  """
  cmd.rotate('x',vect[0],objName,state=state)
  cmd.rotate('y',vect[1],objName,state=state)
  cmd.rotate('z',vect[2],objName,state=state)
  return 1

def rpcGetNames(what='selections',enabledOnly=1):
  """ returns the results of cmd.get_names(what) """
  return cmd.get_names(what,enabled_only=enabledOnly)

def rpcIdAtom(what='all',mode=0):
  """ returns the results of cmd.id_atom(what) """
  return cmd.id_atom(what,mode=mode)

def rpcGetAtomCoords(what='all',state=0):
  """ returns the results of cmd.get_atom_coords(what,state) """
  return cmd.get_atom_coords(what,state=state)


def rpcHelp(what=''):
  """ returns general help text or help on a particular command """
  global serv
  res = 'Command Not Found'
  if not what:
    res = list(serv.funcs.keys())
  else:
    funcs = serv.funcs
    if what in funcs:
      fn = funcs[what]
      res = "Function: %s("%what
      defs = fn.__defaults__
      if defs:
        code = fn.__code__
        nDefs = len(defs)
        args = []
        i = -1
        for i in range(code.co_argcount - nDefs):
          args.append(code.co_varnames[i])
        for j in range(nDefs):
          vName = code.co_varnames[j+i+1]
          args.append("%s=%s"%(vName,repr(defs[j])))
        res += ','.join(args)
      res += ')\n'
      if fn.__doc__:
        res += fn.__doc__
  return res


def launch_XMLRPC(hostname='',port=_xmlPort,nToTry=_nPortsToTry):
  """ launches the xmlrpc server into a separate thread
 
    Arguments:
      hostname: (OPTIONAL) name of the host for the server
         (defaults to be the name of the localhost)
      port: (OPTIONAL) the first port to try for the server
      nToTry: (OPTIONAL) the number of possible ports to try
        (in case the first can't be opened)
 
  """
  if not hostname:
    import os
    hostname = os.environ.get('PYMOL_RPCHOST', 'localhost')

  global cgoDict,serv
  cgoDict = {}
  for i in range(nToTry):
    try:
      serv = SimpleXMLRPCServer.SimpleXMLRPCServer((hostname,port+i),logRequests=0,
          allow_none=True)
    except:
      serv = None
    else:
      break
  if serv:
    print('xml-rpc server running on host %s, port %d'%(hostname,port+i))

    # import PyMOL API
    from pymol import api
    serv.register_instance(cmd)

    # legacy stuff with unique names
    serv.register_function(rpcPing,'ping')
    serv.register_function(rpcResetCGO,'resetCGO')
    serv.register_function(rpcRenderCGO,'renderCGO')
    serv.register_function(rpcSphere,'sphere')
    serv.register_function(rpcSpheres,'spheres')
    serv.register_function(rpcCylinder,'cylinder')
    serv.register_function(rpcDeleteObject,'deleteObject')
    serv.register_function(rpcDeleteAll,'deleteAll')
    serv.register_function(rpcLoadPDB,'loadPDB')
    serv.register_function(rpcLoadMolBlock,'loadMolBlock')
    serv.register_function(rpcLoadSurface,'loadSurface')
    serv.register_function(rpcLoadSurfaceData,'loadSurfaceData')
    serv.register_function(rpcLoadFile,'loadFile')
    serv.register_function(rpcGetNames,'getNames')
    serv.register_function(api.count_atoms,'countAtoms')
    serv.register_function(rpcIdAtom,'idAtom')
    serv.register_function(rpcHelp,'help')
    serv.register_function(rpcGetAtomCoords,'getAtomCoords')

    # legacy stuff, should be removed because overwrites API names!
    serv.register_function(rpcLabel,'label')   # pseudoatom
    serv.register_function(rpcRotate,'rotate')

    serv.register_introspection_functions()
    t = threading.Thread(target=serv.serve_forever)
    t.setDaemon(1)
    t.start()
  else:
    print('xml-rpc server could not be started')

# vi:expandtab:smarttab:sw=2