//
//  2019, Daniel Probst, Reymond Group @ University of Bern
//
//   @@ 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 <boost/python.hpp>
#include <RDBoost/Wrap.h>
#include <GraphMol/Fingerprints/MHFP.h>
#include <vector>

namespace python = boost::python;
using RDKit::MHFPFingerprints::MHFPEncoder;


namespace RDKit {
namespace MHFPWrapper {

typedef std::vector<std::vector<uint32_t>> VectMinHashVect;
typedef std::vector<ExplicitBitVect> VectExplicitBitVect;

template<typename T>
std::vector<T> ListToVector(const python::object& obj)
{
  return std::vector<T>(python::stl_input_iterator<T>(obj),
              python::stl_input_iterator<T>());
}

UINT_VECT
FromStringArray(MHFPEncoder* mhfpEnc, python::list& vec) {
  std::vector<std::string> vec_tmp = ListToVector<std::string>(vec);
  return mhfpEnc->FromStringArray(vec_tmp);
}

UINT_VECT
FromArray(MHFPEncoder* mhfpEnc, python::list& vec) {
  std::vector<uint32_t> vec_tmp = ListToVector<uint32_t>(vec);
  return mhfpEnc->FromArray(vec_tmp);
}

STR_VECT
CreateShinglingFromSmiles(MHFPEncoder* mhfpEnc, 
              std::string smiles, 
              unsigned char radius = 3,
              bool rings = true,
              bool isomeric = false,
              bool kekulize = true,
              unsigned char min_radius = 1) {
  return mhfpEnc->CreateShingling(smiles, radius, rings, isomeric, kekulize, min_radius);
}

STR_VECT
CreateShinglingFromMol(MHFPEncoder* mhfpEnc, 
             ROMol mol, 
             unsigned char radius = 3,
             bool rings = true,
             bool isomeric = false,
             bool kekulize = true,
             unsigned char min_radius = 1) {
  return mhfpEnc->CreateShingling(mol, radius, rings, isomeric, kekulize, min_radius);
}

UINT_VECT
EncodeSmiles(MHFPEncoder* mhfpEnc,
       std::string smiles,
       unsigned char radius = 3,
       bool rings = true,
       bool isomeric = false,
       bool kekulize = true,
       unsigned char min_radius = 1) {
  return mhfpEnc->Encode(smiles, radius, rings, isomeric, kekulize, min_radius);
}

UINT_VECT
EncodeMol(MHFPEncoder* mhfpEnc,
      ROMol mol,
      unsigned char radius = 3,
      bool rings = true,
      bool isomeric = false,
      bool kekulize = true,
      unsigned char min_radius = 1) {
  return mhfpEnc->Encode(mol, radius, rings, isomeric, kekulize, min_radius);
}

VectMinHashVect
EncodeSmilesBulk(MHFPEncoder* mhfpEnc,
         python::list& smiles,
         unsigned char radius = 3,
         bool rings = true,
         bool isomeric = false,
         bool kekulize = true,
         unsigned char min_radius = 1) {
  std::vector<std::string> vec = ListToVector<std::string>(smiles);
  return mhfpEnc->Encode(vec, radius, rings, isomeric, kekulize, min_radius);
}


// There are access problems for vector_indexing_suite and std::vector<ROMol>.
// So let's fallback to a unefficient Python list. 
VectMinHashVect
EncodeMolsBulk(MHFPEncoder* mhfpEnc,
         python::list& mols,
         unsigned char radius = 3,
         bool rings = true,
         bool isomeric = false,
         bool kekulize = true,
         unsigned char min_radius = 1) {
  std::vector<ROMol> vec = ListToVector<ROMol>(mols);
  return mhfpEnc->Encode(vec, radius, rings, isomeric, kekulize, min_radius);
}

// SECFP

ExplicitBitVect
EncodeSECFPSmiles(MHFPEncoder* mhfpEnc,
          std::string smiles,
          unsigned char radius = 3,
          bool rings = true,
          bool isomeric = false,
          bool kekulize = true,
          unsigned char min_radius = 1,
          size_t length = 2048) {
  return mhfpEnc->EncodeSECFP(smiles, radius, rings, isomeric, kekulize, min_radius, length);
}

ExplicitBitVect
EncodeSECFPMol(MHFPEncoder* mhfpEnc,
         ROMol mol,
         unsigned char radius = 3,
         bool rings = true,
         bool isomeric = false,
         bool kekulize = true,
         unsigned char min_radius = 1,
         size_t length = 2048) {
  return mhfpEnc->EncodeSECFP(mol, radius, rings, isomeric, kekulize, min_radius, length);
}

VectExplicitBitVect
EncodeSECFPSmilesBulk(MHFPEncoder* mhfpEnc,
            python::list& smiles,
            unsigned char radius = 3,
            bool rings = true,
            bool isomeric = false,
            bool kekulize = true,
            unsigned char min_radius = 1,
            size_t length = 2048) {
  std::vector<std::string> vec = ListToVector<std::string>(smiles);
  return mhfpEnc->EncodeSECFP(vec, radius, rings, isomeric, kekulize, min_radius, length);
}

VectExplicitBitVect
EncodeSECFPMolsBulk(MHFPEncoder* mhfpEnc,
          python::list& mols,
          unsigned char radius = 3,
          bool rings = true,
          bool isomeric = false,
          bool kekulize = true,
          unsigned char min_radius = 1,
          size_t length = 2048) {
  std::vector<ROMol> vec = ListToVector<ROMol>(mols);
  return mhfpEnc->EncodeSECFP(vec, radius, rings, isomeric, kekulize, min_radius, length);
}

BOOST_PYTHON_FUNCTION_OVERLOADS(CreateShinglingFromSmilesOverloads, CreateShinglingFromSmiles, 2, 7)
BOOST_PYTHON_FUNCTION_OVERLOADS(CreateShinglingFromMolOverloads, CreateShinglingFromMol, 2, 7)

BOOST_PYTHON_FUNCTION_OVERLOADS(EncodeSmilesOverloads, EncodeSmiles, 2, 7)
BOOST_PYTHON_FUNCTION_OVERLOADS(EncodeMolOverloads, EncodeMol, 2, 7)
BOOST_PYTHON_FUNCTION_OVERLOADS(EncodeSmilesBulkOverloads, EncodeSmilesBulk, 2, 7)
BOOST_PYTHON_FUNCTION_OVERLOADS(EncodeMolsBulkOverloads, EncodeMolsBulk, 2, 7)

BOOST_PYTHON_FUNCTION_OVERLOADS(EncodeSECFPSmilesOverloads, EncodeSECFPSmiles, 2, 8)
BOOST_PYTHON_FUNCTION_OVERLOADS(EncodeSECFPMolOverloads, EncodeSECFPMol, 2, 8)
BOOST_PYTHON_FUNCTION_OVERLOADS(EncodeSECFPSmilesBulkOverloads, EncodeSECFPSmilesBulk, 2, 8)
BOOST_PYTHON_FUNCTION_OVERLOADS(EncodeSECFPMolsBulkOverloads, EncodeSECFPMolsBulk, 2, 8)

BOOST_PYTHON_MODULE(rdMHFPFingerprint) {
  python::class_<MHFPEncoder>("MHFPEncoder", python::init<python::optional<unsigned int, unsigned int>>())
    .def("FromStringArray", &FromStringArray, (
      (
        python::arg("vec")
      ),
      "Creates a MHFP vector from a list of arbitrary strings."
     ))
     .def("FromArray", &FromArray, (
      (
        python::arg("vec")
      ),
      "Creates a MHFP vector from a list of unsigned integers."
     ))
    .def("CreateShinglingFromSmiles", CreateShinglingFromSmiles, CreateShinglingFromSmilesOverloads(
      (
        python::arg("smiles"),
        python::arg("radius") = 3,
        python::arg("rings") = true,
        python::arg("isomeric") = false,
        python::arg("kekulize") = false,
        python::arg("min_radius") = 1
      ),
      "Creates a shingling (a list of circular n-grams / substructures) from a SMILES string."
     ))
    .def("CreateShinglingFromMol", CreateShinglingFromMol, CreateShinglingFromMolOverloads(
      (
        python::arg("mol"),
        python::arg("radius") = 3,
        python::arg("rings") = true,
        python::arg("isomeric") = false,
        python::arg("kekulize") = false,
        python::arg("min_radius") = 1
      ),
      "Creates a shingling (a list of circular n-grams / substructures) from a RDKit Mol instance."
     ))
    .def("EncodeSmiles", EncodeSmiles, EncodeSmilesOverloads(
      (
        python::arg("smiles"),
        python::arg("radius") = 3,
        python::arg("rings") = true,
        python::arg("isomeric") = false,
        python::arg("kekulize") = false,
        python::arg("min_radius") = 1
      ),
      "Creates a MHFP vector from a SMILES string."
     ))
    .def("EncodeMol", EncodeMol, EncodeMolOverloads(
      (
        python::arg("mol"),
        python::arg("radius") = 3,
        python::arg("rings") = true,
        python::arg("isomeric") = false,
        python::arg("kekulize") = false,
        python::arg("min_radius") = 1
      ),
      "Creates a MHFP vector from an RDKit Mol instance."
     ))
    .def("EncodeSmilesBulk", EncodeSmilesBulk, EncodeSmilesBulkOverloads(
      (
        python::arg("smiles"),
        python::arg("radius") = 3,
        python::arg("rings") = true,
        python::arg("isomeric") = false,
        python::arg("kekulize") = false,
        python::arg("min_radius") = 1
      ),
      "Creates a MHFP vector from a list of SMILES strings."
     ))
    .def("EncodeMolsBulk", EncodeMolsBulk, EncodeMolsBulkOverloads(
      (
        python::arg("mols"),
        python::arg("radius") = 3,
        python::arg("rings") = true,
        python::arg("isomeric") = false,
        python::arg("kekulize") = false,
        python::arg("min_radius") = 1
      ),
      "Creates a MHFP vector from a list of RDKit Mol instances."
     ))
    .def("EncodeSECFPSmiles", EncodeSECFPSmiles, EncodeSECFPSmilesOverloads(
      (
        python::arg("smiles"),
        python::arg("radius") = 3,
        python::arg("rings") = true,
        python::arg("isomeric") = false,
        python::arg("kekulize") = false,
        python::arg("min_radius") = 1,
        python::arg("length") = 2048
      ),
      "Creates a SECFP binary vector from a SMILES string."
     ))
    .def("EncodeSECFPMol", EncodeSECFPMol, EncodeSECFPMolOverloads(
      (
        python::arg("smiles"),
        python::arg("radius") = 3,
        python::arg("rings") = true,
        python::arg("isomeric") = false,
        python::arg("kekulize") = false,
        python::arg("min_radius") = 1,
        python::arg("length") = 2048
      ),
      "Creates a SECFP binary vector from an RDKit Mol instance."
     ))
    .def("EncodeSECFPSmilesBulk", EncodeSECFPSmilesBulk, EncodeSECFPSmilesBulkOverloads(
      (
        python::arg("smiles"),
        python::arg("radius") = 3,
        python::arg("rings") = true,
        python::arg("isomeric") = false,
        python::arg("kekulize") = false,
        python::arg("min_radius") = 1,
        python::arg("length") = 2048
      ),
      "Creates a SECFP binary vector from a list of SMILES strings."
     ))
    .def("EncodeSECFPMolsBulk", EncodeSECFPMolsBulk, EncodeSECFPMolsBulkOverloads(
      (
        python::arg("smiles"),
        python::arg("radius") = 3,
        python::arg("rings") = true,
        python::arg("isomeric") = false,
        python::arg("kekulize") = false,
        python::arg("min_radius") = 1,
        python::arg("length") = 2048
      ),
      "Creates a SECFP binary vector from a list of RDKit Mol instances."
     ))
    .def("Distance", &MHFPEncoder::Distance)
    .staticmethod("Distance");
}



}  // namespace MHFPWrapper
}  // namespace RDKit
