/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2012 University of Bonn. All rights reserved.
 * Copyright (C)  2013 Frederik Heber. All rights reserved.
 * Please see the COPYING file or "Copyright notice" in builder.cpp for details.
 * 
 *
 *   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/>. 
 */

/*
 * ManyBodyPotential_Tersoff.cpp
 *
 *  Created on: Sep 26, 2012
 *      Author: heber
 */


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

#include "CodePatterns/MemDebug.hpp"

#include "ManyBodyPotential_Tersoff.hpp"

#include <boost/assign/list_of.hpp> // for 'map_list_of()'
#include <boost/bind.hpp>
#include <cmath>
#include <string>

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

#include "FunctionApproximation/Extractors.hpp"
#include "FunctionApproximation/TrainingData.hpp"
#include "Potentials/helpers.hpp"
#include "Potentials/ParticleTypeCheckers.hpp"

class Fragment;

// static definitions
const ManyBodyPotential_Tersoff::ParameterNames_t
ManyBodyPotential_Tersoff::ParameterNames =
      boost::assign::list_of<std::string>
      ("A")
      ("B")
      ("lambda")
      ("mu")
      ("beta")
      ("n")
      ("c")
      ("d")
      ("h")
//      ("R")
//      ("S")
//      ("lambda3")
//      ("alpha")
//      ("chi")
//      ("omega")
    ;
const std::string ManyBodyPotential_Tersoff::potential_token("tersoff");

ManyBodyPotential_Tersoff::ManyBodyPotential_Tersoff() :
    EmpiricalPotential(),
    params(parameters_t(MAXPARAMS, 0.)),
    R(3.2),
    S(3.5),
    lambda3(0.),
    alpha(0.),
    chi(1.),
    omega(1.),
    triplefunction(&Helpers::NoOp_Triplefunction)
{}

ManyBodyPotential_Tersoff::ManyBodyPotential_Tersoff(
    const ParticleTypes_t &_ParticleTypes
    ) :
    EmpiricalPotential(_ParticleTypes),
    params(parameters_t(MAXPARAMS, 0.)),
    R(3.2),
    S(3.5),
    lambda3(0.),
    alpha(0.),
    chi(1.),
    omega(1.),
    triplefunction(&Helpers::NoOp_Triplefunction)
{
  // have some decent defaults for parameter_derivative checking
  params[A] = 3000.;
  params[B] = 300.;
  params[lambda] = 5.;
  params[mu] = 3.;
  params[beta] = 2.;
  params[n] = 1.;
  params[c] = 0.01;
  params[d] = 1.;
  params[h] = 0.01;
}

ManyBodyPotential_Tersoff::ManyBodyPotential_Tersoff(
    const ParticleTypes_t &_ParticleTypes,
    const double &_R,
    const double &_S,
    const double &_A,
    const double &_B,
    const double &_lambda,
    const double &_mu,
    const double &_lambda3,
    const double &_alpha,
    const double &_beta,
    const double &_chi,
    const double &_omega,
    const double &_n,
    const double &_c,
    const double &_d,
    const double &_h) :
  EmpiricalPotential(_ParticleTypes),
  params(parameters_t(MAXPARAMS, 0.)),
  R(_R),
  S(_S),
  lambda3(_lambda3),
  alpha(_alpha),
  chi(_chi),
  omega(_mu),
  triplefunction(&Helpers::NoOp_Triplefunction)
{
//  Info info(__func__);
//  R = _R;
//  S = _S;
  params[A] = _A;
  params[B] = _B;
  params[lambda] = _lambda;
  params[mu] = _mu;
//  lambda3 = _lambda3;
//  alpha = _alpha;
  params[beta] = _beta;
//  chi = _chi;
//  omega = _omega;
  params[n] = _n;
  params[c] = _c;
  params[d] = _d;
  params[h] = _h;
}

