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

#include "base_c/bitarray.h"
#include "base_cpp/scanner.h"
#include "molecule/elements.h"
#include "molecule/molecule.h"
#include "molecule/smiles_loader.h"

#include "indigo.h"

#include "indigo_internal.h"
#include "indigo_molecule.h"

#include "layout/molecule_layout.h"

#include <algorithm>

namespace indigo
{
namespace abbreviations
{

enum Direction { LEFT, RIGHT };

// Single abbreviation token
// Each abbreviation is tokenized first: (C2H5)2N- ->  N (C*2 H*5)*2
struct Token;
typedef std::vector<Token> TokenChain;
struct Token
{
   enum Type { Pattern, Element, Branch };
   Type type;

   int index; // Index in the abbreviations array or element index
   TokenChain branch; // Branch if brackets are present

   int multiplier;
};

// Attachment point with order
struct AttPoint
{
   explicit AttPoint (int index_, int order_) : index(index_), order(order_)
   {
   }
   int index, order;
};

//
// AbbreviationExpander
//
class AbbreviationExpander
{
public:
   AbbreviationExpander (ObjArray<Abbreviation> &abbreviations) : abbreviations(abbreviations)
   {
      ignore_case = false;
      tokenize_level = 0;
      tokenize_direction = expand_direction = RIGHT;
   }

   // Settings
   ObjArray<Abbreviation> &abbreviations;
   Direction tokenize_direction, expand_direction;
   int tokenize_level;

   bool ignore_case;

   // Results
   int input_index, output_index;

   Array<int> added_atoms;

   bool expandAtomAbbreviation (Molecule &mol, int v);
private:
   bool expand (const char *label, int input_order, int output_order, Molecule &dest);

   int tokensizeSubExpression (const char *label, TokenChain &tokens);
   int scanSinlgeToken (const char *label, Token &dest);
   bool tokensizeAbbreviation (const char *label, TokenChain &tokens);

   bool expandParsedTokens (TokenChain &tokens, Molecule &m, AttPoint &attach_to);
   bool tryApplyExpansions (TokenChain &tokens, size_t &offset, Molecule &m, AttPoint &attach_to);
   bool tryCarbonChain (TokenChain &tokens, size_t &offset, Molecule &m, AttPoint &attach_to);
   bool tryRepetition (TokenChain &tokens, size_t &offset, Molecule &m, AttPoint &attach_to);
   bool tryExpandToken (TokenChain &tokens, size_t &offset, Molecule &m, AttPoint &attach_to);
   bool expandParsedTokensWithRev (TokenChain &tokens, Molecule &m, AttPoint &attach_to);
   void attachBond (Molecule &m, AttPoint &attach_to, int index);

