# Copyright (C) 2009 Greg Landrum
#  All Rights Reserved
#

import pickle

from rdkit import Chem, DataStructs

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


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()
    return [int(txt[x:x + 32], 2) for x in range(0, len(txt), 32)]

  @staticmethod
  def GetQueryText(mol, query=True):
    words = LayeredOptions.GetWords(mol, query=query)
    colqs = []
    for idx, word in enumerate(words):
      if not word:
        continue

      colqs.append(f'{word}&Col_{idx + 1}={word}')
    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):
  return Chem.RDKFingerprint(mol, nBitsPerHash=1)


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 BuildAvalonFP(mol, smiles=None):
  from rdkit.Avalon import pyAvalonTools
  if smiles is None:
    return pyAvalonTools.GetAvalonFP(mol)
  return pyAvalonTools.GetAvalonFP(smiles, True)


def DepickleFP(pkl, similarityMethod):
  if not isinstance(pkl, (bytes, str)):
    pkl = str(pkl)
  try:
    klass = similarityMethods[similarityMethod]
    fp = klass(pkl)
  except Exception:
    import traceback
    traceback.print_exc()
    fp = pickle.loads(pkl)
  return fp