void ManyBodyPotential_Tersoff::setParameters(const parameters_t &_params)
{
  const size_t paramsDim = _params.size();
  ASSERT( paramsDim <= getParameterDimension(),
      "ManyBodyPotential_Tersoff::setParameters() - we need not more than "
      +toString(getParameterDimension())+" parameters.");
  for (size_t i=0; i< paramsDim; ++i)
    params[i] = _params[i];

#ifndef NDEBUG
  parameters_t check_params(getParameters());
  check_params.resize(paramsDim); // truncate to same size
  ASSERT( check_params == _params,
      "ManyBodyPotential_Tersoff::setParameters() - failed, mismatch in to be set "
      +toString(_params)+" and set "+toString(check_params)+" params.");
#endif
}

ManyBodyPotential_Tersoff::results_t
ManyBodyPotential_Tersoff::operator()(
    const arguments_t &arguments
    ) const
{
//  Info info(__func__);
  double result = 0.;
  for(arguments_t::const_iterator argiter = arguments.begin();
      argiter != arguments.end();
      ++argiter) {
    const argument_t &r_ij = *argiter;
    ASSERT( ParticleTypeChecker::checkArgumentsAgainstParticleTypes(
        arguments_t(1, r_ij), getParticleTypes()),
        "ManyBodyPotential_Tersoff::operator() - types don't match with ones in arguments.");
 
    const double cutoff = function_cutoff(r_ij.distance);
    const double temp = (cutoff == 0.) ?
        0. :
        cutoff * (
            function_prefactor(
                alpha,
                function_eta(r_ij))
            * function_smoother(
                params[A],
                params[lambda],
                r_ij.distance)
            +
            function_prefactor(
                params[beta],
                function_zeta(r_ij))
            * function_smoother(
                -params[B],
                params[mu],
                r_ij.distance)
        );
    result += temp;
  }
//  LOG(2, "DEBUG: operator()(" << r_ij.distance << ") = " << result);
  return std::vector<result_t>(1, result);
}

ManyBodyPotential_Tersoff::derivative_components_t
ManyBodyPotential_Tersoff::derivative(
    const arguments_t &arguments
    ) const
{
//  Info info(__func__);
  return ManyBodyPotential_Tersoff::derivative_components_t();
}

