/*
 * Copyright 2009-2019 The VOTCA Development Team (http://www.votca.org)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

#include "../../../include/votca/csg/potentialfunctions/potentialfunctioncbspl.h"
#include <boost/lexical_cast.hpp>
#include <iostream>
#include <votca/tools/table.h>

using namespace std;
using namespace votca::tools;

namespace votca {
namespace csg {

PotentialFunctionCBSPL::PotentialFunctionCBSPL(const string &name, Index nlam,
                                               double min, double max)
    : PotentialFunction(name, nlam, min, max) {

  /* Here nlam_ is the total number of coeff values that are to be optimized
   * To ensure that potential and force go to zero smoothly near cut-off,
   * as suggested in Ref. PCCP, 11, 1901, 2009, coeff values leading up to
   * cut-off and beyond take a value of zero.
   *
   * Since region less than rmin is not sampled sufficiently for stability
   * first  nexcl_ coefficients are not optimized instead their values are
   * extrapolated from first statistically significant knot values near rmin
   */

  Index nknots;

  nknots = lam_.size();

  nbreak_ = nknots - 2;

  dr_ = (cut_off_) / (double(nbreak_ - 1));

  // break point locations
  // since ncoeff = nbreak +2 , r values for last two coefficients are also
  // computed
  rbreak_ = Eigen::VectorXd::Zero(nknots);

  for (Index i = 0; i < nknots; i++) {
    rbreak_(i) = double(i) * dr_;
  }

  // exclude knots corresponding to r <=  min_
  nexcl_ = std::min((Index)(min_ / dr_), nbreak_ - 2) + 1;

  // account for finite numerical division of  min_/ dr_
  // e.g. 0.24/0.02 may result in 11.99999999999999
  if (rbreak_(nexcl_) == min_) {
    nexcl_++;
  }

  // fixing last 4 knots to zeros is reasonable
  ncutcoeff_ = 4;

  // check if we have enough parameters to optimize
  if ((Index(lam_.size()) - nexcl_ - ncutcoeff_) < 1) {
    throw std::runtime_error(
        "In potential " + name_ +
        ": no parameters to optimize!\n"
        "All the knot values fall in the range of either excluded (due to high "
        "repulsive region) or cut-off region.\n"
        "This issue can be resolved by one or combination of following steps:\n"
        "1. Make sure you are using large-enough cut-off for this CG "
        "potential.\n"
        "2. Make sure the CG-MD runs are sufficiently Index and CG-MD RDF are "
        "statistically reliable.\n"
        "3. Use more knot values.\n");
  }

  M_ = Eigen::MatrixXd::Zero(4, 4);
  M_(0, 0) = 1.0;
  M_(0, 1) = 4.0;
  M_(0, 2) = 1.0;
  M_(0, 3) = 0.0;
  M_(1, 0) = -3.0;
  M_(1, 1) = 0.0;
  M_(1, 2) = 3.0;
  M_(1, 3) = 0.0;
  M_(2, 0) = 3.0;
  M_(2, 1) = -6.0;
  M_(2, 2) = 3.0;
  M_(2, 3) = 0.0;
  M_(3, 0) = -1.0;
  M_(3, 1) = 3.0;
  M_(3, 2) = -3.0;
  M_(3, 3) = 1.0;
  M_ /= 6.0;
}

Index PotentialFunctionCBSPL::getOptParamSize() const {

  return lam_.size() - nexcl_ - ncutcoeff_;
}

void PotentialFunctionCBSPL::setParam(string filename) {

  Table param;
  param.Load(filename);
  lam_.setZero();

  if (param.size() != lam_.size()) {

    throw std::runtime_error("In potential " + name_ +
                             ": parameters size mismatch!\n"
                             "Check input parameter file \"" +
                             filename + "\" \nThere should be " +
                             boost::lexical_cast<string>(lam_.size()) +
                             " parameters");
  } else {
    // force last  ncutcoeff_ to zero
    Index nonzero = lam_.size() - ncutcoeff_;
    lam_.head(nonzero) = param.y().head(nonzero);
  }
}

