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

/*
 * BondVectors.cpp
 *
 *  Created on: Jun 13, 2017
 *      Author: heber
 */


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

//#include "CodePatterns/MemDebug.hpp"

#include "BondVectors.hpp"

#include <algorithm>
#include <functional>
#include <iterator>
#include <numeric>

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

#include <levmar.h>

#include "Atom/atom.hpp"
#include "Bond/bond.hpp"
#include "Helpers/defs.hpp"

void BondVectors::recalculateBondVectorsAtStep(
    const size_t &_step) const
{
  current_mapped_vectors.clear();

  ASSERT( !container.empty(),
      "BondVectors::getBondVectors() - container empty, not set properly?");
  for (container_t::const_iterator iter = container.begin();
      iter != container.end(); ++iter) {
    const bond::ptr &current_bond = *iter;
    Vector BondVector = current_bond->leftatom->getPositionAtStep(_step)
        - current_bond->rightatom->getPositionAtStep(_step);
    BondVector.Normalize();
    current_mapped_vectors.insert( std::make_pair(current_bond, BondVector) );
  }
  ASSERT( current_mapped_vectors.size() == container.size(),
      "BondVectors::getBondVectors() - not same amount of bond vectors as bonds?");

  map_is_dirty = false;
  current_step_for_map = _step;
}

size_t BondVectors::getIndexForBond(const bond::ptr &_bond) const
{
  std::pair<
      container_t::const_iterator,
      container_t::const_iterator> iters =
          std::equal_range(container.begin(), container.end(), _bond);
  if (iters.first != container.end())
    return std::distance(container.begin(), iters.first);
  else
    return (size_t)-1;
}

std::vector<Vector> BondVectors::getAtomsBondVectorsAtStep(
    const atom &_walker,
    const size_t &_step) const
{
  if (map_is_dirty || (current_step_for_map != _step))
    recalculateBondVectorsAtStep(_step);

  std::vector<Vector> BondVectors;
  // gather subset of BondVectors for the current atom
  const BondList& ListOfBonds = _walker.getListOfBonds();
  for(BondList::const_iterator bonditer = ListOfBonds.begin();
      bonditer != ListOfBonds.end(); ++bonditer) {
    const bond::ptr &current_bond = *bonditer;
    const BondVectors::mapped_t::const_iterator bviter =
        current_mapped_vectors.find(current_bond);
    ASSERT( bviter != current_mapped_vectors.end(),
        "ForceAnnealing() - cannot find current_bond ?");
    ASSERT( bviter != current_mapped_vectors.end(),
        "ForceAnnealing - cannot find current bond "+toString(*current_bond)
        +" in bonds.");
    BondVectors.push_back(bviter->second);
  }
  LOG(4, "DEBUG: BondVectors for atom #" << _walker.getId() << ": " << BondVectors);

  return BondVectors;
}

struct WeightsMinimization {

  static void evaluate(double *p, double *x, int m, int n, void *data)
  {
    // current weights in p, output to x
    double *matrix = static_cast<double*>(data);
    for(size_t i=0;i<(size_t)n;++i) {
      x[i] = 0.;
      for(size_t j=0;j<(size_t)n;++j)
        x[i] += p[j]*matrix[i*n+j];
//      x[i] = .5*x[i]*x[i];
    }
  }

  static void evaluate_derivative(double *p, double *jac, int m, int n, void *data)
  {
    // current weights in p, output to x
    double *matrix = static_cast<double*>(data);
//    // tmp = (Bx - 1)
//    double *tmp = new double[n];
//    for(size_t i=0;i<(size_t)n;++i) {
//      tmp[i] = -1.;
//      for(size_t j=0;j<(size_t)n;++j)
//        tmp[i] += p[j]*matrix[i*n+j];
//    }
//    // tmp(i) * B_(ij)
//    for(size_t i=0;i<(size_t)n;++i)
//      for(size_t j=0;j<(size_t)n;++j)
//        jac[i*n+j] = tmp[i]*matrix[i*n+j];
//    delete[] tmp ;
    for(size_t i=0;i<(size_t)n;++i)
      for(size_t j=0;j<(size_t)n;++j)
        jac[i*n+j] = matrix[i*n+j];
  }

  static double* getMatrixFromBondVectors(
      const std::vector<Vector> &_bondvectors
      )
  {
    const size_t n = _bondvectors.size();
    double *matrix = new double[n*n];
    size_t i=0;
    for (std::vector<Vector>::const_iterator iter = _bondvectors.begin();
        iter != _bondvectors.end(); ++iter, ++i) {
      size_t j=0;
      for (std::vector<Vector>::const_iterator otheriter = _bondvectors.begin();
          otheriter != _bondvectors.end(); ++otheriter, ++j) {
        // only magnitude is important not the sign
        matrix[i*n+j] = fabs((*iter).ScalarProduct(*otheriter));
      }
    }

    return matrix;
  }
};


