#include <openbabel/mol.h>
#include <openbabel/isomorphism.h>
#include <openbabel/obiter.h>
#include <openbabel/oberror.h>
#include <openbabel/query.h>
#include <openbabel/graphsym.h>
#include <ctime>
#include <cassert>
#include <algorithm>

#define DEBUG 0

using namespace std;

namespace OpenBabel {

  template<typename T>
    void print_vector(const std::string &label, const std::vector<T> &v)
    {
      std::cout << label << ": ";
      for (std::size_t i = 0; i < v.size(); ++i)
        if (v[i] < 10)
          std::cout << " " << v[i] << " ";
        else
          std::cout << v[i] << " ";

      std::cout << endl;
    }

  static const char *red    = "\033[1;31m";
  static const char *green  = "\033[1;32m";
  static const char *yellow = "\033[1;33m";
  static const char *blue   = "\033[1;34m";
  static const char *normal = "\033[0m";

  class MapAllFunctor : public OBIsomorphismMapper::Functor
  {
    private:
      OBIsomorphismMapper::Mappings &m_maps;
      std::size_t m_memory, m_maxMemory;
    public:
      MapAllFunctor(OBIsomorphismMapper::Mappings &maps, std::size_t maxMemory)
        : m_maps(maps), m_memory(0), m_maxMemory(maxMemory)
      {
      }
      bool operator()(OBIsomorphismMapper::Mapping &map) override
      {
        m_maps.push_back(map);
        m_memory += 2 * sizeof(unsigned int) * map.size();

        if (m_memory > m_maxMemory) {
          obErrorLog.ThrowError(__FUNCTION__, "memory limit exceeded...", obError);
          // stop mapping
          return true;
        }

        // continue mapping
        return false;
      }
  };


  class VF2Mapper : public OBIsomorphismMapper
  {
    time_t m_startTime;

    public:
      VF2Mapper(OBQuery *query) : OBIsomorphismMapper(query)
      {
      }

      struct Candidate {
        Candidate() : queryAtom(nullptr), queriedAtom(nullptr) {}
        Candidate(OBQueryAtom *_queryAtom, OBAtom *_queriedAtom)
            : queryAtom(_queryAtom), queriedAtom(_queriedAtom) {}

        bool operator==(const Candidate &other)
        {
          if (queryAtom != other.queryAtom)
            return false;
          if (queriedAtom != other.queriedAtom)
            return false;
          return true;
        }

        OBQueryAtom *queryAtom;
        OBAtom *queriedAtom;
      };

      struct State {
        State(Functor &_functor, const OBQuery *_query, const OBMol *_queried, const OBBitVec &mask)
            : functor(_functor), query(_query), queried(_queried), queriedMask(mask)
        {
          abort = false;
          mapping.resize(query->NumAtoms(), nullptr);
          queryDepths.resize(query->NumAtoms(), 0);
          queriedDepths.resize(queried->NumAtoms(), 0);
        }
        bool abort;
        Functor &functor;
        const OBQuery *query; // the query
        const OBMol *queried; // the queried molecule
        OBBitVec queriedMask; // the queriedMask
        std::vector<unsigned int> queryPath; // the path in the query
        std::vector<unsigned int> queriedPath; // the path in the queried molecule

        std::vector<OBAtom*> mapping;

        OBBitVec queryPathBits, queriedPathBits; // the terminal sets
        std::vector<unsigned int> queryDepths, queriedDepths; // the terminal sets
      };

      bool isInTerminalSet(const std::vector<unsigned int> &depths,
          const OBBitVec &path, std::size_t i)
      {
        if (!depths[i])
          return false;

        if (path.BitIsSet(i))
          return false;

        return true;
      }

