/****************************************************************************
 * Copyright (C) 2009-2015 EPAM Systems
 *
 * This file is part of Indigo toolkit.
 *
 * This file may be distributed and/or modified under the terms of the
 * GNU General Public License version 3 as published by the Free Software
 * Foundation and appearing in the file LICENSE.GPL included in the
 * packaging of this file.
 *
 * This file is provided AS IS with NO WARRANTY OF ANY KIND, INCLUDING THE
 * WARRANTY OF DESIGN, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 ***************************************************************************/

#include "reaction/reaction.h"
#include "reaction/query_reaction.h"
#include "molecule/molecule_arom_match.h"
#include "molecule/molecule_3d_constraints.h"
#include "molecule/molecule_substructure_matcher.h"
#include "reaction/base_reaction_substructure_matcher.h"
#include "reaction/reaction_neighborhood_counters.h"
#include "molecule/molecule_neighbourhood_counters.h"
#include "reaction/base_reaction_substructure_matcher.h"

using namespace indigo;

IMPL_ERROR(BaseReactionSubstructureMatcher, "reaction substructure matcher");

CP_DEF(BaseReactionSubstructureMatcher);

BaseReactionSubstructureMatcher::BaseReactionSubstructureMatcher (Reaction &target) :
_target(target),
CP_INIT,
TL_CP_GET(_matchers),
TL_CP_GET(_aam_to_second_side_1),
TL_CP_GET(_aam_to_second_side_2),
TL_CP_GET(_molecule_core_1),
TL_CP_GET(_molecule_core_2),
TL_CP_GET(_aam_core_first_side)
{
   use_aromaticity_matcher = true;
   highlight = false;
   
   match_atoms = 0;
   match_bonds = 0;
   context = 0;
   remove_atom = 0;
   add_bond = 0;
   prepare = 0;
   _match_stereo = true;

   _query_nei_counters = 0;
   _target_nei_counters = 0;

   _query = 0;

   _matchers.clear();
   _matchers.add(new _Matcher(*this));
}

void BaseReactionSubstructureMatcher::setQuery (BaseReaction &query)
{
   _query = &query;
}

void BaseReactionSubstructureMatcher::setNeiCounters (const ReactionAtomNeighbourhoodCounters *query_counters,
                                                      const ReactionAtomNeighbourhoodCounters *target_counters)
{
   _query_nei_counters = query_counters;
   _target_nei_counters = target_counters;
}

bool BaseReactionSubstructureMatcher::find ()
{
   if (_query == 0)
      throw Error("no query");

   if (prepare != 0 && !prepare(*_query, _target, context))
      return false;

   if (_query->reactantsCount() > _target.reactantsCount() || _query->productsCount() > _target.productsCount())
      return false;

   if (_query->reactantsCount() * _target.reactantsCount() < _query->productsCount() * _target.productsCount())
      _first_side = Reaction::REACTANT, _second_side = Reaction::PRODUCT;
   else
      _first_side = Reaction::PRODUCT, _second_side = Reaction::REACTANT;

   _initMap(*_query, _second_side, _aam_to_second_side_1);
   _initMap(_target, _second_side, _aam_to_second_side_2);

   _molecule_core_1.resize(_query->end());
   _molecule_core_1.fffill();
   _molecule_core_2.resize(_target.end());
   _molecule_core_2.fffill();
   _aam_core_first_side.clear();

   _matchers.top()->match_stereo = _match_stereo;

   while (1)
   {
      int command = _matchers.top()->nextPair();

      if (command == _CONTINUE)
         continue;

      if (command == _RETURN)
      {
         if (_checkAAM())
         {
            _highlight();
            return true;
         }
         command = _NO_WAY;
      } else if (command != _NO_WAY)
      {
         int mol1 = _matchers.top()->_current_molecule_1;
         int mol2 = _matchers.top()->_current_molecule_2;
         Array<int> &core1 = _matchers.top()->_current_core_1;
         Array<int> &core2 = _matchers.top()->_current_core_2;
         int mode = _matchers.top()->getMode();

         //_matchers.reserve(_matchers.size() + 1);
         AutoPtr<_Matcher> top_matcher(new _Matcher(*_matchers.top()));
         _matchers.add(top_matcher.release());
         _matchers.top()->setMode(command);
         if (!_matchers.top()->addPair(mol1, mol2, core1, core2, mode == _FIRST_SIDE))
            _matchers.removeLast();
      }

      if (command == _NO_WAY)
      {
         if (_matchers.size() > 1)
         {
            _matchers.top()->restore();
            _matchers.removeLast();
         }
         else
            return false;
      }
   }
}

