#
# 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 DataStructs
from rdkit import six
from rdkit.VLib.Node import VLibNode
from rdkit.six.moves import cPickle


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 Exception:
      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

  >>> from rdkit import RDConfig
  >>> 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 = next(self._dataIter)
    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 import RDConfig
  >>> 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 _runDoctests(verbose=None):  # pragma: nocover
  import sys
  import doctest
  failed, _ = doctest.testmod(optionflags=doctest.ELLIPSIS, verbose=verbose)
  sys.exit(failed)


if __name__ == '__main__':  # pragma: nocover
  _runDoctests()
