/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2010-2012 University of Bonn. 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/>.
 */

/*
 * 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/atom.hpp"
#include "Element/element.hpp"
#include "CodePatterns/enumeration.hpp"
#include "CodePatterns/Info.hpp"
#include "CodePatterns/Verbose.hpp"
#include "CodePatterns/Log.hpp"
#include "Fragmentation/ForceMatrix.hpp"
#include "Helpers/helpers.hpp"
#include "molecule.hpp"
#include "LinearAlgebra/Plane.hpp"
#include "World.hpp"

#include "Dynamics/MinimiseConstrainedPotential.hpp"


MinimiseConstrainedPotential::MinimiseConstrainedPotential(
    World::AtomComposite &_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 (World::AtomComposite::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
  LOG(1, "Allocating, initializting and filling the distance list ... ");
  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
  LOG(1, "Making the PermutationMap injective ... ");
  MakeInjectivePermutation();
  DoubleList.clear();

  // argument minimise the constrained potential in this injective PermutationMap
  LOG(1, "Argument minimising the PermutationMap.");
  OldPotential = 1e+10;
  round = 0;
  do {
    LOG(2, "Starting round " << ++round << ", at current potential " << OldPotential << " ... ");
    OlderPotential = OldPotential;
    World::AtomComposite::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(2, "Current Walker: " << *(*iter) << " with old/next candidate " << *Sprinter << "/" << *DistanceIterators[(*iter)]->second << ".");
        // find source of the new target
        World::AtomComposite::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(2, "Found the corresponding owner " << *(*runner) << " to " << *PermutationMap[(*runner)] << ".");
            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(2, "Current Other: " << *(*runner) << " with old/next candidate " << *PermutationMap[(*runner)] << "/" << *Rider->second << ".");
            // 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(2, "Checking new potential ...");
            Potential = ConstrainedPotential();
            if (Potential > OldPotential) { // we made everything worse! Undo ...
              //LOG(3, "Nay, made the potential worse: " << Potential << " vs. " << OldPotential << "!");
              //LOG(3, "Setting " << *(*runner) << "'s source to " << *DistanceIterators[(*runner)]->second << ".");
              // 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(3, "Setting " << *(*iter) << "'s source to " << *DistanceIterators[(*iter)]->second << ".");
              PermutationMap[(*iter)] = DistanceIterators[(*iter)]->second;
            } else {
              DistanceIterators[(*runner)] = Rider;  // if successful also move the pointer in the iterator list
              LOG(3, "Found a better permutation, new potential is " << Potential << " vs." << OldPotential << ".");
              OldPotential = Potential;
            }
            if (Potential > PenaltyConstants[2]) {
              ELOG(1, "The two-step permutation procedure did not maintain injectivity!");
              exit(255);
            }
          } else {
            ELOG(1, **runner << " was not the owner of " << *Sprinter << "!");
            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);
  LOG(1, "done.");


  return ConstrainedPotential();
};

void MinimiseConstrainedPotential::FillDistanceList()
{
  for (World::AtomComposite::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    for (World::AtomComposite::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 (World::AtomComposite::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
    LOG(2, **iter << " starts with distance " << DistanceList[(*iter)].begin()->first << ".");
  }
};

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

  if (atoms.empty()) {
    ELOG(1, "Molecule is empty.");
    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 (World::AtomComposite::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    // now each single entry in the DoubleList should be <=1
    if (DoubleList[*iter] > 1) {
      ELOG(0, "Failed to create an injective PermutationMap!");
      performCriticalExit();
    }
  }
  LOG(1, "done.");
};

unsigned int MinimiseConstrainedPotential::CalculateDoubleList()
{
  for (World::AtomComposite::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter)
    DoubleList[*iter] = 0;
  unsigned int doubles = 0;
  for (World::AtomComposite::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter)
      DoubleList[ PermutationMap[*iter] ]++;
  for (World::AtomComposite::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter)
    if (DoubleList[*iter] > 1)
      doubles++;
  if (doubles >0)
    LOG(2, "Found " << doubles << " Doubles.");
  return doubles;
};

void MinimiseConstrainedPotential::PrintPermutationMap() const
{
  stringstream zeile1, zeile2;
  int doubles = 0;
  zeile1 << "PermutationMap: ";
  zeile2 << "                ";
  for (World::AtomComposite::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    zeile1 << (*iter)->getName() << " ";
    zeile2 << (PermutationMap[*iter])->getName() << " ";
  }
  for (World::AtomComposite::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)
    LOG(2, "Found " << doubles << " Doubles.");
//  LOG(2, zeile1.str() << endl << zeile2.str());
};

double MinimiseConstrainedPotential::ConstrainedPotential()
{
  double tmp = 0.;
  double result = 0.;
  // go through every atom
  atom *Runner = NULL;
  for (World::AtomComposite::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(4, "Adding " << tmp*constants[0] << ".");

    // 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;
      LOG(3, "Found a new permutation, new potential is " << OldPotential << ".");
    }
  }
  return Potential;
};

double MinimiseConstrainedPotential::PenalizeEqualTargets(atom *Walker)
{
  double result = 0.;
  for (World::AtomComposite::const_iterator iter = atoms.begin(); iter != atoms.end(); ++iter) {
    if ((PermutationMap[Walker] == PermutationMap[(*iter)]) && (Walker < (*iter))) {
//    atom *Sprinter = PermutationMap[Walker->nr];
//    if (DoLog(0)) {
//        std::stringstream output;
//        output << *Walker << " and " << *(*iter) << " are heading to the same target at ";
//        output << Sprinter->getPosition(endstep);
//        output << ", penalting.";
//        LOG(0, output.str());
//    }
      result += PenaltyConstants[2];
      //LOG(4, "INFO: Adding " << constants[2] << ".");
    }
  }
  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 (World::AtomComposite::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
//      std::stringstream output;
//      output << "Both trajectories of " << *Walker << " and " << *Runner << " are linear dependent: ";
//      output << trajectory1 << " and " << trajectory2;
//      LOG(3, output.str());
      tmp = Walker->getPositionAtStep(startstep).distance((*iter)->getPositionAtStep(startstep));
//        LOG(0, " with distance " << tmp << ".");
    } else { // determine distance by finding minimum distance
//      std::stringstream output;
//      output "Both trajectories of " << *Walker << " and " << *(*iter) << " are linear independent -- ";
//      output "First Trajectory: " << trajectory1 << ". Second Trajectory: " << trajectory2);
//      LOG(3, output.str());
      // determine normal vector for both
      normal = Plane(trajectory1, trajectory2,0).getNormal();
      // print all vectors for debugging
//        LOG(3, "INFO: Normal vector in between: " << normal);
      // 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(0, " with distance " << tmp << ".");
      // 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(2, "Test: ok.\tDistance of " << tmp << " is correct.");
      } else {
//        LOG(2, "Test: failed.\tIntersection is off by " << TestVector << ".");
      }
    }
    // add up
    tmp *= IsAngstroem ? 1. : 1./AtomicLengthToAngstroem;
    if (fabs(tmp) > MYEPSILON) {
      result += PenaltyConstants[1] * 1./tmp;
      //LOG(4, "Adding " << 1./tmp*constants[1] << ".");
    }
  }
  return result;
};

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

  /// evaluate forces (only the distance to target dependent part) with the final PermutationMap
  LOG(1, "Calculating forces and adding onto ForceMatrix ... ");
  for(World::AtomComposite::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)));
  }
  LOG(1, "done.");
};