// Init data for reaction substructure search
void BaseReactionSubstructureMatcher::_initMap (BaseReaction &reaction, int side, RedBlackMap<int, int> &aam_map)
{
   int i, j;
   int *val;

   aam_map.clear();

   // collect aam-to-molecule index mapping for reaction second side
   for (i = reaction.sideBegin(side); i < reaction.sideEnd(); i = reaction.sideNext(side, i))
   {
      BaseMolecule &i_mol = reaction.getBaseMolecule(i);

      for (j = i_mol.vertexBegin(); j < i_mol.vertexEnd(); j = i_mol.vertexNext(j))
      {
         int aam_number = reaction.getAAM(i, j);
                    
         if (aam_number != 0)
         {
            if ((val = aam_map.at2(aam_number)) == 0)
               aam_map.insert(aam_number, i);
            else if (*val < 0)
               (*val)--;
            else
               (*val) = -1;
         }
      }
   }
}

// Check correct AAM relationship between query and target reaction
bool BaseReactionSubstructureMatcher::_checkAAM ()
{
   int *aam, aam1, aam2;
   int i, j;

   for (i = 1; i < _matchers.size() - 1; i++)
   {
      if (_matchers[i]->getMode() == _FIRST_SIDE)
         continue;

      BaseMolecule &mol1 = _query->getBaseMolecule(_matchers[i]->_current_molecule_1);

      for (j = mol1.vertexBegin(); j < mol1.vertexEnd(); j = mol1.vertexNext(j))
      {
         int k = _matchers[i]->_current_core_1[j];

         if (k < 0)
            continue;

         aam1 = _query->getAAM(_matchers[i]->_current_molecule_1, j);
         aam2 = _target.getAAM(_matchers[i]->_current_molecule_2, k);

         if (aam1 > 0 && aam2 > 0)
            if ((aam = _aam_core_first_side.at2(aam1)) != 0 && *aam != aam2)
               return false;
      }
   }

   return true;
}

int BaseReactionSubstructureMatcher::getTargetMoleculeIndex (int query_mol_idx)
{
   // can be optimized, but as the number of molecules
   // seldom exceeds 5, the linear search is acceptable
   for (int i = 0; i < _matchers.size() - 1; i++)
      if (_matchers[i]->_current_molecule_1 == query_mol_idx)
         return _matchers[i]->_current_molecule_2;

   throw Error("getTargetMoleculeIndex(): can not find mapping for query molecule %d", query_mol_idx);
}

const int * BaseReactionSubstructureMatcher::getQueryMoleculeMapping (int query_mol_idx)
{
   for (int i = 0; i < _matchers.size() - 1; i++)
      if (_matchers[i]->_current_molecule_1 == query_mol_idx)
         return _matchers[i]->_current_core_1.ptr();

   throw Error("getQueryMoleculeMapping(): can not find mapping for query molecule %d", query_mol_idx);
}


void BaseReactionSubstructureMatcher::_highlight ()
{
   if (!highlight)
      return;

   int i;

   for (i = 0; i < _matchers.size() - 1; i++)
      _target.getBaseMolecule(_matchers[i]->_current_molecule_2).highlightSubmolecule(
         _query->getBaseMolecule(_matchers[i]->_current_molecule_1), _matchers[i]->_current_core_1.ptr(), true);
}