   bool compareStrings (const char *s1, const char *s2, size_t len);
};

bool AbbreviationExpander::compareStrings (const char *s1, const char *s2, size_t len)
{
   if (ignore_case)
      return strncasecmp(s1, s2, len) == 0;
   else
      return strncmp(s1, s2, len) == 0;
}

int AbbreviationExpander::tokensizeSubExpression (const char *label, TokenChain &tokens)
{
   size_t size = strlen(label);
   size_t offset = 0;
   while (offset < size)
   {
      Token t;
      int offset2 = scanSinlgeToken(label + offset, t);
      if (offset2 < 0)
         return -1;
      tokens.push_back(t);
      offset += offset2;
      if (label[offset] == ')')
         break;
   }
   return offset;
}

int AbbreviationExpander::scanSinlgeToken (const char *label, Token &dest)
{
   size_t size = strlen(label);
   int offset = 0;

   if (label[0] == '(')
   {
      dest.type = Token::Branch;
      int offset2 = tokensizeSubExpression(label + 1, dest.branch);
      if (offset2 == -1 || label[offset2 + 1] != ')')
         return -1;
      offset = offset2 + 2;
   }
   else
   {
      // Trying to apply geedy algorithm
      size_t best_matched_length = 0;
      int best_abbreviation = -1;
      for (int i = 0; i < abbreviations.size(); i++)
      {
         Abbreviation &cur = abbreviations[i];

         std::vector<std::string> *aliases[2] = {NULL, NULL};
         if (tokenize_direction == LEFT)
         {
            aliases[0] = &cur.left_aliases;
            if (tokenize_level == 2)
               aliases[1] = &cur.left_aliases2;
         }
         else if (tokenize_direction == RIGHT)
         {
            aliases[0] = &cur.right_aliases;
            if (tokenize_level == 2)
               aliases[1] = &cur.right_aliases2;
         }

         for (int k = 0; k < 2; k++)
         {
            if (aliases[k] == NULL)
               break;
            for (size_t j = 0; j < aliases[k]->size(); j++)
            {
               std::string &alias = aliases[k]->at(j);

               if (alias.length() < best_matched_length)
                  continue;

               if (compareStrings(label, alias.c_str(), alias.length()))
               {
                  best_matched_length = alias.length();
                  best_abbreviation = i;
               }
            }
         }
      }

      if (best_abbreviation != -1)
      {
         dest.type = Token::Pattern;
         dest.index = best_abbreviation;
         offset += best_matched_length;
      }
      else
      {
         // Try to match the tail with any known element: H, O, C, N, Sn, Mg
         // Try to read until hydrogen and digit or other captial letter
         size_t len = 1;
         while (len < size)
         {
            char ch = label[len];
            if (ch == 'H' || !isalpha(ch) || (ch >= 'A' && ch <= 'Z'))
               break;
            len++;
         }
         char symb[4] = {0};
         strncpy(symb, label, __min(NELEM(symb) - 1, len));
         int id = Element::fromString2(symb);
         if (id < 0)
            return -1;

         dest.type = Token::Element;
         dest.index = id;
         offset += len;
      }
   }

   // Try to read multiplier
   if (isdigit(label[offset]))
   {
      int shift;
      sscanf(label + offset, "%d%n", &dest.multiplier, &shift);
      offset += shift;
      
   }
   else
      dest.multiplier = 1;
   return offset;
}

bool AbbreviationExpander::tokensizeAbbreviation (const char *label, TokenChain &tokens)
{
   size_t size = strlen(label);
   size_t offset = 0;
   while (offset < size)
   {
      Token t;
      int offset2 = scanSinlgeToken(label + offset, t);
      if (offset2 < 0)
         return false;
      tokens.push_back(t);
      offset += offset2;
   }
   return true;
}

bool AbbreviationExpander::tryApplyExpansions (TokenChain &tokens, size_t &offset, Molecule &m, AttPoint &attach_to)
{
   if (tryCarbonChain(tokens, offset, m, attach_to))
      return true;
   if (tryRepetition(tokens, offset, m, attach_to))
      return true;
   if (tryExpandToken(tokens, offset, m, attach_to))
      return true;
   return false;
}

void AbbreviationExpander::attachBond (Molecule &m, AttPoint &attach_to, int idx)
{
   if (attach_to.index != -1)
      m.addBond(attach_to.index, idx, attach_to.order);
   else
      input_index = idx;
}

bool AbbreviationExpander::tryCarbonChain (TokenChain &tokens, size_t &offset, Molecule &m, AttPoint &attach_to)
{
   if (attach_to.order != 1)
      return false;

   Token &cur = tokens[offset];
   if (cur.type != Token::Element)
      return false;

   if (cur.multiplier == 1 || cur.index != ELEM_C || offset + 1 == tokens.size())
      return false;

   Token &next = tokens[offset + 1];
   // Check CnH(2n+1), CnH2n pattern
   if (next.multiplier > 1 && next.index == ELEM_H)
   {
      bool tail;
      if (next.multiplier == cur.multiplier * 2)
      {
         // Intermediate carbon chain
         tail = false;
      }
      else if (next.multiplier == cur.multiplier * 2 + 1)
      {
         // Terminator carbon chain
         tail = true;
      }
      else
         return false;

      for (int i = 0; i < cur.multiplier; i++)
      {
         int idx = m.addAtom(ELEM_C);
         attachBond(m, attach_to, idx);

         attach_to = AttPoint(idx, 1);
      }
      if (tail)
         attach_to = AttPoint(-1, 0);
      offset += 2;
      return true;
   }
   return false;
}

bool AbbreviationExpander::expandParsedTokensWithRev (TokenChain &tokens, Molecule &m, AttPoint &attach_to)
{
   if (expandParsedTokens(tokens, m, attach_to))
      return true;
   if (expand_direction != LEFT)
      return false;

   std::reverse(tokens.begin(), tokens.end());
   bool ret = expandParsedTokens(tokens, m, attach_to);
   std::reverse(tokens.begin(), tokens.end());
   return ret;
}

bool AbbreviationExpander::tryRepetition (TokenChain &tokens, size_t &offset, Molecule &m, AttPoint &attach_to)
{
   if (attach_to.order != 1)
      return false;

   Token &cur = tokens[offset];
   if (cur.type != Token::Branch || cur.multiplier == 1)
      return false;
   
   int atom_bound = m.vertexCount();

   AttPoint cur_attach_to = attach_to;
   int i;
   for (i = 0; i < cur.multiplier; i++)
   {
      if (!expandParsedTokensWithRev(cur.branch, m, cur_attach_to))
         break;
      if (cur_attach_to.index == -1)
         break;
   }

   if (i != cur.multiplier)
   {
      // Rollback
      Array<int> new_atoms;
      for (int v = m.vertexBegin(); v != m.vertexEnd(); v = m.vertexNext(v))
         if (v >= atom_bound)
            new_atoms.push(v);
      m.removeAtoms(new_atoms);
      return false;
   }

   offset++;
   attach_to = cur_attach_to;
   return true;
}

bool AbbreviationExpander::tryExpandToken (TokenChain &tokens, size_t &offset, Molecule &m, AttPoint &attach_to)
{
   Token &cur = tokens[offset];

   if (cur.multiplier != 1)
      return false;

   Array<int> connection_points;
   if (cur.type == Token::Element)
   {
      if (cur.index == ELEM_H)
      {
         offset++;
         attach_to = AttPoint(-1, 0);
         return true;
      }
      int added = m.addAtom(cur.index);

      // Get the number of bonds to connect
      int valence, hyd;
      int conn = attach_to.order;
      if (offset + 1 < tokens.size())
      {
         Token &next = tokens[offset + 1];
         conn += next.multiplier;
      }

      if (!Element::calcValence(cur.index, 0, 0, conn, valence, hyd, false))
      {
         // Ignore next atom
         // Appear in the OH3C case when H3 is belong to C
         conn = attach_to.order;
         if (!Element::calcValence(cur.index, 0, 0, conn, valence, hyd, false))
            return false;
      }

      for (int i = 0; i < hyd + conn; i++)
         connection_points.push(added);
   }
   else if (cur.type == Token::Pattern)
   {
      // Add pattern
      BufferScanner scanner(abbreviations[cur.index].expansion.c_str());
      SmilesLoader loader(scanner);

      Molecule abbr;
      loader.loadMolecule(abbr);

      Array<int> mapping;
      Array<int> rsites;
      m.mergeWithMolecule(abbr, &mapping);
      for (int v = abbr.vertexBegin(); v != abbr.vertexEnd(); v = abbr.vertexNext(v))
      {
         int mapped = mapping[v];
         if (m.isRSite(mapped))
         {
            dword bits = m.getRSiteBits(mapped);
            int id1 = bitGetOneHOIndex(bits);
            int id2 = bitGetOneHOIndex(bits);
            if (id1 != id2)
               throw Exception("Invalid abbreviations specification: %s", 
                  abbreviations[cur.index].expansion.c_str());
            if (id1 != 0)
               id1--; // R == R1

            const Vertex &vertex = m.getVertex(mapped);
            int nei = vertex.neiBegin();

            connection_points.expandFill(id1 + 1, -1);
            connection_points[id1] = vertex.neiVertex(nei); // Point connected to the RSite

            rsites.push(mapped);
         }
      }
      m.removeAtoms(rsites);
   }
   else
      return false;

   bool rollback = false;
   int atom_bound = m.vertexCount();
   size_t offset2 = offset + 1;

   attachBond(m, attach_to, connection_points[0]);
   int i = attach_to.order;
   while (i < connection_points.size() - 1 && !rollback)
   {
      if (offset2 >= tokens.size())
      {
         // If we are at the end then there can be an implicit double bond
         // Example: -CH2CH=
         // When we read C H there are no more tokens
         break;
      }

      Token &next = tokens[offset2];
      for (int j = 0; j < next.multiplier; j++)
      {
         if (i >= connection_points.size())
         {
            rollback = true;
            break;
         }

         if (next.type == Token::Branch)
         {
            AttPoint point(connection_points[i], 1);

            if (!expandParsedTokensWithRev(next.branch, m, point) || point.index != -1)
            {
               rollback = true;
               break;
            }
         }
         else
         {
            TokenChain chain;
            chain.push_back(next);
            chain[0].multiplier = 1;
            size_t local_offset = 0;
            AttPoint point(connection_points[i], 1);
            if (!tryExpandToken(chain, local_offset, m, point) || point.index != -1)
            {
               rollback = true;
               break;
            }
         }
         i++;
      }
      offset2++;
   }

   if (i > connection_points.size())
      rollback = true;
   if (!rollback)
   {
      if (i == connection_points.size())
      {
         // This is terminal
         attach_to = AttPoint(-1, 0);
      }
      else if (i == connection_points.size() - 1)
         attach_to = AttPoint(connection_points[i], 1); // Last attachment point
      else
      {
         // Number of tokens are incomlete means that there are double bonds after
         attach_to = AttPoint(connection_points[i], connection_points.size() - i);
      }
   }

   if (rollback)
   {
      // Rollback
      Array<int> new_atoms;
      for (int v = m.vertexBegin(); v != m.vertexEnd(); v = m.vertexNext(v))
         if (v >= atom_bound)
            new_atoms.push(v);
      m.removeAtoms(new_atoms);
      return false;
   }
   offset = offset2;
   return true;
}

bool AbbreviationExpander::expandParsedTokens (TokenChain &tokens, Molecule &m, AttPoint &attach_to)
{
   AttPoint attach_to_saved = attach_to;
   size_t index = 0;
   bool try_swapped = false;

   int atom_bound = m.vertexCount();
   bool first = true;

   while (index < tokens.size())
   {
      if (!first && attach_to.index == -1)
      {
         attach_to = attach_to_saved;
         // Rollback
         Array<int> new_atoms;
         for (int v = m.vertexBegin(); v != m.vertexEnd(); v = m.vertexNext(v))
            if (v >= atom_bound)
               new_atoms.push(v);
         m.removeAtoms(new_atoms);
         return false;
      }

      if (tryApplyExpansions(tokens, index, m, attach_to))
      {
         first = false;
         try_swapped = false;
         continue;
      }

      if (try_swapped)
         return false;

      if (expand_direction == LEFT && index + 1 < tokens.size())
      {
         // Try to swap two tokens if H is the first
         // Happens when left labels are written with right hydrogens: CHO-, C2H5-, etc
         Token &cur = tokens[index];
         Token &next = tokens[index + 1];
         if (cur.type == Token::Element && next.type == Token::Element)
            std::swap(tokens[index], tokens[index + 1]);

         try_swapped = true;
         continue;
      }

      return false;
   }
   return true;
}

bool AbbreviationExpander::expand (const char *label, int input_order, int output_order, Molecule &m)
{
   TokenChain tokens;
   if (!tokensizeAbbreviation(label, tokens))
      return false;

   if (expand_direction == LEFT)
      std::reverse(tokens.begin(), tokens.end());

   m.clear();

   AttPoint begin_att_point(-1, input_order);
   AttPoint end_att_point = begin_att_point;

   if (!expandParsedTokens(tokens, m, end_att_point))
      return false;
   
   if (end_att_point.order != output_order)
      return false;

   output_index = end_att_point.index;

   return true;
}

struct Options
{
   Options (Direction expand_direction, Direction tokenize_direction, int tokenize_level = 1, bool ignore_case = false) :
      expand_direction(expand_direction), 
      tokenize_direction(tokenize_direction), tokenize_level(tokenize_level) {}

