// $Id$
//
// Copyright (C) 2007,2008 Greg Landrum
//
// @@ 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.
//

#include <sqlite3ext.h>
SQLITE_EXTENSION_INIT1
#include <GraphMol/RDKitBase.h>
#include <GraphMol/MolPickler.h>
#include <GraphMol/SmilesParse/SmilesParse.h>
#include <GraphMol/Substruct/SubstructMatch.h>
#include <DataStructs/BitVects.h>
#include <DataStructs/BitOps.h>
#include <DataStructs/SparseIntVect.h>
#include <GraphMol/Fingerprints/Fingerprints.h>
#include <GraphMol/Descriptors/MolDescriptors.h>
#include <boost/cstdint.hpp>
#include <string>
#include <map>

std::string stringFromTextArg(sqlite3_value *arg) {
  const unsigned char *text = sqlite3_value_text(arg);
  int nBytes = sqlite3_value_bytes(arg);
  std::string res((const char *)text, nBytes);
  return res;
}

std::string stringFromBlobArg(sqlite3_value *arg) {
  const void *blob = sqlite3_value_blob(arg);
  int nBytes = sqlite3_value_bytes(arg);
  std::string res((const char *)blob, nBytes);
  return res;
}

RDKit::ROMol *molFromBlobArg(sqlite3_value *arg) {
  std::string pkl = stringFromBlobArg(arg);
  RDKit::ROMol *m;
  try {
    m = new RDKit::ROMol(pkl);
  } catch (RDKit::MolPicklerException &) {
    m = 0;
  }
  return m;
}

ExplicitBitVect *ebvFromBlobArg(sqlite3_value *arg) {
  std::string pkl = stringFromBlobArg(arg);
  ExplicitBitVect *ebv;
  try {
    ebv = new ExplicitBitVect(pkl);
  } catch (ValueErrorException &) {
    ebv = 0;
  }
  return ebv;
}

template <typename T>
RDKit::SparseIntVect<T> *sivFromBlobArg(sqlite3_value *arg) {
  std::string pkl = stringFromBlobArg(arg);
  RDKit::SparseIntVect<T> *siv;
  try {
    siv = new RDKit::SparseIntVect<T>(pkl);
  } catch (ValueErrorException &) {
    siv = 0;
  }
  return siv;
}

/* ---------------------------------

  Benchmarking results.

    Database: 65385 pubchem compounds

    Simple access: select count(*) from molecules where length(molpkl)>40;
                   0.3s
    depickle     : select count(*) from molecules where
 rdk_molNumAtoms(molpkl)>40;
                   11.3s

    substruct1   : select count(*) from molecules where
                   rdk_molHasSubstruct(molpkl,'c1ncncn1');
                   18.0s

    substruct2   : select count(*) from molecules where
                   rdk_molHasSubstruct(molpkl,'[#6;r10]');
                   15.8

    3 Oct 2007:
    depickle     : select count(*) from molecules where
 rdk_molNumAtoms(molpkl)>40;
                   9.4s
    mw           : select count(*) from molecules where rdk_molAMW(molpkl)<200;
                   9.7s



 --------------------------------- */

static void numAtomsFunc(sqlite3_context *context, int argc,
                         sqlite3_value **argv) {
  RDKit::ROMol *m = molFromBlobArg(argv[0]);
  if (m) {
    int res = m->getNumAtoms();
    delete m;
    sqlite3_result_int(context, res);
  } else {
    std::string errorMsg = "BLOB could not be converted into a molecule";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
  }
}

static void molWtFunc(sqlite3_context *context, int argc,
                      sqlite3_value **argv) {
  RDKit::ROMol *m = molFromBlobArg(argv[0]);
  if (m) {
    double res = RDKit::Descriptors::CalcAMW(*m);
    delete m;
    sqlite3_result_double(context, res);
  } else {
    std::string errorMsg = "BLOB could not be converted into a molecule";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
  }
}