CP_DEF(BaseReactionSubstructureMatcher::_Matcher);

BaseReactionSubstructureMatcher::_Matcher::_Matcher (BaseReactionSubstructureMatcher &context) :
CP_INIT,
TL_CP_GET(_current_core_1),
TL_CP_GET(_current_core_2),
_context(context),
TL_CP_GET(_mapped_aams)
{
   _mode = _FIRST_SIDE;
   _selected_molecule_1 = -1;
   _selected_molecule_2 = -1;
   _current_molecule_1 = -1;
   _current_molecule_2 = -1;
   _mapped_aams.clear();
   match_stereo = true;
   _current_core_1.clear();
   _current_core_2.clear();
}

BaseReactionSubstructureMatcher::_Matcher::_Matcher (const BaseReactionSubstructureMatcher::_Matcher &other) :
CP_INIT,
TL_CP_GET(_current_core_1),
TL_CP_GET(_current_core_2),
_context(other._context),
TL_CP_GET(_mapped_aams)
{
   _current_molecule_1 = -1;
   _current_molecule_2 = -1;
   _mapped_aams.clear();
   match_stereo = other.match_stereo;
   _current_core_1.clear();
   _current_core_2.clear();
   _selected_molecule_1 = -1;
   _selected_molecule_2 = -1;
}

int BaseReactionSubstructureMatcher::_Matcher::_nextPair ()
{
   int side;

   if (_mode == _FIRST_SIDE)
      side = _context._first_side;
   else // _SECOND_SIDE_REST
      side = _context._second_side;

   if (_enumerator.get() == 0 || !_enumerator->processNext())
   {
      do
      {
         while (1)
         {
            if (_current_molecule_1 == -1)
            {
               for (_current_molecule_1 = _context._query->sideBegin(side);
                  _current_molecule_1 < _context._query->sideEnd();
                  _current_molecule_1 = _context._query->sideNext(side, _current_molecule_1))
                  if (_context._molecule_core_1[_current_molecule_1] < 0)
                     break;
               if (_current_molecule_1 == _context._query->sideEnd())
               {
                  if (_mode == _FIRST_SIDE)
                  {
                     _mode = _SECOND_SIDE_REST;
                     _current_molecule_1 = -1;
                     return _nextPair();
                  }

                  return _RETURN;
               }
            }

            if (_current_molecule_2 == -1)
               _current_molecule_2 = _context._target.sideBegin(side);
            else
               _current_molecule_2 = _context._target.sideNext(side, _current_molecule_2);

            for ( ; _current_molecule_2 < _context._target.sideEnd(); _current_molecule_2 = _context._target.sideNext(side, _current_molecule_2))
               if (_context._molecule_core_2[_current_molecule_2] < 0)
                  break;

            if (_current_molecule_2 == _context._target.sideEnd())
               return _NO_WAY;

            _enumerator.free();

            BaseMolecule &mol_1 = _context._query->getBaseMolecule(_current_molecule_1);
            Molecule &mol_2 = _context._target.getMolecule(_current_molecule_2);

            if (!_initEnumerator(mol_1, mol_2))
            {
               _enumerator.free();
               continue;
            }
            break;
         } 

         _enumerator->processStart();

      } while (!_enumerator->processNext());
   }

   return _mode == _FIRST_SIDE ? _SECOND_SIDE : _SECOND_SIDE_REST;
}

