// $Id$
//
//  Copyright (C) 2006 Greg Landrum
//       All Rights Reserved
//
// The tests in this file demonstrate the use of RDKit with
// the boost iostreams library. Particularly the compression pieces
// of that library.
//

#include <RDGeneral/test.h>
#include <RDGeneral/RDLog.h>
#include <GraphMol/RDKitBase.h>

#include <GraphMol/FileParsers/FileParsers.h>
#include <GraphMol/SmilesParse/SmilesParse.h>
#include <GraphMol/Depictor/RDDepictor.h>
#include <GraphMol/SmilesParse/SmilesWrite.h>
#include <GraphMol/Substruct/SubstructMatch.h>

#include <fstream>
#include <iostream>
#ifdef SUPPORT_COMPRESSED_IO
#include <boost/iostreams/device/file.hpp>
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <boost/iostreams/filter/bzip2.hpp>
#include <boost/iostreams/operations.hpp>
#include <boost/range/iterator_range.hpp>
#endif
#include <string>

using namespace RDKit;
namespace io = boost::iostreams;

void test1() {
#ifdef SUPPORT_COMPRESSED_IO
  BOOST_LOG(rdInfoLog) << "testing basic reading from compressed streams"
                       << std::endl;

  std::string rdbase = getenv("RDBASE");
  std::string fname2 =
      rdbase + "/Code/Demos/RDKit/BinaryIO/test_data/triazine.mol.gz";
  RWMol *m;

  io::filtering_istream inStrm;
  inStrm.push(io::gzip_decompressor());
  inStrm.push(io::file_source(fname2));
  TEST_ASSERT(inStrm.is_complete());
  unsigned int lineNo = 0;
  m = MolDataStreamToMol(inStrm, lineNo);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(6));

  BOOST_LOG(rdInfoLog) << "done" << std::endl;
#else
  BOOST_LOG(rdInfoLog) << "Compressed IO disabled; test skipped" << std::endl;
#endif
}

void test2() {
#ifdef SUPPORT_COMPRESSED_IO
  BOOST_LOG(rdInfoLog)
      << "testing writing to, then reading from compressed streams"
      << std::endl;

  std::string smiles = "C1CCCC1";
  std::string buff, molBlock;
  RWMol *m;
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(5));
  m->setProp(common_properties::_Name, "monkey");
  io::filtering_ostream outStrm;
  outStrm.push(io::gzip_compressor());
  outStrm.push(io::back_inserter(buff));
  TEST_ASSERT(outStrm.is_complete());

  molBlock = MolToMolBlock(*m);
  outStrm << molBlock;
  outStrm.reset();

  delete m;
  io::filtering_istream inStrm;
  inStrm.push(io::gzip_decompressor());
  inStrm.push(boost::make_iterator_range(buff));
  TEST_ASSERT(inStrm.is_complete());

  unsigned int lineNo = 0;
  m = MolDataStreamToMol(inStrm, lineNo);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(5));

  BOOST_LOG(rdInfoLog) << "done" << std::endl;
#else
  BOOST_LOG(rdInfoLog) << "Compressed IO disabled; test skipped" << std::endl;
#endif
}

void test3() {
#ifdef SUPPORT_COMPRESSED_IO
  BOOST_LOG(rdInfoLog)
      << "testing writing pickles to a file then reading them back"
      << std::endl;
  std::string rdbase = getenv("RDBASE");
  std::string fname2 = rdbase + "/Code/Demos/RDKit/BinaryIO/test_data/mols.rdb";

  std::string smiles, buff;
  RWMol *m;

  std::vector<unsigned int> filePs;
  io::filtering_ostream outStrm;
  outStrm.push(
      io::file_sink(fname2, std::ios_base::out | std::ios_base::binary));
  TEST_ASSERT(outStrm.is_complete());

  smiles = "C1CCC1";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(4));

  MolPickler::pickleMol(*m, outStrm);
  delete m;

  smiles = "C1CCC1";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(4));
  RDDepict::compute2DCoords(*m);
  filePs.push_back(0);
  MolPickler::pickleMol(*m, outStrm);
  delete m;

  smiles = "C1CCCC1";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(5));
  RDDepict::compute2DCoords(*m);
  filePs.push_back(outStrm.tellp());
  MolPickler::pickleMol(*m, outStrm);
  delete m;

  smiles = "c1ccccc1";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(6));
  RDDepict::compute2DCoords(*m);
  filePs.push_back(outStrm.tellp());
  MolPickler::pickleMol(*m, outStrm);
  delete m;

  smiles = "c1ccccc1CC(=O)O";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(9));
  RDDepict::compute2DCoords(*m);
  filePs.push_back(outStrm.tellp());
  MolPickler::pickleMol(*m, outStrm);
  delete m;

  outStrm.flush();
  outStrm.reset();

  io::filtering_istream inStrm;
  inStrm.push(
      io::file_source(fname2, std::ios_base::in | std::ios_base::binary));
  TEST_ASSERT(inStrm.is_complete());

  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(inStrm, *m);
  TEST_ASSERT(m->getNumAtoms(4));
  delete m;

  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(inStrm, *m);
  TEST_ASSERT(m->getNumAtoms(4));
  delete m;

  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(inStrm, *m);
  TEST_ASSERT(m->getNumAtoms(5));
  delete m;

  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(inStrm, *m);
  TEST_ASSERT(m->getNumAtoms(6));
  delete m;

  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(inStrm, *m);
  TEST_ASSERT(m->getNumAtoms(9));
  delete m;

  BOOST_LOG(rdInfoLog) << "done" << std::endl;
