/* ----------------------------------------------------------------------------

 * GTSAM Copyright 2010, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)

 * See LICENSE for the license information

 * -------------------------------------------------------------------------- */

/**
 * @file    HessianFactor.cpp
 * @author  Richard Roberts
 * @date    Dec 8, 2010
 */

#include <gtsam/linear/HessianFactor.h>

#include <gtsam/linear/GaussianConditional.h>
#include <gtsam/linear/GaussianFactor.h>
#include <gtsam/linear/GaussianFactorGraph.h>
#include <gtsam/linear/JacobianFactor.h>
#include <gtsam/linear/linearExceptions.h>
#include <gtsam/base/cholesky.h>
#include <gtsam/base/debug.h>
#include <gtsam/base/FastMap.h>
#include <gtsam/base/Matrix.h>
#include <gtsam/base/ThreadsafeException.h>
#include <gtsam/base/timing.h>

#include <boost/format.hpp>
#include <boost/make_shared.hpp>
#include <boost/tuple/tuple.hpp>
#ifdef __GNUC__
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-variable"
#endif
#include <boost/bind.hpp>
#ifdef __GNUC__
#pragma GCC diagnostic pop
#endif
#include <boost/assign/list_of.hpp>
#include <boost/range/adaptor/transformed.hpp>
#include <boost/range/adaptor/map.hpp>
#include <boost/range/algorithm/copy.hpp>

#include <sstream>
#include <limits>

using namespace std;
using namespace boost::assign;
namespace br {
using namespace boost::range;
using namespace boost::adaptors;
}

namespace gtsam {

/* ************************************************************************* */
void HessianFactor::Allocate(const Scatter& scatter) {
  gttic(HessianFactor_Allocate);

  // Allocate with dimensions for each variable plus 1 at the end for the information vector
  const size_t n = scatter.size();
  keys_.resize(n);
  FastVector<DenseIndex> dims(n + 1);
  DenseIndex slot = 0;
  for(const SlotEntry& slotentry: scatter) {
    keys_[slot] = slotentry.key;
    dims[slot] = slotentry.dimension;
    ++slot;
  }
  dims.back() = 1;
  info_ = SymmetricBlockMatrix(dims);
}

/* ************************************************************************* */
HessianFactor::HessianFactor(const Scatter& scatter) {
  Allocate(scatter);
}

/* ************************************************************************* */
HessianFactor::HessianFactor() :
    info_(cref_list_of<1>(1)) {
  assert(info_.rows() == 1);
  constantTerm() = 0.0;
}

/* ************************************************************************* */
HessianFactor::HessianFactor(Key j, const Matrix& G, const Vector& g, double f) :
    GaussianFactor(cref_list_of<1>(j)), info_(cref_list_of<2>(G.cols())(1)) {
  if (G.rows() != G.cols() || G.rows() != g.size())
    throw invalid_argument(
        "Attempting to construct HessianFactor with inconsistent matrix and/or vector dimensions");
  info_.setDiagonalBlock(0, G);
  linearTerm() = g;
  constantTerm() = f;
}

/* ************************************************************************* */
// error is 0.5*(x-mu)'*inv(Sigma)*(x-mu) = 0.5*(x'*G*x - 2*x'*G*mu + mu'*G*mu)
// where G = inv(Sigma), g = G*mu, f = mu'*G*mu = mu'*g
HessianFactor::HessianFactor(Key j, const Vector& mu, const Matrix& Sigma) :
    GaussianFactor(cref_list_of<1>(j)), info_(cref_list_of<2>(Sigma.cols())(1)) {
  if (Sigma.rows() != Sigma.cols() || Sigma.rows() != mu.size())
    throw invalid_argument(
        "Attempting to construct HessianFactor with inconsistent matrix and/or vector dimensions");
  info_.setDiagonalBlock(0, Sigma.inverse()); // G
  linearTerm() = info_.diagonalBlock(0) * mu; // g
  constantTerm() = mu.dot(linearTerm().col(0)); // f
}

/* ************************************************************************* */
HessianFactor::HessianFactor(Key j1, Key j2, const Matrix& G11,
    const Matrix& G12, const Vector& g1, const Matrix& G22, const Vector& g2,
    double f) :
    GaussianFactor(cref_list_of<2>(j1)(j2)), info_(
        cref_list_of<3>(G11.cols())(G22.cols())(1)) {
  info_.setDiagonalBlock(0, G11);
  info_.setOffDiagonalBlock(0, 1, G12);
  info_.setDiagonalBlock(1, G22);
  linearTerm() << g1, g2;
  constantTerm() = f;
}

/* ************************************************************************* */
HessianFactor::HessianFactor(Key j1, Key j2, Key j3, const Matrix& G11,
    const Matrix& G12, const Matrix& G13, const Vector& g1, const Matrix& G22,
    const Matrix& G23, const Vector& g2, const Matrix& G33, const Vector& g3,
    double f) :
    GaussianFactor(cref_list_of<3>(j1)(j2)(j3)), info_(
        cref_list_of<4>(G11.cols())(G22.cols())(G33.cols())(1)) {
  if (G11.rows() != G11.cols() || G11.rows() != G12.rows()
      || G11.rows() != G13.rows() || G11.rows() != g1.size()
      || G22.cols() != G12.cols() || G33.cols() != G13.cols()
      || G22.cols() != g2.size() || G33.cols() != g3.size())
    throw invalid_argument(
        "Inconsistent matrix and/or vector dimensions in HessianFactor constructor");
  info_.setDiagonalBlock(0, G11);
  info_.setOffDiagonalBlock(0, 1, G12);
  info_.setOffDiagonalBlock(0, 2, G13);
  info_.setDiagonalBlock(1, G22);
  info_.setOffDiagonalBlock(1, 2, G23);
  info_.setDiagonalBlock(2, G33);
  linearTerm() << g1, g2, g3;
  constantTerm() = f;
}

/* ************************************************************************* */
namespace {
DenseIndex _getSizeHF(const Vector& m) {
  return m.size();
}
}

/* ************************************************************************* */
HessianFactor::HessianFactor(const std::vector<Key>& js,
    const std::vector<Matrix>& Gs, const std::vector<Vector>& gs, double f) :
    GaussianFactor(js), info_(gs | br::transformed(&_getSizeHF), true) {
  // Get the number of variables
  size_t variable_count = js.size();

  // Verify the provided number of entries in the vectors are consistent
  if (gs.size() != variable_count
      || Gs.size() != (variable_count * (variable_count + 1)) / 2)
    throw invalid_argument(
        "Inconsistent number of entries between js, Gs, and gs in HessianFactor constructor.\nThe number of keys provided \
        in js must match the number of linear vector pieces in gs. The number of upper-diagonal blocks in Gs must be n*(n+1)/2");

  // Verify the dimensions of each provided matrix are consistent
  // Note: equations for calculating the indices derived from the "sum of an arithmetic sequence" formula
  for (size_t i = 0; i < variable_count; ++i) {
    DenseIndex block_size = gs[i].size();
    // Check rows
    for (size_t j = 0; j < variable_count - i; ++j) {
      size_t index = i * (2 * variable_count - i + 1) / 2 + j;
      if (Gs[index].rows() != block_size) {
        throw invalid_argument(
            "Inconsistent matrix and/or vector dimensions in HessianFactor constructor");
      }
    }
    // Check cols
    for (size_t j = 0; j <= i; ++j) {
      size_t index = j * (2 * variable_count - j + 1) / 2 + (i - j);
      if (Gs[index].cols() != block_size) {
        throw invalid_argument(
            "Inconsistent matrix and/or vector dimensions in HessianFactor constructor");
      }
    }
  }

  // Fill in the blocks
  size_t index = 0;
  for (size_t i = 0; i < variable_count; ++i) {
    for (size_t j = i; j < variable_count; ++j) {
      if (i == j) {
        info_.setDiagonalBlock(i, Gs[index]);
      } else {
        info_.setOffDiagonalBlock(i, j, Gs[index]);
      }
      index++;
    }
    info_.setOffDiagonalBlock(i, variable_count, gs[i]);
  }
  constantTerm() = f;
}

/* ************************************************************************* */
namespace {
void _FromJacobianHelper(const JacobianFactor& jf, SymmetricBlockMatrix& info) {
  gttic(HessianFactor_fromJacobian);
  const SharedDiagonal& jfModel = jf.get_model();
  auto A = jf.matrixObject().full();
  if (jfModel) {
    if (jf.get_model()->isConstrained())
      throw invalid_argument(
          "Cannot construct HessianFactor from JacobianFactor with constrained noise model");

    auto D = (jfModel->invsigmas().array() * jfModel->invsigmas().array()).matrix().asDiagonal();

    info.setFullMatrix(A.transpose() * D * A);
  } else {
    info.setFullMatrix(A.transpose() * A);
  }
}
}

/* ************************************************************************* */
HessianFactor::HessianFactor(const JacobianFactor& jf) :
    GaussianFactor(jf), info_(
        SymmetricBlockMatrix::LikeActiveViewOf(jf.matrixObject())) {
  _FromJacobianHelper(jf, info_);
}

/* ************************************************************************* */
HessianFactor::HessianFactor(const GaussianFactor& gf) :
    GaussianFactor(gf) {
  // Copy the matrix data depending on what type of factor we're copying from
  if (const JacobianFactor* jf = dynamic_cast<const JacobianFactor*>(&gf)) {
    info_ = SymmetricBlockMatrix::LikeActiveViewOf(jf->matrixObject());
    _FromJacobianHelper(*jf, info_);
  } else if (const HessianFactor* hf = dynamic_cast<const HessianFactor*>(&gf)) {
    info_ = hf->info_;
  } else {
    throw std::invalid_argument(
        "In HessianFactor(const GaussianFactor& gf), gf is neither a JacobianFactor nor a HessianFactor");
  }
}

/* ************************************************************************* */
HessianFactor::HessianFactor(const GaussianFactorGraph& factors,
    boost::optional<const Scatter&> scatter) {
  gttic(HessianFactor_MergeConstructor);

  Allocate(scatter ? *scatter : Scatter(factors));

  // Form A' * A
  gttic(update);
  info_.setZero();
  for(const auto& factor: factors)
    if (factor)
      factor->updateHessian(keys_, &info_);
  gttoc(update);
}

/* ************************************************************************* */
void HessianFactor::print(const std::string& s,
    const KeyFormatter& formatter) const {
  cout << s << "\n";
  cout << " keys: ";
  for (const_iterator key = begin(); key != end(); ++key)
    cout << formatter(*key) << "(" << getDim(key) << ") ";
  cout << "\n";
  gtsam::print(Matrix(info_.selfadjointView()),
      "Augmented information matrix: ");
}

/* ************************************************************************* */
bool HessianFactor::equals(const GaussianFactor& lf, double tol) const {
  const HessianFactor* rhs = dynamic_cast<const HessianFactor*>(&lf);
  if (!rhs || !Factor::equals(lf, tol))
    return false;
  return equal_with_abs_tol(augmentedInformation(), rhs->augmentedInformation(),
      tol);
}

/* ************************************************************************* */
Matrix HessianFactor::augmentedInformation() const {
  return info_.selfadjointView();
}

/* ************************************************************************* */
Eigen::SelfAdjointView<SymmetricBlockMatrix::constBlock, Eigen::Upper>
HessianFactor::informationView() const {
  return info_.selfadjointView(0, size());
}

/* ************************************************************************* */
Matrix HessianFactor::information() const {
  return informationView();
}

/* ************************************************************************* */
VectorValues HessianFactor::hessianDiagonal() const {
  VectorValues d;
  for (DenseIndex j = 0; j < (DenseIndex)size(); ++j) {
    d.insert(keys_[j], info_.diagonal(j));
  }
  return d;
}

/* ************************************************************************* */
// Raw memory access version should be called in Regular Factors only currently
void HessianFactor::hessianDiagonal(double* d) const {
  throw std::runtime_error(
      "HessianFactor::hessianDiagonal raw memory access is allowed for Regular Factors only");
}

/* ************************************************************************* */
map<Key, Matrix> HessianFactor::hessianBlockDiagonal() const {
  map<Key, Matrix> blocks;
  // Loop over all variables
  for (DenseIndex j = 0; j < (DenseIndex) size(); ++j) {
    // Get the diagonal block, and insert it
    blocks.emplace(keys_[j], info_.diagonalBlock(j));
  }
  return blocks;
}

/* ************************************************************************* */
Matrix HessianFactor::augmentedJacobian() const {
  return JacobianFactor(*this).augmentedJacobian();
}

/* ************************************************************************* */
std::pair<Matrix, Vector> HessianFactor::jacobian() const {
  return JacobianFactor(*this).jacobian();
}

/* ************************************************************************* */
double HessianFactor::error(const VectorValues& c) const {
  // error 0.5*(f - 2*x'*g + x'*G*x)
  const double f = constantTerm();
  if (empty()) {
    return 0.5 * f;
  }
  double xtg = 0, xGx = 0;
  // extract the relevant subset of the VectorValues
  // NOTE may not be as efficient
  const Vector x = c.vector(keys());
  xtg = x.dot(linearTerm().col(0));
  auto AtA = informationView();
  xGx = x.transpose() * AtA * x;
  return 0.5 * (f - 2.0 * xtg + xGx);
}

/* ************************************************************************* */
void HessianFactor::updateHessian(const FastVector<Key>& infoKeys,
                                  SymmetricBlockMatrix* info) const {
  gttic(updateHessian_HessianFactor);
  assert(info);
  // Apply updates to the upper triangle
  DenseIndex nrVariablesInThisFactor = size(), nrBlocksInInfo = info->nBlocks() - 1;
  vector<DenseIndex> slots(nrVariablesInThisFactor + 1);
  // Loop over this factor's blocks with indices (i,j)
  // For every block (i,j), we determine the block (I,J) in info.
  for (DenseIndex j = 0; j <= nrVariablesInThisFactor; ++j) {
    const bool rhs = (j == nrVariablesInThisFactor);
    const DenseIndex J = rhs ? nrBlocksInInfo : Slot(infoKeys, keys_[j]);
    slots[j] = J;
    for (DenseIndex i = 0; i <= j; ++i) {
      const DenseIndex I = slots[i];  // because i<=j, slots[i] is valid.

      if (i == j) {
        assert(I == J);
        info->updateDiagonalBlock(I, info_.diagonalBlock(i));
      } else {
        assert(i < j);
        assert(I != J);
        info->updateOffDiagonalBlock(I, J, info_.aboveDiagonalBlock(i, j));
      }
    }
  }
}

/* ************************************************************************* */
GaussianFactor::shared_ptr HessianFactor::negate() const {
  shared_ptr result = boost::make_shared<This>(*this);
  // Negate the information matrix of the result
  result->info_.negate();
  return result;
}

/* ************************************************************************* */
void HessianFactor::multiplyHessianAdd(double alpha, const VectorValues& x,
    VectorValues& yvalues) const {

  // Create a vector of temporary y values, corresponding to rows i
  vector<Vector> y;
  y.reserve(size());
  for (const_iterator it = begin(); it != end(); it++)
    y.push_back(Vector::Zero(getDim(it)));

  // Accessing the VectorValues one by one is expensive
  // So we will loop over columns to access x only once per column
  // And fill the above temporary y values, to be added into yvalues after
  for (DenseIndex j = 0; j < (DenseIndex) size(); ++j) {
    // xj is the input vector
    Vector xj = x.at(keys_[j]);
    DenseIndex i = 0;
    for (; i < j; ++i)
      y[i] += info_.aboveDiagonalBlock(i, j) * xj;

    // blocks on the diagonal are only half
    y[i] += info_.diagonalBlock(j) * xj;
    // for below diagonal, we take transpose block from upper triangular part
    for (i = j + 1; i < (DenseIndex) size(); ++i)
      y[i] += info_.aboveDiagonalBlock(j, i).transpose() * xj;
  }

  // copy to yvalues
  for (DenseIndex i = 0; i < (DenseIndex) size(); ++i) {
    bool didNotExist;
    VectorValues::iterator it;
    boost::tie(it, didNotExist) = yvalues.tryInsert(keys_[i], Vector());
    if (didNotExist)
      it->second = alpha * y[i]; // init
    else
      it->second += alpha * y[i]; // add
  }
}

/* ************************************************************************* */
VectorValues HessianFactor::gradientAtZero() const {
  VectorValues g;
  size_t n = size();
  for (size_t j = 0; j < n; ++j)
    g.insert(keys_[j], -info_.aboveDiagonalBlock(j, n));
  return g;
}

/* ************************************************************************* */
// Raw memory access version should be called in Regular Factors only currently
void HessianFactor::gradientAtZero(double* d) const {
  throw std::runtime_error(
      "HessianFactor::gradientAtZero raw memory access is allowed for Regular Factors only");
}

/* ************************************************************************* */
Vector HessianFactor::gradient(Key key, const VectorValues& x) const {
  const Factor::const_iterator it_i = find(key);
  const DenseIndex I = std::distance(begin(), it_i);

  // Sum over G_ij*xj for all xj connecting to xi
  Vector b = Vector::Zero(x.at(key).size());
  for (Factor::const_iterator it_j = begin(); it_j != end(); ++it_j) {
    const DenseIndex J = std::distance(begin(), it_j);

    // Obtain Gij from the Hessian factor
    // Hessian factor only stores an upper triangular matrix, so be careful when i>j
    const Matrix Gij = info_.block(I, J);
    // Accumulate Gij*xj to gradf
    b += Gij * x.at(*it_j);
  }
  // Subtract the linear term gi
  b += -linearTerm(it_i);
  return b;
}

/* ************************************************************************* */
boost::shared_ptr<GaussianConditional> HessianFactor::eliminateCholesky(const Ordering& keys) {
  gttic(HessianFactor_eliminateCholesky);

  GaussianConditional::shared_ptr conditional;

  try {
    // Do dense elimination
    size_t nFrontals = keys.size();
    assert(nFrontals <= size());
    info_.choleskyPartial(nFrontals);

    // TODO(frank): pre-allocate GaussianConditional and write into it
    const VerticalBlockMatrix Ab = info_.split(nFrontals);
    conditional = boost::make_shared<GaussianConditional>(keys_, nFrontals, Ab);

    // Erase the eliminated keys in this factor
    keys_.erase(begin(), begin() + nFrontals);
  } catch (const CholeskyFailed& e) {
    throw IndeterminantLinearSystemException(keys.front());
  }

  // Return result
  return conditional;
}

/* ************************************************************************* */
VectorValues HessianFactor::solve() {
  gttic(HessianFactor_solve);

  // Do Cholesky Factorization
  const size_t n = size();
  assert(info_.nBlocks() == n + 1);
  info_.choleskyPartial(n);
  auto R = info_.triangularView(0, n);
  auto eta = linearTerm();

  // Solve
  const Vector solution = R.solve(eta);

  // Parse into VectorValues
  VectorValues delta;
  std::size_t offset = 0;
  for (DenseIndex j = 0; j < (DenseIndex)n; ++j) {
    const DenseIndex dim = info_.getDim(j);
    delta.insert(keys_[j], solution.segment(offset, dim));
    offset += dim;
  }

  return delta;
}

/* ************************************************************************* */
std::pair<boost::shared_ptr<GaussianConditional>, boost::shared_ptr<HessianFactor> >
EliminateCholesky(const GaussianFactorGraph& factors, const Ordering& keys) {
  gttic(EliminateCholesky);

  // Build joint factor
  HessianFactor::shared_ptr jointFactor;
  try {
    Scatter scatter(factors, keys);
    jointFactor = boost::make_shared<HessianFactor>(factors, scatter);
  } catch (std::invalid_argument&) {
    throw InvalidDenseElimination(
        "EliminateCholesky was called with a request to eliminate variables that are not\n"
        "involved in the provided factors.");
  }

  // Do dense elimination
  auto conditional = jointFactor->eliminateCholesky(keys);

  // Return result
  return make_pair(conditional, jointFactor);
}

/* ************************************************************************* */
std::pair<boost::shared_ptr<GaussianConditional>,
    boost::shared_ptr<GaussianFactor> > EliminatePreferCholesky(
    const GaussianFactorGraph& factors, const Ordering& keys) {
  gttic(EliminatePreferCholesky);

  // If any JacobianFactors have constrained noise models, we have to convert
  // all factors to JacobianFactors.  Otherwise, we can convert all factors
  // to HessianFactors.  This is because QR can handle constrained noise
  // models but Cholesky cannot.
  if (hasConstraints(factors))
    return EliminateQR(factors, keys);
  else
    return EliminateCholesky(factors, keys);
}

} // gtsam


