/*
 * ForceAnnealing.hpp
 *
 *  Created on: Aug 02, 2014
 *      Author: heber
 */

#ifndef FORCEANNEALING_HPP_
#define FORCEANNEALING_HPP_

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

#include "Atom/atom.hpp"
#include "Atom/AtomSet.hpp"
#include "CodePatterns/Assert.hpp"
#include "CodePatterns/Info.hpp"
#include "CodePatterns/Log.hpp"
#include "CodePatterns/Verbose.hpp"
#include "Descriptors/AtomIdDescriptor.hpp"
#include "Dynamics/AtomicForceManipulator.hpp"
#include "Fragmentation/ForceMatrix.hpp"
#include "Graph/BoostGraphCreator.hpp"
#include "Graph/BoostGraphHelpers.hpp"
#include "Graph/BreadthFirstSearchGatherer.hpp"
#include "Helpers/helpers.hpp"
#include "Helpers/defs.hpp"
#include "LinearAlgebra/Vector.hpp"
#include "Thermostats/ThermoStatContainer.hpp"
#include "Thermostats/Thermostat.hpp"
#include "World.hpp"

/** This class is the essential build block for performing structural optimization.
 *
 * Sadly, we have to use some static instances as so far values cannot be passed
 * between actions. Hence, we need to store the current step and the adaptive-
 * step width (we cannot perform a line search, as we have no control over the
 * calculation of the forces).
 *
 * However, we do use the bond graph, i.e. if a single atom needs to be shifted
 * to the left, then the whole molecule left of it is shifted, too. This is
 * controlled by the \a max_distance parameter.
 */
template <class T>
class ForceAnnealing : public AtomicForceManipulator<T>
{
public:
  /** Constructor of class ForceAnnealing.
   *
   * \note We use a fixed delta t of 1.
   *
   * \param _atoms set of atoms to integrate
   * \param _Deltat time step width in atomic units
   * \param _IsAngstroem whether length units are in angstroem or bohr radii
   * \param _maxSteps number of optimization steps to perform
   * \param _max_distance up to this bond order is bond graph taken into account.
   */
  ForceAnnealing(
      AtomSetMixin<T> &_atoms,
      const double _Deltat,
      bool _IsAngstroem,
      const size_t _maxSteps,
      const int _max_distance,
      const double _damping_factor) :
    AtomicForceManipulator<T>(_atoms, _Deltat, _IsAngstroem),
    maxSteps(_maxSteps),
    max_distance(_max_distance),
    damping_factor(_damping_factor)
  {}

  /** Destructor of class ForceAnnealing.
   *
   */
  ~ForceAnnealing()
  {}

  /** Performs Gradient optimization.
   *
   * We assume that forces have just been calculated.
   *
   *
   * \param CurrentTimeStep current time step (i.e. \f$ t + \Delta t \f$ in the sense of the velocity verlet)
   * \param offset offset in matrix file to the first force component
   * \todo This is not yet checked if it is correctly working with DoConstrainedMD set >0.
   */
  void operator()(
      const int _CurrentTimeStep,
      const size_t _offset,
      const bool _UseBondgraph)
  {
    // make sum of forces equal zero
    AtomicForceManipulator<T>::correctForceMatrixForFixedCenterOfMass(_offset, _CurrentTimeStep);

    // are we in initial step? Then set static entities
    Vector maxComponents(zeroVec);
    if (currentStep == 0) {
      currentDeltat = AtomicForceManipulator<T>::Deltat;
      currentStep = 1;
      LOG(2, "DEBUG: Initial step, setting values, current step is #" << currentStep);

      // always use atomic annealing on first step
      anneal(_CurrentTimeStep, _offset, maxComponents);
    } else {
      ++currentStep;
      LOG(2, "DEBUG: current step is #" << currentStep);

      if (_UseBondgraph)
        annealWithBondGraph(_CurrentTimeStep, _offset, maxComponents);
      else
        anneal(_CurrentTimeStep, _offset, maxComponents);
    }

    LOG(1, "STATUS: Largest remaining force components at step #"
        << currentStep << " are " << maxComponents);

    // are we in final step? Remember to reset static entities
    if (currentStep == maxSteps) {
      LOG(2, "DEBUG: Final step, resetting values");
      reset();
    }
  }

