/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2010 University of Bonn. All rights reserved.
 * Please see the LICENSE file or "Copyright notice" in builder.cpp for details.
 */

/*
 * MinimiseConstrainedPotential.cpp
 *
 *  Created on: Feb 23, 2011
 *      Author: heber
 */

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

#include "CodePatterns/MemDebug.hpp"

#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_linalg.h>

#include "atom.hpp"
#include "config.hpp"
#include "element.hpp"
#include "CodePatterns/enumeration.hpp"
#include "CodePatterns/Info.hpp"
#include "CodePatterns/Verbose.hpp"
#include "CodePatterns/Log.hpp"
#include "Helpers/helpers.hpp"
#include "molecule.hpp"
#include "parser.hpp"
#include "LinearAlgebra/Plane.hpp"
#include "World.hpp"

#include "Dynamics/MinimiseConstrainedPotential.hpp"


MinimiseConstrainedPotential::MinimiseConstrainedPotential(
    molecule::atomSet &_atoms,
    std::map<atom*, atom *> &_PermutationMap) :
  atoms(_atoms),
  PermutationMap(_PermutationMap)
{}

MinimiseConstrainedPotential::~MinimiseConstrainedPotential()
{}

double MinimiseConstrainedPotential::operator()(int _startstep, int _endstep, bool IsAngstroem)
{
  double Potential, OldPotential, OlderPotential;
  int round;
  atom *Sprinter = NULL;
  DistanceMap::iterator Rider, Strider;

  // set to zero
  PermutationMap.clear();
  DoubleList.clear();
  for (molecule::atomSet::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    DistanceList[*iter].clear();
  }
  DistanceList.clear();
  DistanceIterators.clear();
  DistanceIterators.clear();

  /// Minimise the potential
  // set Lagrange multiplier constants
  PenaltyConstants[0] = 10.;
  PenaltyConstants[1] = 1.;
  PenaltyConstants[2] = 1e+7;    // just a huge penalty
  // generate the distance list
  DoLog(1) && (Log() << Verbose(1) << "Allocating, initializting and filling the distance list ... " << endl);
  FillDistanceList();

  // create the initial PermutationMap (source -> target)
  CreateInitialLists();

  // make the PermutationMap injective by checking whether we have a non-zero constants[2] term in it
  DoLog(1) && (Log() << Verbose(1) << "Making the PermutationMap injective ... " << endl);
  MakeInjectivePermutation();
  DoubleList.clear();

  // argument minimise the constrained potential in this injective PermutationMap
  DoLog(1) && (Log() << Verbose(1) << "Argument minimising the PermutationMap." << endl);
  OldPotential = 1e+10;
  round = 0;
  do {
    DoLog(2) && (Log() << Verbose(2) << "Starting round " << ++round << ", at current potential " << OldPotential << " ... " << endl);
    OlderPotential = OldPotential;
    molecule::atomSet::const_iterator iter;
    do {
      iter = atoms.begin();
      for (; iter != atoms.end(); ++iter) {
        CalculateDoubleList();
        PrintPermutationMap();
        Sprinter = DistanceIterators[(*iter)]->second;   // store initial partner
        Strider = DistanceIterators[(*iter)];  //remember old iterator
        DistanceIterators[(*iter)] = StepList[(*iter)];
        if (DistanceIterators[(*iter)] == DistanceList[(*iter)].end()) {// stop, before we run through the list and still on
          DistanceIterators[(*iter)] == DistanceList[(*iter)].begin();
          break;
        }
        //Log() << Verbose(2) << "Current Walker: " << *(*iter) << " with old/next candidate " << *Sprinter << "/" << *DistanceIterators[(*iter)]->second << "." << endl;
        // find source of the new target
        molecule::atomSet::const_iterator runner = atoms.begin();
        for (; runner != atoms.end(); ++runner) { // find the source whose toes we might be stepping on (Walker's new target should be in use by another already)
          if (PermutationMap[(*runner)] == DistanceIterators[(*iter)]->second) {
            //Log() << Verbose(2) << "Found the corresponding owner " << *(*runner) << " to " << *PermutationMap[(*runner)] << "." << endl;
            break;
          }
        }
        if (runner != atoms.end()) { // we found the other source
          // then look in its distance list for Sprinter
          Rider = DistanceList[(*runner)].begin();
          for (; Rider != DistanceList[(*runner)].end(); Rider++)
            if (Rider->second == Sprinter)
              break;
          if (Rider != DistanceList[(*runner)].end()) { // if we have found one
            //Log() << Verbose(2) << "Current Other: " << *(*runner) << " with old/next candidate " << *PermutationMap[(*runner)] << "/" << *Rider->second << "." << endl;
            // exchange both
            PermutationMap[(*iter)] = DistanceIterators[(*iter)]->second; // put next farther distance into PermutationMap
            PermutationMap[(*runner)] = Sprinter;  // and hand the old target to its respective owner
            CalculateDoubleList();
            PrintPermutationMap();
            // calculate the new potential
            //Log() << Verbose(2) << "Checking new potential ..." << endl;
            Potential = ConstrainedPotential();
            if (Potential > OldPotential) { // we made everything worse! Undo ...
              //Log() << Verbose(3) << "Nay, made the potential worse: " << Potential << " vs. " << OldPotential << "!" << endl;
              //Log() << Verbose(3) << "Setting " << *(*runner) << "'s source to " << *DistanceIterators[(*runner)]->second << "." << endl;
              // Undo for Runner (note, we haven't moved the iteration yet, we may use this)
              PermutationMap[(*runner)] = DistanceIterators[(*runner)]->second;
              // Undo for Walker
              DistanceIterators[(*iter)] = Strider;  // take next farther distance target
              //Log() << Verbose(3) << "Setting " << *(*iter) << "'s source to " << *DistanceIterators[(*iter)]->second << "." << endl;
              PermutationMap[(*iter)] = DistanceIterators[(*iter)]->second;
            } else {
              DistanceIterators[(*runner)] = Rider;  // if successful also move the pointer in the iterator list
              DoLog(3) && (Log() << Verbose(3) << "Found a better permutation, new potential is " << Potential << " vs." << OldPotential << "." << endl);
              OldPotential = Potential;
            }
            if (Potential > PenaltyConstants[2]) {
              DoeLog(1) && (eLog()<< Verbose(1) << "The two-step permutation procedure did not maintain injectivity!" << endl);
              exit(255);
            }
            //Log() << Verbose(0) << endl;
          } else {
            DoeLog(1) && (eLog()<< Verbose(1) << **runner << " was not the owner of " << *Sprinter << "!" << endl);
            exit(255);
          }
        } else {
          PermutationMap[(*iter)] = DistanceIterators[(*iter)]->second; // new target has no source!
        }
        StepList[(*iter)]++; // take next farther distance target
      }
    } while (++iter != atoms.end());
  } while ((OlderPotential - OldPotential) > 1e-3);
  DoLog(1) && (Log() << Verbose(1) << "done." << endl);


  return ConstrainedPotential();
};

