# $Id$
#
# Copyright (C) 2009 Greg Landrum
#  All Rights Reserved
#
import cPickle
from rdkit import DataStructs,Chem
from rdkit import Chem

similarityMethods={'RDK':DataStructs.ExplicitBitVect,
                   'AtomPairs':DataStructs.IntSparseIntVect,
                   'TopologicalTorsions':DataStructs.LongSparseIntVect,
                   'Pharm2D':DataStructs.SparseBitVect,
                   'Gobbi2D':DataStructs.SparseBitVect,
                   'Morgan':DataStructs.UIntSparseIntVect
                   }
supportedSimilarityMethods=similarityMethods.keys()


class LayeredOptions:
  loadLayerFlags=0xFFFFFFFF
  searchLayerFlags=0x7
  minPath=1
  maxPath=6
  fpSize=1024
  wordSize=32
  nWords=fpSize//wordSize
  @staticmethod
  def GetFingerprint(mol,query=True):
    if query:
      flags=LayeredOptions.searchLayerFlags
    else:
      flags=LayeredOptions.loadLayerFlags
    return Chem.LayeredFingerprint(mol,layerFlags=flags,
                                   minPath=LayeredOptions.minPath,maxPath=LayeredOptions.maxPath,
                                   fpSize=LayeredOptions.fpSize)
  @staticmethod
  def GetWords(mol,query=True):
    txt = LayeredOptions.GetFingerprint(mol,query=query).ToBitString()
    words = [int(txt[x:x+32],2) for x in range(0,len(txt),32)]
    return  words

  @staticmethod
  def GetQueryText(mol,query=True):
    words = LayeredOptions.GetWords(mol,query=query)
    colqs = []
    for idx,word in enumerate(words):
      if not word:
        continue
      idx = idx+1
      colqs.append('%(word)d&Col_%(idx)d=%(word)d'%locals())
    return  ' and '.join(colqs)



def BuildSigFactory(options=None,fdefFile=None,
                    bins=[(2,3),(3,4),(4,5),(5,6),(6,7),(7,8),(8,100)],
                    skipFeats=('LumpedHydrophobe','ZnBinder')):
  if options:
    fdefFile = options.fdefFile
  if not fdefFile:
    raise ValueError,'bad fdef file'
  from rdkit.Chem import ChemicalFeatures
  from rdkit.Chem.Pharm2D import SigFactory
  featFactory = ChemicalFeatures.BuildFeatureFactory(fdefFile)
  sigFactory = SigFactory.SigFactory(featFactory,
                                     skipFeats=skipFeats,
                                     trianglePruneBins=False)
  sigFactory.SetBins(bins)
  return sigFactory

def BuildAtomPairFP(mol):
  from rdkit.Chem.AtomPairs import Pairs
  fp=Pairs.GetAtomPairFingerprintAsIntVect(mol)
  fp._sumCache = fp.GetTotalVal()
  return fp
def BuildTorsionsFP(mol):
  from rdkit.Chem.AtomPairs import Torsions
  fp=Torsions.GetTopologicalTorsionFingerprintAsIntVect(mol)
  fp._sumCache = fp.GetTotalVal()
  return fp
def BuildRDKitFP(mol):
  fp=Chem.RDKFingerprint(mol,nBitsPerHash=1)
  return fp
def BuildPharm2DFP(mol):
  global sigFactory
  from rdkit.Chem.Pharm2D import Generate
  try:
    fp=Generate.Gen2DFingerprint(mol,sigFactory)
  except IndexError:
    print 'FAIL:',Chem.MolToSmiles(mol,True)
    raise
  return fp
def BuildMorganFP(mol):
  from rdkit.Chem import rdMolDescriptors
  fp = rdMolDescriptors.GetMorganFingerprint(mol,2)
  fp._sumCache = fp.GetTotalVal()
  return fp

def DepickleFP(pkl,similarityMethod):
    try:
        klass=similarityMethods[similarityMethod]
        fp = klass(pkl)
    except:
        import traceback
        traceback.print_exc()
        fp = cPickle.loads(pkl)
    return fp