      /**
       * Check bonds around newly mapped atom.
       */
      bool checkBonds(State &state, OBQueryAtom *queryAtom)
      {
        const std::vector<OBQueryBond*> &bonds = queryAtom->GetBonds();
        for (unsigned int i = 0; i < bonds.size(); ++i) {
          OBQueryBond *qbond = bonds[i];
          unsigned int beginIndex = qbond->GetBeginAtom()->GetIndex();
          unsigned int endIndex = qbond->GetEndAtom()->GetIndex();

          OBAtom *begin = state.mapping[beginIndex];
          OBAtom *end = state.mapping[endIndex];
          if (!begin || !end)
            continue;
          OBBond *bond = state.queried->GetBond(begin, end);
          if (!bond)
            return false;
          if (!qbond->Matches(bond))
            return false;
        }
        return true;
      }


      /**
       * Check if the current state is a full mapping of the query.
       */
      bool checkForMap(State &state)
      {
        // store the mapping if all atoms are mapped
        if (state.queryPath.size() != state.query->NumAtoms())
          return false;

        if (DEBUG)
          std::cout << green << "-----------------> MATCH" << normal << std::endl;

        // create the map
        Mapping map;
        map.reserve(state.queryPath.size());
        for (unsigned int k = 0; k < state.queryPath.size(); ++k)
          map.push_back(std::make_pair(state.queryPath[k], state.queriedPath[k]));

        return state.functor(map);
      }

      /**
       * Match the candidate atoms and bonds.
       */
      bool matchCandidate(State &state, OBQueryAtom *queryAtom, OBAtom *queriedAtom)
      {
        if (!queryAtom->Matches(queriedAtom))
          return false;

        // add the neighbors to the paths
        state.queryPath.push_back(queryAtom->GetIndex());
        state.queriedPath.push_back(queriedAtom->GetIndex());
        // update the terminal sets
        state.queryPathBits.SetBitOn(queryAtom->GetIndex());
        state.queriedPathBits.SetBitOn(queriedAtom->GetIndex());
        // update mapping
        state.mapping[queryAtom->GetIndex()] = queriedAtom;

        //
        // update queryDepths
        //
        if (!state.queryDepths[queryAtom->GetIndex()])
          state.queryDepths[queryAtom->GetIndex()] = state.queryPath.size();

        std::vector<OBQueryAtom*> queryNbrs = queryAtom->GetNbrs();
        for (unsigned int i = 0; i < queryNbrs.size(); ++i) {
          unsigned int index = queryNbrs[i]->GetIndex();
          if (!state.queryDepths[index])
            state.queryDepths[index] = state.queryPath.size();
        }

        //
        // update queriedDepths
        //
        if (!state.queriedDepths[queriedAtom->GetIndex()])
          state.queriedDepths[queriedAtom->GetIndex()] = state.queriedPath.size();

        FOR_NBORS_OF_ATOM (nbr, queriedAtom) {
          unsigned int index = nbr->GetIndex();
          // skip atoms not in the mask
          if (!state.queriedMask.BitIsSet(index + 1))
            continue;
          if (!state.queriedDepths[index])
            state.queriedDepths[index] = state.queriedPath.size();
        }

        // check if the bonds match
        if (!checkBonds(state, queryAtom)) {
          if (DEBUG)
            cout << "    bonds do not match..." << endl;
          Backtrack(state);
          return false;
        }

        if (DEBUG) {
          cout << "FOUND:  " << queryAtom->GetIndex() << " -> " << queriedAtom->GetIndex() << "       " << state.queryPath.size() << endl;
          cout << "queryDepths:   ";
          for (unsigned int i = 0; i < state.query->NumAtoms(); ++i)
            cout << state.queryDepths[i] << " ";
          cout <<endl;
          cout << "queriedDepths: ";
          for (unsigned int i = 0; i < state.queried->NumAtoms(); ++i)
            cout << state.queriedDepths[i] << " ";
          cout <<endl;
        }

        //
        // Feasibility rules for the VF2 algorithm:
        //
        //  size( T1(s) ) <= size( T2(s) )
        //
        //  size( N1 - M1(s) - T1(s) ) <= size( N2 - M2(s) - T2(s) )
        //

        // compute T1(s) size
        unsigned int numT1 = 0;
        for (unsigned int i = 0; i < state.query->NumAtoms(); ++i)
          if (isInTerminalSet(state.queryDepths, state.queryPathBits, i))
              numT1++;
        // compute T2(s) size
        unsigned int numT2 = 0;
        for (unsigned int i = 0; i < state.queried->NumAtoms(); ++i)
          if (isInTerminalSet(state.queriedDepths, state.queriedPathBits, i))
              numT2++;

        // T1(s) > T()
        if (numT1 > numT2) {
          Backtrack(state);
          return false;
        }
        //  N1 - M1(s) - T1(s) > N2 - M2(s) - T2(s)
        if ((state.query->NumAtoms() - state.queryPath.size() - numT1) > (state.queried->NumAtoms() - state.queriedPath.size() - numT2)) {
          Backtrack(state);
          return false;
        }

        // Check if there is a mapping found
        state.abort = checkForMap(state);

        return true;
      }