#include <boost/test/unit_test.hpp>
#include "EnergyCalculator.hpp"

#include <cmath>
#include <vector>

#include "Atom.hpp"
#include "common.hpp"
#include "Molecule.hpp"

BOOST_AUTO_TEST_CASE(CH_CH_interaction){
  using namespace fragdock;
  EnergyCalculator calc(1.0, 0.0);

  const Atom atom1(0,Vector3d(0,0,0), 0);
  std::vector<Atom> atoms1(1, atom1);
  const Molecule mol1(atoms1, "mol1", "testtest");

  const Atom atom2(0,Vector3d(4,0,0), 0);
  std::vector<Atom> atoms2(1, atom2);
  const Molecule mol2(atoms2, "mol2", "testtest");
  fltype expected = -0.035579*std::exp(-0.16)
    + (-0.005156) * std::exp(-1.96)
    + (-0.035069) * 1;

  fltype actual = calc.getEnergy(mol1, mol2);
  BOOST_CHECK_CLOSE(expected, actual, 2e-5);
}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS
  @copyright 2016 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_ROUND_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_ROUND_HPP_INCLUDED
#include <boost/simd/detail/overload.hpp>

#include <boost/simd/meta/hierarchy/simd.hpp>
#include <boost/simd/constant/half.hpp>
#include <boost/simd/function/simd/abs.hpp>
#include <boost/simd/function/simd/ceil.hpp>
#include <boost/simd/function/simd/logical_andnot.hpp>
#include <boost/simd/function/simd/copysign.hpp>
#include <boost/simd/function/simd/divides.hpp>
#include <boost/simd/function/simd/is_flint.hpp>
#include <boost/simd/function/simd/is_greater.hpp>
#include <boost/simd/function/simd/is_ltz.hpp>
#include <boost/simd/function/simd/minus.hpp>
#include <boost/simd/function/simd/multiplies.hpp>
#include <boost/simd/function/simd/none.hpp>
#include <boost/simd/function/simd/round2even.hpp>
#include <boost/simd/function/simd/tenpower.hpp>
#include <boost/simd/function/simd/trunc.hpp>
#include <boost/simd/function/simd/unary_minus.hpp>
#include <boost/simd/function/splat.hpp>
#include <boost/mpl/equal_to.hpp>

namespace boost { namespace simd { namespace ext
{
   namespace bd = boost::dispatch;
   namespace bs = boost::simd;
   BOOST_DISPATCH_OVERLOAD(round_
                          , (typename A0, typename X)
                          , bd::cpu_
                          , bs::pack_<bd::integer_<A0>, X>
                          )
   {
      BOOST_FORCEINLINE A0 operator()( const A0& a0) const BOOST_NOEXCEPT
      {
        return a0;
      }
   };

   BOOST_DISPATCH_OVERLOAD(round_
                          , (typename A0, typename X)
                          , bd::cpu_
                          , bs::pack_<bd::floating_<A0>, X>
                          )
   {
      BOOST_FORCEINLINE A0 operator()( const A0& a0) const BOOST_NOEXCEPT
      {
        auto isf = logical_andnot(is_flint(a0+a0), is_flint(a0));
        auto r1 = round2even(a0);
//        if (none(isf)) return r1; // TODO none
        return if_else(isf, -trunc(-a0), r1);
      }
   };