void MinimiseConstrainedPotential::FillDistanceList()
{
  for (molecule::atomSet::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    for (molecule::atomSet::const_iterator runner = atoms.begin(); runner != atoms.end(); ++runner) {
      DistanceList[(*iter)].insert( DistancePair((*iter)->getPositionAtStep(startstep).distance((*runner)->getPositionAtStep(endstep)), (*runner)) );
    }
  }
};

void MinimiseConstrainedPotential::CreateInitialLists()
{
  for (molecule::atomSet::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    StepList[(*iter)] = DistanceList[(*iter)].begin();    // stores the step to the next iterator that could be a possible next target
    PermutationMap[(*iter)] = DistanceList[(*iter)].begin()->second;   // always pick target with the smallest distance
    DoubleList[DistanceList[(*iter)].begin()->second]++;            // increase this target's source count (>1? not injective)
    DistanceIterators[(*iter)] = DistanceList[(*iter)].begin();    // and remember which one we picked
    DoLog(2) && (Log() << Verbose(2) << **iter << " starts with distance " << DistanceList[(*iter)].begin()->first << "." << endl);
  }
};

void MinimiseConstrainedPotential::MakeInjectivePermutation()
{
  molecule::atomSet::const_iterator iter = atoms.begin();
  DistanceMap::iterator NewBase;
  double Potential = fabs(ConstrainedPotential());

  if (atoms.empty()) {
    eLog() << Verbose(1) << "Molecule is empty." << endl;
    return;
  }
  while ((Potential) > PenaltyConstants[2]) {
    CalculateDoubleList();
    PrintPermutationMap();
    iter++;
    if (iter == atoms.end()) // round-robin at the end
      iter = atoms.begin();
    if (DoubleList[DistanceIterators[(*iter)]->second] <= 1)  // no need to make those injective that aren't
      continue;
    // now, try finding a new one
    Potential = TryNextNearestNeighbourForInjectivePermutation((*iter), Potential);
  }
  for (molecule::atomSet::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    // now each single entry in the DoubleList should be <=1
    if (DoubleList[*iter] > 1) {
      DoeLog(0) && (eLog()<< Verbose(0) << "Failed to create an injective PermutationMap!" << endl);
      performCriticalExit();
    }
  }
  DoLog(1) && (Log() << Verbose(1) << "done." << endl);
};

