// $Id$
//
//  Copyright (C) 2008-2011 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.
//
//  Can be built with:
//   g++ -o sample.exe sample.cpp -I$RDBASE/Code -I$RDBASE/Extern \ 
//       -L$RDBASE/lib -lChemReactions -lFileParsers -lSmilesParse -lDepictor \ 
//       -lSubstructMatch -lGraphMol -lDataStructs -lRDGeometryLib -lRDGeneral
//

#include <RDGeneral/Invariant.h>
#include <GraphMol/RDKitBase.h>
#include <GraphMol/SmilesParse/SmilesParse.h>
#include <GraphMol/SmilesParse/SmilesWrite.h>
#include <GraphMol/Substruct/SubstructMatch.h>
#include <GraphMol/Depictor/RDDepictor.h>
#include <GraphMol/FileParsers/FileParsers.h>
#include <GraphMol/ChemReactions/Reaction.h>
#include <GraphMol/ChemReactions/ReactionParser.h>
#include <GraphMol/ChemReactions/ReactionPickler.h>

#include <RDGeneral/RDLog.h>
#include <vector>
#include <algorithm>

using namespace RDKit;

void BuildSimpleMolecule() {
  // build the molecule: C/C=C\C
  RWMol *mol = new RWMol();

  // add atoms and bonds:
  mol->addAtom(new Atom(6));         // atom 0
  mol->addAtom(new Atom(6));         // atom 1
  mol->addAtom(new Atom(6));         // atom 2
  mol->addAtom(new Atom(6));         // atom 3
  mol->addBond(0, 1, Bond::SINGLE);  // bond 0
  mol->addBond(1, 2, Bond::DOUBLE);  // bond 1
  mol->addBond(2, 3, Bond::SINGLE);  // bond 2
  // setup the stereochem:
  mol->getBondWithIdx(0)->setBondDir(Bond::ENDUPRIGHT);
  mol->getBondWithIdx(2)->setBondDir(Bond::ENDDOWNRIGHT);

  // do the chemistry perception:
  RDKit::MolOps::sanitizeMol(*mol);

  // Get the canonical SMILES, include stereochemistry:
  std::string smiles;
  smiles = MolToSmiles(*(static_cast<ROMol *>(mol)), true);
  BOOST_LOG(rdInfoLog) << " sample 1 SMILES: " << smiles << std::endl;
}