int BaseReactionSubstructureMatcher::_Matcher::nextPair ()
{
   if (_mode != _SECOND_SIDE)
   {
      int next = _nextPair();

      if (next != _SECOND_SIDE)
         return next;

      // Switch to _SECOND_SIDE
      BaseMolecule &mol_1 = _context._query->getBaseMolecule(_current_molecule_1);

      int first_aam_1 = 0;
      int first_aam_2 = 0;

      int i;

      for (i = mol_1.vertexBegin(); i < mol_1.vertexEnd(); i = mol_1.vertexNext(i))
         if (_current_core_1[i] >= 0)
         {
            first_aam_1 = _context._query->getAAM(_current_molecule_1, i);
            first_aam_2 = _context._target.getAAM(_current_molecule_2, _current_core_1[i]);
            break;
         }


      if (first_aam_1 > 0 && first_aam_2 > 0)
      {
         // Check the other side if needed
         int* mol_1_idx_ss_ptr = _context._aam_to_second_side_1.at2(first_aam_1);
         int* mol_2_idx_ss_ptr = _context._aam_to_second_side_2.at2(first_aam_2);

         if (mol_1_idx_ss_ptr == 0 && mol_2_idx_ss_ptr == 0)
            // There is no pair for both atom
            return _FIRST_SIDE; 

         if (mol_1_idx_ss_ptr == 0 || mol_2_idx_ss_ptr == 0)
            // One atom has a pair atom while other hasn't one
            return _CONTINUE; 

         int mol_1_idx_ss = *mol_1_idx_ss_ptr;
         int mol_2_idx_ss = *mol_2_idx_ss_ptr;
         if ((mol_1_idx_ss < 0 && mol_1_idx_ss < mol_2_idx_ss))
            return _CONTINUE; // subreactions equal AAM-numbers more than superreaction

         if (mol_2_idx_ss < 0)
            return _FIRST_SIDE; // check this molecules in the completion phase

         if (_context._molecule_core_1[mol_1_idx_ss] >= 0)
         {
            if (_context._molecule_core_1[mol_1_idx_ss] != mol_2_idx_ss)
               return _CONTINUE;

            int first_idx_1_ss = _context._query->findAtomByAAM(mol_1_idx_ss, first_aam_1);
            int first_idx_2_ss = _context._target.findAtomByAAM(mol_2_idx_ss, first_aam_2);
            int i;

            for (i = 0; i < _context._matchers.size(); i++)
               if (_context._matchers[i]->_current_molecule_1 == mol_1_idx_ss)
               {
                  if (_context._matchers[i]->_current_core_1[first_idx_1_ss] != first_idx_2_ss)
                     return _CONTINUE;
                  return _FIRST_SIDE;
               }
         }

         return _SECOND_SIDE;
      }

      return _FIRST_SIDE;
   }

   // _SECOND_SIDE
   if (_enumerator.get() == 0)
   {
      BaseMolecule &src_mol_1 = _context._query->getBaseMolecule(_selected_molecule_1);
      Molecule &src_mol_2 = _context._target.getMolecule(_selected_molecule_2);

      int src_aam_1 = 0;
      int src_aam_2 = 0;

      Array<int> &prev_core_1 = _context._matchers[_context._matchers.size() - 2]->_current_core_1;
      for (int i = src_mol_1.vertexBegin(); i < src_mol_1.vertexEnd(); i = src_mol_1.vertexNext(i))
         if (prev_core_1[i] >= 0)
         {
            src_aam_1 = _context._query->getAAM(_selected_molecule_1, i);
            src_aam_2 = _context._target.getAAM(_selected_molecule_2, prev_core_1[i]);
            break;
         }

      BaseMolecule &mol_1 = _context._query->getBaseMolecule(_current_molecule_1);
      Molecule &mol_2 = _context._target.getMolecule(_current_molecule_2);

      int first_idx_1 = _context._query->findAtomByAAM(_current_molecule_1, src_aam_1);
      int first_idx_2 = _context._target.findAtomByAAM(_current_molecule_2, src_aam_2);

      // init embedding enumerator context
      _initEnumerator(mol_1, mol_2);

      if (!_enumerator->fix(first_idx_1, first_idx_2))
         return _NO_WAY;

      _enumerator->processStart();
   }

   if (!_enumerator->processNext())
      return _NO_WAY;

   return _FIRST_SIDE;
}