unsigned int MinimiseConstrainedPotential::CalculateDoubleList()
{
  for (molecule::atomSet::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter)
    DoubleList[*iter] = 0;
  unsigned int doubles = 0;
  for (molecule::atomSet::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter)
      DoubleList[ PermutationMap[*iter] ]++;
  for (molecule::atomSet::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter)
    if (DoubleList[*iter] > 1)
      doubles++;
  if (doubles >0)
    DoLog(2) && (Log() << Verbose(2) << "Found " << doubles << " Doubles." << endl);
  return doubles;
};

void MinimiseConstrainedPotential::PrintPermutationMap() const
{
  stringstream zeile1, zeile2;
  int doubles = 0;
  zeile1 << "PermutationMap: ";
  zeile2 << "                ";
  for (molecule::atomSet::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    zeile1 << (*iter)->getName() << " ";
    zeile2 << (PermutationMap[*iter])->getName() << " ";
  }
  for (molecule::atomSet::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    std::map<atom *, unsigned int>::const_iterator value_iter = DoubleList.find(*iter);
    if (value_iter->second > (unsigned int)1)
      doubles++;
  }
  if (doubles >0)
    DoLog(2) && (Log() << Verbose(2) << "Found " << doubles << " Doubles." << endl);
//  Log() << Verbose(2) << zeile1.str() << endl << zeile2.str() << endl;
};

double MinimiseConstrainedPotential::ConstrainedPotential()
{
  double tmp = 0.;
  double result = 0.;
  // go through every atom
  atom *Runner = NULL;
  for (molecule::atomSet::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    // first term: distance to target
    Runner = PermutationMap[(*iter)];   // find target point
    tmp = ((*iter)->getPositionAtStep(startstep).distance(Runner->getPositionAtStep(endstep)));
    tmp *= IsAngstroem ? 1. : 1./AtomicLengthToAngstroem;
    result += PenaltyConstants[0] * tmp;
    //Log() << Verbose(4) << "Adding " << tmp*constants[0] << "." << endl;

    // second term: sum of distances to other trajectories
    result += SumDistanceOfTrajectories((*iter));

    // third term: penalty for equal targets
    result += PenalizeEqualTargets((*iter));
  }

  return result;
};