void WorkWithRingInfo() {
  // use a more complicated molecule to demonstrate querying about
  // ring information
  ROMol *mol = SmilesToMol("OC1CCC2C1CCCC2");
  // the molecule from SmilesToMol is already sanitized, so we don't
  // need to worry about that.

  // work with ring information
  RingInfo *ringInfo = mol->getRingInfo();

  TEST_ASSERT(ringInfo->numRings() == 2);

  // can ask how many rings an atom is in:
  TEST_ASSERT(ringInfo->numAtomRings(0) == 0);
  TEST_ASSERT(ringInfo->numAtomRings(1) == 1);
  TEST_ASSERT(ringInfo->numAtomRings(4) == 2);
  // same with bonds:
  TEST_ASSERT(ringInfo->numBondRings(0) == 0);
  TEST_ASSERT(ringInfo->numBondRings(1) == 1);

  // can check if an atom is in a ring of a particular size:
  TEST_ASSERT(!ringInfo->isAtomInRingOfSize(0, 5));
  TEST_ASSERT(ringInfo->isAtomInRingOfSize(1, 5));
  TEST_ASSERT(ringInfo->isAtomInRingOfSize(4, 5));
  TEST_ASSERT(ringInfo->isAtomInRingOfSize(4, 6));
  // same with bonds:
  TEST_ASSERT(!ringInfo->isBondInRingOfSize(0, 5));
  TEST_ASSERT(ringInfo->isBondInRingOfSize(1, 5));

  // can also get the full list of rings as atom indices:
  VECT_INT_VECT atomRings;  // VECT_INT_VECT is vector< vector<int> >
  atomRings = ringInfo->atomRings();
  TEST_ASSERT(atomRings.size() == 2);
  TEST_ASSERT(atomRings[0].size() == 5);
  TEST_ASSERT(atomRings[1].size() == 6);
  // this sort is just here for test/demo purposes:
  std::sort(atomRings[0].begin(), atomRings[0].end());
  TEST_ASSERT(atomRings[0][0] == 1);
  TEST_ASSERT(atomRings[0][1] == 2);
  TEST_ASSERT(atomRings[0][2] == 3);
  TEST_ASSERT(atomRings[0][3] == 4);
  TEST_ASSERT(atomRings[0][4] == 5);
  // same with bonds:
  VECT_INT_VECT bondRings;  // VECT_INT_VECT is vector< vector<int> >
  bondRings = ringInfo->bondRings();
  TEST_ASSERT(bondRings.size() == 2);
  TEST_ASSERT(bondRings[0].size() == 5);
  TEST_ASSERT(bondRings[1].size() == 6);
  // the same trick played above with the contents of each ring
  // can be played, but we won't

  // count the number of rings of size 5:
  unsigned int nRingsSize5 = 0;
  for (VECT_INT_VECT_CI ringIt = atomRings.begin(); ringIt != atomRings.end();
       ++ringIt) {
    if (ringIt->size() == 5) nRingsSize5++;
  }
  TEST_ASSERT(nRingsSize5 == 1);
  delete mol;

  // count the number of atoms in 5-rings where all the atoms
  // are aromatic:
  mol = SmilesToMol("C1CC2=C(C1)C1=C(NC3=C1C=CC=C3)C=C2");
  ringInfo = mol->getRingInfo();
  atomRings = ringInfo->atomRings();

  unsigned int nMatchingAtoms = 0;
  for (VECT_INT_VECT_CI ringIt = atomRings.begin(); ringIt != atomRings.end();
       ++ringIt) {
    if (ringIt->size() != 5) {
      continue;
    }
    bool isAromatic = true;
    for (INT_VECT_CI atomIt = ringIt->begin(); atomIt != ringIt->end();
         ++atomIt) {
      if (!mol->getAtomWithIdx(*atomIt)->getIsAromatic()) {
        isAromatic = false;
        break;
      }
    }
    if (isAromatic) {
      nMatchingAtoms += 5;
    }
  }
  TEST_ASSERT(nMatchingAtoms == 5);
  delete mol;

  // count the number of rings where all the bonds
  // are aromatic.
  mol = SmilesToMol("c1cccc2c1CCCC2");
  ringInfo = mol->getRingInfo();
  bondRings = ringInfo->bondRings();

  unsigned int nAromaticRings = 0;
  for (VECT_INT_VECT_CI ringIt = bondRings.begin(); ringIt != bondRings.end();
       ++ringIt) {
    bool isAromatic = true;
    for (INT_VECT_CI bondIt = ringIt->begin(); bondIt != ringIt->end();
         ++bondIt) {
      if (!mol->getBondWithIdx(*bondIt)->getIsAromatic()) {
        isAromatic = false;
        break;
      }
    }
    if (isAromatic) nAromaticRings++;
  }
  TEST_ASSERT(nAromaticRings == 1);
  delete mol;
}

void WorkWithSmarts() {
  // demonstrate the use of substructure searching
  ROMol *mol = SmilesToMol("ClCC=CCC");
  // a simple SMARTS pattern for rotatable bonds:
  ROMol *pattern = SmartsToMol("[!$(*#*)&!D1]-&!@[!$(*#*)&!D1]");

  std::vector<MatchVectType> matches;
  unsigned int nMatches;
  nMatches = SubstructMatch(*mol, *pattern, matches);
  TEST_ASSERT(nMatches == 2);
  TEST_ASSERT(matches.size() == 2);  // <- there are two rotatable bonds

  // a MatchVect is a vector of std::pairs with (patternIdx, molIdx):
  TEST_ASSERT(matches[0].size() == 2);
  TEST_ASSERT(matches[0][0].first == 0);
  TEST_ASSERT(matches[0][0].second == 1);
  TEST_ASSERT(matches[0][1].first == 1);
  TEST_ASSERT(matches[0][1].second == 2);

  delete pattern;
  delete mol;
}

void DepictDemo() {
  // demonstrate the use of the depiction-generation code2D coordinates:
  ROMol *mol = SmilesToMol("ClCC=CCC");

  // generate the 2D coordinates:
  RDDepict::compute2DCoords(*mol);

  // generate a mol block (could also go to a file):
  std::string molBlock = MolToMolBlock(*mol);
  BOOST_LOG(rdInfoLog) << molBlock;

  delete mol;
}

