/*
 * VerletForceIntegration.hpp
 *
 *  Created on: Feb 23, 2011
 *      Author: heber
 */

#ifndef VERLETFORCEINTEGRATION_HPP_
#define VERLETFORCEINTEGRATION_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 "Dynamics/MinimiseConstrainedPotential.hpp"
#include "Fragmentation/ForceMatrix.hpp"
#include "Helpers/helpers.hpp"
#include "Helpers/defs.hpp"
#include "LinearAlgebra/Vector.hpp"
#include "Thermostats/ThermoStatContainer.hpp"
#include "Thermostats/Thermostat.hpp"
#include "World.hpp"

template <class T>
class VerletForceIntegration
{
public:
  /** Constructor of class VerletForceIntegration.
   *
   * \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
   */
  VerletForceIntegration(AtomSetMixin<T> &_atoms, double _Deltat, bool _IsAngstroem) :
    Deltat(_Deltat),
    IsAngstroem(_IsAngstroem),
    atoms(_atoms)
  {}
  /** Destructor of class VerletForceIntegration.
   *
   */
  ~VerletForceIntegration()
  {}

  /** Parses nuclear forces from file.
   * Forces are stored in the time step \a TimeStep in the atomicForces in \a atoms.
   * \param *file filename
   * \param TimeStep time step to parse forces file into
   * \return true - file parsed, false - file not found or imparsable
   */
  bool parseForcesFile(const char *file, const int TimeStep)
  {
    Info FunctionInfo(__func__);
    ForceMatrix Force;

    // parse file into ForceMatrix
    std::ifstream input(file);
    if ((input.good()) && (!Force.ParseMatrix(input, 0,0,0))) {
      ELOG(0, "Could not parse Force Matrix file " << file << ".");
      return false;
    }
    input.close();
    if (Force.RowCounter[0] != (int)atoms.size()) {
      ELOG(0, "Mismatch between number of atoms in file " << Force.RowCounter[0] << " and in molecule " << atoms.size() << ".");
      return false;
    }

    addForceMatrixToAtomicForce(Force, TimeStep, 1);
    return true;
  }

  /** Performs Verlet integration.
   *
   * We assume that forces have just been calculated. Then, we perform the velocity
   * and the position calculation for \f$ t + \Delta t \f$, such that forces may be
   * again calculated with respect to the new position.
   *
   * \param NextStep 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 DoConstrainedMD whether a constrained MD shall be done
   * \param FixedCenterOfMass whether forces and velocities are correct to have fixed center of mass
   * \return true - file found and parsed, false - no atoms, file not found or imparsable
   * \todo This is not yet checked if it is correctly working with DoConstrainedMD set >0.
   */
  bool operator()(const int NextStep, const size_t offset, const int DoConstrainedMD, const bool FixedCenterOfMass)
  {
    Info FunctionInfo(__func__);

    // check that atoms are present at all
    if (atoms.size() == 0) {
      ELOG(2, "VerletForceIntegration::operator() - no atoms to integrate.");
      return false;
    }

    // make sum of forces equal zero
    if (FixedCenterOfMass)
      correctForceMatrixForFixedCenterOfMass(offset,NextStep-1);

    // solve a constrained potential if we are meant to
    if (DoConstrainedMD)
      performConstraintMinimization(DoConstrainedMD,NextStep-1);

    if (NextStep > 0) {
      for(typename AtomSetMixin<T>::iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
        //std::cout << "Id of atom is " << (*iter)->getId() << std::endl;
        (*iter)->VelocityVerletUpdateU((*iter)->getId(), NextStep-1, Deltat, IsAngstroem);
      }

      // make sum of velocities equal zero
      if (FixedCenterOfMass)
        correctVelocitiesForFixedCenterOfMass(NextStep-1);

      // thermostat
      performThermostatControl(NextStep-1);
    }

    //std::cout << "Force before velocity verlet, " << Force << std::endl;
    // and perform Verlet integration for each atom with position, velocity and force vector
    // check size of vectors
    for(typename AtomSetMixin<T>::iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
      //std::cout << "Id of atom is " << (*iter)->getId() << std::endl;
      (*iter)->VelocityVerletUpdateX((*iter)->getId(), NextStep, Deltat, IsAngstroem);
    }

    // exit
    return true;
  };

private:
  void addForceMatrixToAtomicForce(const ForceMatrix &Force, const int &TimeStep, const int offset)
  {
    // place forces from matrix into atoms
    Vector tempVector;
    size_t i=0;
    for(typename AtomSetMixin<T>::iterator iter = atoms.begin(); iter != atoms.end(); ++iter,++i) {
      for(size_t d=0;d<NDIM;d++) {
        tempVector[d] = Force.Matrix[0][i][d+offset]*(IsAngstroem ? AtomicLengthToAngstroem : 1.);
      }
      tempVector += (*iter)->getAtomicForceAtStep(TimeStep);
      (*iter)->setAtomicForceAtStep(TimeStep, tempVector);
    }
  }