double MinimiseConstrainedPotential::TryNextNearestNeighbourForInjectivePermutation(atom *Walker, double &OldPotential)
{
  double Potential = 0;
  DistanceMap::iterator NewBase = DistanceIterators[Walker];  // store old base
  do {
    NewBase++;  // take next further distance in distance to targets list that's a target of no one
  } while ((DoubleList[NewBase->second] != 0) && (NewBase != DistanceList[Walker].end()));
  if (NewBase != DistanceList[Walker].end()) {
    PermutationMap[Walker] = NewBase->second;
    Potential = fabs(ConstrainedPotential());
    if (Potential > OldPotential) { // undo
      PermutationMap[Walker] = DistanceIterators[Walker]->second;
    } else {  // do
      DoubleList[DistanceIterators[Walker]->second]--;  // decrease the old entry in the doubles list
      DoubleList[NewBase->second]++;    // increase the old entry in the doubles list
      DistanceIterators[Walker] = NewBase;
      OldPotential = Potential;
      DoLog(3) && (Log() << Verbose(3) << "Found a new permutation, new potential is " << OldPotential << "." << endl);
    }
  }
  return Potential;
};

double MinimiseConstrainedPotential::PenalizeEqualTargets(atom *Walker)
{
  double result = 0.;
  for (molecule::atomSet::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    if ((PermutationMap[Walker] == PermutationMap[(*iter)]) && (Walker < (*iter))) {
  //    atom *Sprinter = PermutationMap[Walker->nr];
  //        Log() << Verbose(0) << *Walker << " and " << *(*iter) << " are heading to the same target at ";
  //        Log() << Verbose(0) << Sprinter->getPosition(endstep);
  //        Log() << Verbose(0) << ", penalting." << endl;
      result += PenaltyConstants[2];
      //Log() << Verbose(4) << "Adding " << constants[2] << "." << endl;
    }
  }
  return result;
};

