/****************************************************************************
 * 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 "molecule/molecule_pi_systems_matcher.h"

#include "graph/filter.h"
#include "graph/spanning_tree.h"
#include "molecule/elements.h"

using namespace indigo;

IMPL_ERROR(MoleculePiSystemsMatcher, "Pi-system matcher");

CP_DEF(MoleculePiSystemsMatcher);

MoleculePiSystemsMatcher::MoleculePiSystemsMatcher (Molecule &target) 
   : _target(target),
   CP_INIT,
   TL_CP_GET(_atom_pi_system_idx),
   TL_CP_GET(_pi_systems),
   TL_CP_GET(_connectivity)
{
   _calcConnectivity(_target, _connectivity);

   _atom_pi_system_idx.clear_resize(target.vertexEnd());
   int n_pi_systems = _initMarks();
   _pi_systems.clear();
   _pi_systems.resize(n_pi_systems);
}

void MoleculePiSystemsMatcher::_calcConnectivity (Molecule &mol, Array<int> &conn)
{
   conn.clear_resize(mol.vertexEnd());
   conn.zerofill();
   for (int e = mol.edgeBegin(); e != mol.edgeEnd(); e = mol.edgeNext(e))
   {
      int bond_order = mol.getBondOrder(e);

      const Edge &edge = mol.getEdge(e);
      conn[edge.beg] += bond_order;
      conn[edge.end] += bond_order;
   }
   for (int v = mol.vertexBegin(); v != mol.vertexEnd(); v = mol.vertexNext(v))
      if (!mol.isPseudoAtom(v) && !mol.isRSite(v) && !mol.isTemplateAtom(v))
         conn[v] += mol.getImplicitH(v);
}

int MoleculePiSystemsMatcher::_initMarks (void)
{
   _markAtomsFirst();

   // Decompose molecule into pi-systems
   Filter filter(_atom_pi_system_idx.ptr(), Filter::NEQ, _NOT_IN_PI_SYSTEM);

   // Decompose 'pi_systems' into connected components
   _decomposer.create(_target);
   int n_comp = _decomposer->decompose(&filter);

   // Copy pi-system indices
   _copyPiSystemsIdFromDecomposer();

   QS_DEF(Array<bool>, pi_system_used);
   pi_system_used.clear_resize(n_comp);
   for (int i = 0; i < n_comp; i++)
      pi_system_used[i] = false;

   _markUnstablePiSystems(pi_system_used);
   
   bool use_any_pi_system = false;
   for (int i = 0; i < n_comp; i++)
   {
      use_any_pi_system |= pi_system_used[i];
      if (use_any_pi_system)
         break;
   }
   _markVerticesInUnusedPiSystems(pi_system_used);
   if (!use_any_pi_system)
      return 0;

   _markVerticesInSingleAtomPiSystem(n_comp);
   _markVerticesInPiSystemsWithCycles();

   // Decompose again because cycles could split pi-system
   n_comp = _decomposer->decompose(&filter);

   _copyPiSystemsIdFromDecomposer();
   _markVerticesInSingleAtomPiSystem(n_comp);

   return n_comp;
}

void MoleculePiSystemsMatcher::_calculatePiSystemsSizes (int n_pi_systems, 
                                                         Array<int> &sizes)
{
   sizes.clear_resize(n_pi_systems);
   sizes.zerofill();

   // Collect sizes
   for (int v = _target.vertexBegin(); 
            v != _target.vertexEnd(); 
            v = _target.vertexNext(v))
   {
      int p = _atom_pi_system_idx[v];
      if (p != _NOT_IN_PI_SYSTEM)
         sizes[p]++;
   }
}

void MoleculePiSystemsMatcher::_markVerticesInSingleAtomPiSystem (int n_pi_systems)
{
   QS_DEF(Array<int>, pi_system_size);
   _calculatePiSystemsSizes(n_pi_systems, pi_system_size);

   // Exclude single atoms
   for (int v = _target.vertexBegin(); 
            v != _target.vertexEnd(); 
            v = _target.vertexNext(v))
   {
      int p = _atom_pi_system_idx[v];
      if (p != _NOT_IN_PI_SYSTEM && pi_system_size[p] == 1)
         _atom_pi_system_idx[v] = _NOT_IN_PI_SYSTEM;
   }
}


void MoleculePiSystemsMatcher::_copyPiSystemsIdFromDecomposer ()
{
   const Array<int>& pi_system_per_vertex = _decomposer->getDecomposition();
   for (int v = _target.vertexBegin(); 
            v != _target.vertexEnd(); 
            v = _target.vertexNext(v))
   {
      int pi_system = pi_system_per_vertex[v];
      if (pi_system == -1)
         pi_system = _NOT_IN_PI_SYSTEM;

      _atom_pi_system_idx[v] = pi_system;
   }
}

void MoleculePiSystemsMatcher::_markUnstablePiSystems (Array<bool> &pi_system_used)
{
   // Mark pi-systems that contain atoms with partail octet or charge
   for (int v = _target.vertexBegin(); 
            v != _target.vertexEnd(); 
            v = _target.vertexNext(v))
   {
      int pi_system = _atom_pi_system_idx[v];
      if (pi_system < 0)
         continue;
      if (_connectivity[v] == -1)
         continue; // This means that atom is in aromatic ring

      if (_target.getAtomCharge(v) != 0 || _target.getAtomRadical(v) != 0)
         pi_system_used[pi_system] = true;


      int lonepairs;
      int vac = _target.getVacantPiOrbitals(v, _connectivity[v], &lonepairs);
      if (vac != 0)
         pi_system_used[pi_system] = true;
   }
}

void MoleculePiSystemsMatcher::_markVerticesInUnusedPiSystems (
   Array<bool> &pi_system_used)
{
   for (int v = _target.vertexBegin(); 
            v != _target.vertexEnd(); 
            v = _target.vertexNext(v))
   {
      int pi_system = _atom_pi_system_idx[v];
      if (pi_system == _NOT_IN_PI_SYSTEM)
         continue;

      if (!pi_system_used[pi_system])
         _atom_pi_system_idx[v] = _NOT_IN_PI_SYSTEM;
   }
}

void MoleculePiSystemsMatcher::_markVerticesInPiSystemsWithCycles ()
{
   // Developed algorithm doesn't operate with pi-systems with cycles
   Filter filter(_atom_pi_system_idx.ptr(), Filter::NEQ, _NOT_IN_PI_SYSTEM);
   SpanningTree sp_tree(_target, &filter);

   QS_DEF(Array<int>, edge_in_cycle);
   edge_in_cycle.clear_resize(_target.edgeEnd());
   edge_in_cycle.zerofill();
   sp_tree.markAllEdgesInCycles(edge_in_cycle.ptr(), 1);

   for (int e = _target.edgeBegin(); 
            e != _target.edgeEnd(); 
            e = _target.edgeNext(e))
   {
      if (edge_in_cycle[e])
      {
         const Edge &edge = _target.getEdge(e);

         _atom_pi_system_idx[edge.beg] = _NOT_IN_PI_SYSTEM;
         _atom_pi_system_idx[edge.end] = _NOT_IN_PI_SYSTEM;
      }
   }
}


void MoleculePiSystemsMatcher::_markAtomsFirst ()
{
   for (int i = 0; i < _atom_pi_system_idx.size(); i++)
      _atom_pi_system_idx[i] = _UNKNOWN;

   // Aromatic bonds matches by aromaticity matcher
   for (int e = _target.edgeBegin(); 
            e != _target.edgeEnd(); 
            e = _target.edgeNext(e))
   {
      if (_target.getBondOrder(e) == BOND_AROMATIC)
      {
         const Edge &edge = _target.getEdge(e);
         _atom_pi_system_idx[edge.beg] = _IN_AROMATIC;
         _atom_pi_system_idx[edge.end] = _IN_AROMATIC;

         _connectivity[edge.beg] = -1;
         _connectivity[edge.end] = -1;
      }
   }

   // Mark atoms that can be in pi system
   for (int v = _target.vertexBegin(); 
            v != _target.vertexEnd(); 
            v = _target.vertexNext(v))
   {
      if (_atom_pi_system_idx[v] == _NOT_IN_PI_SYSTEM)
         continue;
      // Atoms in aromatic rings will be exclused later 
      // because they are in rings
      if (_atom_pi_system_idx[v] == _IN_AROMATIC)
         continue;

      if (!_canAtomBeInPiSystem(v))
         _atom_pi_system_idx[v] = _NOT_IN_PI_SYSTEM;
   }
}

bool MoleculePiSystemsMatcher::_canAtomBeInPiSystem (int v)
{
   if (_target.isPseudoAtom(v) || _target.isRSite(v) || _target.isTemplateAtom(v))
      return false;

   int label = _target.getAtomNumber(v);
   if (!Element::canBeAromatic(label))
      return false;
   if (label == ELEM_Pb || label == ELEM_Tl)
      return false; // Skip this elements because of lonepairs behavior

   // Check if d-orbital is used (not supported now)
   int electrons = Element::electrons(label, _target.getAtomCharge(v));
   int conn = _connectivity[v];
   int radical = _target.getAtomRadical(v);
   int used = electrons + conn + 2 * Element::radicalOrbitals(radical);
   if (used > 8 || conn > 4)
      return false;

   int skeleton_conn = _target.getVertex(v).degree() + _target.getImplicitH(v);
   if (skeleton_conn == conn && conn == 4)
      return false;

   // Check if atom is correct: charge corresponds to connectivity
   // For examples there are molecules like Cl-Ga(v2)-Cl, that are incorrect
   // but stored in database
   int total_electrons = electrons + conn + Element::radicalElectrons(radical);
   if (total_electrons % 2 != 0)
      return false;
   return true;
}

bool MoleculePiSystemsMatcher::isAtomInPiSystem (int atom)
{
   return _atom_pi_system_idx[atom] != _NOT_IN_PI_SYSTEM;
}

bool MoleculePiSystemsMatcher::isBondInPiSystem (int bond)
{
   const Edge &edge = _target.getEdge(bond);
   int p1 = _atom_pi_system_idx[edge.beg];
   int p2 = _atom_pi_system_idx[edge.end];
   if (p1 != p2)
      return false;
   return p1 != _NOT_IN_PI_SYSTEM;
}

void MoleculePiSystemsMatcher::_extractPiSystem (int pi_system_index)
{
   _Pi_System &pi_system = _pi_systems[pi_system_index];
   pi_system.initialized = true;

   Molecule &ps = pi_system.pi_system;
   Filter filt(_decomposer->getDecomposition().ptr(), Filter::EQ, pi_system_index);
   ps.makeSubmolecule(_target, filt, &pi_system.mapping, &pi_system.inv_mapping);

   // Replace bonds outside pi-system to implicit hydrogens
   QS_DEF(Array<int>, conn);
   _calcConnectivity(ps, conn);
   for(int v = ps.vertexBegin(); v != ps.vertexEnd(); v = ps.vertexNext(v))
   {
      int original_v = pi_system.mapping[v];
      int delta = _connectivity[original_v] - conn[v];
      if (delta < 0)
         throw Error("Internal error: delta must be >= 0");
      ps.setImplicitH(v, ps.getImplicitH(v) + delta);
   }

   pi_system.localizations.clear();
   pi_system.localizer.create(pi_system.pi_system);
   _findPiSystemLocalization(pi_system_index);
}

void MoleculePiSystemsMatcher::_findPiSystemLocalization (int pi_system_index)
{
   // This function should be implemented outside this class
   // for finding pi-systems localizations during databse index 
   // creation.

   _Pi_System &pi_system = _pi_systems[pi_system_index];

   // Calculate number of electrons and electrons for zero-charge 
   int electrons = 0, zc_electrons = 0, octet_electron_pairs = 0;

   Molecule &mol = pi_system.pi_system;
   for(int v = mol.vertexBegin(); v != mol.vertexEnd(); v = mol.vertexNext(v))
   {
      int degree = mol.getVertex(v).degree();
      int min_connectivity = mol.getImplicitH(v) + degree;

      int label = mol.getAtomNumber(v);
      electrons += Element::electrons(label, mol.getAtomCharge(v)) - min_connectivity;
      zc_electrons += __max(Element::electrons(label, 0) - min_connectivity, 0);

      octet_electron_pairs += 4 - min_connectivity - 
         Element::radicalOrbitals(mol.getAtomRadical(v));
   }

   // Find localizations
   if (electrons % 2 != 0)
      throw Error("Electrons number must be even");
   int electron_pairs = electrons / 2;
   int zc_lonepairs = pi_system.localizer->getZeroChargeLonepairs();

   int double_bonds = electron_pairs;
   if (double_bonds > octet_electron_pairs / 2)
      double_bonds = octet_electron_pairs / 2;

   bool was_found = false;

   int other_lp = 0, best_other_lp = -1;
   bool double_bonds_checks = false;
   int octet_charges_count = -1;
   while (double_bonds >= 0)
   {
      int zero_lp = electron_pairs - double_bonds - other_lp;
      if (zero_lp > zc_lonepairs)
      {
         other_lp = zero_lp - zc_lonepairs;
         zero_lp = zc_lonepairs;
      }
      if (zero_lp < 0)
      {
         other_lp = 0;
         double_bonds--;
         double_bonds_checks = false;
         continue;
      }

      if (best_other_lp != -1 && best_other_lp < other_lp)
      {
         // Such localization has more atoms with 
         // charges then in previous localization
         break;
      }

      pi_system.localizer->setParameters(double_bonds, zero_lp, other_lp);

      bool found;
      if (!double_bonds_checks)
      {
         found = pi_system.localizer->localize(true);
         if (!found)
         {
            // Such number of double bonds can't be
            other_lp = 0;
            double_bonds--;
            continue;
         }
         double_bonds_checks = true;
      }

      found = pi_system.localizer->localize();

      if (found)
      {
         bool octet = pi_system.localizer->isAllAtomsHaveOctet();
         int charges_count = pi_system.localizer->getLocalizationChargesCount();

         if (octet)
            octet_charges_count = charges_count;
         else
         {
            if (octet_charges_count != -1 && octet_charges_count <= charges_count)
               break;
         }

         best_other_lp = other_lp;
         was_found = true;
      }
      else
      {
         other_lp++;
         continue;
      }

      // Add localization
      _Pi_System::Localizations &loc = pi_system.localizations.push();
      loc.double_bonds = double_bonds;
      loc.primary_lp = zero_lp;
      loc.seconary_lp = other_lp;

      if (zero_lp == zc_lonepairs)
         break;
      
      other_lp = 0;
      double_bonds--;
      double_bonds_checks = false;
   }
}

bool MoleculePiSystemsMatcher::checkEmbedding (QueryMolecule &query, const int *mapping)
{
   for (int i = 0; i < _pi_systems.size(); i++)
      _pi_systems[i].pi_system_mapped = false;

   _markMappedPiSystems(query, mapping);

   bool ret;
   
   // Fix bonds
   ret = _fixBonds(query, mapping);
   if (!ret)
      return false;

   // Fix charges
   ret = _fixAtoms(query, mapping);
   if (!ret)
      return false;
         
   return _findMatching();
}

bool MoleculePiSystemsMatcher::_findMatching ()
{
   for (int i = 0; i < _pi_systems.size(); i++)
   {
      if (!_pi_systems[i].initialized)
         continue;

      _Pi_System &pi_system = _pi_systems[i];
      if (!pi_system.pi_system_mapped)
         continue;

      if (!_findMatchingForPiSystem(i))
         return false;
   }

   return true;
}

bool MoleculePiSystemsMatcher::_findMatchingForPiSystem (int pi_system_index)
{
   _Pi_System &pi_system = _pi_systems[pi_system_index];

   for (int i = 0; i < pi_system.localizations.size(); i++)
   {
      _Pi_System::Localizations &loc = pi_system.localizations[i];
      pi_system.localizer->setParameters(loc.double_bonds, 
         loc.primary_lp, loc.seconary_lp);

      if (pi_system.localizer->localize())
         return true;
   }

   return false;
}

void MoleculePiSystemsMatcher::_markMappedPiSystems (QueryMolecule &query, 
      const int *mapping)
{
   for (int qv = query.vertexBegin(); 
            qv != query.vertexEnd(); 
            qv = query.vertexNext(qv))
   {
      int v = mapping[qv];

      if (v < 0)
         continue; // Such vertex must be ignored

      int pi_system_idx = _atom_pi_system_idx[v];
      if (pi_system_idx == _NOT_IN_PI_SYSTEM)
         continue;

      if (!_pi_systems[pi_system_idx].initialized)
         _extractPiSystem(pi_system_idx);
      
      _Pi_System &pi_system = _pi_systems[pi_system_idx];
      if (!pi_system.pi_system_mapped)
      {
         pi_system.pi_system_mapped = true;
         pi_system.localizer->unfixAll();
      }
   }
}

bool MoleculePiSystemsMatcher::_fixAtoms (QueryMolecule &query, const int *mapping)
{
   // Fix charges
   for (int qv = query.vertexBegin(); 
            qv != query.vertexEnd(); 
            qv = query.vertexNext(qv))
   {
      int v = mapping[qv];
      if (v < 0)
         continue; // Such vertex must be ignored

      int pi_system_idx = _atom_pi_system_idx[v];
      if (pi_system_idx == _NOT_IN_PI_SYSTEM)
         continue;

      _Pi_System &pi_system = _pi_systems[pi_system_idx];

      QueryMolecule::Atom &qatom = query.getAtom(qv);
      int pv = pi_system.inv_mapping[v];

      int charge = query.getAtomCharge(qv);
      if (charge != CHARGE_UNKNOWN)
      {
         bool ret = pi_system.localizer->fixAtomCharge(pv, charge);
         if (!ret)
            return false;
      } 
      else if (qatom.hasConstraint(QueryMolecule::ATOM_CHARGE))
         throw Error("Unsupported atom charge specified");

      int valence = query.getExplicitValence(qv);
      if (valence != -1)
      {
         bool ret = pi_system.localizer->fixAtomConnectivity(pv, valence);
         if (!ret)
            return false;
      }
      else if (qatom.hasConstraint(QueryMolecule::ATOM_VALENCE))
         throw Error("Unsupported atom charge specified");
   }
   return true;
}

bool MoleculePiSystemsMatcher::_fixBonds (QueryMolecule &query, const int *mapping)
{
   for (int e = query.edgeBegin(); 
           e != query.edgeEnd(); 
           e = query.edgeNext(e))
   {
      const Edge &query_edge = query.getEdge(e);
      if (mapping[query_edge.beg] < 0 || mapping[query_edge.end] < 0)
         continue; // Edges connected with ignored vertices

      int target_edge = Graph::findMappedEdge(query, _target, e, mapping);
      const Edge &edge = _target.getEdge(target_edge);

      int p1_idx = _atom_pi_system_idx[edge.beg];
      int p2_idx = _atom_pi_system_idx[edge.end];
      if (p1_idx == _NOT_IN_PI_SYSTEM || p2_idx == _NOT_IN_PI_SYSTEM || p1_idx != p2_idx)
         continue;

      if (!_pi_systems[p1_idx].initialized)
         throw Error("pi-system must be initialized here");

      _Pi_System &pi_system = _pi_systems[p1_idx];

      int pi_sys_edge = Graph::findMappedEdge(_target, pi_system.pi_system, 
         target_edge, pi_system.inv_mapping.ptr());

      // Get target topology
      int topology = _target.getBondTopology(target_edge);

      QueryMolecule::Bond &qbond = query.getBond(e);

      bool can_be_single = qbond.possibleValuePair(
         QueryMolecule::BOND_ORDER, BOND_SINGLE,
         QueryMolecule::BOND_TOPOLOGY, topology);
      bool can_be_double = qbond.possibleValuePair(
         QueryMolecule::BOND_ORDER, BOND_DOUBLE,
         QueryMolecule::BOND_TOPOLOGY, topology);
      bool can_be_triple = qbond.possibleValuePair(
         QueryMolecule::BOND_ORDER, BOND_TRIPLE,
         QueryMolecule::BOND_TOPOLOGY, topology);

      if (!can_be_single && !can_be_double && !can_be_triple)
         return false;
      if (can_be_single && can_be_double && can_be_triple)
         continue;

      bool ret = false; // initializing to avoid compiler warning
      if (can_be_single && can_be_double)
         // Here can_be_triple = false because of previous check
         ret = pi_system.localizer->fixBondSingleDouble(pi_sys_edge);
      else
      {
         if (can_be_triple)
         {
            if (can_be_single)
               throw Error("Unsupported bond order specified (can be single or triple)");
            else if (can_be_double)
               throw Error("Unsupported bond order specified (can be double or triple)");
            ret = pi_system.localizer->fixBond(pi_sys_edge, BOND_TRIPLE);
         }
         if (can_be_single)
            ret = pi_system.localizer->fixBond(pi_sys_edge, BOND_SINGLE);
         if (can_be_double)
            ret = pi_system.localizer->fixBond(pi_sys_edge, BOND_DOUBLE);
      }

      if (!ret)
         return false;
   }

   return true;
}

void MoleculePiSystemsMatcher::copyLocalization (Molecule &target)
{
   for (int i = 0; i < _pi_systems.size(); i++)
   {
      if (!_pi_systems[i].initialized)
         continue;

      _Pi_System &system = _pi_systems[i];
      system.localizer->copyBondsAndCharges(target, system.mapping);
   }
}

void MoleculePiSystemsMatcher::_Pi_System::clear ()
{
   initialized = false;
   localizer.free();
   pi_system.clear();
   inv_mapping.clear();
   mapping.clear();
   localizations.clear();
}