//
// file RDKitSV.cc
// David Cosgrove
// AstraZeneca
//
// 20th June 2014
//

#include "RDKitSVMainWindow.H"
#include "RDKitSVPanel.H"
#include "RDKitSVSettings.H"
#include "QTGet2Strings.H"
#include "QT4SelectItems.H"

#include <QAction>
#include <QApplication>
#include <QFileDialog>
#include <QFileInfo>
#include <QHBoxLayout>
#include <QMenu>
#include <QMenuBar>
#include <QMessageBox>
#include <QStatusBar>
#include <QString>

#include <boost/algorithm/string.hpp>
#include <boost/algorithm/string/split.hpp>
#include <boost/bind.hpp>
#include <boost/foreach.hpp>
#include <boost/iostreams/device/file.hpp>
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <boost/iostreams/stream.hpp>
#include <boost/lexical_cast.hpp>

#include <fstream>
#include <list>

#include <GraphMol/FileParsers/MolSupplier.h>
#include <GraphMol/FileParsers/MolWriters.h>
#include <GraphMol/SmilesParse/SmilesParse.h>
#include <GraphMol/Substruct/SubstructMatch.h>

using namespace boost;
using namespace std;

namespace RDKitSV {

// ****************************************************************************
RDKitSVMainWindow::RDKitSVMainWindow(int argc, char **argv) : QMainWindow() {
  build_actions();
  build_menubar();
  build_widget();

  parse_args(argc, argv);
}

// ****************************************************************************
void RDKitSVMainWindow::build_actions() {
  build_file_actions();
  build_smarts_actions();
}

// ****************************************************************************
void RDKitSVMainWindow::build_file_actions() {
  file_exit_ = new QAction(tr("Exit"), this);
  file_exit_->setShortcut(QString("Ctrl+Q"));
  connect(file_exit_, &QAction::triggered, this, &RDKitSVMainWindow::slot_exit);

  file_read_mols_ = new QAction(tr("Molecules"), this);
  connect(file_read_mols_, &QAction::triggered, this,
          &RDKitSVMainWindow::slot_read_mols);

  file_read_smarts_ = new QAction(tr("SMARTS"), this);
  connect(file_read_smarts_, &QAction::triggered, this,
          &RDKitSVMainWindow::slot_read_smarts);

  file_write_left_ = new QAction(tr("Molecules (Left Panel)"), this);
  connect(file_write_left_, &QAction::triggered, this,
          &RDKitSVMainWindow::slot_write_left_molecules);

  file_write_right_ = new QAction(tr("Molecules (Right Panel)"), this);
  connect(file_write_right_, &QAction::triggered, this,
          &RDKitSVMainWindow::slot_write_right_molecules);

  file_write_smarts_ = new QAction(tr("SMARTS"), this);
  connect(file_write_smarts_, &QAction::triggered, this,
          &RDKitSVMainWindow::slot_write_smarts);
}

// ****************************************************************************
void RDKitSVMainWindow::build_smarts_actions() {
  smarts_match_ = new QAction(tr("Match"), this);
  smarts_match_->setShortcut(QString("Ctrl+M"));
  connect(smarts_match_, &QAction::triggered, this,
          &RDKitSVMainWindow::slot_match_smarts);

  smarts_edit_ = new QAction(tr("Edit"), this);
  connect(smarts_edit_, &QAction::triggered, this,
          &RDKitSVMainWindow::slot_edit_smarts);

  smarts_new_ = new QAction(tr("New"), this);
  connect(smarts_new_, &QAction::triggered, this,
          &RDKitSVMainWindow::slot_new_smarts);
}

// ****************************************************************************
void RDKitSVMainWindow::build_menubar() {
  QMenu *file_menu = menuBar()->addMenu(tr("File"));
  QMenu *read_menu = file_menu->addMenu(tr("Read"));
  read_menu->addAction(file_read_mols_);
  read_menu->addAction(file_read_smarts_);

  QMenu *write_menu = file_menu->addMenu(tr("Write"));
  write_menu->addAction(file_write_left_);
  write_menu->addAction(file_write_right_);
  write_menu->addAction(file_write_smarts_);

  file_menu->addSeparator();
  file_menu->addAction(file_exit_);

  QMenu *smarts_menu = menuBar()->addMenu(tr("SMARTS"));
  smarts_menu->addAction(smarts_match_);
  smarts_menu->addAction(smarts_edit_);
  smarts_menu->addAction(smarts_new_);
}

// ****************************************************************************
void RDKitSVMainWindow::build_widget() {
  left_panel_ = new RDKitSVPanel;
  right_panel_ = new RDKitSVPanel(false);

  QWidget *cen_wid = new QWidget;
  QHBoxLayout *hbox = new QHBoxLayout;
  hbox->addWidget(left_panel_);
  hbox->addWidget(right_panel_);
  right_panel_->hide();

  cen_wid->setLayout(hbox);
  setCentralWidget(cen_wid);
}

// ****************************************************************************
void RDKitSVMainWindow::parse_args(int argc, char **argv) {
  RDKitSVSettings s(argc, argv);

  vector<string> mol_files = s.mol_files();
  if (!mol_files.empty()) {
    BOOST_FOREACH (string mf, mol_files) {
      read_mols(mf);
    }
  }

  string smarts_file = s.smarts_file();
  if (!smarts_file.empty()) {
    read_smarts(smarts_file);
  }
}

// ****************************************************************************
RDKitSVMainWindow::FILE_TYPE RDKitSVMainWindow::get_filetype(
    const string &filename, bool &is_compressed) {
  is_compressed = false;
  if (filename.substr(filename.length() - 3) == string(".gz")) {
    is_compressed = true;
  }

  if (string(".smi") == filename.substr(filename.length() - 4) ||
      string(".smi.gz") == filename.substr(filename.length() - 7) ||
      string(".ism") == filename.substr(filename.length() - 4) ||
      string(".ism.gz") == filename.substr(filename.length() - 7)) {
    return SMILES;
  } else if (string(".sdf") == filename.substr(filename.length() - 4) ||
             string(".sdf.gz") == filename.substr(filename.length() - 7)) {
    return SDF;
  }

  return UNKNOWN;
}

// ****************************************************************************
void RDKitSVMainWindow::read_mols(const string &filename) {
  QApplication::setOverrideCursor(Qt::WaitCursor);
  bool is_compressed;
  FILE_TYPE file_type = get_filetype(filename, is_compressed);

  boost::iostreams::filtering_istream ins;
  if (is_compressed) {
    ins.push(boost::iostreams::gzip_decompressor());
  }
  ins.push(boost::iostreams::file_source(filename.c_str()));
  if (!ins || !ins.good()) {
    QMessageBox::warning(this, tr("File Read Error"),
                         QString("%1 %2 %3 for reading.")
                             .arg(tr("Couldn't open"))
                             .arg(filename.c_str())
                             .arg(tr("for reading.")));
    return;
  }

  vector<RDKit::ROMOL_SPTR> new_mols;
  if (SDF == file_type) {
    RDKit::MolSupplier *mol_supplier =
        new RDKit::ForwardSDMolSupplier(&ins, true);
    while (!mol_supplier->atEnd()) {
      new_mols.push_back(RDKit::ROMOL_SPTR(mol_supplier->next()));
      if (!new_mols.back()) {
        new_mols.pop_back();
        break;
      }
    }
  } else if (SMILES == file_type) {
    // SmilesMolSupplier doesn't seem to work with compressed streams. It might
    // be because of the use of getline.
    // For now, assume the file is a simple one with no header and just SMILES
    // followed by optional name
    istreambuf_iterator<char> isb(ins), eos;

    vector<char> nextline;

    while (isb != eos) {
      nextline.clear();
      while (isb != eos && *isb != '\n') {
        nextline.push_back(*isb);
        ++isb;
      }
      if ('#' != nextline[0]) {
        vector<string> split_line;
        split(split_line, nextline, is_any_of(" \t"));
        string smi(split_line[0]);
        string smi_name;
        if (1 == split_line.size()) {
          smi_name = string("MOL_") + lexical_cast<string>(mols_.size() + 1);
        } else {
          smi_name = string(nextline.begin() + split_line[0].length() + 1,
                            nextline.end());
        }
        boost::trim(smi);
        boost::trim(smi_name);
        RDKit::ROMol *mol = RDKit::SmilesToMol(smi);
        if (mol) {
          mol->setProp("_Name", smi_name);
          mols_.push_back(RDKit::ROMOL_SPTR(mol));
        }
      }
      ++isb;  // get past '\n'
    }
  } else {
    QMessageBox::warning(this, tr("File Read Error"),
                         QString("%1 %2.")
                             .arg(tr("Unrecognised file type for"))
                             .arg(filename.c_str()));
    return;
  }

  mols_.insert(mols_.end(), new_mols.begin(), new_mols.end());
  left_panel_->set_molecules(mols_);

  QApplication::restoreOverrideCursor();
}

// ****************************************************************************
void RDKitSVMainWindow::read_smarts(const string &filename) {
#ifdef NOTYET
  cout << "reading SMARTS from " << filename << endl;
#endif
  ifstream ifs(filename.c_str());
  if (!ifs || !ifs.good()) {
    QMessageBox::warning(this, tr("File Read Error"),
                         QString("%1 %2 %3 for reading.")
                             .arg(tr("Couldn't open"))
                             .arg(filename.c_str())
                             .arg(tr("for reading.")));
    return;
  }

  while (1) {
    string next_line;
    getline(ifs, next_line);
#ifdef NOTYET
    cout << next_line << "XX" << endl;
#endif
    if (!ifs.good() || ifs.eof()) {
      break;
    }
    boost::trim(next_line);
    if (next_line.empty()) {
      continue;
    }
    if ('#' == next_line[0]) {
      continue;  // it's a comment
    }
    list<string> splits;
    boost::split(splits, next_line, boost::is_any_of(" \t"));
    string smt, smt_name;
    if (splits.size() >= 2) {
      smt = splits.front();
      splits.pop_front();
      boost::trim(smt);
      smt_name = boost::join(splits, " ");
      boost::trim(smt_name);
#ifdef NOTYET
      cout << "smt = " << smt << " name = " << smt_name << endl;
#endif
    } else if (1 == splits.size()) {
      smt = splits.front();
      boost::trim(smt);
      smt_name = string("SMARTS_") + lexical_cast<string>(smarts_.size() + 1);
    }
    smarts_.push_back(make_pair(smt_name, smt));
  }
  statusBar()->showMessage(
      QString("%1 %2.").arg(tr("Number of SMARTS now")).arg(smarts_.size()),
      2000);
}

// ****************************************************************************
void RDKitSVMainWindow::write_mols(RDKitSVPanel &panel,
                                   const string &filename) {
  vector<RDKit::ROMOL_SPTR> mols = panel.get_molecules();
  if (mols.empty()) {
    QMessageBox::information(
        this, tr("No molecules."),
        tr("Panel contains no molecules, so nothing to write."));
    return;
  }

  bool is_compressed;
  FILE_TYPE file_type = get_filetype(filename, is_compressed);

  boost::iostreams::filtering_ostream ons;
  if (is_compressed) {
    ons.push(boost::iostreams::gzip_compressor());
  }
  ons.push(boost::iostreams::file_sink(filename.c_str()));
  if (!ons || !ons.good()) {
    QMessageBox::warning(
        this, tr("Can't write to file"),
        QString("%1 %2.").arg(
            tr("Couldn't write to file").arg(filename.c_str())));
    return;
  }

  RDKit::MolWriter *mw;
  if (SMILES == file_type) {
    bool ism = filename.substr(filename.length() - 4) == ".ism" ||
                       filename.substr(filename.length() - 7) == ".ism.gz"
                   ? true
                   : false;
    mw = new RDKit::SmilesWriter(&ons, " ", "Name", true, false, ism);
  } else if (SDF == file_type) {
    mw = new RDKit::SDWriter(&ons, false);
  } else {
    QMessageBox::warning(this, tr("File Write Error"),
                         QString("%1 %2.")
                             .arg(tr("Unrecognised file type for"))
                             .arg(filename.c_str()));
    return;
  }

  BOOST_FOREACH (RDKit::ROMOL_SPTR mol, mols) {
    mw->write(*mol);
  }

  delete mw;
}

// ****************************************************************************
void RDKitSVMainWindow::match_smarts(const vector<pair<string, string>> &smts) {
  cout << "Matching " << smts.size() << " SMARTS" << endl;
  vector<vector<int>> smarts_match(mols_.size(), vector<int>());

  QString smts_label;
  for (int i = 0, is = smts.size(); i < is; ++i) {
    cout << "Matching " << i << " : " << smts[i].second << endl;
    smts_label += QString("%1 (%2)")
                      .arg(smts[i].first.c_str())
                      .arg(smts[i].second.c_str());
    if (i < is - 1) {
      smts_label += "|";
    }
    RDKit::RWMol *q = RDKit::SmartsToMol(smts[i].second);
    for (int j = 0, js = mols_.size(); j < js; ++j) {
      vector<RDKit::MatchVectType> hits_vect;
      if (RDKit::SubstructMatch(*mols_[j], *q, hits_vect)) {
        // store the numbers of the atoms that matched
        typedef pair<int, int> INTPAIR;
        BOOST_FOREACH (RDKit::MatchVectType hit, hits_vect) {
          BOOST_FOREACH (INTPAIR ip, hit) {
            smarts_match[j].push_back(ip.second);
          }
        }
        std::sort(smarts_match[j].begin(), smarts_match[j].end());
        smarts_match[j].erase(
            unique(smarts_match[j].begin(), smarts_match[j].end()),
            smarts_match[j].end());
      }
    }
    delete q;
  }

  vector<RDKit::ROMOL_SPTR> match_mols, miss_mols;
  vector<vector<int>> match_atom_matches;
  for (int i = 0, is = mols_.size(); i < is; ++i) {
    if (smarts_match[i].empty()) {
      miss_mols.push_back(mols_[i]);
    } else {
      match_mols.push_back(mols_[i]);
      match_atom_matches.push_back(smarts_match[i]);
    }
  }

  right_panel_->show();
  left_panel_->set_molecules(match_mols, match_atom_matches);
  // Details, details!!
  if (1 == match_mols.size()) {
    left_panel_->set_label(
        QString("%1 %2.").arg(tr("1 molecule matched.").arg(smts_label)));
  } else {
    left_panel_->set_label(QString("%1 %2 %3.")
                               .arg(match_mols.size())
                               .arg(tr("molecules matched"))
                               .arg(smts_label));
  }
  right_panel_->set_molecules(miss_mols);
  if (1 == miss_mols.size()) {
    right_panel_->set_label(tr("1 molecule didn't match"));
  } else {
    right_panel_->set_label(QString("%1 %2")
                                .arg(miss_mols.size())
                                .arg(tr("molecules didn't match.")));
  }
}

// ****************************************************************************
vector<pair<string, string>> RDKitSVMainWindow::select_smarts(
    bool multi_select) {
  vector<pair<string, string>> smts_to_use;

  if (smarts_.empty()) {
    return smts_to_use;
  }

  vector<char> selected_smarts(smarts_.size(), 0);
  vector<QString> smts;
  for (int i = 0, is = smarts_.size(); i < is; ++i) {
#ifdef NOTYET
    cout << i << " : " << smarts_[i].first << " :: " << smarts_[i].second
         << endl;
#endif
    smts.push_back(QString("%1 (%2)")
                       .arg(smarts_[i].first.c_str())
                       .arg(smarts_[i].second.c_str()));
  }

  DACLIB::QT4SelectItems si(string("Select SMARTS to use."), smts,
                            selected_smarts, !multi_select, this);
  if (QDialog::Accepted != si.exec()) {
    return smts_to_use;
  }
  si.get_results(selected_smarts);

  for (int i = 0, is = smarts_.size(); i < is; ++i) {
    if (selected_smarts[i]) {
      smts_to_use.push_back(smarts_[i]);
    }
  }

  return smts_to_use;
}

// ****************************************************************************
void RDKitSVMainWindow::update_smarts(const string &new_name,
                                      const string &new_val) {
  vector<pair<string, string>>::iterator p =
      std::find_if(smarts_.begin(), smarts_.end(),
                   bind(equal_to<string>(), new_name,
                        bind(&pair<string, string>::first, _1)));
  if (p == smarts_.end()) {
    smarts_.push_back(make_pair(new_name, new_val));
  } else {
    p->second = new_val;
  }
}

// ****************************************************************************
void RDKitSVMainWindow::slot_read_mols() {
  QString filename = QFileDialog::getOpenFileName(
      this, "Select molecule file", last_dir_,
      "Mol. files (*.smi *.smi.gz *.ism *.ism.gz *.sdf *.sdf.gz)");
  if (filename.isEmpty() || filename.isNull()) {
    return;
  }

  QFileInfo fi(filename);
  last_dir_ = fi.absolutePath();

  read_mols(filename.toLocal8Bit().data());
}

// ****************************************************************************
void RDKitSVMainWindow::slot_read_smarts() {
  QString filename =
      QFileDialog::getOpenFileName(this, "Select SMARTS file", last_dir_,
                                   "SMARTS files (*.smt);;Any file (*)");
  if (filename.isEmpty() || filename.isNull()) {
    return;
  }

  QFileInfo fi(filename);
  last_dir_ = fi.absolutePath();

  read_smarts(filename.toLocal8Bit().data());
}

// ****************************************************************************
void RDKitSVMainWindow::slot_exit() { exit(0); }

// ****************************************************************************
void RDKitSVMainWindow::slot_match_smarts() {
  vector<pair<string, string>> smts_to_use = select_smarts(true);
  if (smts_to_use.empty()) {
    return;
  }

  match_smarts(smts_to_use);
}

// ****************************************************************************
void RDKitSVMainWindow::slot_edit_smarts() {
  if (smarts_.empty()) {
    slot_new_smarts();
  }

  vector<pair<string, string>> smts_to_use = select_smarts(false);
  if (smts_to_use.empty()) {
    return;
  }

  DACLIB::QTGet2Strings g2s(
      tr("SMARTS name"), smts_to_use.front().first.c_str(), tr("SMARTS value"),
      smts_to_use.front().second.c_str(), this);
  if (QDialog::Accepted != g2s.exec()) {
    return;
  }

  QString new_name, new_smarts;
  g2s.get_values(new_name, new_smarts);

  update_smarts(new_name.toLocal8Bit().data(), new_smarts.toLocal8Bit().data());
}

// ****************************************************************************
void RDKitSVMainWindow::slot_new_smarts() {
  DACLIB::QTGet2Strings g2s(tr("SMARTS name"),
                            QString("SMARTS_%1").arg(smarts_.size() + 1),
                            tr("SMARTS value"), QString(""), this);
  if (QDialog::Accepted != g2s.exec()) {
    return;
  }

  QString new_name, new_smarts;
  g2s.get_values(new_name, new_smarts);

  update_smarts(new_name.toLocal8Bit().data(), new_smarts.toLocal8Bit().data());
}

// ****************************************************************************
void RDKitSVMainWindow::slot_write_left_molecules() {
  QString filename = QFileDialog::getSaveFileName(
      this, "File for left panel molecules", last_dir_,
      "Mol. files (*.smi *.smi.gz *.ism *.ism.gz *.sdf *.sdf.gz)");
  if (filename.isEmpty()) {
    return;
  }

  QFileInfo fi(filename);
  last_dir_ = fi.absolutePath();

  write_mols(*left_panel_, filename.toLocal8Bit().data());
}

// ****************************************************************************
void RDKitSVMainWindow::slot_write_right_molecules() {
  if (right_panel_->isHidden()) {
    // it's empty, so nothing to do.
    QMessageBox::information(
        this, tr("No molecules."),
        tr("Right panel is inactive, so nothing to write."));
    return;
  }

  QString filename = QFileDialog::getSaveFileName(
      this, "File for left panel molecules", last_dir_,
      "Mol. files (*.smi *.smi.gz *.ism *.ism.gz *.sdf *.sdf.gz)");
  if (filename.isEmpty()) {
    return;
  }

  QFileInfo fi(filename);
  last_dir_ = fi.absolutePath();

  write_mols(*right_panel_, filename.toLocal8Bit().data());
}

// ****************************************************************************
void RDKitSVMainWindow::slot_write_smarts() {
  QString filename =
      QFileDialog::getSaveFileName(this, "File for SMARTS.", last_dir_,
                                   "SMARTS files (*.smt);;Any file (*)");
  if (filename.isEmpty()) {
    return;
  }

  QFileInfo fi(filename);
  last_dir_ = fi.absolutePath();

  ofstream ofs(filename.toLocal8Bit().data());
  if (!ofs || !ofs.good()) {
    QMessageBox::warning(
        this, tr("Can't write to file"),
        QString("%1 %2.").arg(tr("Couldn't write to file").arg(filename)));
    return;
  }

  vector<pair<string, string>> smts_to_use = select_smarts(true);
  if (smts_to_use.empty()) {
    return;
  }

  typedef pair<string, string> STRINGPAIR;
  BOOST_FOREACH (STRINGPAIR smt, smts_to_use) {
    ofs << smt.first << " " << smt.second << endl;
  }
}

}  // namespace RDKitSV