double MinimiseConstrainedPotential::SumDistanceOfTrajectories(atom *Walker)
{
  gsl_matrix *A = gsl_matrix_alloc(NDIM,NDIM);
  gsl_vector *x = gsl_vector_alloc(NDIM);
  atom *Sprinter = NULL;
  Vector trajectory1, trajectory2, normal, TestVector;
  double Norm1, Norm2, tmp, result = 0.;

  for (molecule::atomSet::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    if ((*iter) == Walker) // hence, we only go up to the Walker, not beyond (similar to i=0; i<j; i++)
      break;
    // determine normalized trajectories direction vector (n1, n2)
    Sprinter = PermutationMap[Walker];   // find first target point
    trajectory1 = Sprinter->getPositionAtStep(endstep) - Walker->getPositionAtStep(startstep);
    trajectory1.Normalize();
    Norm1 = trajectory1.Norm();
    Sprinter = PermutationMap[(*iter)];   // find second target point
    trajectory2 = Sprinter->getPositionAtStep(endstep) - (*iter)->getPositionAtStep(startstep);
    trajectory2.Normalize();
    Norm2 = trajectory1.Norm();
    // check whether either is zero()
    if ((Norm1 < MYEPSILON) && (Norm2 < MYEPSILON)) {
      tmp = Walker->getPositionAtStep(startstep).distance((*iter)->getPositionAtStep(startstep));
    } else if (Norm1 < MYEPSILON) {
      Sprinter = PermutationMap[Walker];   // find first target point
      trajectory1 = Sprinter->getPositionAtStep(endstep) - (*iter)->getPositionAtStep(startstep);
      trajectory2 *= trajectory1.ScalarProduct(trajectory2); // trajectory2 is scaled to unity, hence we don't need to divide by anything
      trajectory1 -= trajectory2;   // project the part in norm direction away
      tmp = trajectory1.Norm();  // remaining norm is distance
    } else if (Norm2 < MYEPSILON) {
      Sprinter = PermutationMap[(*iter)];   // find second target point
      trajectory2 = Sprinter->getPositionAtStep(endstep) - Walker->getPositionAtStep(startstep);  // copy second offset
      trajectory1 *= trajectory2.ScalarProduct(trajectory1); // trajectory1 is scaled to unity, hence we don't need to divide by anything
      trajectory2 -= trajectory1;   // project the part in norm direction away
      tmp = trajectory2.Norm();  // remaining norm is distance
    } else if ((fabs(trajectory1.ScalarProduct(trajectory2)/Norm1/Norm2) - 1.) < MYEPSILON) { // check whether they're linear dependent
  //        Log() << Verbose(3) << "Both trajectories of " << *Walker << " and " << *Runner << " are linear dependent: ";
  //        Log() << Verbose(0) << trajectory1;
  //        Log() << Verbose(0) << " and ";
  //        Log() << Verbose(0) << trajectory2;
      tmp = Walker->getPositionAtStep(startstep).distance((*iter)->getPositionAtStep(startstep));
  //        Log() << Verbose(0) << " with distance " << tmp << "." << endl;
    } else { // determine distance by finding minimum distance
  //        Log() << Verbose(3) << "Both trajectories of " << *Walker << " and " << *(*iter) << " are linear independent ";
  //        Log() << Verbose(0) << endl;
  //        Log() << Verbose(0) << "First Trajectory: ";
  //        Log() << Verbose(0) << trajectory1 << endl;
  //        Log() << Verbose(0) << "Second Trajectory: ";
  //        Log() << Verbose(0) << trajectory2 << endl;
      // determine normal vector for both
      normal = Plane(trajectory1, trajectory2,0).getNormal();
      // print all vectors for debugging
  //        Log() << Verbose(0) << "Normal vector in between: ";
  //        Log() << Verbose(0) << normal << endl;
      // setup matrix
      for (int i=NDIM;i--;) {
        gsl_matrix_set(A, 0, i, trajectory1[i]);
        gsl_matrix_set(A, 1, i, trajectory2[i]);
        gsl_matrix_set(A, 2, i, normal[i]);
        gsl_vector_set(x,i, (Walker->getPositionAtStep(startstep)[i] - (*iter)->getPositionAtStep(startstep)[i]));
      }
      // solve the linear system by Householder transformations
      gsl_linalg_HH_svx(A, x);
      // distance from last component
      tmp = gsl_vector_get(x,2);
  //        Log() << Verbose(0) << " with distance " << tmp << "." << endl;
      // test whether we really have the intersection (by checking on c_1 and c_2)
      trajectory1.Scale(gsl_vector_get(x,0));
      trajectory2.Scale(gsl_vector_get(x,1));
      normal.Scale(gsl_vector_get(x,2));
      TestVector = (*iter)->getPositionAtStep(startstep) + trajectory2 + normal
                   - (Walker->getPositionAtStep(startstep) + trajectory1);
      if (TestVector.Norm() < MYEPSILON) {
  //          Log() << Verbose(2) << "Test: ok.\tDistance of " << tmp << " is correct." << endl;
      } else {
  //          Log() << Verbose(2) << "Test: failed.\tIntersection is off by ";
  //          Log() << Verbose(0) << TestVector;
  //          Log() << Verbose(0) << "." << endl;
      }
    }
    // add up
    tmp *= IsAngstroem ? 1. : 1./AtomicLengthToAngstroem;
    if (fabs(tmp) > MYEPSILON) {
      result += PenaltyConstants[1] * 1./tmp;
      //Log() << Verbose(4) << "Adding " << 1./tmp*constants[1] << "." << endl;
    }
  }
  return result;
};

void MinimiseConstrainedPotential::EvaluateConstrainedForces(ForceMatrix *Force)
{
  double constant = 10.;

  /// evaluate forces (only the distance to target dependent part) with the final PermutationMap
  DoLog(1) && (Log() << Verbose(1) << "Calculating forces and adding onto ForceMatrix ... " << endl);
  for(molecule::atomSet::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    atom *Sprinter = PermutationMap[(*iter)];
    // set forces
    for (int i=NDIM;i++;)
      Force->Matrix[0][(*iter)->getNr()][5+i] += 2.*constant*sqrt((*iter)->getPositionAtStep(startstep).distance(Sprinter->getPositionAtStep(endstep)));
  }
  DoLog(1) && (Log() << Verbose(1) << "done." << endl);
};