bool BaseReactionSubstructureMatcher::_Matcher::_initEnumerator (BaseMolecule &mol_1, Molecule &mol_2)
{
   // init embedding enumerator context
   _enumerator.create(mol_2);

   _enumerator->cb_match_edge = _matchBonds;
   _enumerator->cb_match_vertex = _matchAtoms;
   _enumerator->cb_edge_add = _addBond;
   _enumerator->cb_vertex_remove = _removeAtom;
   _enumerator->cb_embedding = _embedding;

   if (mol_1.isQueryMolecule() && _context.use_aromaticity_matcher &&
       AromaticityMatcher::isNecessary(mol_1.asQueryMolecule()))
      _am.reset(new AromaticityMatcher(mol_1.asQueryMolecule(), mol_2, _context.arom_options));
   else
      _am.reset(0);

   _enumerator->userdata = this;
   _enumerator->setSubgraph(mol_1);

   if (_context.prepare_ee != 0)
   {
      if (!_context.prepare_ee(_enumerator.ref(), mol_1, mol_2, _context.context))
         return false;
   }

   return true;
}

bool BaseReactionSubstructureMatcher::_Matcher::addPair (int mol1_idx, int mol2_idx, const Array<int> &core1, const Array<int> &core2, bool from_first_side)
{
   _selected_molecule_1 = mol1_idx;
   _selected_molecule_2 = mol2_idx;

   _mapped_aams.clear();

   BaseMolecule &mol1 = _context._query->getBaseMolecule(mol1_idx);

   if (from_first_side)
   {
      int i;

      for (i = mol1.vertexBegin(); i < mol1.vertexEnd(); i = mol1.vertexNext(i))
         if (core1[i] >= 0)
         {
            int aam1 = _context._query->getAAM(mol1_idx, i);
            int aam2 = _context._target.getAAM(mol2_idx, core1[i]);
            int *aam;

            if (aam1 > 0 && aam2 > 0)
            {
               if ((aam = _context._aam_core_first_side.at2(aam1)) == 0)
               {
                  _context._aam_core_first_side.insert(aam1, aam2);
                  _mapped_aams.push(aam1);
               } else if (*aam != aam2)
               {
                  while (_mapped_aams.size() > 0)
                     _context._aam_core_first_side.remove(_mapped_aams.pop());
                  return false;
               }
            }
         }
   }

   if (_mode == _SECOND_SIDE)
   {
      int first_aam_1 = 0;
      int first_aam_2 = 0;

      int i;

      for (i = mol1.vertexBegin(); i < mol1.vertexEnd(); i = mol1.vertexNext(i))
         if (core1[i] >= 0)
         {
            first_aam_1 = _context._query->getAAM(mol1_idx, i);
            first_aam_2 = _context._target.getAAM(mol2_idx, core1[i]);
            break;
         }

      _current_molecule_1 = _context._aam_to_second_side_1.at(first_aam_1);
      _current_molecule_2 = _context._aam_to_second_side_2.at(first_aam_2);
   }

   _context._molecule_core_1[mol1_idx] = mol2_idx;
   _context._molecule_core_2[mol2_idx] = mol1_idx;

   return true;
}

void BaseReactionSubstructureMatcher::_Matcher::restore()
{
   _context._molecule_core_1[_selected_molecule_1] = -1;
   _context._molecule_core_2[_selected_molecule_2] = -1;

   while (_mapped_aams.size() > 0)
      _context._aam_core_first_side.remove(_mapped_aams.pop());
}

int BaseReactionSubstructureMatcher::_Matcher::_embedding (Graph &subgraph, Graph &supergraph, int *core_sub, int *core_super, void *userdata)
{
   BaseReactionSubstructureMatcher::_Matcher &self = *(BaseReactionSubstructureMatcher::_Matcher *)userdata;

   QueryMolecule &query = (QueryMolecule &)subgraph;
   Molecule &target = (Molecule &)supergraph;

   if (self.match_stereo)
   {
      if (!MoleculeStereocenters::checkSub(query.stereocenters, target.stereocenters, core_sub, false))
         return 1;

      if (!MoleculeCisTrans::checkSub(query, target, core_sub))
         return 1;
   }

   // Check possible aromatic configuration
   if (self._am.get() != 0)
   {
      if (!self._am->match(core_sub, core_super))
         return 1;
   }

   self._current_core_1.copy(core_sub, subgraph.vertexEnd());
   self._current_core_2.copy(core_super, supergraph.vertexEnd());

   return 0;
}