void CleanupMolecule() {
  // an example of doing some cleaning up of a molecule before
  // calling the sanitizeMol function()

  // build: C1CC1C(:O):O
  RWMol *mol = new RWMol();

  // add atoms and bonds:
  mol->addAtom(new Atom(6));           // atom 0
  mol->addAtom(new Atom(6));           // atom 1
  mol->addAtom(new Atom(6));           // atom 2
  mol->addAtom(new Atom(6));           // atom 3
  mol->addAtom(new Atom(8));           // atom 4
  mol->addAtom(new Atom(8));           // atom 5
  mol->addBond(3, 4, Bond::AROMATIC);  // bond 0
  mol->addBond(3, 5, Bond::AROMATIC);  // bond 1
  mol->addBond(3, 2, Bond::SINGLE);    // bond 2
  mol->addBond(2, 1, Bond::SINGLE);    // bond 3
  mol->addBond(1, 0, Bond::SINGLE);    // bond 4
  mol->addBond(0, 2, Bond::SINGLE);    // bond 5

  // instead of calling sanitize mol, which would generate an error,
  // we'll perceive the rings, then take care of aromatic bonds
  // that aren't in a ring, then sanitize:
  MolOps::findSSSR(*mol);
  for (ROMol::BondIterator bondIt = mol->beginBonds();
       bondIt != mol->endBonds(); ++bondIt) {
    if (((*bondIt)->getIsAromatic() ||
         (*bondIt)->getBondType() == Bond::AROMATIC) &&
        !mol->getRingInfo()->numBondRings((*bondIt)->getIdx())) {
      // remove the aromatic flag on the bond:
      (*bondIt)->setIsAromatic(false);
      // and cleanup its attached atoms as well (they were
      // also marked aromatic when the bond was added)
      (*bondIt)->getBeginAtom()->setIsAromatic(false);
      (*bondIt)->getEndAtom()->setIsAromatic(false);

      // NOTE: this isn't really reasonable:
      (*bondIt)->setBondType(Bond::SINGLE);
    }
  }

  // now it's safe to sanitize:
  RDKit::MolOps::sanitizeMol(*mol);

  // Get the canonical SMILES, include stereochemistry:
  std::string smiles;
  smiles = MolToSmiles(*(static_cast<ROMol *>(mol)), true);
  BOOST_LOG(rdInfoLog) << " fixed SMILES: " << smiles << std::endl;
}

void ReactionDemo() {
  // reaction smarts for a crude amide-bond formation definition:
  std::string sma = "[C:1](=[O:2])[OH].[N:3]>>[O:2]=[C:1][N:3]";
  // construct the reaction:
  ChemicalReaction *rxn = RxnSmartsToChemicalReaction(sma);
  // now initialize it and check for errors:
  rxn->initReactantMatchers();
  unsigned int nWarn, nError;
  rxn->validate(nWarn, nError);

  ROMol *mol;
  MOL_SPTR_VECT reacts;

  // build the list of reactants:
  ROMOL_SPTR react1(SmilesToMol("CC(=O)O"));
  ROMOL_SPTR react2(SmilesToMol("CCNCC1NC1"));
  reacts.push_back(react1);
  reacts.push_back(react2);

  // run the reaction, it returns a vector of vectors of product molecules:
  std::vector<MOL_SPTR_VECT> prods;
  prods = rxn->runReactants(reacts);

  // for each of the possible applications of the reaction to the reactants:
  for (unsigned int i = 0; i < prods.size(); ++i) {
    BOOST_LOG(rdInfoLog) << " product set: " << i << std::endl;
    // for each product of that application:
    for (unsigned int j = 0; j < prods[i].size(); ++j) {
      std::string psmiles = MolToSmiles(*prods[i][j], true);
      BOOST_LOG(rdInfoLog) << "   product : " << j << " " << psmiles
                           << std::endl;
    }
  }
}

int main(int argc, char *argv[]) {
  RDLog::InitLogs();
  BuildSimpleMolecule();
  WorkWithRingInfo();
  WorkWithSmarts();
  DepictDemo();
  CleanupMolecule();
  ReactionDemo();
}