ManyBodyPotential_Tersoff::results_t
ManyBodyPotential_Tersoff::parameter_derivative(
    const arguments_t &arguments,
    const size_t index
    ) const
{
//  Info info(__func__);
//  ASSERT( arguments.size() == 1,
//      "ManyBodyPotential_Tersoff::parameter_derivative() - requires exactly one argument.");

  double result = 0.;
  for(arguments_t::const_iterator argiter = arguments.begin();
      argiter != arguments.end();
      ++argiter) {
    const argument_t &r_ij = *argiter;
    ASSERT( ParticleTypeChecker::checkArgumentsAgainstParticleTypes(
        arguments_t(1, r_ij), getParticleTypes()),
        "ManyBodyPotential_Tersoff::operator() - types don't match with ones in arguments.");

  switch (index) {
//    case R:
//    {
//      result += 0.;
//      break;
//    }
//    case S:
//    {
//      result += 0.;
//      break;
//    }
    case A:
    {
      const double cutoff = function_cutoff(r_ij.distance);
      result += (cutoff == 0.) ?
          0. :
          cutoff *
              function_prefactor(
                  alpha,
                  function_eta(r_ij))
              * function_smoother(
                  1.,
                  params[lambda],
                  r_ij.distance);
//          cutoff * function_prefactor(
//              alpha,
//              function_eta(r_ij))
//          * function_smoother(
//              1.,
//              params[mu],
//              r_ij.distance);
      break;
    }
    case B:
    {
      const double cutoff = function_cutoff(r_ij.distance);
      result += (cutoff == 0.) ?
          0. :
          cutoff * function_prefactor(
              params[beta],
              function_zeta(r_ij))
          * function_smoother(
              -1.,
              params[mu],
              r_ij.distance);
//          cutoff * function_prefactor(
//              beta,
//              function_zeta(r_ij))
//          * function_smoother(
//              -params[B],
//              params[mu],
//              r_ij.distance)/params[B];
      break;
    }
    case lambda:
    {
      const double cutoff = function_cutoff(r_ij.distance);
      result += (cutoff == 0.) ?
          0. :
          -r_ij.distance * cutoff *
              function_prefactor(
                  alpha,
                  function_eta(r_ij))
              * function_smoother(
                  params[A],
                  params[lambda],
                  r_ij.distance);
      break;
    }
    case mu:
    {
      const double cutoff = function_cutoff(r_ij.distance);
      result += (cutoff == 0.) ?
          0. :
          -r_ij.distance * cutoff *(
          function_prefactor(
              params[beta],
              function_zeta(r_ij))
          * function_smoother(
              -params[B],
              params[mu],
              r_ij.distance)
      );
      break;
    }
//    case lambda3:
//    {
//      result += 0.;
//      break;
//    }
//    case alpha:
//    {
//      const double temp =
//          pow(alpha*function_eta(r_ij), params[n]);
//      const double cutoff = function_cutoff(r_ij.distance);
//      result += (cutoff == 0.) || (alpha == 0. )?
//          0. :
//          function_smoother(
//              params[A],
//              params[lambda],
//              r_ij.distance)
//          * (-.5) * alpha * (temp/alpha)
//          / (1. + temp)
//          ;
//      break;
//    }
//    case chi:
//    {
//      result += 0.;
//      break;
//    }
//    case omega:
//    {
//      result += 0.;
//      break;
//    }
    case beta:
    {
      const double temp =
          pow(params[beta]*function_zeta(r_ij), params[n]);
      const double cutoff = function_cutoff(r_ij.distance);
      result += (cutoff == 0.) || (params[beta] == 0. )?
          0. : cutoff *
          function_smoother(
              -params[B],
              params[mu],
              r_ij.distance)
          * (-.5)
          * function_prefactor(
              params[beta],
              function_zeta(r_ij))
          * (temp/params[beta])
          / (1. + temp)
          ;
      break;
    }
    case n:
    {
      const double zeta = function_zeta(r_ij);
      const double temp = pow( params[beta]*zeta , params[n]);
      const double cutoff = function_cutoff(r_ij.distance);
      const double tempres = ((cutoff == 0.) || (zeta == 0.)) ? // zeta must be caught if zero due to log
          0. : .5 * cutoff *
          function_smoother(
              -params[B],
              params[mu],
              r_ij.distance)
          * function_prefactor(
              params[beta],
              function_zeta(r_ij))
          * ( log(1.+temp)/(params[n]*params[n]) - temp
              * (log(function_zeta(r_ij)) + log(params[beta]))
              /(params[n]*(1.+temp)))
          ;
//      if (tempres != tempres)
//	LOG(2, "DEBUG: tempres is NaN."); 
//      LOG(2, "DEBUG: Adding " << tempres << " for p.d. w.r.t n, temp=" << temp << ", cutoff=" << cutoff);
      result += tempres;
      break;
    }
    case c:
    {
      const double zeta = function_zeta(r_ij);
      if (zeta == 0.)
        break;
      const double temp =
          pow(zeta, params[n]-1.) * pow(params[beta],params[n]);
      const double cutoff = function_cutoff(r_ij.distance);
      const double tempres = (cutoff == 0.) ?
          0. : cutoff *
          function_smoother(
              -params[B],
              params[mu],
              r_ij.distance)
          * function_prefactor(
              params[beta],
              zeta)
           * (-1.) * temp / (1.+temp*zeta);
      double factor = function_derivative_c(r_ij);
      result += tempres*factor;
      if (result != result)
        ELOG(1, "result is NaN, zeta=" << zeta << ", temp=" << temp << ", cutoff=" << cutoff << ", tempres=" << tempres << ", factor=" << factor);
      break;
    }
    case d:
    {
      const double zeta = function_zeta(r_ij);
      const double temp =
          pow(zeta, params[n]-1.) * pow(params[beta],params[n]);
      const double cutoff = function_cutoff(r_ij.distance);
      const double tempres = (cutoff == 0.) ?
          0. : cutoff *
          function_smoother(
              -params[B],
              params[mu],
              r_ij.distance)
          * function_prefactor(
              params[beta],
              zeta)
           * (-1.) * temp / (1.+temp*zeta);
      double factor = function_derivative_d(r_ij);
      result += tempres*factor;
      if (result != result)
        ELOG(1, "result is NaN, zeta=" << zeta << ", temp=" << temp << ", cutoff=" << cutoff << ", tempres=" << tempres << ", factor=" << factor);
      break;
    }
    case h:
    {
      const double zeta = function_zeta(r_ij);
      const double temp =
          pow(zeta, params[n]-1.) * pow(params[beta],params[n]);
      const double cutoff = function_cutoff(r_ij.distance);
      const double tempres = (cutoff == 0.) ?
          0. : cutoff *
          function_smoother(
              -params[B],
              params[mu],
              r_ij.distance)
          * function_prefactor(
              params[beta],
              zeta)
           * (-1.) * temp / (1.+temp*zeta);
      double factor = function_derivative_h(r_ij);
      result += tempres*factor;
      if (result != result)
        ELOG(1, "result is NaN, zeta=" << zeta << ", temp=" << temp << ", cutoff=" << cutoff << ", tempres=" << tempres << ", factor=" << factor);
      break;
    }
    default:
      ASSERT(0, "ManyBodyPotential_Tersoff::parameter_derivative() - derivative to unknown parameter desired.");
      break;
  }
  if (result != result)
    ELOG(1, "result is NaN.");
  }
  return results_t(1,-result);
}