  /** Performs Gradient optimization on the atoms.
   *
   * We assume that forces have just been calculated.
   *
   * \param CurrentTimeStep current time step (i.e. \f$ t + \Delta t \f$ in the sense of the velocity verlet)
   * \param offset offset in matrix file to the first force component
   * \param maxComponents to be filled with maximum force component over all atoms
   */
  void anneal(
      const int CurrentTimeStep,
      const size_t offset,
      Vector &maxComponents)
  {
    for(typename AtomSetMixin<T>::iterator iter = AtomicForceManipulator<T>::atoms.begin();
        iter != AtomicForceManipulator<T>::atoms.end(); ++iter) {
      // atom's force vector gives steepest descent direction
      const Vector oldPosition = (*iter)->getPositionAtStep(CurrentTimeStep-2 >= 0 ? CurrentTimeStep - 2 : 0);
      const Vector currentPosition = (*iter)->getPosition();
      const Vector oldGradient = (*iter)->getAtomicForceAtStep(CurrentTimeStep-2 >= 0 ? CurrentTimeStep - 2 : 0);
      const Vector currentGradient = (*iter)->getAtomicForce();
      LOG(4, "DEBUG: oldPosition for atom " << **iter << " is " << oldPosition);
      LOG(4, "DEBUG: currentPosition for atom " << **iter << " is " << currentPosition);
      LOG(4, "DEBUG: oldGradient for atom " << **iter << " is " << oldGradient);
      LOG(4, "DEBUG: currentGradient for atom " << **iter << " is " << currentGradient);
//      LOG(4, "DEBUG: Force for atom " << **iter << " is " << currentGradient);

      // we use Barzilai-Borwein update with position reversed to get descent
      const Vector PositionDifference = currentPosition - oldPosition;
      const Vector GradientDifference = (currentGradient - oldGradient);
      double stepwidth = 0.;
      if (GradientDifference.Norm() > MYEPSILON)
        stepwidth = fabs(PositionDifference.ScalarProduct(GradientDifference))/
            GradientDifference.NormSquared();
      if (fabs(stepwidth) < 1e-10) {
        // dont' warn in first step, deltat usage normal
        if (currentStep != 1)
          ELOG(1, "INFO: Barzilai-Borwein stepwidth is zero, using deltat " << currentDeltat << " instead.");
        stepwidth = currentDeltat;
      }
      Vector PositionUpdate = stepwidth * currentGradient;
      LOG(3, "DEBUG: Update would be " << stepwidth << "*" << currentGradient << " = " << PositionUpdate);

      // extract largest components for showing progress of annealing
      for(size_t i=0;i<NDIM;++i)
        if (currentGradient[i] > maxComponents[i])
          maxComponents[i] = currentGradient[i];

      // are we in initial step? Then don't check against velocity
      if ((currentStep > 1) && (!(*iter)->getAtomicVelocity().IsZero()))
        // update with currentDelta tells us how the current gradient relates to
        // the last one: If it has become larger, reduce currentDelta
        if ((PositionUpdate.ScalarProduct((*iter)->getAtomicVelocity()) < 0)
            && (currentDeltat > MinimumDeltat)) {
          currentDeltat = .5*currentDeltat;
          LOG(2, "DEBUG: Upgrade in other direction: " << PositionUpdate.NormSquared()
              << " > " << (*iter)->getAtomicVelocity().NormSquared()
              << ", decreasing deltat: " << currentDeltat);
          PositionUpdate = currentDeltat * currentGradient;
      }
      // finally set new values
      (*iter)->setPosition(currentPosition + PositionUpdate);
      (*iter)->setAtomicVelocity(PositionUpdate);
      //std::cout << "Id of atom is " << (*iter)->getId() << std::endl;
//        (*iter)->VelocityVerletUpdateU((*iter)->getId(), CurrentTimeStep-1, Deltat, IsAngstroem);
    }
  }