BondVectors::weights_t BondVectors::getWeightsForAtomAtStep(
    const atom &_walker,
    const std::vector<Vector> &_bondvectors,
    const size_t &_step) const
{
  // let levmar optimize
  register int i, j;
  int ret;
  double *p;
  double *x;
  int n=_bondvectors.size();
  double opts[LM_OPTS_SZ], info[LM_INFO_SZ];

  double *matrix = WeightsMinimization::getMatrixFromBondVectors(_bondvectors);

  weights_t weights(n, 0.);

  // minim. options [\tau, \epsilon1, \epsilon2, \epsilon3]. Respectively the scale factor for initial \mu,
  // * stopping thresholds for ||J^T e||_inf, ||Dp||_2 and ||e||_2.
  opts[0]=LM_INIT_MU; opts[1]=1E-15; opts[2]=1E-15; opts[3]=1E-20;
  opts[4]= LM_DIFF_DELTA; // relevant only if the Jacobian is approximated using finite differences; specifies forward differencing
  //opts[4]=-LM_DIFF_DELTA; // specifies central differencing to approximate Jacobian; more accurate but more expensive to compute!

  // prepare initial values for weights
  p = new double[n];
  for (i=0;i<n;++i)
    p[i] = 1.;
  // prepare output value, i.e. the row sums
  x = new double[n];
  for (i=0;i<n;++i)
    x[i] = 1.;

  {
    double *work, *covar;
    work=(double *)malloc((LM_DIF_WORKSZ(n, n)+n*n)*sizeof(double));
    if(!work){
      ELOG(0, "BondVectors::getWeightsForAtomAtStep() - memory allocation request failed.");
      return weights;
    }
    covar=work+LM_DIF_WORKSZ(n, n);

    // give this pointer as additional data to construct function pointer in
    // LevMarCallback and call
    double *lb = new double[n];
    double *ub = new double[n];
    for (i=0;i<n;++i) {
      lb[i] = 1./(double)n;
      ub[i] = 1.;
    }
    ret=dlevmar_bc_der(
        &WeightsMinimization::evaluate,
        &WeightsMinimization::evaluate_derivative,
        p, x, n, n, lb, ub, NULL, 1000, opts, info, work, covar, matrix); // no Jacobian, caller allocates work memory, covariance estimated
    delete[] lb;
    delete[] ub;

    if (0)
    {
      std::stringstream covar_msg;
      covar_msg << "Covariance of the fit:\n";
      for(i=0; i<n; ++i){
        for(j=0; j<n; ++j)
          covar_msg << covar[i*n+j] << " ";
        covar_msg << std::endl;
      }
      covar_msg << std::endl;
      LOG(1, "INFO: " << covar_msg.str());
    }

    free(work);
  }

  if (DoLog(4)) {
   std::stringstream result_msg;
   result_msg << "Levenberg-Marquardt returned " << ret << " in " << info[5] << " iter, reason " << info[6] << "\nSolution: ";
   for(i=0; i<n; ++i)
     result_msg << p[i] << " ";
   result_msg << "\n\nMinimization info:\n";
   std::vector<std::string> infonames(LM_INFO_SZ);
   infonames[0] = std::string("||e||_2 at initial p");
   infonames[1] = std::string("||e||_2");
   infonames[2] = std::string("||J^T e||_inf");
   infonames[3] = std::string("||Dp||_2");
   infonames[4] = std::string("mu/max[J^T J]_ii");
   infonames[5] = std::string("# iterations");
   infonames[6] = std::string("reason for termination");
   infonames[7] = std::string(" # function evaluations");
   infonames[8] = std::string(" # Jacobian evaluations");
   infonames[9] = std::string(" # linear systems solved");
   for(i=0; i<LM_INFO_SZ; ++i)
      result_msg << infonames[i] << ": " << info[i] << " ";
    result_msg << std::endl;
    LOG(4, "DEBUG: " << result_msg.str());
  }

  std::copy(p, p+n, weights.begin());
  LOG(4, "DEBUG: Weights for atom #" << _walker.getId() << ": " << weights);

  delete[] p;
  delete[] x;
  delete[] matrix;

  return weights;
}

BondVectors::weights_t BondVectors::getWeightsForAtomAtStep(
    const atom &_walker,
    const size_t &_step) const
{
  const std::vector<Vector> BondVectors =
      getAtomsBondVectorsAtStep(_walker, _step);

  return getWeightsForAtomAtStep(_walker, BondVectors, _step);
}