ManyBodyPotential_Tersoff::result_t
ManyBodyPotential_Tersoff::function_derivative_c(
    const argument_t &r_ij
  ) const
{
  double result = 0.;
  std::vector<arguments_t> triples = triplefunction(r_ij, S);
  for (std::vector<arguments_t>::const_iterator iter = triples.begin();
      iter != triples.end(); ++iter) {
    ASSERT( iter->size() == 2,
        "ManyBodyPotential_Tersoff::function_derivative_c() - the triples result must contain exactly two distances.");
    const argument_t &r_ik = (*iter)[0];
    const argument_t &r_jk = (*iter)[1];
    const double tempangle = params[h] - function_theta(r_ij.distance, r_ik.distance, r_jk.distance);
    const double cutoff = function_cutoff(r_ik.distance);
    result += (cutoff == 0.) ?
        0. : cutoff * omega * exp( Helpers::pow(lambda3 * (r_ij.distance - r_ik.distance) ,3)) * (
            params[c]/Helpers::pow(params[d],2)
            - params[c] / ( Helpers::pow(params[d],2) + Helpers::pow(tempangle,2) )
        );
  }
  return result;
}

ManyBodyPotential_Tersoff::result_t
ManyBodyPotential_Tersoff::function_derivative_d(
    const argument_t &r_ij
  ) const
{
  double result = 0.;
  std::vector<arguments_t> triples = triplefunction(r_ij, S);
  for (std::vector<arguments_t>::const_iterator iter = triples.begin();
      iter != triples.end(); ++iter) {
    ASSERT( iter->size() == 2,
        "ManyBodyPotential_Tersoff::function_derivative_d() - the triples result must contain exactly two distances.");
    const argument_t &r_ik = (*iter)[0];
    const argument_t &r_jk = (*iter)[1];
    const double tempangle = params[h] - function_theta(r_ij.distance, r_ik.distance, r_jk.distance);
    const double cutoff = function_cutoff(r_ik.distance);
    result += (cutoff == 0.) ?
        0. : cutoff * omega  * exp( Helpers::pow(lambda3 * (r_ij.distance - r_ik.distance) ,3)) * (
          - Helpers::pow(params[c],2)/Helpers::pow(params[d],3)
          + Helpers::pow(params[c],2) * params[d]
            / Helpers::pow(Helpers::pow(params[d],2) + Helpers::pow(tempangle,2),2)
        );
  }
  return result;
}