  void correctForceMatrixForFixedCenterOfMass(const size_t offset, const int &TimeStep) {
    Vector ForceVector;
    // correct Forces
    //std::cout << "Force before correction, " << Force << std::endl;
    ForceVector.Zero();
    for(typename AtomSetMixin<T>::iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
      ForceVector += (*iter)->getAtomicForceAtStep(TimeStep);
    }
    ForceVector.Scale(1./(double)atoms.size());
    //std::cout << "Force before second correction, " << Force << std::endl;
    for(typename AtomSetMixin<T>::iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
      const Vector tempVector = (*iter)->getAtomicForceAtStep(TimeStep) - ForceVector;
      (*iter)->setAtomicForceAtStep(TimeStep, tempVector);
    }
    LOG(3, "INFO: forces correct by " << ForceVector << "each.");
  }

  void correctVelocitiesForFixedCenterOfMass(const int &TimeStep) {
    Vector Velocity;
    double IonMass;
    // correct velocities (rather momenta) so that center of mass remains motionless
    Velocity = atoms.totalMomentumAtStep(TimeStep);
    IonMass = atoms.totalMass();

    // correct velocities (rather momenta) so that center of mass remains motionless
    Velocity *= 1./IonMass;
    atoms.addVelocityAtStep(-1.*Velocity,TimeStep);

    LOG(3, "INFO: Velocities corrected by " << Velocity << " each.");
  }

  void performConstraintMinimization(const int &DoConstrainedMD, const int &TimeStep) {
    // calculate forces and potential
    ForceMatrix Force;
    std::map<atom *, atom*> PermutationMap;
    MinimiseConstrainedPotential Minimiser(atoms, PermutationMap);
    //double ConstrainedPotentialEnergy =
    Minimiser(DoConstrainedMD, 0, IsAngstroem);
    Minimiser.EvaluateConstrainedForces(&Force);
    addForceMatrixToAtomicForce(Force, TimeStep, 1);
  }

  void performThermostatControl(const int &TimeStep) {
    double ActualTemp;

    // calculate current temperature
    ActualTemp = atoms.totalTemperatureAtStep(TimeStep);
    LOG(3, "INFO: Current temperature is " << ActualTemp);

    // rescale to desired value
    double ekin = World::getInstance().getThermostats()->getActive()->scaleAtoms(TimeStep,ActualTemp,atoms);
    ActualTemp = atoms.totalTemperatureAtStep(TimeStep);
    LOG(3, "INFO: New temperature after thermostat is " << ActualTemp);
    LOG(1, "Kinetic energy is " << ekin << ".");
  }

private:
  double Deltat;
  bool IsAngstroem;
  AtomSetMixin<T> atoms;
};

#endif /* VERLETFORCEINTEGRATION_HPP_ */