   Direction expand_direction, tokenize_direction;
   int tokenize_level;
   bool ignore_case;

   void set (AbbreviationExpander &expander)
   {
      expander.expand_direction = expand_direction;
      expander.tokenize_direction = tokenize_direction;
      expander.tokenize_level = tokenize_level;
      expander.ignore_case = ignore_case;
   }
};

bool AbbreviationExpander::expandAtomAbbreviation (Molecule &mol, int v)
{
   // Determine bonds configuration: number of bonds on the left and on the right
   Vec3f &pos = mol.getAtomXyz(v);
   const Vertex &vertex = mol.getVertex(v);
   int count_left = 0, count_right = 0, count_middle = 0;
   Array<int> left_atoms, right_atoms;
   for (int nei = vertex.neiBegin(); nei != vertex.neiEnd(); nei = vertex.neiNext(nei))
   {
      int nei_atom = vertex.neiVertex(nei);
      Vec3f &nei_pos = mol.getAtomXyz(nei_atom);
      int order = mol.getBondOrder(vertex.neiEdge(nei));
      if (nei_pos.x < pos.x)
      {
         count_left += order;
         left_atoms.push(nei_atom);
      }
      else
      {
         count_right += order;
         right_atoms.push(nei_atom);
      }

      Vec3f diff = nei_pos;
      diff.sub(pos);
      if (fabs(diff.y) > fabs(diff.x))
         count_middle += order;
   }

   int input_order, output_order;

   bool on_left = false, on_right = false, is_single = false;
   if (vertex.degree() == 1)
   {
      if (count_middle)
         on_left = on_right = true;
      if (count_left)
         on_right = true;
      if (count_right)
         on_left = true;
      input_order = std::max(count_left, count_right);
      output_order = 0;
      is_single = true;
   }
   else
   {
      on_right = true;
      input_order = count_left;
      output_order = count_right;
   }

   // Try to expand according to the directions
   Array<Options> options;

   if (on_right && !on_left)
   {
      options.push(Options(RIGHT, RIGHT));
      options.push(Options(LEFT, LEFT));
   }
   else
   {
      if (on_right)
         options.push(Options(RIGHT, RIGHT));
      if (on_left)
      {
         options.push(Options(LEFT, LEFT));
         options.push(Options(LEFT, RIGHT));
         options.push(Options(RIGHT, RIGHT));
         options.push(Options(LEFT, LEFT, 2));
      }
   }

   // Duplicate all the options with ignore case flag
   int opt_cnt = options.size();
   for (int i = 0; i < opt_cnt; i++)
   {
      Options opt = options[i];
      opt.ignore_case = true;
      options.push(opt);
   }

   bool found = false;
   Molecule expanded;
   for (int i = 0; i < options.size(); i++)
   {
      options[i].set(*this);

      if (expand(mol.getPseudoAtom(v), input_order, output_order, expanded))
      {
         found = true;
         break;
      }
   }

   if (!found)
      return false;

   // Merge expanded abbreviation with the source molecule and connect bonds
   Array<int> mapping;
   mol.mergeWithMolecule(expanded, &mapping);

   //if (right_atoms.size() == 0)
   //   right_atoms.swap(left_atoms);

   for (int i = 0; i < left_atoms.size(); i++)
      mol.flipBond(left_atoms[i], v, mapping[input_index]);

   int target_end = output_index;
   if (is_single == 1)
      target_end = input_index;

   for (int i = 0; i < right_atoms.size(); i++)
      mol.flipBond(right_atoms[i], v, mapping[target_end]);

   // Collect added atoms and set their coordinate to the initial atom
   // Layout procedure will find correct atom coordinates, but neighbours 
   // should know correct relative position
   for (int ve = expanded.vertexBegin(); ve != expanded.vertexEnd(); ve = expanded.vertexNext(ve))
   {
      int idx = mapping[ve];
      mol.setAtomXyz(idx, mol.getAtomXyz(v));
      added_atoms.push(mapping[ve]);
   }

   int sid = mol.sgroups.addSGroup(SGroup::SG_TYPE_SUP);
   Superatom &super = (Superatom &)mol.sgroups.getSGroup(sid);
   super.subscript.readString(mol.getPseudoAtom(v), true);
   for (int ve = expanded.vertexBegin(); ve != expanded.vertexEnd(); ve = expanded.vertexNext(ve))
      super.atoms.push(mapping[ve]);

   mol.removeAtom(v);

   return true;
}

//
// Interface functions
//
CEXPORT int indigoExpandAbbreviations (int molecule)
{
   INDIGO_BEGIN
   {
      AbbreviationExpander expander(indigoGetAbbreviationsInstance().abbreviations);

      IndigoObject &obj = self.getObject(molecule);
      Molecule &mol = obj.getMolecule();

      float avg_bond_length = 0;
      if (Molecule::hasCoord(mol))
      {
         // Detect average bond length
         for (int e = mol.edgeBegin(); e != mol.edgeEnd(); e = mol.edgeNext(e))
         {
            const Edge &edge = mol.getEdge(e);
            Vec3f p1 = mol.getAtomXyz(edge.beg);
            Vec3f p2 = mol.getAtomXyz(edge.end);
            Vec3f diff;
            diff.diff(p1, p2);
            avg_bond_length += diff.length();
         }
         avg_bond_length /= mol.edgeCount();
      }

      // Collect pseudoatoms first because abbreviation can change atom ordering because 
      // it deletes atoms and adds new ones
      Array<int> pseudoatoms;
      for (int v = mol.vertexBegin(); v != mol.vertexEnd(); v = mol.vertexNext(v))
      {
         if (mol.isPseudoAtom(v))
            pseudoatoms.push(v);
      }
      int count = 0;
      for (int i = 0; i < pseudoatoms.size(); i++)
      {
         // Try to expand this pseudoatom
         count += expander.expandAtomAbbreviation(mol, pseudoatoms[i]);
      }

      if (count > 0 && Molecule::hasCoord(mol))
      {
         // Layout expanded parts
         MoleculeLayout ml(mol, self.smart_layout);
         ml.max_iterations = self.layout_max_iterations;
         ml.respect_existing_layout = true;
         // Check if there are not bonds at all or if the coordinates are invalid
         if (fabs(avg_bond_length) > 0.001)
            ml.bond_length = avg_bond_length;

         Filter f;
         f.initNone(mol.vertexEnd());
         for (int i = 0; i < expander.added_atoms.size(); i++)
            f.unhide(expander.added_atoms[i]);
         ml.filter = &f;

         ml.make();
         RedBlackSet<int> atoms_with_nei;

         for (int i = 0; i < expander.added_atoms.size(); i++)
         {
            const Vertex &vertex = mol.getVertex(expander.added_atoms[i]);
            for (int nei = vertex.neiBegin(); nei != vertex.neiEnd(); nei = vertex.neiNext(nei))
            {
               int nei_atom = vertex.neiVertex(nei);
               atoms_with_nei.find_or_insert(nei_atom);
            }
            atoms_with_nei.find_or_insert(expander.added_atoms[i]);
         }
         for (int i = atoms_with_nei.begin(); i != atoms_with_nei.end(); i = atoms_with_nei.next(i))
            mol.stereocenters.markBond(atoms_with_nei.key(i));
      }

      return count;
   }
   INDIGO_END(-1)
}

}} // namespace indigo::abbreviations