ManyBodyPotential_Tersoff::result_t
ManyBodyPotential_Tersoff::function_derivative_h(
    const argument_t &r_ij
  ) const
{
  double result = 0.;
  std::vector<arguments_t> triples = triplefunction(r_ij, S);
  for (std::vector<arguments_t>::const_iterator iter = triples.begin();
      iter != triples.end(); ++iter) {
    ASSERT( iter->size() == 2,
        "ManyBodyPotential_Tersoff::function_derivative_h() - the triples result must contain exactly two distances.");
    const argument_t &r_ik = (*iter)[0];
    const argument_t &r_jk = (*iter)[1];
    const double tempangle = params[h] - function_theta(r_ij.distance, r_ik.distance, r_jk.distance);
    const double cutoff = function_cutoff(r_ik.distance);
    result += (cutoff == 0.) ?
        0. : cutoff * omega  * exp( Helpers::pow(lambda3 * (r_ij.distance - r_ik.distance) ,3)) * (
          ( Helpers::pow(params[c],2)*tempangle )
            / Helpers::pow(Helpers::pow(params[d],2) + Helpers::pow(tempangle,2),2)
        );
  }
  return result;
}

ManyBodyPotential_Tersoff::result_t
ManyBodyPotential_Tersoff::function_cutoff(
    const double &distance
  ) const
{
//  Info info(__func__);
  double result = 0.;
  if (distance < R)
    result = 1.;
  else if (distance > S)
    result = 0.;
  else {
    result = (0.5 + 0.5 * cos( M_PI * (distance - R)/(S-R)));
  }
//  LOG(2, "DEBUG: function_cutoff(" << distance << ") = " << result);
  return result;
}

ManyBodyPotential_Tersoff::result_t
ManyBodyPotential_Tersoff::function_prefactor(
    const double &alpha,
    const double &eta
  ) const
{
//  Info info(__func__);
  const double result = chi * pow(
      (1. + pow(alpha * eta, params[n])),
      -1./(2.*params[n]));
//  LOG(2, "DEBUG: function_prefactor(" << alpha << "," << eta << ") = " << result);
  return result;
}

ManyBodyPotential_Tersoff::result_t
ManyBodyPotential_Tersoff::function_smoother(
    const double &prefactor,
    const double &lambda,
    const double &distance
    ) const
{
//  Info info(__func__);
  const double result = prefactor * exp(-lambda * distance);
//  LOG(2, "DEBUG: function_smoother(" << prefactor << "," << lambda << "," << distance << ") = " << result);
  return result;
}

ManyBodyPotential_Tersoff::result_t
ManyBodyPotential_Tersoff::function_eta(
    const argument_t &r_ij
  ) const
{
//  Info info(__func__);
  result_t result = 0.;

  // get all triples within the cutoff
  std::vector<arguments_t> triples = triplefunction(r_ij, S);
  for (std::vector<arguments_t>::const_iterator iter = triples.begin();
      iter != triples.end(); ++iter) {
    ASSERT( iter->size() == 2,
        "ManyBodyPotential_Tersoff::function_zeta() - the triples result must contain of exactly two distances.");
    const argument_t &r_ik = (*iter)[0];
    result += function_cutoff(r_ik.distance)
        * exp( Helpers::pow(lambda3 * (r_ij.distance - r_ik.distance) ,3));
  }

//  LOG(2, "DEBUG: function_eta(" << r_ij.distance << ") = " << result);
  return result;
}

ManyBodyPotential_Tersoff::result_t
ManyBodyPotential_Tersoff::function_zeta(
    const argument_t &r_ij
  ) const
{
//  Info info(__func__);
  result_t result = 0.;

  // get all triples within the cutoff
  std::vector<arguments_t> triples = triplefunction(r_ij, S);
  for (std::vector<arguments_t>::const_iterator iter = triples.begin();
      iter != triples.end(); ++iter) {
    ASSERT( iter->size() == 2,
        "ManyBodyPotential_Tersoff::function_zeta() - the triples result must contain exactly two distances.");
    const argument_t &r_ik = (*iter)[0];
    const argument_t &r_jk = (*iter)[1];
    result +=
        function_cutoff(r_ik.distance)
        * omega
        * function_angle(r_ij.distance, r_ik.distance, r_jk.distance)
        * exp( Helpers::pow(lambda3 * (r_ij.distance - r_ik.distance) ,3));
  }

//  LOG(2, "DEBUG: function_zeta(" << r_ij.distance << ") = " << result);
  return result;
}

