# $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 fingerprints from databases
#DOC 

"""
from rdkit import RDConfig
from rdkit.VLib.Node import VLibNode
from rdkit import DataStructs
from rdkit import six
from rdkit.six.moves import cPickle
import sys
def warning(msg,dest=sys.stderr):
  dest.write(msg)
  
class DbFpSupplier(VLibNode):
  """
    new fps come back with all additional fields from the
    database set in a "_fieldsFromDb" data member
  
  """
  def __init__(self,dbResults,fpColName='AutoFragmentFp',usePickles=True):
    """

      DbResults should be a subclass of Dbase.DbResultSet.DbResultBase

    """
    VLibNode.__init__(self)
    self._usePickles = usePickles
    self._data = dbResults
    self._fpColName = fpColName.upper()
    self._colNames = [x.upper() for x in self._data.GetColumnNames()]
    if self._fpColName not in self._colNames:
      raise ValueError('fp column name "%s" not found in result set: %s'%(self._fpColName,str(self._colNames)))
    self.fpCol = self._colNames.index(self._fpColName)
    del self._colNames[self.fpCol]
    self._colNames = tuple(self._colNames)
    self._numProcessed=0


  def GetColumnNames(self):
    return self._colNames

  def _BuildFp(self,data):
    data = list(data)
    if six.PY3:
      pkl = bytes(data[self.fpCol],encoding='Latin1')
    else:
      pkl = str(data[self.fpCol])
    del data[self.fpCol]
    self._numProcessed+=1;
    try:
      if self._usePickles:
        newFp = cPickle.loads(pkl,encoding='bytes')          
      else:
        newFp = DataStructs.ExplicitBitVect(pkl)
    except:
      import traceback
      traceback.print_exc()
      newFp = None
    if newFp:
      newFp._fieldsFromDb = data
    return newFp

  def next(self):
    itm = self.NextItem()
    if itm is None:
      raise StopIteration
    return itm

  __next__ = next # py3

  
class ForwardDbFpSupplier(DbFpSupplier):
  """ DbFp supplier supporting only forward iteration

  >>> import os.path
  >>> from rdkit.Dbase.DbConnection import DbConnect
  >>> fName = RDConfig.RDTestDatabase
  >>> conn = DbConnect(fName,'simple_combined')
  >>> suppl = ForwardDbFpSupplier(conn.GetData())

  we can loop over the supplied fingerprints:
  >>> fps = []
  >>> for fp in suppl:
  ...   fps.append(fp)
  >>> len(fps)
  12

  """
  def __init__(self,*args,**kwargs):
    DbFpSupplier.__init__(self,*args,**kwargs)
    self.reset()

  def reset(self):
    DbFpSupplier.reset(self)
    self._dataIter = iter(self._data)

  def NextItem(self):
    """ 

      NOTE: this has side effects

    """
    try:
      d = self._dataIter.next()
    except StopIteration:
      d = None
    if d is not None:
      newFp = self._BuildFp(d)
    else:
      newFp = None
    return newFp

class RandomAccessDbFpSupplier(DbFpSupplier):  
  """ DbFp supplier supporting random access:
  >>> import os.path
  >>> from rdkit.Dbase.DbConnection import DbConnect
  >>> fName = RDConfig.RDTestDatabase
  >>> conn = DbConnect(fName,'simple_combined')
  >>> suppl = RandomAccessDbFpSupplier(conn.GetData())
  >>> len(suppl)
  12

  we can pull individual fingerprints:
  >>> fp = suppl[5]
  >>> fp.GetNumBits()
  128
  >>> fp.GetNumOnBits()
  54
  
  a standard loop over the fingerprints:
  >>> fps = []
  >>> for fp in suppl:
  ...   fps.append(fp)
  >>> len(fps)
  12

  or we can use an indexed loop:
  >>> fps = [None]*len(suppl)
  >>> for i in range(len(suppl)):
  ...   fps[i] = suppl[i]
  >>> len(fps)
  12

  """
  def __init__(self,*args,**kwargs):
    DbFpSupplier.__init__(self,*args,**kwargs)
    self.reset()
    
  def __len__(self):
    return len(self._data)
  
  def __getitem__(self,idx):
    newD = self._data[idx]
    return self._BuildFp(newD)

  def reset(self):
    self._pos = -1

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



#------------------------------------
#
#  doctest boilerplate
#
def _test():
  import doctest,sys
  return doctest.testmod(sys.modules["__main__"])

if __name__ == '__main__':
  import sys
  failed,tried = _test()
  sys.exit(failed)
  

