/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2020 Frederik Heber. All rights reserved.
 * 
 *
 *   This file is part of MoleCuilder.
 *
 *    MoleCuilder is free software: you can redistribute it and/or modify
 *    it under the terms of the GNU General Public License as published by
 *    the Free Software Foundation, either version 2 of the License, or
 *    (at your option) any later version.
 *
 *    MoleCuilder is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU General Public License for more details.
 *
 *    You should have received a copy of the GNU General Public License
 *    along with MoleCuilder.  If not, see <http://www.gnu.org/licenses/>.
 */

/*
 * BondifyAction.cpp
 *
 *  Created on: Oct 07, 2020
 *      Author: heber
 */

// include config.h
#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

//#include "CodePatterns/MemDebug.hpp"

#include "CodePatterns/Assert.hpp"
#include "CodePatterns/Log.hpp"

#include <string>

#include "Actions/AtomAction/BondifyAction.hpp"
#include "Actions/UndoRedoHelpers.hpp"
#include "Atom/atom.hpp"
#include "Atom/AtomicInfo.hpp"
#include "Bond/bond.hpp"
#include "Bond/BondInfo.hpp"
#include "Descriptors/AtomsWithinDistanceOfDescriptor.hpp"
#include "Graph/BondGraph.hpp"
#include "RandomNumbers/RandomNumberDistributionFactory.hpp"
#include "RandomNumbers/RandomNumberDistribution.hpp"
#include "RandomNumbers/RandomNumberGenerator.hpp"
#include "RandomNumbers/RandomNumberGeneratorFactory.hpp"
#include "World.hpp"

using namespace MoleCuilder;

// and construct the stuff
#include "BondifyAction.def"
#include "Action_impl_pre.hpp"

typedef std::vector< std::pair<atom *, int> > candidates_t;

const double MAX_DISTANCE = 5.0;

static int getNumberOfHydrogenAtoms(
    const atom *Walker)
{
  int num_hydrogens = 0;
  const BondList& ListOfBonds = Walker->getListOfBonds();
  for (BondList::const_iterator bonditer = ListOfBonds.begin();
      bonditer != ListOfBonds.end(); ++bonditer) {
    num_hydrogens += (int)((*bonditer)->GetOtherAtom(Walker)->getElementNo() == 1);
  }
  return num_hydrogens;
}

static candidates_t getCandidatesInVicinity(
    const BondGraph &BG,
    const atom * const Walker)
{
  const std::vector< atom *> atoms_in_vicinity = World::getInstance().getAllAtoms(
      AtomsWithinDistanceOf(MAX_DISTANCE, Walker->getPosition()));
  candidates_t candidates;
  for (std::vector<atom *>::const_iterator iter = atoms_in_vicinity.begin();
      iter != atoms_in_vicinity.end(); ++iter) {
    const atom *OtherWalker = *iter;
    // skip hydrogens
    if (OtherWalker->getElementNo() == 1)
      continue;
    // skip atoms that are bonded already
    if (OtherWalker->IsBondedTo(Walker))
      continue;
    const double distance = OtherWalker->getPosition().distance(Walker->getPosition());
    range<double> typical_distance = BG.getMinMaxDistance(Walker, OtherWalker);
    if (typical_distance.isInRange(distance)) {
      int num_hydrogens = getNumberOfHydrogenAtoms(OtherWalker);
      candidates.push_back( std::make_pair(const_cast<atom *>(OtherWalker), num_hydrogens));
    }
  }
  return candidates;
}

static int getTotalBondDegree(
    const atom * Walker)
{
  int num_degreed_bonds = 0;
  const BondList& ListOfBonds = Walker->getListOfBonds();
  for (BondList::const_iterator bonditer = ListOfBonds.begin();
      bonditer != ListOfBonds.end(); ++bonditer) {
    num_degreed_bonds += (*bonditer)->getDegree();
  }
  return num_degreed_bonds;
}

class CandidatesAtValence {
public:
  typedef std::multimap<int, atom*> num_hydrogens_atom_map_t;
  typedef std::vector< std::pair<atom *, int> > candidates_hydrogens_to_remove_t;
  typedef std::vector<candidates_hydrogens_to_remove_t> all_candidates_t;

  int max_hydrogens;

  CandidatesAtValence(const int _max_hydrogens) : max_hydrogens(_max_hydrogens) {}

  void setMaxHydrogens(const int _max_hydrogens) {
    max_hydrogens = _max_hydrogens;
  }