void PotentialFunctionCBSPL::SaveParam(const string &filename) {

  extrapolExclParam();

  Table param;
  param.SetHasYErr(false);
  param.resize(lam_.size());

  // write extrapolated knots with flag 'o'
  // points close to rmin can also be stastically not reliable
  // so flag 3 more points next to rmin as 'o'
  for (Index i = 0; i < nexcl_ + 3; i++) {
    param.set(i, rbreak_(i), lam_(i), 'o');
  }

  for (Index i = nexcl_ + 3; i < lam_.size(); i++) {
    param.set(i, rbreak_(i), lam_(i), 'i');
  }

  param.Save(filename);
}

void PotentialFunctionCBSPL::SavePotTab(const string &filename, double step,
                                        double rmin, double rcut) {
  extrapolExclParam();
  PotentialFunction::SavePotTab(filename, step, rmin, rcut);
}

void PotentialFunctionCBSPL::SavePotTab(const string &filename, double step) {
  extrapolExclParam();
  PotentialFunction::SavePotTab(filename, step);
}

void PotentialFunctionCBSPL::extrapolExclParam() {

  double u0 = lam_(nexcl_);
  double m = (lam_(nexcl_ + 1) - lam_(nexcl_)) /
             (rbreak_(nexcl_ + 1) - rbreak_(nexcl_));
  double r0 = rbreak_(nexcl_);

  /* If the slope m is positive then the potential core
   * will be attractive. So, artificially forcing core to be
   * repulsive by setting m = -m
   */
  if (m > 0) {
    cout << name_ << " potential's extrapolated core is attractive!" << endl;
    cout << "Artifically enforcing repulsive core.\n" << endl;
    m *= -1.0;
  }
  // using linear extrapolation
  // u(r) = ar + b
  // a = m
  // b = - m*r0 + u0
  // m = (u1-u0)/(r1-r0)

  double a = m;
  double b = -1.0 * m * r0 + u0;
  for (Index i = 0; i < nexcl_; i++) {
    lam_(i) = a * rbreak_(i) + b;
  }
}

void PotentialFunctionCBSPL::setOptParam(Index i, double val) {

  lam_(i + nexcl_) = val;
}

double PotentialFunctionCBSPL::getOptParam(Index i) const {

  return lam_(i + nexcl_);
}

double PotentialFunctionCBSPL::CalculateF(double r) const {

  if (r <= cut_off_) {

    double u = 0.0;
    Index indx = std::min((Index)(r / dr_), nbreak_ - 2);
    double rk = (double)indx * dr_;
    double t = (r - rk) / dr_;

    Eigen::Vector4d R = Eigen::Vector4d::Zero();
    R(0) = 1.0;
    R(1) = t;
    R(2) = t * t;
    R(3) = t * t * t;
    Eigen::Vector4d B = lam_.segment<4>(indx);
    u += ((R.transpose() * M_) * B).value();
    return u;

  } else {
    return 0.0;
  }
}

// calculate first derivative w.r.t. ith parameter
double PotentialFunctionCBSPL::CalculateDF(Index i, double r) const {

  // since first  nexcl_ parameters are not optimized for stability reasons

  if (r <= cut_off_) {

    Index i_opt = i + nexcl_;
    Index indx;
    double rk;

    indx = std::min((Index)(r / dr_), nbreak_ - 2);
    rk = (double)indx * dr_;

    if (i_opt >= indx && i_opt <= indx + 3) {

      Eigen::Vector4d R = Eigen::Vector4d::Zero();

      double t = (r - rk) / dr_;

      R(0) = 1.0;
      R(1) = t;
      R(2) = t * t;
      R(3) = t * t * t;

      Eigen::Vector4d RM = R.transpose() * M_;

      return RM(i_opt - indx);

    } else {
      return 0.0;
    }

  } else {
    return 0.0;
  }
}

// calculate second derivative w.r.t. ith parameter
double PotentialFunctionCBSPL::CalculateD2F(Index, Index, double) const {

  return 0.0;
}

}  // namespace csg
}  // namespace votca