ManyBodyPotential_Tersoff::result_t
ManyBodyPotential_Tersoff::function_theta(
    const double &r_ij,
    const double &r_ik,
    const double &r_jk
  ) const
{
  const double angle = Helpers::pow(r_ij,2) + Helpers::pow(r_ik,2) - Helpers::pow(r_jk,2);
  const double divisor = 2.* r_ij * r_ik;
  if (divisor != 0.) {
    LOG(2, "DEBUG: cos(theta)= " << angle/divisor);
    return angle/divisor;
  } else
    return 0.;
}

ManyBodyPotential_Tersoff::result_t
ManyBodyPotential_Tersoff::function_angle(
    const double &r_ij,
    const double &r_ik,
    const double &r_jk
  ) const
{
//  Info info(__func__);
  const double result =
      1.
      + (Helpers::pow(params[c]/params[d], 2))
      - Helpers::pow(params[c], 2)/(Helpers::pow(params[d], 2) +
          Helpers::pow(params[h] - function_theta(r_ij, r_ik, r_jk),2));

//  LOG(2, "DEBUG: function_angle(" << r_ij << "," << r_ik << "," << r_jk << ") = " << result);
  return result;
}

FunctionModel::extractor_t
ManyBodyPotential_Tersoff::getSpecificExtractor() const
{
  FunctionModel::extractor_t returnfunction =
      boost::bind(&Extractors::gatherAllDistances,
          boost::bind(&Fragment::getPositions, _1),
          boost::bind(&Fragment::getCharges, _1),
          _2);
  return returnfunction;
}

FunctionModel::filter_t ManyBodyPotential_Tersoff::getSpecificFilter() const
{
  FunctionModel::filter_t returnfunction =
      boost::bind(&Extractors::filterArgumentsByParticleTypes,
          _1,
          getParticleTypes());
  return returnfunction;
}

void
ManyBodyPotential_Tersoff::setParametersToRandomInitialValues(
    const TrainingData &data)
{
//  params[ManyBodyPotential_Tersoff::R] = 1./AtomicLengthToAngstroem;
//  params[ManyBodyPotential_Tersoff::S] = 2./AtomicLengthToAngstroem;
  params[ManyBodyPotential_Tersoff::A] = 1e+4*rand()/(double)RAND_MAX;//1.393600e+03;
  params[ManyBodyPotential_Tersoff::B] = 1e+4*rand()/(double)RAND_MAX;//3.467000e+02;
  params[ManyBodyPotential_Tersoff::lambda] = 1e+1*rand()/(double)RAND_MAX;//3.487900e+00;
  params[ManyBodyPotential_Tersoff::mu] = 1e+1*rand()/(double)RAND_MAX;//2.211900e+00;
  //  params[ManyBodyPotential_Tersoff::lambda3] = 0.;
  //  params[ManyBodyPotential_Tersoff::alpha] = 0.;
  params[ManyBodyPotential_Tersoff::beta] = 1e-1*rand()/(double)RAND_MAX;//1.572400e-07;
  //  params[ManyBodyPotential_Tersoff::chi] = 1.;
  //  params[ManyBodyPotential_Tersoff::omega] = 1.;
  params[ManyBodyPotential_Tersoff::n] = 1e+1*rand()/(double)RAND_MAX;//7.275100e-01;
  params[ManyBodyPotential_Tersoff::c] = 1e+1*rand()/(double)RAND_MAX;//3.804900e+04;
  params[ManyBodyPotential_Tersoff::d] = 1e+1*rand()/(double)RAND_MAX;//4.384000e+00;
  params[ManyBodyPotential_Tersoff::h] = 1e+1*rand()/(double)RAND_MAX;//-5.705800e-01;
}