static void molLogPFunc(sqlite3_context *context, int argc,
                        sqlite3_value **argv) {
  RDKit::ROMol *m = molFromBlobArg(argv[0]);
  if (m) {
    double res, tmp;
    RDKit::Descriptors::CalcCrippenDescriptors(*m, res, tmp);
    delete m;
    sqlite3_result_double(context, res);
  } else {
    std::string errorMsg = "BLOB could not be converted into a molecule";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
  }
}

static void smilesToBlob(sqlite3_context *context, int argc,
                         sqlite3_value **argv) {
  std::string smiles = stringFromTextArg(argv[0]);
  RDKit::ROMol *m = 0;
  try {
    m = RDKit::SmilesToMol(smiles);
  } catch (RDKit::MolSanitizeException &) {
    m = 0;
  }
  if (m) {
    std::string text;
    RDKit::MolPickler::pickleMol(*m, text);
    delete m;
    sqlite3_result_blob(context, text.c_str(), text.length(), SQLITE_TRANSIENT);
  } else {
    std::string errorMsg = "SMILES could not be converted into a molecule";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
  }
}

static void molHasSubstruct(sqlite3_context *context, int argc,
                            sqlite3_value **argv) {
  RDKit::ROMol *m = molFromBlobArg(argv[0]);
  if (!m) {
    std::string errorMsg =
        "BLOB (argument 1) could not be converted into a molecule";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }

  std::string smarts = stringFromTextArg(argv[1]);

  std::map<std::string, boost::any> &molMap =
      *static_cast<std::map<std::string, boost::any> *>(
          sqlite3_user_data(context));
  RDKit::ROMol *patt = 0;
  if (molMap.find(smarts) != molMap.end()) {
    patt = boost::any_cast<RDKit::ROMOL_SPTR>(molMap[smarts]).get();
  } else {
    patt = static_cast<RDKit::ROMol *>(RDKit::SmartsToMol(smarts));
    molMap[smarts] = boost::any(RDKit::ROMOL_SPTR(patt));
  }
  if (!patt) {
    std::string errorMsg =
        "SMARTS (argument 2) could not be converted into a molecule";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }
  RDKit::MatchVectType match;
  int res = RDKit::SubstructMatch(*m, *patt, match, true, false, true);
  delete m;
  sqlite3_result_int(context, res);
}

static void molSubstructCount(sqlite3_context *context, int argc,
                              sqlite3_value **argv) {
  RDKit::ROMol *m = molFromBlobArg(argv[0]);
  if (!m) {
    std::string errorMsg =
        "BLOB (argument 1) could not be converted into a molecule";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }

  std::string smarts = stringFromTextArg(argv[1]);

  std::map<std::string, boost::any> &molMap =
      *static_cast<std::map<std::string, boost::any> *>(
          sqlite3_user_data(context));
  RDKit::ROMol *patt = 0;
  if (molMap.find(smarts) != molMap.end()) {
    patt = boost::any_cast<RDKit::ROMOL_SPTR>(molMap[smarts]).get();
  } else {
    patt = static_cast<RDKit::ROMol *>(RDKit::SmartsToMol(smarts));
    molMap[smarts] = boost::any(RDKit::ROMOL_SPTR(patt));
  }
  if (!patt) {
    std::string errorMsg =
        "SMARTS (argument 2) could not be converted into a molecule";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }
  std::vector<RDKit::MatchVectType> matches;
  int res = RDKit::SubstructMatch(*m, *patt, matches, true, true, false);
  delete m;
  sqlite3_result_int(context, res);
}

static void blobToRDKitFingerprint(sqlite3_context *context, int argc,
                                   sqlite3_value **argv) {
  RDKit::ROMol *m = molFromBlobArg(argv[0]);
  if (!m) {
    std::string errorMsg =
        "BLOB (argument 1) could not be converted into a molecule";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }
  ExplicitBitVect *fp =
      RDKit::DaylightFingerprintMol(*m, 1, 7, 2048, 4, true, 0.3, 128);
  std::string text = fp->toString();
  delete fp;
  delete m;
  sqlite3_result_text(context, text.c_str(), text.length(), SQLITE_TRANSIENT);
}