      Candidate NextCandidate(State &state, const Candidate &lastCandidate)
      {
        std::size_t lastQueryAtom = lastCandidate.queryAtom ? lastCandidate.queryAtom->GetIndex() : 0;
        std::size_t lastQueriedAtom = lastCandidate.queriedAtom ? lastCandidate.queriedAtom->GetIndex() + 1 : 0;

        std::size_t querySize = state.query->NumAtoms();
        std::size_t queriedSize = state.queried->NumAtoms();

        std::size_t queryTerminalSize = state.queryDepths.size() - std::count(state.queryDepths.begin(), state.queryDepths.end(), 0);
        std::size_t queriedTerminalSize = state.queriedDepths.size() - std::count(state.queriedDepths.begin(), state.queriedDepths.end(), 0);

        std::size_t mappingSize = state.queryPath.size();

        if (queryTerminalSize > mappingSize && queriedTerminalSize > mappingSize) {
          while (lastQueryAtom < querySize && (state.queryPathBits.BitIsSet(lastQueryAtom) || !state.queryDepths[lastQueryAtom])) {
            lastQueryAtom++;
            lastQueriedAtom = 0;
          }
        } else {
          while(lastQueryAtom < querySize && state.queryPathBits.BitIsSet(lastQueryAtom)) {
            lastQueryAtom++;
            lastQueriedAtom = 0;
          }
        }

        if (queryTerminalSize > mappingSize && queriedTerminalSize > mappingSize) {
          while (lastQueriedAtom < queriedSize && (state.queriedPathBits.BitIsSet(lastQueriedAtom) || !state.queriedDepths[lastQueriedAtom]))
            lastQueriedAtom++;
        } else {
          while(lastQueriedAtom < queriedSize && state.queriedPathBits[lastQueriedAtom])
            lastQueriedAtom++;
        }

        if (lastQueryAtom < querySize && lastQueriedAtom < queriedSize)
          return Candidate(state.query->GetAtoms()[lastQueryAtom], state.queried->GetAtom(lastQueriedAtom + 1));

        return Candidate();
      }

      /**
       * The depth-first isomorphism algorithm.
       */
      void MapNext(State &state, OBQueryAtom * /*queryAtom*/, OBAtom * /*queriedAtom*/)
      {
        if (time(nullptr) - m_startTime > m_timeout)
          return;
        if (state.abort)
          return;

        Candidate candidate;
        while (!state.abort) {
          candidate = NextCandidate(state, candidate);

          if (!candidate.queryAtom)
            return;

          if (DEBUG)
            cout << yellow << "candidate: " << candidate.queryAtom->GetIndex() << " -> " << candidate.queriedAtom->GetIndex() << normal << endl;


          if (matchCandidate(state, candidate.queryAtom, candidate.queriedAtom)) {
            MapNext(state, candidate.queryAtom, candidate.queriedAtom);
            Backtrack(state);
          }
        }

      }