  all_candidates_t operator()(
      const int open_valence,
      const num_hydrogens_atom_map_t& sorted_candidates) const
  {
    all_candidates_t all_candidates;
    if (sorted_candidates.empty())
      return all_candidates;


   /* pick from beginning (i.e. smallest valence first) and gather all
    * possible candidate sets.
    */
   int remaining_valence = open_valence;
   num_hydrogens_atom_map_t::const_iterator iter = sorted_candidates.begin();
   candidates_hydrogens_to_remove_t candidate_set;
   exploreCandidateSetRecursively(
       iter, sorted_candidates.end(),
       candidate_set, remaining_valence,
       all_candidates);

   for (all_candidates_t::const_iterator all_iter = all_candidates.begin();
       all_iter != all_candidates.end(); ++all_iter) {
     LOG(2, "DEBUG: Candidate set is " << candidateToString(*all_iter));
   }

   return all_candidates;
  }

  void exploreCandidateSetRecursively(
      num_hydrogens_atom_map_t::const_iterator iter,
      num_hydrogens_atom_map_t::const_iterator iter_end,
      candidates_hydrogens_to_remove_t candidate_set,
      int remaining_valence,
      all_candidates_t& all_candidates
      ) const
  {
    const int & num_hydrogens = iter->first;
    atom * const walker = iter->second;
    ++iter;
    const int usable_valence =
        std::min(
            std::min(num_hydrogens, remaining_valence),
            max_hydrogens
        );
    candidate_set.push_back(
        std::make_pair(walker, 0)
    );
    for (int i=0;i<= usable_valence; ++i) {
      candidate_set.back().second = i;
      // check whether we are complete with the candidate_set
      if (remaining_valence - i <= 0) {
        // we have a complete set
        ASSERT(remaining_valence -i == 0,
            "exploreCandidateSetRecursively() - usable valence "+toString(i)
            +" must always be less or equal remaining valence "+toString(remaining_valence));
        all_candidates.push_back(candidate_set);
      } else {
        // not complete, need to recurse further if we can
        if (iter != iter_end)
          exploreCandidateSetRecursively(
              iter,
              iter_end,
              candidate_set,
              remaining_valence-i,
              all_candidates);
      }
    }
  }

  static std::string candidateToString(
      const candidates_hydrogens_to_remove_t &candidates)
  {
    std::stringstream output;
    for (candidates_hydrogens_to_remove_t::const_iterator iter = candidates.begin();
        iter != candidates.end(); ++iter)
      output << "[" << *iter->first << "," << iter->second << "], ";
    return output.str();
  }


};