   BOOST_DISPATCH_OVERLOAD(round_
                             , (typename A0, typename A1, typename X)
                             , bd::cpu_
                             , bs::pack_<bd::floating_<A0>, X>
                             , bs::pack_<bd::integer_<A1>, X>
                             )
   {
      BOOST_FORCEINLINE A0 operator()( const A0& a0, const  A1&  a1) const BOOST_NOEXCEPT
      {
        A0 fac = tenpower(a1);
        A0 tmp1 = round(a0*fac)/fac;
        return if_else(is_ltz(a1), round(tmp1), tmp1);
      }
   };

   BOOST_DISPATCH_OVERLOAD(round_
                          , (typename A0, typename A1, typename X)
                          , bd::cpu_
                          , bs::pack_<bd::floating_<A0>, X>
                          , bd::scalar_<bd::unsigned_<A1>>
                          )
   {
     BOOST_FORCEINLINE A0 operator()(A0 const & a0,  A1 const & a1) const
      {
        using itype_t = bd::as_integer_t<A0, unsigned>;
        return round(a0, splat<itype_t>(a1));
      }
   };

   BOOST_DISPATCH_OVERLOAD(round_
                          , (typename A0, typename X, typename A1)
                          , bd::cpu_
                          , bs::pack_<bd::floating_<A0>, X>
                          , bd::scalar_<bd::integer_<A1>>
                          )
   {
     BOOST_FORCEINLINE A0 operator()(A0 const & a0,  A1 const & a1) const
      {
        using itype_t = bd::as_integer_t<A0, unsigned>;
        return round(a0, splat<itype_t>(a1));
      }
   };
} } }

#endif


#include <boost/geometry/formulas/gnomonic_spheroid.hpp>


#include "mex.h"
#include <Eigen/Dense>
#include <iostream>
#include "drakeUtil.h"
#include "RigidBodyManipulator.h"

using namespace Eigen;
using namespace std;

/*
 * A C version of the HandC function from the Featherstone library
 *
 * To set up the model, use HandCpmex(model[,grav_accn])
 * Then to evaluate the dynamics use
 *   [H,C] = HandCmex(q,qd[,f_ext]);
 */

void mexFunction( int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[] ) {

  if (nrhs<1) {
    mexErrMsgIdAndTxt("Drake:HandCmex:NotEnoughInputs","Usage [H,C,dH,dC] = HandCmex(model_ptr,q,qd[,f_ext,df_ext]).");
  }

  if (nrhs==4 && nlhs>2) {
    mexErrMsgIdAndTxt("Drake:HandCmex:NotEnoughInputs","You need to provide df_ext if you request dH or dC while supplying f_ext");
  }

  // first get the model_ptr back from matlab
  RigidBodyManipulator *model= (RigidBodyManipulator*) getDrakeMexPointer(prhs[0]);

  Map<MatrixXd> *f_ext=NULL;
  Map<MatrixXd> *df_ext=NULL;
  if (static_cast<int>(mxGetNumberOfElements(prhs[1]))!=model->num_positions || static_cast<int>(mxGetNumberOfElements(prhs[2]))!=model->num_positions)
    mexErrMsgIdAndTxt("Drake:HandCmex:BadInputs","q and qd must be size %d x 1",model->num_positions);
  Map<VectorXd> q(mxGetPrSafe(prhs[1]), model->num_positions);
  Map<VectorXd> qd(mxGetPrSafe(prhs[2]), model->num_velocities);

  if (nrhs>3) {
    if (!mxIsEmpty(prhs[3])) {
      f_ext = new Map<MatrixXd>(mxGetPrSafe(prhs[3]),6,model->NB);
    }
  }
  if (nrhs>4) {
    if (!mxIsEmpty(prhs[4])) {
      df_ext = new Map<MatrixXd>(mxGetPrSafe(prhs[4]),6*model->NB,2*model->num_positions);
    }
  }

  Map<MatrixXd> *dH=NULL, *dC=NULL;

  plhs[0] = mxCreateDoubleMatrix(model->num_positions,model->num_positions,mxREAL);
  Map<MatrixXd> H(mxGetPrSafe(plhs[0]),model->num_positions,model->num_positions);

  plhs[1] = mxCreateDoubleMatrix(model->num_positions,1,mxREAL);
  Map<VectorXd> C(mxGetPrSafe(plhs[1]),model->num_positions);

  if (nlhs>2) {
    plhs[2] = mxCreateDoubleMatrix(model->num_positions*model->num_positions,model->num_positions,mxREAL);
    dH = new Map<MatrixXd>(mxGetPrSafe(plhs[2]),model->num_positions*model->num_positions,model->num_positions);
  }
  if (nlhs>3) {
    plhs[3] = mxCreateDoubleMatrix(model->num_positions,2*model->num_positions,mxREAL);
    dC = new Map<MatrixXd>(mxGetPrSafe(plhs[3]),model->num_positions,2*model->num_positions);
  }

  model->HandC(q,qd,f_ext,H,C,dH,dC,df_ext);

  // destroy dynamically allocated Map<MatrixXd> (but not the underlying data!)
  if (f_ext) delete f_ext;
  if (df_ext) delete df_ext;
  if (nlhs>2) delete dH;
  if (nlhs>3) delete dC;
}


// Distributed under the MIT License.
// See LICENSE.txt for details.

#include "PointwiseFunctions/AnalyticSolutions/GeneralRelativity/Tov.hpp"

// Need Boost MultiArray because it is used internally by ODEINT
#include "DataStructures/BoostMultiArray.hpp"  // IWYU pragma: keep

#include <algorithm>
#include <array>
#include <boost/numeric/odeint.hpp>  // IWYU pragma: keep
#include <cmath>
#include <cstddef>
#include <functional>
#include <ostream>
#include <pup.h>
#include <vector>

#include "DataStructures/DataVector.hpp"
#include "DataStructures/Tensor/Tensor.hpp"
#include "PointwiseFunctions/Hydro/EquationsOfState/EquationOfState.hpp"
#include "Utilities/ConstantExpressions.hpp"
#include "Utilities/ContainerHelpers.hpp"
#include "Utilities/ErrorHandling/Assert.hpp"
#include "Utilities/GenerateInstantiations.hpp"
#include "Utilities/Gsl.hpp"
#include "Utilities/MakeWithValue.hpp"

// IWYU pragma: no_include <boost/numeric/odeint/integrate/integrate_adaptive.hpp>
// IWYU pragma: no_include <boost/numeric/odeint/stepper/controlled_runge_kutta.hpp>
// IWYU pragma: no_include <boost/numeric/odeint/stepper/dense_output_runge_kutta.hpp>
// IWYU pragma: no_include <boost/numeric/odeint/stepper/generation/make_dense_output.hpp>
// IWYU pragma: no_include <boost/numeric/odeint/stepper/runge_kutta_dopri5.hpp>
// IWYU pragma: no_include <complex>

// IWYU pragma: no_forward_declare boost::numeric::odeint::controlled_runge_kutta
// IWYU pragma: no_forward_declare EquationsOfState::EquationOfState
// IWYU pragma: no_forward_declare Tensor

namespace {

void lindblom_rhs(
    const gsl::not_null<std::array<double, 2>*> dvars,
    const std::array<double, 2>& vars, const double log_enthalpy,
    const EquationsOfState::EquationOfState<true, 1>& equation_of_state) {
  const double& radius_squared = vars[0];
  const double& mass_over_radius = vars[1];
  double& d_radius_squared = (*dvars)[0];
  double& d_mass_over_radius = (*dvars)[1];
  const double specific_enthalpy = std::exp(log_enthalpy);
  const double rest_mass_density =
      get(equation_of_state.rest_mass_density_from_enthalpy(
          Scalar<double>{specific_enthalpy}));
  const double pressure = get(equation_of_state.pressure_from_density(
      Scalar<double>{rest_mass_density}));
  const double energy_density =
      specific_enthalpy * rest_mass_density - pressure;

  // At the center of the star: (u,v) = (0,0)
  if (UNLIKELY((radius_squared == 0.0) and (mass_over_radius == 0.0))) {
    d_radius_squared = -3.0 / (2.0 * M_PI * (energy_density + 3.0 * pressure));
    d_mass_over_radius =
        -2.0 * energy_density / (energy_density + 3.0 * pressure);
  } else {
    const double common_factor =
        (1.0 - 2.0 * mass_over_radius) /
        (4.0 * M_PI * radius_squared * pressure + mass_over_radius);
    d_radius_squared = -2.0 * radius_squared * common_factor;
    d_mass_over_radius =
        -(4.0 * M_PI * radius_squared * energy_density - mass_over_radius) *
        common_factor;
  }
}

class Observer {
 public:
  void operator()(const std::array<double, 2>& vars,
                  const double current_log_enthalpy) {
    radius.push_back(std::sqrt(vars[0]));
    mass_over_radius.push_back(vars[1]);
    log_enthalpy.push_back(current_log_enthalpy);
  }
  std::vector<double> radius;
  std::vector<double> mass_over_radius;
  std::vector<double> log_enthalpy;
};

}  // namespace

namespace gr::Solutions {

TovSolution::TovSolution(
    const EquationsOfState::EquationOfState<true, 1>& equation_of_state,
    const double central_mass_density,
    const double log_enthalpy_at_outer_radius, const double absolute_tolerance,
    const double relative_tolerance) {
  std::array<double, 2> u_and_v = {{0.0, 0.0}};
  std::array<double, 2> dudh_and_dvdh{};
  const double central_log_enthalpy =
      std::log(get(equation_of_state.specific_enthalpy_from_density(
          Scalar<double>{central_mass_density})));
  lindblom_rhs(&dudh_and_dvdh, u_and_v, central_log_enthalpy,
               equation_of_state);
  const double initial_step = -std::min(std::abs(1.0 / dudh_and_dvdh[0]),
                                        std::abs(1.0 / dudh_and_dvdh[1]));
  using StateDopri5 =
      boost::numeric::odeint::runge_kutta_dopri5<std::array<double, 2>>;
  boost::numeric::odeint::dense_output_runge_kutta<
      boost::numeric::odeint::controlled_runge_kutta<StateDopri5>>
      dopri5 = make_dense_output(absolute_tolerance, relative_tolerance,
                                 StateDopri5{});
  Observer observer{};
  boost::numeric::odeint::integrate_adaptive(
      dopri5,
      [&equation_of_state](const std::array<double, 2>& lindblom_u_and_v,
                           std::array<double, 2>& lindblom_dudh_and_dvdh,
                           const double lindblom_enthalpy) {
        return lindblom_rhs(&lindblom_dudh_and_dvdh, lindblom_u_and_v,
                            lindblom_enthalpy, equation_of_state);
      },
      u_and_v, central_log_enthalpy, log_enthalpy_at_outer_radius, initial_step,
      std::ref(observer));
  outer_radius_ = observer.radius.back();
  const double total_mass_over_radius = observer.mass_over_radius.back();
  total_mass_ = total_mass_over_radius * outer_radius_;
  injection_energy_ = sqrt(1. - 2. * total_mass_ / outer_radius_);
  mass_over_radius_interpolant_ =
      intrp::BarycentricRational(observer.radius, observer.mass_over_radius, 5);
  // log_enthalpy(radius) is almost linear so an interpolant of order 3
  // maximizes precision
  log_enthalpy_interpolant_ =
      intrp::BarycentricRational(observer.radius, observer.log_enthalpy, 3);
}

template <typename DataType>
DataType TovSolution::mass_over_radius(const DataType& r) const {
  // Possible optimization: Support DataVector in intrp::BarycentricRational
  auto result = make_with_value<DataType>(r, 0.);
  for (size_t i = 0; i < get_size(r); ++i) {
    ASSERT(
        get_element(r, i) >= 0.0 and get_element(r, i) <= outer_radius_,
        "Invalid radius: " << r << " not in [0.0, " << outer_radius_ << "]\n");
    get_element(result, i) = mass_over_radius_interpolant_(get_element(r, i));
  }
  return result;
}

template <typename DataType>
DataType TovSolution::log_specific_enthalpy(const DataType& r) const {
  // Possible optimization: Support DataVector in intrp::BarycentricRational
  auto result = make_with_value<DataType>(r, 0.);
  for (size_t i = 0; i < get_size(r); ++i) {
    ASSERT(
        get_element(r, i) >= 0.0 and get_element(r, i) <= outer_radius_,
        "Invalid radius: " << r << " not in [0.0, " << outer_radius_ << "]\n");
    get_element(result, i) = log_enthalpy_interpolant_(get_element(r, i));
  }
  return result;
}

void TovSolution::pup(PUP::er& p) {  // NOLINT
  p | outer_radius_;
  p | total_mass_;
  p | injection_energy_;
  p | mass_over_radius_interpolant_;
  p | log_enthalpy_interpolant_;
}

#define DTYPE(data) BOOST_PP_TUPLE_ELEM(0, data)

#define INSTANTIATE(_, data)                                                \
  template DTYPE(data) TovSolution::mass_over_radius(const DTYPE(data) & r) \
      const;                                                                \
  template DTYPE(data)                                                      \
      TovSolution::log_specific_enthalpy(const DTYPE(data) & r) const;

GENERATE_INSTANTIATIONS(INSTANTIATE, (double, DataVector))

#undef DTYPE

}  // namespace gr::Solutions