      void Backtrack(State &state)
      {
        if (DEBUG)
          cout << red << "backtrack... " << normal << state.queryPath.size()-1 << endl;
        // remove last atoms from the mapping
        if (state.queryPath.size()) {
          state.mapping[state.queryPath.back()] = nullptr;
          state.queryPathBits.SetBitOff(state.queryPath.back());
          state.queryPath.pop_back();
        }
        if (state.queriedPath.size()) {
          state.queriedPathBits.SetBitOff(state.queriedPath.back());
          state.queriedPath.pop_back();
        }
        // restore queryDepths and queriedDepths
        unsigned int depth = state.queryPath.size() + 1;
        std::replace(state.queryDepths.begin(), state.queryDepths.end(), depth, static_cast<unsigned int>(0));
        std::replace(state.queriedDepths.begin(), state.queriedDepths.end(), depth, static_cast<unsigned int>(0)); // O(n)  n = # vertices in the queried
      }

      /**
       * Get the first mappings of the query on the queried molecule.
       * @param queried The queried molecule.
       * @return The mapping.
       */
      void MapFirst(const OBMol *queried, Mapping &map, const OBBitVec &mask) override
      {
        class MapFirstFunctor : public Functor
        {
          private:
            Mapping &m_map;
          public:
            MapFirstFunctor(Mapping &map) : m_map(map)
            {
            }
            bool operator()(Mapping &map) override
            {
              m_map = map;
              // stop mapping
              return true;
            }
        };

        MapFirstFunctor functor(map);
        MapGeneric(functor, queried, mask);
      }

      /**
       * Get all unique mappings of the query on the queried molecule.
       * @param queried The queried molecule.
       * @return The unique mappings
       */
      void MapUnique(const OBMol *queried, Mappings &maps, const OBBitVec &mask) override
      {
        class MapUniqueFunctor : public OBIsomorphismMapper::Functor
        {
          private:
            OBIsomorphismMapper::Mappings &m_maps;
          public:
            MapUniqueFunctor(OBIsomorphismMapper::Mappings &maps) : m_maps(maps)
            {
            }
            bool operator()(OBIsomorphismMapper::Mapping &map) override
            {
              // get the values from the map
              std::vector<unsigned int> values;
              for (OBIsomorphismMapper::Mapping::const_iterator it = map.begin(); it != map.end(); ++it)
                values.push_back(it->second);
              std::sort(values.begin(), values.end());
               // print_vector("values ", values);

              bool isUnique = true;
              for (unsigned int k = 0; k < m_maps.size(); ++k) {
                std::vector<unsigned int> kValues;
                for (OBIsomorphismMapper::Mapping::iterator it = m_maps[k].begin(); it != m_maps[k].end(); ++it)
                  kValues.push_back(it->second);
                std::sort(kValues.begin(), kValues.end());

              //  print_vector("kValues", kValues);
                if (values == kValues)
                  isUnique = false;
              }

              if (isUnique)
                m_maps.push_back(map);

              // continue mapping
              return false;
            }
        };


        maps.clear();
        MapUniqueFunctor functor(maps);
        MapGeneric(functor, queried, mask);

        if (DEBUG)
          for (unsigned int i =0; i < maps.size(); ++i) {
            cout << "mapping:" << endl;
            for (Mapping::iterator it = maps[i].begin(); it != maps[i].end(); ++it)
              cout << "    " << it->first << " -> " << it->second << endl;
          }
      }

      /**
       * Get all mappings of the query on the queried molecule. Duplicates are
       * ignored but unlinke MapUnique, multiple mappings of the query on the same
       * part of the queried structure are allowed. This makes it possible to use
       * MapAll for finding the automorphism group.
       * @param queried The queried molecule.
       * @return The mappings.
       */
      void MapAll(const OBMol *queried, Mappings &maps, const OBBitVec &mask, std::size_t maxMemory) override
      {
        maps.clear();
        MapAllFunctor functor(maps, maxMemory);
        MapGeneric(functor, queried, mask);

        if (DEBUG)
          for (unsigned int i =0; i < maps.size(); ++i) {
            cout << "mapping:" << endl;
            for (Mapping::iterator it = maps[i].begin(); it != maps[i].end(); ++it)
              cout << "    " << it->first << " -> " << it->second << endl;
          }

      }

