//  Created by Guillaume GODIN
//  Copyright (C) 2012-2020 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 <RDGeneral/test.h>
#include <RDGeneral/Invariant.h>
#include <GraphMol/RDKitBase.h>
#include <GraphMol/GraphMol.h>
#include <GraphMol/FileParsers/MolSupplier.h>
#include <GraphMol/FileParsers/MolWriters.h>
#include <GraphMol/SmilesParse/SmilesParse.h>
#include <RDGeneral/RDLog.h>
#include <vector>
#include <algorithm>
#include <fstream>
#include <sstream>

#include <chrono>  // for high_resolution_clock

#include <GraphMol/Descriptors/AtomFeat.h>

using namespace RDKit;

void test1() {
  auto m = "CO"_smiles;

  TEST_ASSERT(m);

  std::vector<double> res;

  int atomid = 0;
  RDKit::Descriptors::AtomFeatVect(*m, res, atomid);

  std::vector<double> exp{0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                          0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.,
                          0., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0., 0.,
                          0., 0., 0., 0., 0., 0., 0., 1., 0., 0.5};

  // 49 features
  TEST_ASSERT(res.size() == 49);

  for (std::size_t i = 0; i < res.size(); i++) {
    TEST_ASSERT(fabs(res[i] - exp[i]) < 0.001);
    // std::cout << res[i] << "," ;
  }
}

void test2() {
  auto m = "C1CC=C1[NH3+]"_smiles;

  TEST_ASSERT(m);

  std::vector<double> res;

  std::vector<double> molatf{
      0, 1, 0, 0, 0, 0, 0,   0, 0, 0, 0, 0, 0, 0,   0, 0, 1, 0, 0, 0, 0,
      0, 0, 1, 0, 0, 0, 0,   1, 0, 0, 0, 0, 0, 1,   0, 0, 1, 0, 0, 0, 0,
      0, 0, 0, 1, 0, 0, 0.2, 0, 1, 0, 0, 0, 0, 0,   0, 0, 0, 0, 0, 0, 0,
      0, 0, 1, 0, 0, 0, 0,   0, 0, 1, 0, 0, 0, 0,   1, 0, 0, 0, 0, 0, 1,
      0, 0, 1, 0, 0, 0, 0,   0, 0, 0, 1, 0, 0, 0.2, 0, 1, 0, 0, 0, 0, 0,
      0, 0, 0, 0, 0, 0, 0,   0, 0, 1, 0, 0, 0, 0,   0, 1, 0, 0, 0, 0, 1,
      0, 0, 0, 0, 0, 0, 1,   0, 0, 1, 0, 0, 0, 0,   0, 0, 1, 0, 0, 0, 0.2,
      0, 1, 0, 0, 0, 0, 0,   0, 0, 0, 0, 0, 0, 0,   0, 0, 0, 1, 0, 0, 0,
      0, 1, 0, 0, 0, 1, 0,   0, 0, 0, 0, 0, 0, 1,   0, 0, 1, 0, 0, 0, 0,
      0, 1, 0, 0, 0, 0, 0.2, 0, 0, 1, 0, 0, 0, 0,   0, 0, 0, 0, 0, 0, 0,
      0, 1, 0, 0, 0, 0, 0,   0, 0, 1, 0, 0, 1, 0,   0, 0, 0, 0, 0, 0, 0,
      1, 0, 0, 0, 0, 0, 0,   0, 0, 0, 0, 1, 0, 0.2};

  for (unsigned int j = 0; j < m->getNumAtoms(); j++) {
    RDKit::Descriptors::AtomFeatVect(*m, res, j);

    // 49 features
    TEST_ASSERT(res.size() == 49);

    for (std::size_t i = 0; i < res.size(); i++) {
      TEST_ASSERT(fabs(res[i] - molatf[i + j * 49]) < 0.001);
    }

    res.clear();
    res.resize(49);
  }
}