static void blobToAtomPairFingerprint(sqlite3_context *context, int argc,
                                      sqlite3_value **argv) {
  RDKit::ROMol *m = molFromBlobArg(argv[0]);
  if (!m) {
    std::string errorMsg =
        "BLOB (argument 1) could not be converted into a molecule";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }
  RDKit::SparseIntVect<int> *fp =
      RDKit::Descriptors::AtomPairs::getAtomPairFingerprint(*m);
  std::string text = fp->toString();
  delete fp;
  delete m;
  sqlite3_result_text(context, text.c_str(), text.length(), SQLITE_TRANSIENT);
}

static void bvTanimotoSim(sqlite3_context *context, int argc,
                          sqlite3_value **argv) {
  ExplicitBitVect *bv1 = ebvFromBlobArg(argv[0]);
  if (!bv1) {
    std::string errorMsg =
        "BLOB (argument 1) could not be converted into a bit vector";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }
  ExplicitBitVect *bv2 = ebvFromBlobArg(argv[1]);
  if (!bv2) {
    delete bv1;
    std::string errorMsg =
        "BLOB (argument 2) could not be converted into a bit vector";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }
  double res = SimilarityWrapper(*bv1, *bv2, TanimotoSimilarity);
  delete bv1;
  delete bv2;
  sqlite3_result_double(context, res);
}
static void ucvTanimotoSim(sqlite3_context *context, int argc,
                           sqlite3_value **argv) {
  // table from Andrew Dalke:
  static const unsigned int popCounts[] = {
      0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 1, 2, 2, 3, 2, 3, 3, 4,
      2, 3, 3, 4, 3, 4, 4, 5, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
      2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 1, 2, 2, 3, 2, 3, 3, 4,
      2, 3, 3, 4, 3, 4, 4, 5, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
      2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6,
      4, 5, 5, 6, 5, 6, 6, 7, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
      2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 2, 3, 3, 4, 3, 4, 4, 5,
      3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
      2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6,
      4, 5, 5, 6, 5, 6, 6, 7, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
      4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8,
  };
  const unsigned char *t1 = (const unsigned char *)sqlite3_value_blob(argv[0]);
  int nB1 = sqlite3_value_bytes(argv[0]);
  const unsigned char *t2 = (const unsigned char *)sqlite3_value_blob(argv[1]);
  int nB2 = sqlite3_value_bytes(argv[1]);
  if (nB1 != nB2) {
    std::string errorMsg = "bit vectors not ths same length";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }
  unsigned int x = 0, y = 0, z = 0;
  for (unsigned int i = 0; i < (unsigned int)nB1; ++i) {
    y += popCounts[*t1];
    z += popCounts[*t2];
    x += popCounts[(*t1) & (*t2)];
    ++t1;
    ++t2;
  }
  double res = 0;
  if (y + z - x > 0) {
    res = double(x) / (y + z - x);
  }
  sqlite3_result_double(context, res);
}