      void MapGeneric(Functor &functor, const OBMol *queried, const OBBitVec &mask) override
      {
        m_startTime = time(nullptr);
        if(m_query->NumAtoms() == 0) return;
        // set all atoms to 1 if the mask is empty
        OBBitVec queriedMask = mask;
        if (!queriedMask.CountBits())
          for (unsigned int i = 0; i < queried->NumAtoms(); ++i)
            queriedMask.SetBitOn(i + 1);

        OBQueryAtom *queryAtom = m_query->GetAtoms()[0];
        for (unsigned int i = 0; i < queried->NumAtoms(); ++i) {
          if (!queriedMask.BitIsSet(i + 1)) {
            continue;
          }
          State state(functor, m_query, queried, queriedMask);
          OBAtom *queriedAtom = queried->GetAtom(i+1);
          if (!queryAtom->Matches(queriedAtom)) {
            continue;
          }
          if (DEBUG)
            std::cout << blue << "START: 0 -> " << queriedAtom->GetIndex() << normal << std::endl;

          if (m_query->NumAtoms() > 1) {
            if (matchCandidate(state, queryAtom, queriedAtom))
              MapNext(state, queryAtom, queriedAtom);
          } else {
            Mapping map;
            map.push_back(std::make_pair(queryAtom->GetIndex(), queriedAtom->GetIndex()));
            functor(map);
          }
        }

        if (time(nullptr) - m_startTime > m_timeout)
          obErrorLog.ThrowError(__FUNCTION__, "time limit exceeded...", obError);

      }

  };

  OBIsomorphismMapper::OBIsomorphismMapper(OBQuery *query) : m_query(query), m_timeout(60)
  {
  }

  OBIsomorphismMapper::~OBIsomorphismMapper()
  {
  }

  OBIsomorphismMapper* OBIsomorphismMapper::GetInstance(OBQuery *query, const std::string &algorithm)
  {
    if (algorithm == "VF2")
      return new VF2Mapper(query);
    // return VF2 mapper as default
    return new VF2Mapper(query);
  }







  class OBAutomorphismQueryAtom : public OBQueryAtom
  {
    public:
      OBAutomorphismQueryAtom(unsigned int _symClass, const std::vector<unsigned int> &_symClasses)
          : OBQueryAtom(), symClass(_symClass), symClasses(_symClasses)
      {
      }

      bool Matches(const OBAtom *atom) const override
      {
        return (symClasses[atom->GetIndex()] == symClass);
      }
      unsigned int symClass;
      std::vector<unsigned int> symClasses;
  };

  bool isFerroceneBond(OBBond *bond);

  OBQuery* CompileAutomorphismQuery(OBMol *mol, const OBBitVec &mask, const std::vector<unsigned int> &symClasses)
  {
    OBQuery *query = new OBQuery;
    unsigned int offset = 0;
    std::vector<unsigned int> indexes;
    FOR_ATOMS_OF_MOL (obatom, mol) {
      indexes.push_back(obatom->GetIndex() - offset);
      if (!mask.BitIsSet(obatom->GetIndex() + 1)) {
        offset++;
        continue;
      }
      query->AddAtom(new OBAutomorphismQueryAtom(symClasses[obatom->GetIndex()], symClasses));
    }
    FOR_BONDS_OF_MOL (obbond, mol) {
      if (isFerroceneBond(&*obbond))
        continue;
      unsigned int beginIndex = obbond->GetBeginAtom()->GetIndex();
      unsigned int endIndex = obbond->GetEndAtom()->GetIndex();
      if (!mask.BitIsSet(beginIndex + 1) || !mask.BitIsSet(endIndex + 1))
        continue;
      query->AddBond(new OBQueryBond(query->GetAtoms()[indexes[beginIndex]], query->GetAtoms()[indexes[endIndex]],
            obbond->GetBondOrder(), obbond->IsAromatic()));
    }

    return query;
  }