#include <boost/math/special_functions/next.hpp>
#include <boost/random.hpp>

#include <algorithm>
#include <limits>
#include <vector>

#include "caffe/common.hpp"
#include "caffe/util/math_functions.hpp"
#include "caffe/util/rng.hpp"

namespace caffe {

template<>
void caffe_cpu_gemm<float>(const CBLAS_TRANSPOSE TransA,
    const CBLAS_TRANSPOSE TransB, const int M, const int N, const int K,
    const float alpha, const float* A, const float* B, const float beta,
    float* C) {
  int lda = (TransA == CblasNoTrans) ? K : M;
  int ldb = (TransB == CblasNoTrans) ? N : K;
  cblas_sgemm(CblasRowMajor, TransA, TransB, M, N, K, alpha, A, lda, B,
      ldb, beta, C, N);
}

template<>
void caffe_cpu_gemm<double>(const CBLAS_TRANSPOSE TransA,
    const CBLAS_TRANSPOSE TransB, const int M, const int N, const int K,
    const double alpha, const double* A, const double* B, const double beta,
    double* C) {
  int lda = (TransA == CblasNoTrans) ? K : M;
  int ldb = (TransB == CblasNoTrans) ? N : K;
  cblas_dgemm(CblasRowMajor, TransA, TransB, M, N, K, alpha, A, lda, B,
      ldb, beta, C, N);
}

#ifndef CPU_ONLY
template<>
void caffe_cpu_gemm<float16>(const CBLAS_TRANSPOSE TransA,
    const CBLAS_TRANSPOSE TransB, const int M, const int N, const int K,
    const float16 alpha, const float16* A, const float16* B, const float16 beta,
    float16* C) {
  if (M <= 0 || N <= 0 || K <= 0) {
    return;
  }
  std::vector<float> a(M*K), b(K*N), c(M*N);
  caffe_cpu_convert(a.size(), A, &a.front());
  caffe_cpu_convert(b.size(), B, &b.front());
  caffe_cpu_convert(c.size(), C, &c.front());
  const int lda = (TransA == CblasNoTrans) ? K : M;
  const int ldb = (TransB == CblasNoTrans) ? N : K;
  cblas_sgemm(CblasRowMajor, TransA, TransB, M, N, K, static_cast<float>(alpha), &a.front(),
      lda, &b.front(), ldb, static_cast<float>(beta), &c.front(), N);
  caffe_cpu_convert(c.size(), &c.front(), C);
}
#endif

template <>
void caffe_cpu_gemv<float>(const CBLAS_TRANSPOSE TransA, const int M,
    const int N, const float alpha, const float* A, const float* x,
    const float beta, float* y) {
  cblas_sgemv(CblasRowMajor, TransA, M, N, alpha, A, N, x, 1, beta, y, 1);
}

template <>
void caffe_cpu_gemv<double>(const CBLAS_TRANSPOSE TransA, const int M,
    const int N, const double alpha, const double* A, const double* x,
    const double beta, double* y) {
  cblas_dgemv(CblasRowMajor, TransA, M, N, alpha, A, N, x, 1, beta, y, 1);
}

#ifndef CPU_ONLY
template <>
void caffe_cpu_gemv<float16>(const CBLAS_TRANSPOSE TransA, const int M,
    const int N, const float16 alpha, const float16* A, const float16* x,
    const float16 beta, float16* y) {
  if (M <= 0 || N <= 0) {
    return;
  }
  const int lx = (TransA == CblasNoTrans) ? N : M;
  const int ly = (TransA == CblasNoTrans) ? M : N;
  std::vector<float> a(M*N), xv(lx), yv(ly);
  caffe_cpu_convert(a.size(), A, &a.front());
  caffe_cpu_convert(xv.size(), x, &xv.front());
  caffe_cpu_convert(yv.size(), y, &yv.front());
  cblas_sgemv(CblasRowMajor, TransA, M, N, static_cast<float>(alpha), &a.front(), N,
      &xv.front(), 1, static_cast<float>(beta), &yv.front(), 1);
  caffe_cpu_convert(yv.size(), &yv.front(), y);
}
#endif

template <>
void caffe_axpy<float>(const int N, const float alpha, const float* X,
    float* Y) { cblas_saxpy(N, alpha, X, 1, Y, 1); }

template <>
void caffe_axpy<double>(const int N, const double alpha, const double* X,
    double* Y) { cblas_daxpy(N, alpha, X, 1, Y, 1); }

#ifndef CPU_ONLY
template<>
void caffe_axpy<float16>(const int N, const float16 alpha, const float16* X,
    float16* Y) {
  for (int i = 0; i < N; ++i) {
    Y[i] = alpha * X[i] + Y[i];
  }
}
#endif

template <typename Dtype>
void caffe_set(const int N, const Dtype alpha, Dtype* Y) {
  if (alpha == 0) {
    memset(Y, 0, sizeof(Dtype) * N);  // NOLINT(caffe/alt_fn)
    return;
  }
  for (int i = 0; i < N; ++i) {
    Y[i] = alpha;
  }
}

template void caffe_set<int>(const int N, const int alpha, int* Y);
template void caffe_set<float>(const int N, const float alpha, float* Y);
template void caffe_set<double>(const int N, const double alpha, double* Y);

#ifndef CPU_ONLY
template void caffe_set<float16>(const int N, const float16 alpha, float16* Y);
#endif

template <>
void caffe_add_scalar(const int N, const float alpha, float* Y) {
  for (int i = 0; i < N; ++i) {
    Y[i] += alpha;
  }
}

template <>
void caffe_add_scalar(const int N, const double alpha, double* Y) {
  for (int i = 0; i < N; ++i) {
    Y[i] += alpha;
  }
}

#ifndef CPU_ONLY
template <>
void caffe_add_scalar(const int N, const float16 alpha, float16* Y) {
  for (int i = 0; i < N; ++i) {
    Y[i] = Y[i] + alpha;
  }
}
#endif

template <typename Dtype>
void caffe_copy(const int N, const Dtype* X, Dtype* Y) {
  if (X != Y) {
    if (Caffe::mode() == Caffe::GPU) {
#ifndef CPU_ONLY
      caffe_gpu_memcpy(sizeof(Dtype) * N, X, Y);
#else
      NO_GPU;
#endif
    } else {
      memcpy(Y, X, sizeof(Dtype) * N);  // NOLINT(caffe/alt_fn)
    }
  }
}

template void caffe_copy<int>(const int N, const int* X, int* Y);
template void caffe_copy<unsigned int>(const int N, const unsigned int* X,
    unsigned int* Y);
template void caffe_copy<float>(const int N, const float* X, float* Y);
template void caffe_copy<double>(const int N, const double* X, double* Y);

#ifndef CPU_ONLY
template void caffe_copy<float16>(const int N, const float16* X, float16* Y);
#endif

template <>
void caffe_scal<float>(const int N, const float alpha, float *X) {
  cblas_sscal(N, alpha, X, 1);
}

template <>
void caffe_scal<double>(const int N, const double alpha, double *X) {
  cblas_dscal(N, alpha, X, 1);
}

#ifndef CPU_ONLY
template <>
void caffe_scal<float16>(const int N, const float16 alpha, float16 *X) {
  // cblas_hscal(N, alpha, X, 1); ?
  for (int i = 0; i < N; ++i) {
    X[i] = alpha * X[i];
  }
}
#endif

template <>
void caffe_cpu_axpby<float>(const int N, const float alpha, const float* X,
                            const float beta, float* Y) {
  cblas_saxpby(N, alpha, X, 1, beta, Y, 1);
}

template <>
void caffe_cpu_axpby<double>(const int N, const double alpha, const double* X,
                             const double beta, double* Y) {
  cblas_daxpby(N, alpha, X, 1, beta, Y, 1);
}

#ifndef CPU_ONLY
template <>
void caffe_cpu_axpby<float16>(const int N, const float16 alpha,
    const float16* X, const float16 beta, float16* Y) {
  for (int i = 0; i < N; ++i) {
    Y[i] = alpha * X[i] + beta * Y[i];
  }
}
#endif

template <>
void caffe_add<float>(const int n, const float* a, const float* b,
    float* y) {
  vsAdd(n, a, b, y);
}

template <>
void caffe_add<double>(const int n, const double* a, const double* b,
    double* y) {
  vdAdd(n, a, b, y);
}

#ifndef CPU_ONLY
template <>
void caffe_add<float16>(const int n, const float16* a, const float16* b,
    float16* y) {
  for (int i = 0; i < n; ++i) {
    y[i] = a[i] + b[i];
  }
}
#endif

template <>
void caffe_sub<float>(const int n, const float* a, const float* b,
    float* y) {
  vsSub(n, a, b, y);
}

template <>
void caffe_sub<double>(const int n, const double* a, const double* b,
    double* y) {
  vdSub(n, a, b, y);
}

#ifndef CPU_ONLY
template <>
void caffe_sub<float16>(const int n, const float16* a, const float16* b,
    float16* y) {
  for (int i = 0; i < n; ++i) {
    y[i] = a[i] - b[i];
  }
}
#endif

template <>
void caffe_mul<float>(const int n, const float* a, const float* b,
    float* y) {
  vsMul(n, a, b, y);
}

template <>
void caffe_mul<double>(const int n, const double* a, const double* b,
    double* y) {
  vdMul(n, a, b, y);
}

#ifndef CPU_ONLY
template <>
void caffe_mul<float16>(const int n, const float16* a, const float16* b, float16* y) {
  for (int i = 0; i < n; ++i) {
    //  vhMul(n, a, b, y);
    y[i] = a[i] * b[i];
  }
}
#endif

template <>
void caffe_div<float>(const int n, const float* a, const float* b,
    float* y) {
  vsDiv(n, a, b, y);
}

template <>
void caffe_div<double>(const int n, const double* a, const double* b,
    double* y) {
  vdDiv(n, a, b, y);
}

#ifndef CPU_ONLY
template <>
void caffe_div<float16>(const int n, const float16* a, const float16* b, float16* y) {
  //  vhDiv(n, a, b, y);
  for (int i = 0; i < n; ++i) {
    y[i] = a[i] / b[i];
  }
}
#endif

template <>
void caffe_powx<float>(const int n, const float* a, const float b,
    float* y) {
  vsPowx(n, a, b, y);
}

template <>
void caffe_powx<double>(const int n, const double* a, const double b,
    double* y) {
  vdPowx(n, a, b, y);
}

#ifndef CPU_ONLY
template <>
void caffe_powx<float16>(const int n, const float16* a, const float16 b,
    float16* y) {
  for (int i = 0; i < n; ++i) {
    y[i] = pow(static_cast<float>(a[i]), static_cast<float>(b));
  }
}
#endif

template <>
void caffe_sqr<float>(const int n, const float* a, float* y) {
  vsSqr(n, a, y);
}

template <>
void caffe_sqr<double>(const int n, const double* a, double* y) {
  vdSqr(n, a, y);
}

#ifndef CPU_ONLY
template <>
void caffe_sqr<float16>(const int n, const float16* a, float16* y) {
  vhSqr(n, a, y);
}
#endif

template <>
void caffe_exp<float>(const int n, const float* a, float* y) {
  vsExp(n, a, y);
}

template <>
void caffe_exp<double>(const int n, const double* a, double* y) {
  vdExp(n, a, y);
}

template <>
void caffe_log<float>(const int n, const float* a, float* y) {
  vsLn(n, a, y);
}

template <>
void caffe_log<double>(const int n, const double* a, double* y) {
  vdLn(n, a, y);
}

#ifndef CPU_ONLY
template <>
void caffe_exp<float16>(const int n, const float16* a, float16* y) {
  vhExp(n, a, y);
}
template <>
void caffe_log<float16>(const int n, const float16* a, float16* y) {
  vhLn(n, a, y);
}
#endif  // ! CPU_ONLY

template <>
void caffe_abs<float>(const int n, const float* a, float* y) {
    vsAbs(n, a, y);
}

template <>
void caffe_abs<double>(const int n, const double* a, double* y) {
    vdAbs(n, a, y);
}

#ifndef CPU_ONLY
template <>
void caffe_abs<float16>(const int n, const float16* a, float16* y) {
  for (int i = 0; i < n; ++i) {
    y[i] = fabs(a[i]);
  }
}
#endif

unsigned int caffe_rng_rand() {
  return (*caffe_rng())();
}

template <typename Dtype>
Dtype caffe_nextafter(const Dtype b) {
  return boost::math::nextafter<Dtype>(
      b, std::numeric_limits<Dtype>::max());
}

template
float caffe_nextafter(const float b);

template
double caffe_nextafter(const double b);

template <typename Dtype>
void caffe_rng_uniform(int n, float a, float b, Dtype* r) {
  CHECK_GE(n, 0);
  CHECK(r);
  CHECK_LE(a, b);
  boost::uniform_real<float> random_distribution(a, caffe_nextafter<float>(b));
  boost::variate_generator<caffe::rng_t*, boost::uniform_real<float> >
      variate_generator(caffe_rng(), random_distribution);
  for (int i = 0; i < n; ++i) {
    r[i] = variate_generator();
  }
}

template
void caffe_rng_uniform<float>(int n, float a, float b, float* r);

template
void caffe_rng_uniform<double>(int n, float a, float b, double* r);

#ifndef CPU_ONLY
template
void caffe_rng_uniform<float16>(int n, float a, float b, float16* r);
#endif

template <typename Dtype>
void caffe_rng_gaussian(int n, float a, float sigma, Dtype* r) {
  CHECK_GE(n, 0);
  CHECK(r);
  CHECK_GT(sigma, 0);
  boost::normal_distribution<float> random_distribution(a, sigma);
  boost::variate_generator<caffe::rng_t*, boost::normal_distribution<float> >
      variate_generator(caffe_rng(), random_distribution);
  for (int i = 0; i < n; ++i) {
    r[i] = variate_generator();
  }
}

template
void caffe_rng_gaussian<float>(int n, float mu, float sigma, float* r);

template
void caffe_rng_gaussian<double>(int n, float mu, float sigma, double* r);

#ifndef CPU_ONLY
template
void caffe_rng_gaussian<float16>(const int n, float mu, float sigma, float16* r);
#endif

template <typename Dtype>
void caffe_rng_bernoulli(const int n, const Dtype p, int* r) {
  CHECK_GE(n, 0);
  CHECK(r);
  CHECK_GE(p, 0);
  CHECK_LE(p, 1);
  boost::bernoulli_distribution<Dtype> random_distribution(p);
  boost::variate_generator<caffe::rng_t*, boost::bernoulli_distribution<Dtype> >
      variate_generator(caffe_rng(), random_distribution);
  for (int i = 0; i < n; ++i) {
    r[i] = static_cast<int>(variate_generator());
  }
}

template
void caffe_rng_bernoulli<double>(const int n, const double p, int* r);

template
void caffe_rng_bernoulli<float>(const int n, const float p, int* r);

#ifndef CPU_ONLY
template<>
void caffe_rng_bernoulli<float16>(const int n, const float16 p, int* r) {
  caffe_rng_bernoulli(n, static_cast<const float>(p), r);
}
#endif

template <typename Dtype>
void caffe_rng_bernoulli(const int n, const Dtype p, unsigned int* r) {
  CHECK_GE(n, 0);
  CHECK(r);
  CHECK_GE(p, 0);
  CHECK_LE(p, 1);
  boost::bernoulli_distribution<Dtype> random_distribution(p);
  boost::variate_generator<caffe::rng_t*, boost::bernoulli_distribution<Dtype> >
      variate_generator(caffe_rng(), random_distribution);
  for (int i = 0; i < n; ++i) {
    r[i] = static_cast<unsigned int>(variate_generator());
  }
}

template
void caffe_rng_bernoulli<double>(const int n, const double p, unsigned int* r);

template
void caffe_rng_bernoulli<float>(const int n, const float p, unsigned int* r);

#ifndef CPU_ONLY
template<>
void caffe_rng_bernoulli<float16>(const int n, const float16 p,
    unsigned int* r) {
  caffe_rng_bernoulli(n, static_cast<const float>(p), r);
}
#endif

template <>
float caffe_cpu_strided_dot<float>(const int n, const float* x, const int incx,
    const float* y, const int incy) {
  return cblas_sdot(n, x, incx, y, incy);
}

template <>
double caffe_cpu_strided_dot<double>(const int n, const double* x,
    const int incx, const double* y, const int incy) {
  return cblas_ddot(n, x, incx, y, incy);
}

#ifndef CPU_ONLY
template <>
float16 caffe_cpu_strided_dot<float16>(const int n, const float16* x,
    const int incx, const float16 *y, const int incy) {
  float sum = 0.0f;
  int idx_x, idx_y;
  for (int i = 0; i < n; ++i) {
    idx_x = i*incx;
    idx_y = i*incy;
    sum += x[idx_x] * y[idx_y];
  }
  return float16(sum);
}
#endif

template <typename Dtype>
Dtype caffe_cpu_dot(const int n, const Dtype* x, const Dtype* y) {
  return caffe_cpu_strided_dot(n, x, 1, y, 1);
}

template
float caffe_cpu_dot<float>(const int n, const float* x, const float* y);
template
double caffe_cpu_dot<double>(const int n, const double* x, const double* y);

#ifndef CPU_ONLY
template
float16 caffe_cpu_dot<float16>(const int n, const float16* x, const float16* y);
#endif

template <>
int caffe_cpu_hamming_distance<float>(const int n, const float* x,
                                      const float* y) {
  int dist = 0;
  for (int i = 0; i < n; ++i) {
    dist += __builtin_popcount(static_cast<uint32_t>(x[i]) ^
                               static_cast<uint32_t>(y[i]));
  }
  return dist;
}

template <>
int caffe_cpu_hamming_distance<double>(const int n, const double* x,
                                       const double* y) {
  int dist = 0;
  for (int i = 0; i < n; ++i) {
    dist += __builtin_popcountl(static_cast<uint64_t>(x[i]) ^
                                static_cast<uint64_t>(y[i]));
  }
  return dist;
}

#ifndef CPU_ONLY
template <>
int caffe_cpu_hamming_distance<float16>(const int n, const float16* x,
                                        const float16* y) {
  int dist = 0;
  for (int i = 0; i < n; ++i) {
    dist += __builtin_popcount(x[i].getx() ^ y[i].getx());
  }
  return dist;
}
#endif

template <typename Dtype>
Dtype caffe_cpu_amax(const int n, const Dtype* x) {
  Dtype y = 0.;
  for (int i = 0; i < n; ++i) {
    if (x[i] > 0)
      y = std::max(y, x[i]);
    else
      y = std::max(y, -x[i]);
  }
  return y;
}

template
float caffe_cpu_amax<float>(const int n, const float* x);
template
double caffe_cpu_amax<double>(const int n, const double* x);
#ifndef CPU_ONLY
template
float16 caffe_cpu_amax<float16>(const int n, const float16* x);
#endif

template <>
float caffe_cpu_asum<float>(const int n, const float* x) {
  return cblas_sasum(n, x, 1);
}
template <>
float caffe_cpu_asum<double>(const int n, const double* x) {
  return cblas_dasum(n, x, 1);
}

#ifndef CPU_ONLY
template <>
float caffe_cpu_asum<float16>(const int n, const float16 *x) {
  float sum = 0.0f;
  for (int i = 0; i < n; ++i) {
    sum += fabs(x[i]);
  }
  return sum;
}
#endif

template <>
void caffe_cpu_scale<float>(const int n, const float alpha, const float *x,
                            float* y) {
  cblas_scopy(n, x, 1, y, 1);
  cblas_sscal(n, alpha, y, 1);
}

template <>
void caffe_cpu_scale<double>(const int n, const double alpha, const double *x,
                             double* y) {
  cblas_dcopy(n, x, 1, y, 1);
  cblas_dscal(n, alpha, y, 1);
}

#ifndef CPU_ONLY
template <>
void caffe_cpu_scale<float16>(const int n, const float16 alpha,
    const float16 *x, float16 *y) {
  for (int i = 0; i < n; i++) {
    y[i] = alpha * x[i];
  }
}
#endif

// y[i]= max(a*x[i], b*y[i])
template <typename Dtype>
void caffe_cpu_eltwise_max(const int N, const Dtype alpha, const Dtype* x,
  const Dtype beta, Dtype* y) {
  for (int i = 0; i < N; ++i) {
    y[i] = std::max(alpha * x[i], beta * y[i]);
  }
}
template void caffe_cpu_eltwise_max<float>(const int N,
    const float alpha, const float* x, const float beta, float* y);
template void caffe_cpu_eltwise_max<double>(const int N,
    const double alpha, const double* x, const double beta, double* y);
#ifndef CPU_ONLY
template void caffe_cpu_eltwise_max<float16>(const int N,
    const float16 alpha, const float16* x, const float16 beta, float16* y);
#endif

// y[i]= min(a*x[i], b*y[i])
template <typename Dtype>
void caffe_cpu_eltwise_min(const int N, const Dtype alpha, const Dtype* x,
  const Dtype beta, Dtype* y) {
  for (int i = 0; i < N; ++i) {
    y[i] = std::min(alpha * x[i], beta * y[i]);
  }
}
template void caffe_cpu_eltwise_min<float>(const int N,
    const float alpha, const float* x, const float beta, float* y);
template void caffe_cpu_eltwise_min<double>(const int N,
    const double alpha, const double* x, const double beta, double* y);
#ifndef CPU_ONLY
template void caffe_cpu_eltwise_min<float16>(const int N,
    const float16 alpha, const float16* x, const float16 beta, float16* y);
#endif

}  // namespace caffe