void test3() {
  auto m = "O[C@H](F)[C@H](F)O"_smiles;

  TEST_ASSERT(m);

  std::vector<double> molatf{
      0, 0, 0, 1, 0, 0, 0,        0, 0, 0, 0, 0, 0, 0,
      0, 1, 0, 0, 0, 0, 0,        0, 0, 1, 0, 0, 0, 1,
      0, 0, 0, 0, 0, 0, 1,        0, 0, 0, 0, 0, 0, 0,
      0, 0, 1, 0, 0, 0, 0.166667, 0, 1, 0, 0, 0, 0, 0,
      0, 0, 0, 0, 0, 0, 0,        0, 0, 0, 1, 0, 0, 0,
      0, 0, 1, 0, 0, 1, 0,        0, 0, 0, 0, 0, 0, 1,
      0, 0, 0, 0, 0, 0, 0,        0, 0, 1, 0, 0, 0, 0.166667,
      0, 0, 0, 0, 0, 1, 0,        0, 0, 0, 0, 0, 0, 0,
      0, 1, 0, 0, 0, 0, 0,        0, 0, 1, 0, 0, 1, 0,
      0, 0, 0, 0, 0, 0, 1,        0, 0, 0, 0, 0, 0, 0,
      0, 1, 0, 0, 0, 0, 0.166667, 0, 1, 0, 0, 0, 0, 0,
      0, 0, 0, 0, 0, 0, 0,        0, 0, 0, 1, 0, 0, 0,
      0, 0, 1, 0, 0, 1, 0,        0, 0, 0, 0, 0, 0, 1,
      0, 0, 0, 0, 0, 0, 0,        0, 0, 1, 0, 0, 0, 0.166667,
      0, 0, 0, 0, 0, 1, 0,        0, 0, 0, 0, 0, 0, 0,
      0, 1, 0, 0, 0, 0, 0,        0, 0, 1, 0, 0, 1, 0,
      0, 0, 0, 0, 0, 0, 1,        0, 0, 0, 0, 0, 0, 0,
      0, 1, 0, 0, 0, 0, 0.166667, 0, 0, 0, 1, 0, 0, 0,
      0, 0, 0, 0, 0, 0, 0,        0, 1, 0, 0, 0, 0, 0,
      0, 0, 1, 0, 0, 0, 1,        0, 0, 0, 0, 0, 0, 1,
      0, 0, 0, 0, 0, 0, 0,        0, 0, 1, 0, 0, 0, 0.166667};

  std::vector<double> res;

  for (unsigned int j = 0; j < m->getNumAtoms(); j++) {
    RDKit::Descriptors::AtomFeatVect(*m, res, j);

    // 49 features
    TEST_ASSERT(res.size() == 49);

    for (std::size_t i = 0; i < res.size(); i++) {
      TEST_ASSERT(fabs(res[i] - molatf[i + j * 49]) < 0.001);
    }

    res.clear();
    res.resize(49);
  }
}

void test4() {
  auto m = "O[C@H](F)[C@H](F)O"_smiles;

  TEST_ASSERT(m);

  std::vector<double> res;

  std::vector<double> molatf{0, 0, 0, 1, 0, 0, 0, 0, 0, 0,        0, 0, 0,
                             0, 0, 1, 0, 0, 0, 0, 0, 0, 0,        1, 0, 0,
                             0, 1, 0, 0, 0, 0, 0, 0, 1, 0,        0, 0, 0,
                             0, 0, 0, 0, 0, 1, 0, 0, 0, 0.166667, 0, 0, 0,
                             0, 1, 0, 0, 0, 0, 0, 0, 0, 0,        0, 0, 0,
                             0, 0, 0, 0, 1, 0, 0, 0, 0, 0,        1, 0, 0,
                             1, 0, 0, 0, 0, 0, 0, 0, 1, 0,        0, 0, 0,
                             0, 0, 0, 0, 0, 1, 0, 0, 0, 0.166667, 0, 1, 0,
                             0, 0, 0, 0, 0, 1, 0, 0, 0, 0,        0, 0, 0,
                             0, 0, 1, 0, 0, 0, 0, 0, 0, 0,        1, 0, 0,
                             1, 0, 0, 0, 0, 0, 0, 0, 1, 0,        0, 0, 0,
                             0, 0, 0, 0, 1, 0, 0, 0, 0, 0.166667, 0, 0, 0,
                             0, 1, 0, 0, 0, 0, 0, 0, 0, 0,        0, 0, 0,
                             0, 0, 0, 0, 1, 0, 0, 0, 0, 0,        1, 0, 0,
                             1, 0, 0, 0, 0, 0, 0, 0, 1, 0,        0, 0, 0,
                             0, 0, 0, 0, 0, 1, 0, 0, 0, 0.166667, 0, 1, 0,
                             0, 0, 0, 0, 0, 1, 0, 0, 0, 0,        0, 0, 0,
                             0, 0, 1, 0, 0, 0, 0, 0, 0, 0,        1, 0, 0,
                             1, 0, 0, 0, 0, 0, 0, 0, 1, 0,        0, 0, 0,
                             0, 0, 0, 0, 1, 0, 0, 0, 0, 0.166667, 0, 0, 0,
                             0, 0, 0, 1, 0, 0, 0, 0, 0, 0,        0, 0, 0,
                             0, 0, 1, 0, 0, 0, 0, 0, 0, 0,        1, 0, 0,
                             0, 1, 0, 0, 0, 0, 0, 0, 1, 0,        0, 0, 0,
                             0, 0, 0, 0, 0, 1, 0, 0, 0, 0.166667, 0, 0, 0};

  for (unsigned int j = 0; j < m->getNumAtoms(); j++) {
    RDKit::Descriptors::AtomFeatVect(*m, res, j, true);

    // 52 features
    TEST_ASSERT(res.size() == 52);

    for (std::size_t i = 0; i < res.size(); i++) {
      TEST_ASSERT(fabs(res[i] - molatf[i + j * 52]) < 0.001);
    }

    res.clear();
    res.resize(52);
  }
}