bool BaseReactionSubstructureMatcher::_Matcher::_matchAtoms (Graph &subgraph, Graph &supergraph, const int *core_sub, int sub_idx, int super_idx, void *userdata)
{
   BaseReactionSubstructureMatcher::_Matcher *self = (BaseReactionSubstructureMatcher::_Matcher *)userdata;

   if (self->_context.match_atoms != 0 && !self->_context.match_atoms(
      *self->_context._query, self->_context._target, self->_current_molecule_1, sub_idx,
       self->_current_molecule_2, super_idx, self->_context.context))
      return false;

   if (self->_mode == _SECOND_SIDE)
   {
      int *aam, aam1, aam2;

      aam1 = self->_context._query->getAAM(self->_current_molecule_1, sub_idx);
      if (aam1 != 0)
      {
         aam2 = self->_context._target.getAAM(self->_current_molecule_2, super_idx);
         if (aam2 != 0)
            if ((aam = self->_context._aam_core_first_side.at2(aam1)) != 0 && *aam != aam2)
               return false;
      }
   }

   if (self->_context._query_nei_counters != 0 && self->_context._target_nei_counters != 0)
   {
      const MoleculeAtomNeighbourhoodCounters &mol_count1 = self->_context._query_nei_counters->getCounters(self->_current_molecule_1);
      const MoleculeAtomNeighbourhoodCounters &mol_count2 = self->_context._target_nei_counters->getCounters(self->_current_molecule_2);

      if (!mol_count1.testSubstructure(mol_count2, sub_idx, super_idx, true))
         return false;
   }

   int sub_atom_inv = self->_context._query->getInversion(self->_current_molecule_1, sub_idx);
   int super_atom_inv = self->_context._target.getInversion(self->_current_molecule_2, super_idx);

   if (sub_atom_inv != STEREO_UNMARKED && sub_atom_inv != super_atom_inv)
      return false;

   return true;
}


bool BaseReactionSubstructureMatcher::_Matcher::_matchBonds (Graph &subgraph, Graph &supergraph, int sub_idx, int super_idx, void *userdata)
{
   BaseReactionSubstructureMatcher::_Matcher *self = (BaseReactionSubstructureMatcher::_Matcher *)userdata;

   if (self->_context.match_bonds != 0 && !self->_context.match_bonds(
      *self->_context._query, self->_context._target, self->_current_molecule_1,
       sub_idx, self->_current_molecule_2, super_idx, self->_am.get(), self->_context.context))
      return false;

   return true;
}

void BaseReactionSubstructureMatcher::_Matcher::_removeAtom (Graph &subgraph, int sub_idx, void *userdata)
{
   BaseReactionSubstructureMatcher::_Matcher *self = (BaseReactionSubstructureMatcher::_Matcher *)userdata;

   if (self->_context.remove_atom != 0)
      self->_context.remove_atom((BaseMolecule &)subgraph, sub_idx, self->_am.get());
}

void BaseReactionSubstructureMatcher::_Matcher::_addBond (Graph &subgraph, Graph &supergraph,
                           int sub_idx, int super_idx, void *userdata)
{
   BaseReactionSubstructureMatcher::_Matcher *self = (BaseReactionSubstructureMatcher::_Matcher *)userdata;

   if (self->_context.add_bond != 0)
      self->_context.add_bond((BaseMolecule &)subgraph, (Molecule &)supergraph,
              sub_idx, super_idx, self->_am.get());
}