#else
  BOOST_LOG(rdInfoLog) << "Compressed IO disabled; test skipped" << std::endl;
#endif
}

void test4() {
#ifdef SUPPORT_COMPRESSED_IO
  BOOST_LOG(rdInfoLog) << "testing writing pickles to a single compressed file "
                          "then reading them back"
                       << std::endl;
  std::string rdbase = getenv("RDBASE");
  std::string fname2 = rdbase + "/Code/Demos/RDKit/BinaryIO/test_data/mols.rdz";

  std::string smiles;
  RWMol *m;

  io::filtering_ostream outStrm;
  outStrm.push(io::gzip_compressor());
  outStrm.push(
      io::file_sink(fname2, std::ios_base::out | std::ios_base::binary));
  TEST_ASSERT(outStrm.is_complete());

  smiles = "C1CCC1";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(4));

  MolPickler::pickleMol(*m, outStrm);
  delete m;

  smiles = "C1CCC1";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(4));
  RDDepict::compute2DCoords(*m);

  MolPickler::pickleMol(*m, outStrm);
  delete m;

  smiles = "C1CCCC1";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(5));
  RDDepict::compute2DCoords(*m);
  MolPickler::pickleMol(*m, outStrm);
  delete m;

  smiles = "c1ccccc1";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(6));
  RDDepict::compute2DCoords(*m);
  MolPickler::pickleMol(*m, outStrm);
  delete m;

  smiles = "c1ccccc1CC(=O)O";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(9));
  RDDepict::compute2DCoords(*m);
  MolPickler::pickleMol(*m, outStrm);
  delete m;

  io::flush(outStrm);
  outStrm.pop();

  io::filtering_istream inStrm;
  inStrm.push(io::gzip_decompressor());
  inStrm.push(
      io::file_source(fname2, std::ios_base::in | std::ios_base::binary));
  TEST_ASSERT(inStrm.is_complete());

  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(inStrm, *m);
  TEST_ASSERT(m->getNumAtoms(4));
  delete m;

  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(inStrm, *m);
  TEST_ASSERT(m->getNumAtoms(4));
  delete m;

  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(inStrm, *m);
  TEST_ASSERT(m->getNumAtoms(5));
  delete m;

  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(inStrm, *m);
  TEST_ASSERT(m->getNumAtoms(6));
  delete m;

  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(inStrm, *m);
  TEST_ASSERT(m->getNumAtoms(9));
  delete m;

  BOOST_LOG(rdInfoLog) << "done" << std::endl;
#else
  BOOST_LOG(rdInfoLog) << "Compressed IO disabled; test skipped" << std::endl;
#endif
}