void test5() {
  auto m = "O[C@@H](F)[C@@H](F)O"_smiles;

  TEST_ASSERT(m);

  std::vector<double> res;

  std::vector<double> molatf{0, 0, 0, 1, 0, 0, 0, 0, 0, 0,        0, 0, 0,
                             0, 0, 1, 0, 0, 0, 0, 0, 0, 0,        1, 0, 0,
                             0, 1, 0, 0, 0, 0, 0, 0, 1, 0,        0, 0, 0,
                             0, 0, 0, 0, 0, 1, 0, 0, 0, 0.166667, 0, 0, 0,
                             0, 1, 0, 0, 0, 0, 0, 0, 0, 0,        0, 0, 0,
                             0, 0, 0, 0, 1, 0, 0, 0, 0, 0,        1, 0, 0,
                             1, 0, 0, 0, 0, 0, 0, 0, 1, 0,        0, 0, 0,
                             0, 0, 0, 0, 0, 1, 0, 0, 0, 0.166667, 1, 0, 0,
                             0, 0, 0, 0, 0, 1, 0, 0, 0, 0,        0, 0, 0,
                             0, 0, 1, 0, 0, 0, 0, 0, 0, 0,        1, 0, 0,
                             1, 0, 0, 0, 0, 0, 0, 0, 1, 0,        0, 0, 0,
                             0, 0, 0, 0, 1, 0, 0, 0, 0, 0.166667, 0, 0, 0,
                             0, 1, 0, 0, 0, 0, 0, 0, 0, 0,        0, 0, 0,
                             0, 0, 0, 0, 1, 0, 0, 0, 0, 0,        1, 0, 0,
                             1, 0, 0, 0, 0, 0, 0, 0, 1, 0,        0, 0, 0,
                             0, 0, 0, 0, 0, 1, 0, 0, 0, 0.166667, 1, 0, 0,
                             0, 0, 0, 0, 0, 1, 0, 0, 0, 0,        0, 0, 0,
                             0, 0, 1, 0, 0, 0, 0, 0, 0, 0,        1, 0, 0,
                             1, 0, 0, 0, 0, 0, 0, 0, 1, 0,        0, 0, 0,
                             0, 0, 0, 0, 1, 0, 0, 0, 0, 0.166667, 0, 0, 0,
                             0, 0, 0, 1, 0, 0, 0, 0, 0, 0,        0, 0, 0,
                             0, 0, 1, 0, 0, 0, 0, 0, 0, 0,        1, 0, 0,
                             0, 1, 0, 0, 0, 0, 0, 0, 1, 0,        0, 0, 0,
                             0, 0, 0, 0, 0, 1, 0, 0, 0, 0.166667, 0, 0, 0};

  for (unsigned int j = 0; j < m->getNumAtoms(); j++) {
    RDKit::Descriptors::AtomFeatVect(*m, res, j, true);

    // 52 features
    TEST_ASSERT(res.size() == 52);

    for (std::size_t i = 0; i < res.size(); i++) {
      TEST_ASSERT(fabs(res[i] - molatf[i + j * 52]) < 0.001);
    }

    res.clear();
    res.resize(52);
  }
}

int main() {
  RDLog::InitLogs();
  test1();
  test2();
  test3();
  test4();
  test5();
}