/** =========== define the function ====================== */
ActionState::ptr AtomBondifyAction::performCall() {
  // ensure we have one selected atom
  World &world = World::getInstance();
  const BondGraph &BG = *world.getBondGraph();

  // check precondition
  const std::vector<atom *> atoms = world.getSelectedAtoms();
  if (atoms.size() != 1) {
    STATUS("Exactly one atom must be selected for bondify.");
    return Action::failure;
  }
  atom *Walker = atoms[0];

  // check whether atom has unoccupied valence orbitals
  const int num_valence = Walker->getType()->getNoValenceOrbitals();

  // Look at all bond neighbors
  const int num_degreed_bonds = getTotalBondDegree(Walker);
  int open_valence = num_valence - num_degreed_bonds;
  LOG(1, "INFO: We have " << open_valence << " unoccupied  valence");

  if (open_valence <= 0) {
    STATUS("Nothing to do. Atom is saturated already, we need some unoccupied valence orbitals to have something to work on.");
    return Action::success;
  }

  // gather all atoms in vicinity
  candidates_t candidates = getCandidatesInVicinity(BG, Walker);


  /** We revert the map's key and values, as we want the entries
   * sorted by the number of hydrogens.
   */
  std::multimap<int, atom*> sorted_candidates;
  for (candidates_t::const_iterator iter = candidates.begin();
      iter != candidates.end(); ++iter) {
    if (iter->second > 0) {
      sorted_candidates.insert( std::make_pair(iter->second, iter->first) );
    }
  }

  /**
   * Try to find the maximum number of bonds to use instead of having to
   * add hydrogens in the end.
   */
  CandidatesAtValence::all_candidates_t all_candidates;
  CandidatesAtValence candidateGetter(params.max_hydrogens.get());
  ++open_valence;
  while ((all_candidates.size() == 0) && (--open_valence > 0)) {
    all_candidates = candidateGetter(open_valence, sorted_candidates);
  }
  if (all_candidates.empty()) {
    STATUS("Could not find any suitable candidate sets.");
    return Action::success;
  }

  // pick one set at random
  const std::string oldtype = RandomNumberDistributionFactory::getConstInstance().getCurrentTypeName();
  RandomNumberDistributionFactory::getInstance().setCurrentType("uniform_01");
  const RandomNumberGenerator& rng = RandomNumberGeneratorFactory::getConstInstance().makeRandomNumberGenerator();
  int id = (int)(rng()*(all_candidates.size()-1));
  RandomNumberDistributionFactory::getInstance().setCurrentType(oldtype);
  CandidatesAtValence::candidates_hydrogens_to_remove_t &picked_set = all_candidates[id];
  LOG(1, "INFO: Picked candidate set is " << CandidatesAtValence::candidateToString(picked_set));

  // execute that
  std::vector<AtomicInfo> removedHydrogens;
  std::vector<BondInfo> addedBonds;
  int removed_hydrogens = 0;
  int added_bonds = 0;
  for (CandidatesAtValence::candidates_hydrogens_to_remove_t::iterator iter = picked_set.begin();
      iter != picked_set.end(); ++iter) {
    atom * const other_walker = iter->first;
    const int &num_hydrogens = iter->second;

    // remove the hydrogens closest to Walker
    typedef std::map<double, atom *> closest_map_t;
    closest_map_t closest_map;
    const BondList& ListOfBonds = other_walker->getListOfBonds();
    for (BondList::const_iterator bonditer = ListOfBonds.begin();
        bonditer != ListOfBonds.end(); ++bonditer) {
      atom * const hydrogen = (*bonditer)->GetOtherAtom(other_walker);
      if (hydrogen->getElementNo() != 1)
        continue;
      closest_map.insert(
          std::make_pair(
            Walker->getPosition().distance(hydrogen->getPosition()),
            hydrogen
          )
      );
    }
    ASSERT( closest_map.size() >= (size_t)num_hydrogens,
        "AtomBondifyAction::performCall() - Atom "+toString(*other_walker)
        +" has less hydrogens "+toString(closest_map.size())
        +" than expected "+toString(num_hydrogens));
    closest_map_t::iterator removeiter = closest_map.begin();
    for (int i=0;i<num_hydrogens; ++i) {
      ASSERT( removeiter != closest_map.end(),
          "AtomBondifyAction::performCall() - already at end of closest map?");
      LOG(2, "DEBUG: Removing hydrogen at distance " << removeiter->first
          << " to Walker at " << Walker->getPosition());
      removedHydrogens.push_back(AtomicInfo(*removeiter->second));
      world.destroyAtom(removeiter->second);
      removeiter->second = NULL;
      ++removeiter;
      ++removed_hydrogens;
    }

    // add the bond
    bond::ptr new_bond = Walker->addBond(other_walker);
    addedBonds.push_back(BondInfo(new_bond));
    new_bond->setDegree(num_hydrogens);
    ++added_bonds;
    LOG(2, "DEBUG: Added new bond " << *new_bond);
  }
  ASSERT( removed_hydrogens == open_valence,
      "AtomBondifyAction::performCall() - candidate set has not the number of hydrogens "
      +toString(removed_hydrogens)+" to match the unoccupied valence "+toString(open_valence));

  LOG(1, "INFO: Removed " << removed_hydrogens << " hydrogen atoms and added "
      << added_bonds << " new bonds.");

  return ActionState::ptr(new AtomBondifyState(removedHydrogens, addedBonds, params));
}

ActionState::ptr AtomBondifyAction::performUndo(ActionState::ptr _state) {
  AtomBondifyState *state = assert_cast<AtomBondifyState*>(_state.get());

  AddAtomsFromAtomicInfo(state->removedHydrogens);
  RemoveBondsFromBondInfo(state->addedBonds);

  return ActionState::ptr(_state);
}

ActionState::ptr AtomBondifyAction::performRedo(ActionState::ptr _state){
  AtomBondifyState *state = assert_cast<AtomBondifyState*>(_state.get());

  RemoveAtomsFromAtomicInfo(state->removedHydrogens);
  AddBondsFromBondInfo(state->addedBonds);

  return ActionState::ptr(_state);
}

bool AtomBondifyAction::canUndo() {
  return true;
}

bool AtomBondifyAction::shouldUndo() {
  return true;
}
/** =========== end of function ====================== */