  bool FindAutomorphisms(OBMol *mol, Automorphisms &maps, const OBBitVec &mask, std::size_t maxMemory)
  {
    // set all atoms to 1 if the mask is empty
    OBBitVec queriedMask = mask;
    if (!queriedMask.CountBits())
      for (unsigned int i = 0; i < mol->NumAtoms(); ++i)
        queriedMask.SetBitOn(i + 1);

    // get the symmetry classes
    OBGraphSym gs(mol, &queriedMask);
    std::vector<unsigned int> symClasses;
    gs.GetSymmetry(symClasses);

    return FindAutomorphisms(mol, maps, symClasses, mask, maxMemory);
  }


  bool FindAutomorphisms(OBMol *mol, Automorphisms &maps, const std::vector<unsigned int> &symClasses, const OBBitVec &mask, std::size_t maxMemory)
  {
    maps.clear();
    MapAllFunctor functor(maps, maxMemory);
    FindAutomorphisms(functor, mol, symClasses, mask);
    return maps.size();
  }

  OBBitVec getFragment(OBAtom *atom, const OBBitVec &mask, const std::vector<OBBond*> &metalloceneBonds = std::vector<OBBond*>());

  void FindAutomorphisms(OBIsomorphismMapper::Functor &functor, OBMol *mol,
      const std::vector<unsigned int> &symClasses, const OBBitVec &mask)
  {
    class AutomorphismFunctor : public OBIsomorphismMapper::Functor
    {
      private:
        OBIsomorphismMapper::Functor &m_functor;
        const OBBitVec &m_fragment;
        std::vector<unsigned int> m_indexes;
      public:
        AutomorphismFunctor(OBIsomorphismMapper::Functor &functor, const OBBitVec &fragment, unsigned int numAtoms)
            : m_functor(functor), m_fragment(fragment)
        {
          for (unsigned int j = 0; j < numAtoms; ++j)
            if (m_fragment.BitIsSet(j + 1))
              m_indexes.push_back(j);
        }
        bool operator()(Automorphism &map) override
        {
          // convert the continuous mapping map to a mapping with gaps (considering key values)
          for (Automorphism::iterator it = map.begin(); it != map.end(); ++it)
            it->first = m_indexes[it->first];
          return m_functor(map);
        }
    };

    // set all atoms to 1 if the mask is empty
    OBBitVec queriedMask = mask;
    if (!queriedMask.CountBits())
      for (unsigned int i = 0; i < mol->NumAtoms(); ++i)
        queriedMask.SetBitOn(i + 1);

    if (DEBUG)
    for (unsigned int i = 0; i < symClasses.size(); ++i)
      cout << i << ": " << symClasses[i] << endl;

    // compute the connected fragments
    OBBitVec visited;
    std::vector<OBBitVec> fragments;
    for (std::size_t i = 0; i < mol->NumAtoms(); ++i) {
      if (!queriedMask.BitIsSet(i+1) || visited.BitIsSet(i+1))
        continue;
      fragments.push_back(getFragment(mol->GetAtom(i+1), queriedMask));
      visited |= fragments.back();
    }

    // count the symmetry classes
    std::vector<int> symClassCounts(symClasses.size() + 1, 0);
    for (unsigned int i = 0; i < symClasses.size(); ++i) {
      if (!queriedMask.BitIsSet(i + 1))
        continue;
      unsigned int symClass = symClasses[i];
      symClassCounts[symClass]++;
    }

    for (std::size_t i = 0; i < fragments.size(); ++i) {
      OBQuery *query = CompileAutomorphismQuery(mol, fragments[i], symClasses);
      OBIsomorphismMapper *mapper = OBIsomorphismMapper::GetInstance(query);

      AutomorphismFunctor autFunctor(functor, fragments[i], mol->NumAtoms());
      mapper->MapGeneric(autFunctor, mol, fragments[i]);
      delete mapper;
      delete query;
    }

  }




}

/// @file isomorphism.cpp
/// @brief OBIsomorphismMapper class for finding isomorphisms.