  /** Performs Gradient optimization on the bonds.
   *
   * We assume that forces have just been calculated. These forces are projected
   * onto the bonds and these are annealed subsequently by moving atoms in the
   * bond neighborhood on either side conjunctively.
   *
   *
   * \param CurrentTimeStep current time step (i.e. \f$ t + \Delta t \f$ in the sense of the velocity verlet)
   * \param offset offset in matrix file to the first force component
   * \param maxComponents to be filled with maximum force component over all atoms
   */
  void annealWithBondGraph(
      const int CurrentTimeStep,
      const size_t offset,
      Vector &maxComponents)
  {
    // get nodes on either side of selected bond via BFS discovery
//    std::vector<atomId_t> atomids;
//    for(typename AtomSetMixin<T>::iterator iter = AtomicForceManipulator<T>::atoms.begin();
//        iter != AtomicForceManipulator<T>::atoms.end(); ++iter) {
//      atomids.push_back((*iter)->getId());
//    }
//    ASSERT( atomids.size() == AtomicForceManipulator<T>::atoms.size(),
//        "ForceAnnealing() - could not gather all atomic ids?");
    BoostGraphCreator BGcreator;
    BGcreator.createFromRange(
        AtomicForceManipulator<T>::atoms.begin(),
        AtomicForceManipulator<T>::atoms.end(),
        AtomicForceManipulator<T>::atoms.size(),
        BreadthFirstSearchGatherer::AlwaysTruePredicate);
    BreadthFirstSearchGatherer NodeGatherer(BGcreator);

    std::map<atomId_t, Vector> GatheredUpdates; //!< gathers all updates which are applied at the end
    for(typename AtomSetMixin<T>::iterator iter = AtomicForceManipulator<T>::atoms.begin();
        iter != AtomicForceManipulator<T>::atoms.end(); ++iter) {
      // atom's force vector gives steepest descent direction
      const Vector oldPosition = (*iter)->getPositionAtStep(CurrentTimeStep-2 >= 0 ? CurrentTimeStep - 2 : 0);
      const Vector currentPosition = (*iter)->getPosition();
      const Vector oldGradient = (*iter)->getAtomicForceAtStep(CurrentTimeStep-2 >= 0 ? CurrentTimeStep - 2 : 0);
      const Vector currentGradient = (*iter)->getAtomicForce();
      LOG(4, "DEBUG: Force for atom " << **iter << " is " << currentGradient);

      // we use Barzilai-Borwein update with position reversed to get descent
      const Vector GradientDifference = (currentGradient - oldGradient);
      const double stepwidth =
          fabs((currentPosition - oldPosition).ScalarProduct(GradientDifference))/
          GradientDifference.NormSquared();
      Vector PositionUpdate = stepwidth * currentGradient;
      if (fabs(stepwidth) < 1e-10) {
        // dont' warn in first step, deltat usage normal
        if (currentStep != 1)
          ELOG(1, "INFO: Barzilai-Borwein stepwidth is zero, using deltat " << currentDeltat << " instead.");
        PositionUpdate = currentDeltat * currentGradient;
      }
      LOG(3, "DEBUG: Update would be " << PositionUpdate);

//      // add update to central atom
//      const atomId_t atomid = (*iter)->getId();
//      if (GatheredUpdates.count(atomid)) {
//        GatheredUpdates[atomid] += PositionUpdate;
//      } else
//        GatheredUpdates.insert( std::make_pair(atomid, PositionUpdate) );

      // We assume that a force is local, i.e. a bond is too short yet and hence
      // the atom needs to be moved. However, all the adjacent (bound) atoms might
      // already be at the perfect distance. If we just move the atom alone, we ruin
      // all the other bonds. Hence, it would be sensible to move every atom found
      // through the bond graph in the direction of the force as well by the same
      // PositionUpdate. This is just what we are going to do.

      /// get all nodes from bonds in the direction of the current force

      // remove edges facing in the wrong direction
      std::vector<bond::ptr> removed_bonds;
      const BondList& ListOfBonds = (*iter)->getListOfBonds();
      for(BondList::const_iterator bonditer = ListOfBonds.begin();
          bonditer != ListOfBonds.end(); ++bonditer) {
        const bond &current_bond = *(*bonditer);
        LOG(2, "DEBUG: Looking at bond " << current_bond);
        Vector BondVector = (*iter)->getPosition();
        BondVector -= ((*iter)->getId() == current_bond.rightatom->getId())
            ? current_bond.rightatom->getPosition() : current_bond.leftatom->getPosition();
        BondVector.Normalize();
        if (BondVector.ScalarProduct(currentGradient) < 0) {
          removed_bonds.push_back(*bonditer);
#ifndef NDEBUG
          const bool status =
#endif
              BGcreator.removeEdge(current_bond.leftatom->getId(), current_bond.rightatom->getId());
          ASSERT( status, "ForceAnnealing() - edge to found bond is not present?");
        }
      }
      BoostGraphHelpers::Nodeset_t bondside_set = NodeGatherer((*iter)->getId(), max_distance);
      const BreadthFirstSearchGatherer::distance_map_t &distance_map = NodeGatherer.getDistances();
      std::sort(bondside_set.begin(), bondside_set.end());
      // re-add those edges
      for (std::vector<bond::ptr>::const_iterator bonditer = removed_bonds.begin();
          bonditer != removed_bonds.end(); ++bonditer)
        BGcreator.addEdge((*bonditer)->leftatom->getId(), (*bonditer)->rightatom->getId());

      // apply PositionUpdate to all nodes in the bondside_set
      for (BoostGraphHelpers::Nodeset_t::const_iterator setiter = bondside_set.begin();
          setiter != bondside_set.end(); ++setiter) {
        const BreadthFirstSearchGatherer::distance_map_t::const_iterator diter
          = distance_map.find(*setiter);
        ASSERT( diter != distance_map.end(),
            "ForceAnnealing() - could not find distance to an atom.");
        const double factor = pow(damping_factor, diter->second);
        LOG(3, "DEBUG: Update for atom #" << *setiter << " will be "
            << factor << "*" << PositionUpdate);
        if (GatheredUpdates.count((*setiter))) {
          GatheredUpdates[(*setiter)] += factor*PositionUpdate;
        } else {
          GatheredUpdates.insert(
              std::make_pair(
                  (*setiter),
                  factor*PositionUpdate) );
        }
      }

      // extract largest components for showing progress of annealing
      for(size_t i=0;i<NDIM;++i)
        if (currentGradient[i] > maxComponents[i])
          maxComponents[i] = currentGradient[i];
    }
    // apply the gathered updates
    for (std::map<atomId_t, Vector>::const_iterator iter = GatheredUpdates.begin();
        iter != GatheredUpdates.end(); ++iter) {
      const atomId_t &atomid = iter->first;
      const Vector &update = iter->second;
      atom* const walker = World::getInstance().getAtom(AtomById(atomid));
      ASSERT( walker != NULL,
          "ForceAnnealing() - walker with id "+toString(atomid)+" has suddenly disappeared.");
      LOG(3, "DEBUG: Applying update " << update << " to atom #" << atomid
          << ", namely " << *walker);
      walker->setPosition( walker->getPosition() + update );
    }
  }

  /** Reset function to unset static entities and artificial velocities.
   *
   */
  void reset()
  {
    currentDeltat = 0.;
    currentStep = 0;
  }

private:
  //!> contains the current step in relation to maxsteps
  static size_t currentStep;
  //!> contains the maximum number of steps, determines initial and final step with currentStep
  size_t maxSteps;
  static double currentDeltat;
  //!> minimum deltat for internal while loop (adaptive step width)
  static double MinimumDeltat;
  //!> contains the maximum bond graph distance up to which shifts of a single atom are spread
  const int max_distance;
  //!> the shifted is dampened by this factor with the power of the bond graph distance to the shift causing atom
  const double damping_factor;
};

template <class T>
double ForceAnnealing<T>::currentDeltat = 0.;
template <class T>
size_t ForceAnnealing<T>::currentStep = 0;
template <class T>
double ForceAnnealing<T>::MinimumDeltat = 1e-8;

#endif /* FORCEANNEALING_HPP_ */