#if 0
// Naive approach: actually construct two sparse int vects:
static void sivDiceSim(
  sqlite3_context *context,
  int argc,
  sqlite3_value **argv
){
  RDKit::SparseIntVect<int> *v1=sivFromBlobArg<int>(argv[0]);
  if(!v1){
    std::string errorMsg="BLOB (argument 1) could not be converted into an int vector";
    sqlite3_result_error(context,errorMsg.c_str(),errorMsg.length());
    return;
  }
  RDKit::SparseIntVect<int> *v2=sivFromBlobArg<int>(argv[1]);
  if(!v2){
    delete v1;
    std::string errorMsg="BLOB (argument 2) could not be converted into a bit vector";
    sqlite3_result_error(context,errorMsg.c_str(),errorMsg.length());
    return;
  }
  double res= RDKit::DiceSimilarity(*v1,*v2);
  delete v1;
  delete v2;
  sqlite3_result_double(context, res);
}
#else
// faster, just parse the format directly
static void sivDiceSim(sqlite3_context *context, int argc,
                       sqlite3_value **argv) {
  const unsigned char *t1 = (const unsigned char *)sqlite3_value_blob(argv[0]);
  int nB1 = sqlite3_value_bytes(argv[0]);
  const unsigned char *t2 = (const unsigned char *)sqlite3_value_blob(argv[1]);
  int nB2 = sqlite3_value_bytes(argv[1]);

  // check the version flags:
  boost::uint32_t tmp;
  tmp = *(reinterpret_cast<const boost::uint32_t *>(t1));
  t1 += sizeof(boost::uint32_t);
  if (tmp != ci_SPARSEINTVECT_VERSION) {
    std::string errorMsg =
        "BLOB (argument 1) could not be converted into an int vector";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }
  tmp = *(reinterpret_cast<const boost::uint32_t *>(t2));
  t2 += sizeof(boost::uint32_t);
  if (tmp != ci_SPARSEINTVECT_VERSION) {
    std::string errorMsg =
        "BLOB (argument 2) could not be converted into an int vector";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }

  // check the element size:
  tmp = *(reinterpret_cast<const boost::uint32_t *>(t1));
  t1 += sizeof(boost::uint32_t);
  if (tmp != sizeof(boost::uint32_t)) {
    std::string errorMsg =
        "BLOB (argument 1) could not be converted into an uint32_t vector";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }
  tmp = *(reinterpret_cast<const boost::uint32_t *>(t2));
  t2 += sizeof(boost::uint32_t);
  if (tmp != sizeof(boost::uint32_t)) {
    std::string errorMsg =
        "BLOB (argument 2) could not be converted into an uint32_t vector";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }

  double res = 0.;
  // start reading:
  boost::uint32_t len1, len2;
  len1 = *(reinterpret_cast<const boost::uint32_t *>(t1));
  t1 += sizeof(boost::uint32_t);
  len2 = *(reinterpret_cast<const boost::uint32_t *>(t2));
  t2 += sizeof(boost::uint32_t);
  if (len1 != len2) {
    std::string errorMsg =
        "attempt to compare fingerprints of different length";
    sqlite3_result_error(context, errorMsg.c_str(), errorMsg.length());
    return;
  }

  boost::uint32_t nElem1, nElem2;
  nElem1 = *(reinterpret_cast<const boost::uint32_t *>(t1));
  t1 += sizeof(boost::uint32_t);
  nElem2 = *(reinterpret_cast<const boost::uint32_t *>(t2));
  t2 += sizeof(boost::uint32_t);

  if (!nElem1 || !nElem2) {
    res = 0.0;
    sqlite3_result_double(context, res);
  }

  double v1Sum = 0, v2Sum = 0, numer = 0;
  boost::uint32_t idx1 = 0;
  boost::int32_t v1;
  boost::uint32_t idx2 = 0;
  boost::int32_t v2;
  idx1 = *(reinterpret_cast<const boost::uint32_t *>(t1));
  t1 += sizeof(boost::uint32_t);
  v1 = *(reinterpret_cast<const boost::int32_t *>(t1));
  t1 += sizeof(boost::int32_t);
  nElem1--;
  v1Sum += v1;

  idx2 = *(reinterpret_cast<const boost::uint32_t *>(t2));
  t2 += sizeof(boost::uint32_t);
  v2 = *(reinterpret_cast<const boost::int32_t *>(t2));
  t2 += sizeof(boost::int32_t);
  nElem2--;
  v2Sum += v2;

  while (1) {
    while (nElem2 && idx2 < idx1) {
      idx2 = *(reinterpret_cast<const boost::uint32_t *>(t2));
      t2 += sizeof(boost::uint32_t);
      v2 = *(reinterpret_cast<const boost::int32_t *>(t2));
      t2 += sizeof(boost::int32_t);
      nElem2--;
      v2Sum += v2;
    }
    if (idx2 == idx1) {
      // std::cerr<<"   --- "<<idx1<<" "<<v1<<" - "<<idx2<<" "<<v2<<std::endl;
      numer += std::min(v1, v2);
    }
    if (nElem1) {
      idx1 = *(reinterpret_cast<const boost::uint32_t *>(t1));
      t1 += sizeof(boost::uint32_t);
      v1 = *(reinterpret_cast<const boost::int32_t *>(t1));
      t1 += sizeof(boost::int32_t);
      nElem1--;
      v1Sum += v1;
    } else {
      break;
    }
  }
  while (nElem2) {
    idx2 = *(reinterpret_cast<const boost::uint32_t *>(t2));
    t2 += sizeof(boost::uint32_t);
    v2 = *(reinterpret_cast<const boost::int32_t *>(t2));
    t2 += sizeof(boost::int32_t);
    nElem2--;
    v2Sum += v2;
  }
  double denom = v1Sum + v2Sum;
  if (fabs(denom) < 1e-6) {
    res = 0.0;
  } else {
    res = 2. * numer / denom;
  }
  // std::cerr<<" "<<v1Sum<<" "<<v2Sum<<" "<<numer<<" "<<res<<std::endl;
  sqlite3_result_double(context, res);
}
#endif