void test5() {
#ifdef SUPPORT_COMPRESSED_IO
  BOOST_LOG(rdInfoLog) << "testing writing compressed pickles to a single file "
                          "then reading them back"
                       << std::endl;
  std::string rdbase = getenv("RDBASE");
  std::string fname2 = rdbase + "/Code/Demos/RDKit/BinaryIO/test_data/tmp.rdz";

  std::string smiles, buff;
  RWMol *m;
  std::vector<unsigned int> filePs;

  io::filtering_ostream outStrm;
  outStrm.push(
      io::file_sink(fname2, std::ios_base::out | std::ios_base::binary));
  TEST_ASSERT(outStrm.is_complete());

  smiles = "C1CCC1";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(4));

  io::filtering_ostream *tmpStrm;
  tmpStrm = new io::filtering_ostream();
  tmpStrm->push(io::gzip_compressor());
  tmpStrm->push(io::back_inserter(buff));

  filePs.push_back(0);
  MolPickler::pickleMol(*m, *tmpStrm);
  delete m;
  tmpStrm->reset();

  outStrm << buff.size();
  outStrm << buff;
  std::cerr << "sz: " << buff.size() << " " << outStrm.tellp() << std::endl;
  buff = "";
  tmpStrm->push(io::gzip_compressor());
  tmpStrm->push(io::back_inserter(buff));

  smiles = "C1CCC1";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(4));
  RDDepict::compute2DCoords(*m);

  filePs.push_back(outStrm.tellp());
  MolPickler::pickleMol(*m, *tmpStrm);
  delete m;
  tmpStrm->reset();
  outStrm << buff.size();
  outStrm << buff;
  std::cerr << "sz: " << buff.size() << " " << outStrm.tellp() << std::endl;
  buff = "";
  tmpStrm->push(io::gzip_compressor());
  tmpStrm->push(io::back_inserter(buff));

  smiles = "C1CCCC1";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(5));
  RDDepict::compute2DCoords(*m);

  filePs.push_back(outStrm.tellp());
  MolPickler::pickleMol(*m, *tmpStrm);
  delete m;
  tmpStrm->reset();
  outStrm << buff.size();
  outStrm << buff;
  std::cerr << "sz: " << buff.size() << " " << outStrm.tellp() << std::endl;
  buff = "";
  tmpStrm->push(io::gzip_compressor());
  tmpStrm->push(io::back_inserter(buff));

  smiles = "c1ccccc1";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(6));
  RDDepict::compute2DCoords(*m);
  filePs.push_back(outStrm.tellp());
  MolPickler::pickleMol(*m, *tmpStrm);
  delete m;
  tmpStrm->reset();
  outStrm << buff.size();
  outStrm << buff;
  std::cerr << "sz: " << buff.size() << " " << outStrm.tellp() << std::endl;
  buff = "";
  tmpStrm->push(io::gzip_compressor());
  tmpStrm->push(io::back_inserter(buff));

  smiles = "c1ccccc1CC(=O)O";
  m = SmilesToMol(smiles);
  TEST_ASSERT(m);
  TEST_ASSERT(m->getNumAtoms(9));
  RDDepict::compute2DCoords(*m);
  filePs.push_back(outStrm.tellp());
  MolPickler::pickleMol(*m, *tmpStrm);
  delete m;
  tmpStrm->reset();
  outStrm << buff.size();
  outStrm << buff;
  std::cerr << "sz: " << buff.size() << " " << outStrm.tellp() << std::endl;

  delete tmpStrm;
  io::flush(outStrm);
  outStrm.pop();

  io::filtering_istream tmpIStrm;
  io::filtering_istream inStrm;
  unsigned int sz;
  char *charArr;

  inStrm.push(
      io::file_source(fname2, std::ios_base::in | std::ios_base::binary));
  TEST_ASSERT(inStrm.is_complete());

  inStrm >> sz;
  charArr = new char[sz];
  inStrm.read(charArr, sz);
  buff = "";
  buff.append(charArr, sz);
  tmpIStrm.push(io::gzip_decompressor());
  tmpIStrm.push(boost::make_iterator_range(buff));
  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(tmpIStrm, *m);
  TEST_ASSERT(m->getNumAtoms(4));
  delete m;

  inStrm >> sz;
  delete[] charArr;
  charArr = new char[sz];
  inStrm.read(charArr, sz);
  buff = "";
  buff.append(charArr, sz);
  tmpIStrm.reset();
  tmpIStrm.push(io::gzip_decompressor());
  tmpIStrm.push(boost::make_iterator_range(buff));
  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(tmpIStrm, *m);
  TEST_ASSERT(m->getNumAtoms(4));
  delete m;

  inStrm >> sz;
  delete[] charArr;
  charArr = new char[sz];
  inStrm.read(charArr, sz);
  buff = "";
  buff.append(charArr, sz);
  tmpIStrm.reset();
  tmpIStrm.push(io::gzip_decompressor());
  tmpIStrm.push(boost::make_iterator_range(buff));
  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(tmpIStrm, *m);
  TEST_ASSERT(m->getNumAtoms(5));
  delete m;

  inStrm >> sz;
  delete[] charArr;
  charArr = new char[sz];
  inStrm.read(charArr, sz);
  buff = "";
  buff.append(charArr, sz);
  tmpIStrm.reset();
  tmpIStrm.push(io::gzip_decompressor());
  tmpIStrm.push(boost::make_iterator_range(buff));
  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(tmpIStrm, *m);
  TEST_ASSERT(m->getNumAtoms(6));
  delete m;

  inStrm >> sz;
  delete[] charArr;
  charArr = new char[sz];
  inStrm.read(charArr, sz);
  buff = "";
  buff.append(charArr, sz);
  tmpIStrm.reset();
  tmpIStrm.push(io::gzip_decompressor());
  tmpIStrm.push(boost::make_iterator_range(buff));
  m = new RWMol();
  TEST_ASSERT(m);
  MolPickler::molFromPickle(tmpIStrm, *m);
  TEST_ASSERT(m->getNumAtoms(9));
  delete m;

  BOOST_LOG(rdInfoLog) << "done" << std::endl;
#else
  BOOST_LOG(rdInfoLog) << "Compressed IO disabled; test skipped" << std::endl;
#endif
}

int main(int argc, char *argv[]) {
  RDLog::InitLogs();
  test1();
  test2();
  test3();
  test4();
  test5();
  return 0;
}