/*
 * Copyright (c) 2017 Spotify AB.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 <boost/test/unit_test.hpp>

#include "Knee.h"
#include "types.h"

namespace nativeformat {
namespace plugin {
namespace compressor {

BOOST_AUTO_TEST_SUITE(CompressorPluginTests)

BOOST_AUTO_TEST_CASE(compressor_test_hard_knee) {
  nativeformat::plugin::compressor::HardKnee<CompressorType> knee;
  // tests for <= threshold
  BOOST_CHECK_CLOSE(knee(-30, -15, 0, 10), -30, .001);
  BOOST_CHECK_CLOSE(knee(-15, -15, 0, 10), -15, .001);
  // tests for > threshold
  BOOST_CHECK_CLOSE(knee(-10, -15, 0, 10), -14.5, .001);
  BOOST_CHECK_CLOSE(knee(0, -15, 0, 10), -13.5, .001);
}

BOOST_AUTO_TEST_CASE(compressor_test_soft_knee) {
  nativeformat::plugin::compressor::SoftKnee<CompressorType> knee;
  BOOST_CHECK_CLOSE(knee(-30, -15, 5, 10), -30, .001);
  BOOST_CHECK_CLOSE(knee(-20, -15, 5, 10), -20, .001);
  BOOST_CHECK_CLOSE(knee(-13, -15, 5, 10), -14.8225, .001);
  BOOST_CHECK_CLOSE(knee(0, -15, 5, 10), -13.5, .001);
}

BOOST_AUTO_TEST_CASE(expander_test_hard_knee) {
  nativeformat::plugin::compressor::HardKnee<ExpanderType> knee;
  // tests for <= threshold
  BOOST_CHECK_CLOSE(knee(-30, -15, 0, 10), -165, .001);
  BOOST_CHECK_CLOSE(knee(-15, -15, 0, 10), -15, .001);
  // tests for > threshold
  BOOST_CHECK_CLOSE(knee(-10, -15, 0, 10), -10, .001);
  BOOST_CHECK_CLOSE(knee(0, -15, 0, 10), -0, .001);
}
BOOST_AUTO_TEST_SUITE_END()

}  // namespace compressor
}  // namespace plugin
}  // namespace nativeformat


#include <Eigen/Dense>
#include "control.hpp"
#include "defines.hpp"
#include "fad.hpp"
#include "global_residual.hpp"
#include "J2_small_strain.hpp"
#include "material_params.hpp"

namespace calibr8 {

static ParameterList get_valid_local_residual_params() {
  ParameterList p;
  p.set<std::string>("type", "J2_small_strain");
  p.set<int>("nonlinear max iters", 0);
  p.set<double>("nonlinear absolute tol", 0.);
  p.set<double>("nonlinear relative tol", 0.);
  p.sublist("materials");
  return p;
}
static ParameterList get_valid_material_params() {
  ParameterList p;
  p.set<double>("E", 0.);
  p.set<double>("nu", 0.);
  p.set<double>("K", 0.);
  p.set<double>("Y", 0.);
  p.set<double>("cte", 0.);
  p.set<double>("delta_T", 0.);
  return p;
}

template <typename T>
J2_small_strain<T>::J2_small_strain(ParameterList const& inputs, int ndims) {

  this->m_params_list = inputs;
  this->m_params_list.validateParameters(get_valid_local_residual_params(), 0);

  int const num_residuals = 2;
  int const num_params = 6;

  this->m_num_residuals = num_residuals;
  this->m_num_eqs.resize(num_residuals);
  this->m_var_types.resize(num_residuals);
  this->m_resid_names.resize(num_residuals);

  this->m_resid_names[0] = "pstrain";
  this->m_var_types[0] = SYM_TENSOR;
  this->m_num_eqs[0] = get_num_eqs(SYM_TENSOR, ndims);

  this->m_resid_names[1] = "alpha";
  this->m_var_types[1] = SCALAR;
  this->m_num_eqs[1] = get_num_eqs(SCALAR, ndims);

  m_max_iters = inputs.get<int>("nonlinear max iters");
  m_abs_tol = inputs.get<double>("nonlinear absolute tol");
  m_rel_tol = inputs.get<double>("nonlinear relative tol");

}

template <typename T>
J2_small_strain<T>::~J2_small_strain() {
}

template <typename T>
void J2_small_strain<T>::init_params() {

  int const num_params = 6;
  this->m_params.resize(num_params);
  this->m_param_names.resize(num_params);

  this->m_param_names[0] = "E";
  this->m_param_names[1] = "nu";
  this->m_param_names[2] = "K";
  this->m_param_names[3] = "Y";
  this->m_param_names[4] = "cte";
  this->m_param_names[5] = "delta_T";

  int const num_elem_sets = this->m_elem_set_names.size();
  resize(this->m_param_values, num_elem_sets, num_params);

  ParameterList& all_material_params =
      this->m_params_list.sublist("materials", true);

  for (int es = 0; es < num_elem_sets; ++es) {
    std::string const& elem_set_name = this->m_elem_set_names[es];
    ParameterList& material_params =
        all_material_params.sublist(elem_set_name, true);
    material_params.validateParameters(get_valid_material_params(), 0);
    this->m_param_values[es][0] = material_params.get<double>("E");
    this->m_param_values[es][1] = material_params.get<double>("nu");
    this->m_param_values[es][2] = material_params.get<double>("K");
    this->m_param_values[es][3] = material_params.get<double>("Y");
    this->m_param_values[es][4] = material_params.get<double>("cte");
    this->m_param_values[es][5] = material_params.get<double>("delta_T");
  }

  this->m_active_indices.resize(1);
  this->m_active_indices[0].resize(1);
  this->m_active_indices[0][0] = 0;
}

template <typename T>
void J2_small_strain<T>::init_variables_impl() {

  int const ndims = this->m_num_dims;
  int const pstrain_idx = 0;
  int const alpha_idx = 1;

  T const alpha = 0.0;
  Tensor<T> const pstrain = minitensor::zero<T>(ndims);

  this->set_scalar_xi(alpha_idx, alpha);
  this->set_sym_tensor_xi(pstrain_idx, pstrain);

}

template <>
int J2_small_strain<double>::solve_nonlinear(RCP<GlobalResidual<double>>) {
  return 0;
}

template <>
int J2_small_strain<FADT>::solve_nonlinear(RCP<GlobalResidual<FADT>> global) {

  int path;

  // pick an initial guess for the local variables
  {
    Tensor<FADT> const pstrain_old = this->sym_tensor_xi_prev(0);
    Tensor<FADT> const pstrain = pstrain_old;
    FADT const alpha_old = this->scalar_xi_prev(1);
    FADT const alpha = alpha_old;
    this->set_sym_tensor_xi(0, pstrain);
    this->set_scalar_xi(1, alpha);
    path = ELASTIC;
  }

  // newton iteration until convergence

  int iter = 1;
  double R_norm_0 = 1.;
  bool converged = false;

  while ((iter <= m_max_iters) && (!converged)) {

    path = this->evaluate(global);

    double const R_norm = this->norm_residual();
    if (iter == 1) R_norm_0 = R_norm;
    double const R_norm_rel = R_norm / R_norm_0;
    if ((R_norm_rel < m_rel_tol) || (R_norm < m_abs_tol)) {
      converged = true;
      break;
    }

    EMatrix const J = this->eigen_jacobian();
    EVector const R = this->eigen_residual();
    EVector const dxi = J.fullPivLu().solve(-R);

    this->add_to_sym_tensor_xi(0, dxi);
    this->add_to_scalar_xi(1, dxi);

    iter++;

  }

  // fail if convergence was not achieved
  if ((iter > m_max_iters) && (!converged)) {
    fail("J2_small_strain:solve_nonlinear failed in %d iterations", m_max_iters);
  }

  return path;

}

template <typename T>
int J2_small_strain<T>::evaluate(
    RCP<GlobalResidual<T>> global,
    bool force_path,
    int path_in) {

  int path = ELASTIC;
  int const ndims = this->m_num_dims;
  double const sqrt_23 = std::sqrt(2./3.);
  double const sqrt_32 = std::sqrt(3./2.);

  T const E = this->m_params[0];
  T const nu = this->m_params[1];
  T const K = this->m_params[2];
  T const Y = this->m_params[3];
  T const mu = compute_mu(E, nu);

  Tensor<T> const pstrain_old = this->sym_tensor_xi_prev(0);
  T const alpha_old = this->scalar_xi_prev(1);

  Tensor<T> const pstrain = this->sym_tensor_xi(0);
  T const alpha = this->scalar_xi(1);

  Tensor<T> const I = minitensor::eye<T>(ndims);
  Tensor<T> const s = this->dev_cauchy(global);
  T const s_mag = minitensor::norm(s);
  Tensor<T> const n = s / s_mag;
  T const sigma_yield = Y + K * alpha;
  T const f = s_mag - sqrt_23 * sigma_yield;

  Tensor<T> const grad_u_prev = global->grad_vector_x_prev(0);
  Array2D<int> const& active_indices = this->active_indices();
  T const dummy3 = this->params(active_indices[0][0]);

  Tensor<T> R_pstrain;
  T R_alpha;

  if (!force_path) {
    // plastic step
    if (f > m_abs_tol || std::abs(f) < m_abs_tol) {
      T const dgam = sqrt_32 * (alpha - alpha_old);
      R_pstrain = (0. * dummy3 + 1.) * pstrain - pstrain_old - dgam * n + 0. * grad_u_prev;
      R_alpha = (s_mag - sqrt_23 * sigma_yield) / val(mu) + 0. * dummy3;
      path = PLASTIC;
    }
    // elastic step
    else {
      R_pstrain = (0. * dummy3 + 1.) * pstrain - pstrain_old + 0. * s + 0. * grad_u_prev;
      R_alpha = alpha - alpha_old + 0. * dummy3 + 0. * s_mag;
      path = ELASTIC;
    }
  }

  // force the path
  else {
    path = path_in;
    // plastic step
    if (path == PLASTIC) {
      T const dgam = sqrt_32 * (alpha - alpha_old);
      R_pstrain = (0. * dummy3 + 1.) * pstrain - pstrain_old - dgam * n + 0. * grad_u_prev;
      R_alpha = (s_mag - sqrt_23 * sigma_yield) / val(mu) + 0. * dummy3;
    }
    // elastic step
    else {
      R_pstrain = (0. * dummy3 + 1.) * pstrain - pstrain_old + 0. * s + 0. * grad_u_prev;
      R_alpha = alpha - alpha_old + 0. * dummy3 + 0. * s_mag;
    }
  }

  this->set_sym_tensor_R(0, R_pstrain);
  this->set_scalar_R(1, R_alpha);

  return path;

}

template <typename T>
Tensor<T> J2_small_strain<T>::dev_cauchy(RCP<GlobalResidual<T>> global) {
  int const ndims = global->num_dims();
  Tensor<T> const I = minitensor::eye<T>(ndims);
  T const E = this->m_params[0];
  T const nu = this->m_params[1];
  T const mu = E / (2. * (1. + nu));
  Tensor<T> const pstrain = this->sym_tensor_xi(0);
  Tensor<T> const grad_u = global->grad_vector_x(0);
  Tensor<T> const eps = 0.5 * (grad_u + minitensor::transpose(grad_u));
  Tensor<T> const dev_eps = eps - (minitensor::trace(eps) / 3.) * I;
  return 2. * mu * (dev_eps - pstrain);
}

template <typename T>
Tensor<T> J2_small_strain<T>::cauchy(RCP<GlobalResidual<T>> global, T p) {
  int const ndims = global->num_dims();
  T const E = this->m_params[0];
  T const nu = this->m_params[1];
  T const cte = this->m_params[4];
  T const kappa = compute_kappa(E, nu);
  T const delta_T = this->m_params[5];
  Tensor<T> const I = minitensor::eye<T>(ndims);
  Tensor<T> const dev_sigma = this->dev_cauchy(global);
  Tensor<T> const sigma = dev_sigma - p * I - 3.*kappa*cte*delta_T*I;
  return sigma;
}

template class J2_small_strain<double>;
template class J2_small_strain<FADT>;

}


#include <cmath>
#include <boost/filesystem/fstream.hpp>
#include "array.hpp"
#include "scoring_function.hpp"
#include "receptor.hpp"

receptor::receptor(const path& p, const array<float, 3>& center, const array<float, 3>& size, const float granularity) : center(center), size(size), corner0(center - 0.5f * size), corner1(corner0 + size), granularity(granularity), granularity_inverse(1.0f / granularity), num_probes({static_cast<int>(size[0] * granularity_inverse) + 2, static_cast<int>(size[1] * granularity_inverse) + 2, static_cast<int>(size[2] * granularity_inverse) + 2}), num_probes_product(num_probes[0] * num_probes[1] * num_probes[2]), map_bytes(sizeof(float) * num_probes_product), p_offset(scoring_function::n), maps(scoring_function::n)
{
	// Parse the receptor line by line.
	atoms.reserve(2000); // A receptor typically consists of <= 2,000 atoms within bound.
	string residue = "XXXX"; // Current residue sequence located at 1-based [23, 26], used to track residue change, initialized to a dummy value.
	size_t residue_start; // The starting atom of the current residue.
	string line;
	for (boost::filesystem::ifstream ifs(p); getline(ifs, line);)
	{
		const string record = line.substr(0, 6);
		if (record == "ATOM  " || record == "HETATM")
		{
			// If this line is the start of a new residue, mark the starting index within the atoms vector.
			if (line[25] != residue[3] || line[24] != residue[2] || line[23] != residue[1] || line[22] != residue[0])
			{
				residue[3] = line[25];
				residue[2] = line[24];
				residue[1] = line[23];
				residue[0] = line[22];
				residue_start = atoms.size();
			}

			// Parse the ATOM/HETATM line.
			atom a(line);

			// Skip unsupported atom types.
			if (a.ad_unsupported()) continue;

			// Skip non-polar hydrogens.
			if (a.is_nonpolar_hydrogen()) continue;

			// For a polar hydrogen, the bonded hetero atom must be a hydrogen bond donor.
			if (a.is_polar_hydrogen())
			{
				for (size_t i = atoms.size(); i > residue_start;)
				{
					atom& b = atoms[--i];
					if (b.is_hetero() && b.has_covalent_bond(a))
					{
						b.donorize();
						break;
					}
				}
				continue;
			}

			// For a hetero atom, its connected carbon atoms are no longer hydrophobic.
			if (a.is_hetero())
			{
				for (size_t i = atoms.size(); i > residue_start;)
				{
					atom& b = atoms[--i];
					if (!b.is_hetero() && b.has_covalent_bond(a))
					{
						b.dehydrophobicize();
					}
				}
			}
			// For a carbon atom, it is no longer hydrophobic when connected to a hetero atom.
			else
			{
				for (size_t i = atoms.size(); i > residue_start;)
				{
					const atom& b = atoms[--i];
					if (b.is_hetero() && b.has_covalent_bond(a))
					{
						a.dehydrophobicize();
						break;
					}
				}
			}

			// Save the atom if and only if its distance to its projection point on the box is within cutoff.
			float r2 = 0;
			for (size_t i = 0; i < 3; ++i)
			{
				if (a.coord[i] < corner0[i])
				{
					const float d = a.coord[i] - corner0[i];
					r2 += d * d;
				}
				else if (a.coord[i] > corner1[i])
				{
					const float d = a.coord[i] - corner1[i];
					r2 += d * d;
				}
			}
			if (r2 < scoring_function::cutoff_sqr)
			{
				atoms.push_back(move(a));
			}
		}
		else if (record == "TER   ")
		{
			residue = "XXXX";
		}
	}
}

void receptor::precalculate(const scoring_function& sf, const vector<size_t>& xs)
{
	const size_t nxs = xs.size();
	for (size_t t0 = 0; t0 < sf.n; ++t0)
	{
		vector<size_t>& p = p_offset[t0];
		p.resize(nxs);
		for (size_t i = 0; i < nxs; ++i)
		{
			const size_t t1 = xs[i];
			p[i] = sf.nr * mp(t0, t1);
		}
	}
}

void receptor::populate(const vector<size_t>& xs, const size_t z, const scoring_function& sf)
{
	const size_t n = xs.size();
	const float z_coord = corner0[2] + granularity * z;
	const size_t z_offset = num_probes[0] * num_probes[1] * z;

	for (const auto& a : atoms)
	{
		assert(!a.is_hydrogen());
		const float dz = z_coord - a.coord[2];
		const float dz_sqr = dz * dz;
		const float dydx_sqr_ub = scoring_function::cutoff_sqr - dz_sqr;
		if (dydx_sqr_ub <= 0) continue;
		const float dydx_ub = sqrt(dydx_sqr_ub);
		const float y_lb = a.coord[1] - dydx_ub;
		const float y_ub = a.coord[1] + dydx_ub;
		const size_t y_beg = y_lb > corner0[1] ? (y_lb < corner1[1] ? static_cast<size_t>((y_lb - corner0[1]) * granularity_inverse)     : num_probes[1]) : 0;
		const size_t y_end = y_ub > corner0[1] ? (y_ub < corner1[1] ? static_cast<size_t>((y_ub - corner0[1]) * granularity_inverse) + 1 : num_probes[1]) : 0;
		const vector<size_t>& p = p_offset[a.xs];
		size_t zy_offset = z_offset + num_probes[0] * y_beg;
		float dy = corner0[1] + granularity * y_beg - a.coord[1];
		for (size_t y = y_beg; y < y_end; ++y, zy_offset += num_probes[0], dy += granularity)
		{
			const float dy_sqr = dy * dy;
			const float dx_sqr_ub = dydx_sqr_ub - dy_sqr;
			if (dx_sqr_ub <= 0) continue;
			const float dx_ub = sqrt(dx_sqr_ub);
			const float x_lb = a.coord[0] - dx_ub;
			const float x_ub = a.coord[0] + dx_ub;
			const size_t x_beg = x_lb > corner0[0] ? (x_lb < corner1[0] ? static_cast<size_t>((x_lb - corner0[0]) * granularity_inverse)     : num_probes[0]) : 0;
			const size_t x_end = x_ub > corner0[0] ? (x_ub < corner1[0] ? static_cast<size_t>((x_ub - corner0[0]) * granularity_inverse) + 1 : num_probes[0]) : 0;
			const float dzdy_sqr = dz_sqr + dy_sqr;
			size_t zyx_offset = zy_offset + x_beg;
			float dx = corner0[0] + granularity * x_beg - a.coord[0];
			for (size_t x = x_beg; x < x_end; ++x, ++zyx_offset, dx += granularity)
			{
				const float dx_sqr = dx * dx;
				const float r2 = dzdy_sqr + dx_sqr;
				if (r2 >= scoring_function::cutoff_sqr) continue;
				const size_t r_offset = static_cast<size_t>(sf.ns * r2);
				for (size_t i = 0; i < n; ++i)
				{
					maps[xs[i]][zyx_offset] += sf.e[p[i] + r_offset];
				}
			}
		}
	}
}


#ifndef STATE_HPP
#define STATE_HPP

#include <polyfem/ElementBases.hpp>
#include <polyfem/ElementAssemblyValues.hpp>
#include <polyfem/Problem.hpp>
#include <polyfem/Mesh.hpp>
#include <polyfem/Problem.hpp>
#include <polyfem/LocalBoundary.hpp>
#include <polyfem/InterfaceData.hpp>
#include <polyfem/Common.hpp>

#include <polyfem/Mesh2D.hpp>
#include <polyfem/Mesh3D.hpp>

#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <memory>
#include <string>

namespace polyfem
{
	class SolutionFrame
	{
	public:
		std::string name;
		Eigen::MatrixXd points;
		Eigen::MatrixXi connectivity;
		Eigen::MatrixXd solution;
		Eigen::MatrixXd pressure;
		Eigen::MatrixXd exact;
		Eigen::MatrixXd error;
		Eigen::MatrixXd scalar_value;
		Eigen::MatrixXd scalar_value_avg;
	};

	class State
	{
	public:
		~State() = default;

		State();
		void init_logger(const std::string &log_file, int log_level, const bool is_quiet);

		void init(const json &args);
		void init(const std::string &json);

		json args;

		std::shared_ptr<Problem> problem;

		std::vector< ElementBases >    bases;
		std::vector< ElementBases >    pressure_bases;
		std::vector< ElementBases >    geom_bases;

		std::vector< int >                   boundary_nodes;
		std::vector< LocalBoundary >         local_boundary;
		std::vector< LocalBoundary >         local_neumann_boundary;
		std::map<int, InterfaceData> 		 poly_edge_to_data;

		std::vector<int> flipped_elements;

		std::unique_ptr<Mesh> mesh;

		std::map<int, Eigen::MatrixXd> polys;
		std::map<int, std::pair<Eigen::MatrixXd, Eigen::MatrixXi> > polys_3d;
		std::vector<int> parent_elements;

		StiffnessMatrix stiffness, mass;
		Eigen::MatrixXd rhs, rhs_in;
		Eigen::MatrixXd sol, pressure;

		Eigen::Vector4d spectrum;

		json solver_info;

		int n_bases, n_pressure_bases;
		Eigen::VectorXi disc_orders;

		double mesh_size;
		double min_edge_length;
		double average_edge_length;

		double l2_err, linf_err, lp_err, h1_err, h1_semi_err, grad_max_err;

		long long nn_zero, mat_size, num_dofs;

		double building_basis_time;
		double loading_mesh_time;
		double computing_poly_basis_time;
		double assembling_stiffness_mat_time;
		double assigning_rhs_time;
		double solving_time;
		double computing_errors_time;
		double max_angle;
		double sigma_max, sigma_min, sigma_avg;

		int n_flipped;

		int simplex_count;
		int regular_count;
		int regular_boundary_count;
		int simple_singular_count;
		int multi_singular_count;
		int boundary_count;
		int non_regular_boundary_count;
		int non_regular_count;
		int undefined_count;
		int multi_singular_boundary_count;

		bool solve_export_to_file = true;
		std::vector<SolutionFrame> solution_frames;

		json build_json_params();

		void compute_mesh_size(const Mesh &mesh, const std::vector< ElementBases > &bases, const int n_samples);

		void load_mesh();
		void load_mesh(GEO::Mesh &meshin, const std::function<int(const RowVectorNd&)> &boundary_marker, bool skip_boundary_sideset = false);
		void build_basis();

		void assemble_stiffness_mat();
		void assemble_rhs();
		void solve_problem();
		void compute_errors();
		void export_data();

		void compute_vertex_values(int actual_dim, const std::vector< ElementBases > &basis,
			const MatrixXd &fun, Eigen::MatrixXd &result);
		void compute_stress_at_quadrature_points(const MatrixXd &fun, Eigen::MatrixXd &result, Eigen::VectorXd & von_mises);

		void interpolate_function(const int n_points, const Eigen::MatrixXd &fun, Eigen::MatrixXd &result, const bool boundary_only = false);
		void interpolate_function(const int n_points, const int actual_dim, const std::vector< ElementBases > &basis, const MatrixXd &fun, MatrixXd &result, const bool boundary_only = false);

		void compute_scalar_value(const int n_points, const Eigen::MatrixXd &fun, Eigen::MatrixXd &result, const bool boundary_only = false);
		void compute_tensor_value(const int n_points, const Eigen::MatrixXd &fun, Eigen::MatrixXd &result, const bool boundary_only = false);
		void average_grad_based_function(const int n_points, const MatrixXd &fun, MatrixXd &result_scalar, MatrixXd &result_tensor, const bool boundary_only = false);

		void interpolate_boundary_function(const MatrixXd &pts, const MatrixXi &faces, const MatrixXd &fun, const bool compute_avg, MatrixXd &result);
		void interpolate_boundary_function_at_vertices(const MatrixXd &pts, const MatrixXi &faces, const MatrixXd &fun, MatrixXd &result);
		void interpolate_boundary_tensor_function(const MatrixXd &pts, const MatrixXi &faces, const MatrixXd &fun, const MatrixXd &disp, const bool compute_avg, MatrixXd &result);
		void interpolate_boundary_tensor_function(const MatrixXd &pts, const MatrixXi &faces, const MatrixXd &fun, const bool compute_avg, MatrixXd &result);

		void get_sidesets(Eigen::MatrixXd &pts, Eigen::MatrixXi &faces, Eigen::MatrixXd &sidesets);

		void save_json(std::ostream &out);
		void save_json(nlohmann::json &j);
		void save_json();

		void compute_mesh_stats();

		void build_vis_mesh(Eigen::MatrixXd &points, Eigen::MatrixXi &tets, Eigen::MatrixXd &discr);
		void save_vtu(const std::string &name);
		void save_wire(const std::string &name, bool isolines = false);

		inline std::string mesh_path() const { return args["mesh"]; }

		inline std::string formulation() const { return problem->is_scalar() ? scalar_formulation() : tensor_formulation(); }
		inline bool iso_parametric() const {
			if(non_regular_count > 0 || non_regular_boundary_count > 0 || undefined_count > 0)
				return true;

			if(args["use_spline"])
				return true;

			if(mesh->is_rational())
				return false;

			if(args["use_p_ref"])
				return false;

			if(mesh->orders().size() <= 0){
				if(args["discr_order"] == 1)
					return true;
				else
					return args["iso_parametric"];
			}

			if(mesh->orders().minCoeff() != mesh->orders().maxCoeff())
				return false;

			if(args["discr_order"] == mesh->orders().minCoeff())
				return true;

			if( args["disc_orders"] == 1 && args["force_linear_geometry"])
				return true;

			return args["iso_parametric"];
		}

		inline std::string solver_type() const { return args["solver_type"]; }
		inline std::string precond_type() const { return args["precond_type"]; }
		inline const json &solver_params() const { return args["solver_params"]; }

		inline std::string scalar_formulation() const { return args["scalar_formulation"]; }
		inline std::string tensor_formulation() const { return args["tensor_formulation"]; }
		// inline std::string mixed_formulation() const { return args["mixed_formulation"]; }

		void p_refinement(const Mesh2D &mesh2d);
		void p_refinement(const Mesh3D &mesh3d);

	private:
		void sol_to_pressure();
		void build_polygonal_basis();

	};

}
#endif //STATE_HPP


// Copyright (C) 2011 by the BEM++ Authors
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#ifndef bempp_check_arrays_are_close_hpp
#define bempp_check_arrays_are_close_hpp

#include "fiber/scalar_traits.hpp"
#include "fiber/_2d_array.hpp"
#include "fiber/_3d_array.hpp"
#include "fiber/_4d_array.hpp"

#include "common/armadillo_fwd.hpp"
#include <complex>
#include <iomanip>
#include <boost/test/unit_test.hpp>

template <typename ValueType>
boost::test_tools::predicate_result
check_arrays_are_close(const arma::Mat<ValueType>& left,
                       const arma::Mat<ValueType>& right,
                       typename Fiber::ScalarTraits<ValueType>::RealType tolerance)
{
    typedef typename Fiber::ScalarTraits<ValueType>::RealType RealType;
    const int digits10 = std::numeric_limits<RealType>::digits10;

    boost::test_tools::predicate_result result(true);
    if (left.n_rows != right.n_rows || left.n_cols != right.n_cols) {
        result = false;
        result.message() << "Size mismatch [("
                         << left.n_rows << ", " << left.n_cols << ") != ("
                         << right.n_rows << ", " << right.n_cols << ")]";
        return result;
    }
    for (size_t r = 0; r < left.n_rows; ++r)
        for (size_t c = 0; c < left.n_cols; ++c) {
            RealType diff = std::abs(left(r, c) - right(r, c));
            RealType avg = std::abs(left(r, c) + right(r, c)) / 2.;
            if (diff > tolerance * (1 + avg)) {
                result = false;
                result.message() << std::setprecision(digits10)
                                 << "\n  Mismatch at position ("
                                 << r << ", " << c << ") ["
                                 << left(r, c) << " != " << right(r, c) << "]";
            }
        }

    return result;
}

template <typename ValueType>
boost::test_tools::predicate_result
check_arrays_are_close(const arma::Cube<ValueType>& left,
                       const arma::Cube<ValueType>& right,
                       typename Fiber::ScalarTraits<ValueType>::RealType tolerance)
{
    typedef typename Fiber::ScalarTraits<ValueType>::RealType RealType;
    const int digits10 = std::numeric_limits<RealType>::digits10;

    boost::test_tools::predicate_result result(true);
    if (left.n_rows != right.n_rows ||
            left.n_cols != right.n_cols ||
            left.n_slices != right.n_slices) {
        result = false;
        result.message() << "Size mismatch [("
                         << left.n_rows << ", " << left.n_cols << ", "
                         << left.n_slices << ") != ("
                         << right.n_rows << ", " << right.n_cols << ", "
                         << right.n_slices << ")]";
        return result;
    }
    for (size_t r = 0; r < left.n_rows; ++r)
        for (size_t c = 0; c < left.n_cols; ++c)
            for (size_t s = 0; s < left.n_slices; ++s) {
                RealType diff = std::abs(left(r, c, s) - right(r, c, s));
                RealType avg = std::abs(left(r, c, s) + right(r, c, s)) / 2.;
                if (diff > tolerance * (1 + avg)) {
                    result = false;
                    result.message() << std::setprecision(digits10)
                                     << "\n  Mismatch at position ("
                                     << r << ", " << c << ", " << s << ") ["
                                     << left(r, c, s) << " != "
                                     << right(r, c, s) << "]";
                }
            }

    return result;
}

template <typename ValueType>
boost::test_tools::predicate_result
check_arrays_are_close(const Fiber::_4dArray<ValueType>& left,
                       const Fiber::_4dArray<ValueType>& right,
                       typename Fiber::ScalarTraits<ValueType>::RealType tolerance)
{
    typedef typename Fiber::ScalarTraits<ValueType>::RealType RealType;
    const int digits10 = std::numeric_limits<RealType>::digits10;

    boost::test_tools::predicate_result result(true);
    if (left.extent(0) != right.extent(0) ||
            left.extent(1) != right.extent(1) ||
            left.extent(2) != right.extent(2) ||
            left.extent(3) != right.extent(3)) {
        result = false;
        result.message() << "Size mismatch [("
                         << left.extent(0) << ", " << left.extent(1) << ", "
                         << left.extent(2) << ", " << left.extent(3) << ") != ("
                         << right.extent(0) << ", " << right.extent(1) << ", "
                         << right.extent(2) << ", " << right.extent(3) << "]";
        return result;
    }
    for (size_t r = 0; r < left.extent(0); ++r)
        for (size_t c = 0; c < left.extent(1); ++c)
            for (size_t s = 0; s < left.extent(2); ++s)
                for (size_t u = 0; u < left.extent(3); ++u) {
                    RealType diff =
                            std::abs(left(r, c, s, u) - right(r, c, s, u));
                    RealType avg =
                            std::abs(left(r, c, s, u) + right(r, c, s, u)) / 2.;
                    if (diff > tolerance * (1 + avg)) {
                        result = false;
                        result.message() << std::setprecision(digits10)
                                         << "\n  Mismatch at position ("
                                         << r << ", " << c << ", "
                                         << s << ", " << u << ") ["
                                         << left(r, c, s, u) << " != "
                                         << right(r, c, s, u) << "]";
                    }
                }

    return result;
}

template <typename ValueType>
boost::test_tools::predicate_result
check_arrays_are_close(const Fiber::_3dArray<ValueType>& left,
                       const Fiber::_3dArray<ValueType>& right,
                       typename Fiber::ScalarTraits<ValueType>::RealType tolerance)
{
    typedef typename Fiber::ScalarTraits<ValueType>::RealType RealType;
    const int digits10 = std::numeric_limits<RealType>::digits10;

    boost::test_tools::predicate_result result(true);
    if (left.extent(0) != right.extent(0) ||
            left.extent(1) != right.extent(1) ||
            left.extent(2) != right.extent(2)) {
        result = false;
        result.message() << "Size mismatch [("
                         << left.extent(0) << ", " << left.extent(1) << ", "
                         << left.extent(2) << ") != ("
                         << right.extent(0) << ", " << right.extent(1) << ", "
                         << right.extent(2)  << "]";
        return result;
    }
    for (size_t r = 0; r < left.extent(0); ++r)
        for (size_t c = 0; c < left.extent(1); ++c)
            for (size_t s = 0; s < left.extent(2); ++s) {
                    RealType diff = std::abs(left(r, c, s) - right(r, c, s));
                    RealType avg = std::abs(left(r, c, s) + right(r, c, s)) / 2.;
                    if (diff > tolerance * (1 + avg)) {
                        result = false;
                        result.message() << std::setprecision(digits10)
                                         << "\n  Mismatch at position ("
                                         << r << ", " << c << ", "
                                         << s << ") ["
                                         << left(r, c, s) << " != "
                                         << right(r, c, s) << "]";
                    }
            }

    return result;
}

template <typename ValueType>
boost::test_tools::predicate_result
check_arrays_are_close(const Fiber::_2dArray<arma::Mat<ValueType> >& leftArrays,
                       const Fiber::_2dArray<arma::Mat<ValueType> >& rightArrays,
                       typename Fiber::ScalarTraits<ValueType>::RealType tolerance)
{
    typedef typename Fiber::ScalarTraits<ValueType>::RealType RealType;
    const int digits10 = std::numeric_limits<RealType>::digits10;

    boost::test_tools::predicate_result result(true);
    if (leftArrays.extent(0) != rightArrays.extent(0) ||
            leftArrays.extent(1) != rightArrays.extent(1)) {
        result = false;
        result.message() << "Size mismatch of Fiber::_2dArray [("
                         << leftArrays.extent(0) << ", "
                         << leftArrays.extent(1) << ") != ("
                         << rightArrays.extent(0) << ", "
                         << rightArrays.extent(1) << "]";
        return result;
    }

    for (size_t ra = 0; ra < leftArrays.extent(0); ++ra)
        for (size_t ca = 0; ca < leftArrays.extent(1); ++ca) {
            const arma::Mat<ValueType>& left = leftArrays(ra, ca);
            const arma::Mat<ValueType>& right = rightArrays(ra, ca);
            if (left.n_rows != right.n_rows || left.n_cols != right.n_cols) {
                result = false;
                result.message() << "Size mismatch of matrix ("
                                 << ra << ", " << ca << ") [("
                                 << left.n_rows << ", " << left.n_cols << ") != ("
                                 << right.n_rows << ", " << right.n_cols << ")]";
                return result;
            }
            for (size_t r = 0; r < left.n_rows; ++r)
                for (size_t c = 0; c < left.n_cols; ++c) {
                    RealType diff = std::abs(left(r, c) - right(r, c));
                    RealType avg = std::abs(left(r, c) + right(r, c)) / 2.;
                    if (diff > tolerance * (1 + avg)) {
                        result = false;
                        result.message() << std::setprecision(digits10)
                                         << "\n  Mismatch in matrix ("
                                         << ra << ", " << ca << ") at position ("
                                         << r << ", " << c << ") ["
                                         << left(r, c) << " != " << right(r, c) << "]";
                    }
                }
        }

    return result;
}

template <typename ValueType>
boost::test_tools::predicate_result
check_arrays_are_close(const std::vector<arma::Mat<ValueType> >& leftArrays,
                       const std::vector<arma::Mat<ValueType> >& rightArrays,
                       typename Fiber::ScalarTraits<ValueType>::RealType tolerance)
{
    typedef typename Fiber::ScalarTraits<ValueType>::RealType RealType;
    const int digits10 = std::numeric_limits<RealType>::digits10;

    boost::test_tools::predicate_result result(true);
    if (leftArrays.size() != rightArrays.size()) {
        result = false;
        result.message() << "Size mismatch of std::vector ["
                         << leftArrays.size() << " != "
                         << rightArrays.size() << "]";
        return result;
    }

    for (size_t ra = 0; ra < leftArrays.size(); ++ra) {
        const arma::Mat<ValueType>& left = leftArrays[ra];
        const arma::Mat<ValueType>& right = rightArrays[ra];
        if (left.n_rows != right.n_rows || left.n_cols != right.n_cols) {
            result = false;
            result.message() << "Size mismatch of matrix ("
                             << ra << ") [("
                             << left.n_rows << ", " << left.n_cols << ") != ("
                             << right.n_rows << ", " << right.n_cols << ")]";
            return result;
        }
        for (size_t r = 0; r < left.n_rows; ++r)
            for (size_t c = 0; c < left.n_cols; ++c) {
                RealType diff = std::abs(left(r, c) - right(r, c));
                RealType avg = std::abs(left(r, c) + right(r, c)) / 2.;
                if (diff > tolerance * (1 + avg)) {
                    result = false;
                    result.message() << std::setprecision(digits10)
                                     << "\n  Mismatch in matrix ("
                                     << ra << ") at position ("
                                     << r << ", " << c << ") ["
                                     << left(r, c) << " != " << right(r, c) << "]";
                }
            }
    }

    return result;
}

#endif