/* SQLite invokes this routine once when it loads the extension.
** Create new functions, collating sequences, and virtual table
** modules here.  This is usually the only exported symbol in
** the shared library.
*/
extern "C" int sqlite3_extension_init(sqlite3 *db, char **pzErrMsg,
                                      const sqlite3_api_routines *pApi) {
  SQLITE_EXTENSION_INIT2(pApi);
  std::map<std::string, boost::any> *molMap =
      new std::map<std::string, boost::any>();
  sqlite3_create_function(db, "rdk_molNumAtoms", 1, SQLITE_ANY, 0, numAtomsFunc,
                          0, 0);
  sqlite3_create_function(db, "rdk_molAMW", 1, SQLITE_ANY, 0, molWtFunc, 0, 0);
  sqlite3_create_function(db, "rdk_smilesToBlob", 1, SQLITE_ANY, 0,
                          smilesToBlob, 0, 0);
  sqlite3_create_function(db, "rdk_molToRDKitFP", 1, SQLITE_ANY, 0,
                          blobToRDKitFingerprint, 0, 0);
  sqlite3_create_function(db, "rdk_bvTanimotoSim", 2, SQLITE_ANY, 0,
                          bvTanimotoSim, 0, 0);
  sqlite3_create_function(db, "rdk_ucvTanimotoSim", 2, SQLITE_ANY, 0,
                          ucvTanimotoSim, 0, 0);
  sqlite3_create_function(db, "rdk_molToAtomPairFP", 1, SQLITE_ANY, 0,
                          blobToAtomPairFingerprint, 0, 0);
  sqlite3_create_function(db, "rdk_sivDiceSim", 2, SQLITE_ANY, 0, sivDiceSim, 0,
                          0);
  sqlite3_create_function(db, "rdk_sivDiceSim2", 2, SQLITE_ANY, 0, sivDiceSim2,
                          0, 0);
  sqlite3_create_function(db, "rdk_molHasSubstruct", 2, SQLITE_ANY,
                          static_cast<void *>(molMap), molHasSubstruct, 0, 0);
  sqlite3_create_function(db, "rdk_molSubstructCount", 2, SQLITE_ANY,
                          static_cast<void *>(molMap), molSubstructCount, 0, 0);
  sqlite3_create_function(db, "rdk_molLogP", 1, SQLITE_ANY, 0, molLogPFunc, 0,
                          0);
  return 0;
}