# $Id$
#
# Copyright (C) 2003-2006 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.
#
""" Supplies a class for working with molecules from databases
#DOC 

"""
from rdkit import Chem
from rdkit.Chem.Suppliers.MolSupplier import MolSupplier
import sys
def warning(msg,dest=sys.stderr):
  dest.write(msg)
  
class DbMolSupplier(MolSupplier):
  """
    new molecules come back with all additional fields from the
    database set in a "_fieldsFromDb" data member
  
  """
  def __init__(self,dbResults,
               molColumnFormats={'SMILES':'SMI',
                                 'SMI':'SMI',
                                 'MOLPKL':'PKL'},
               nameCol='',
               transformFunc=None,
               **kwargs):
    """

      DbResults should be a subclass of Dbase.DbResultSet.DbResultBase

    """
    self._data = dbResults
    self._colNames = [x.upper() for x in self._data.GetColumnNames()]
    nameCol = nameCol.upper()
    self.molCol = -1
    self.transformFunc=transformFunc
    try:
      self.nameCol = self._colNames.index(nameCol)
    except ValueError:
      self.nameCol = -1
    for name in molColumnFormats.keys():
      name = name.upper()
      try:
        idx = self._colNames.index(name)
      except ValueError:
        pass
      else:
        self.molCol = idx
        self.molFmt = molColumnFormats[name]
        break
    if self.molCol < 0:
      raise ValueError('DbResultSet has no recognizable molecule column')
    del self._colNames[self.molCol]
    self._colNames  = tuple(self._colNames)
    self._numProcessed=0
  def GetColumnNames(self):
    return self._colNames

  def _BuildMol(self,data):
    data = list(data)
    molD = data[self.molCol]
    del data[self.molCol]
    self._numProcessed+=1;
    try:
      if self.molFmt =='SMI':
        newM = Chem.MolFromSmiles(str(molD))
        if not newM:
          warning('Problems processing mol %d, smiles: %s\n'%(self._numProcessed,molD))
      elif self.molFmt =='PKL':
        newM = Chem.Mol(str(molD))
    except:
      import traceback
      traceback.print_exc()
      newM = None
    else:
      if newM and self.transformFunc:
        try:
          newM = self.transformFunc(newM,data)
        except:
          import traceback
          traceback.print_exc()
          newM = None
      if newM:
        newM._fieldsFromDb = data
        nFields = len(data)
        for i in range(nFields):
          newM.SetProp(self._colNames[i],str(data[i]))
        if self.nameCol >=0 :
          newM.SetProp('_Name',str(data[self.nameCol]))
    return newM
  
class ForwardDbMolSupplier(DbMolSupplier):
  """ DbMol supplier supporting only forward iteration


    new molecules come back with all additional fields from the
    database set in a "_fieldsFromDb" data member
  
  """
  def __init__(self,dbResults,**kwargs):
    """

      DbResults should be an iterator for Dbase.DbResultSet.DbResultBase

    """
    DbMolSupplier.__init__(self,dbResults,**kwargs)
    self.Reset()

  def Reset(self):
    self._dataIter = iter(self._data)
    
  def NextMol(self):
    """ 

      NOTE: this has side effects

    """
    try:
      d = self._dataIter.next()
    except StopIteration:
      d = None
    if d is not None:
      newM = self._BuildMol(d)
    else:
      newM = None

    return newM

class RandomAccessDbMolSupplier(DbMolSupplier):  
  def __init__(self,dbResults,**kwargs):
    """

      DbResults should be a Dbase.DbResultSet.RandomAccessDbResultSet

    """
    DbMolSupplier.__init__(self,dbResults,**kwargs)
    self._pos = -1
    
  def __len__(self):
    return len(self._data)
  
  def __getitem__(self,idx):
    newD = self._data[idx]
    return self._BuildMol(newD)

  def Reset(self):
    self._pos = -1
  def NextMol(self):
    self._pos += 1
    res = None
    if self._pos < len(self):
      res = self[self._pos]
    return res