void BondVectors::getProjectedGradientsForAtomAtStep(
    const atom &_walker,
    const Vector &_walkerGradient,
    const size_t _timestep,
    std::vector< std::vector<double> > &_projected_forces) const
{
  // gather subset of BondVectors for the current atom
  const std::vector<Vector> BondVectors = getAtomsBondVectorsAtStep(_walker, _timestep);
  const BondList& ListOfBonds = _walker.getListOfBonds(); // we always use current bonds

  // go over all its bonds
  std::vector<Vector>::const_iterator vectoriter = BondVectors.begin();
  for(BondList::const_iterator bonditer = ListOfBonds.begin();
      bonditer != ListOfBonds.end(); ++bonditer, ++vectoriter) {
    const bond::ptr &current_bond = *bonditer;
    const Vector &BondVector = *vectoriter;

    // bv goes from rightatom to leftatom
    // as plus sign in force indicates expansion, minus indicates contraction
    // leftatom: from right to left means same sign expansion, opposite contraction
    // rightatom: from right to left means opposite sign expansion, same sign contraction
    const double sign = (&_walker == current_bond->leftatom) ? 1. : -1.;
    const double temp = sign*_walkerGradient.ScalarProduct(BondVector);
    LOG(4, "DEBUG: BondVector " << BondVector << " receives projected force of "
        << sign << "*" << _walkerGradient << "*" << BondVector << " = " << temp);
    const size_t index = getIndexForBond(current_bond);
    ASSERT( index != (size_t)-1,
        "ForceAnnealing() - could not find bond "+toString(*current_bond)
        +" in bondvectors");
    std::vector<double> &forcelist = (&_walker == current_bond->leftatom) ?
        _projected_forces[leftside] : _projected_forces[rightside];
    forcelist[index] = temp;
  }
  ASSERT( vectoriter == BondVectors.end(),
      "BondVectors::getRemnantGradientForAtomAtStep() - vectoriter is not at end when it should be.");
}

Vector BondVectors::getRemnantGradientForAtomAtStep(
    const atom &_walker,
    const Vector &_walkerGradient,
    const std::vector<Vector> _BondVectors,
    const BondVectors::weights_t &_weights,
    const size_t &_step,
    forcestore_t _forcestore) const
{
  BondVectors::weights_t::const_iterator weightiter = _weights.begin();
  std::vector<Vector>::const_iterator vectoriter = _BondVectors.begin();
  const BondList& ListOfBonds = _walker.getListOfBonds();

  Vector forcesum;
  for(BondList::const_iterator bonditer = ListOfBonds.begin();
      bonditer != ListOfBonds.end(); ++bonditer, ++weightiter, ++vectoriter) {
    const bond::ptr &current_bond = *bonditer;
    const Vector &BondVector = *vectoriter;

    const double temp = (*weightiter)*_walkerGradient.ScalarProduct(BondVector);
    _forcestore(_walker, current_bond, _step, temp);
    LOG(4, "DEBUG: BondVector " << BondVector << " receives projected force of "
        << (*weightiter) << "*" << _walkerGradient << "*" << BondVector << " = " << temp);
    forcesum += temp * BondVector;
  }
  ASSERT( weightiter == _weights.end(),
      "BondVectors::getRemnantGradientForAtomAtStep() - weightiter is not at end when it should be.");
  ASSERT( vectoriter == _BondVectors.end(),
      "BondVectors::getRemnantGradientForAtomAtStep() - vectoriter is not at end when it should be.");

  return _walkerGradient-forcesum;
}

bool BondVectors::getCheckWeightSumForAtomAtStep(
    const atom &_walker,
    const std::vector<Vector> _BondVectors,
    const BondVectors::weights_t &_weights,
    const size_t &_step) const
{
  bool status = true;
  for (std::vector<Vector>::const_iterator iter = _BondVectors.begin();
      iter != _BondVectors.end(); ++iter) {
    std::vector<double> scps;
    scps.reserve(_BondVectors.size());
    std::transform(
        _BondVectors.begin(), _BondVectors.end(),
        _weights.begin(),
        std::back_inserter(scps),
        boost::bind(static_cast< double (*)(double) >(&fabs),
            boost::bind(std::multiplies<double>(),
                boost::bind(&Vector::ScalarProduct, boost::cref(*iter), _1),
                _2))
        );
    const double scp_sum = std::accumulate(scps.begin(), scps.end(), 0.);
    if (fabs(scp_sum -1.) > MYEPSILON)
      status = false;
  }

  return status;
}