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

 * 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/base/debug.h>
#include <gtsam/base/timing.h>
#include <gtsam/base/Matrix.h>
#include <gtsam/base/FastMap.h>
#include <gtsam/base/cholesky.h>
#include <gtsam/linear/linearExceptions.h>
#include <gtsam/linear/GaussianConditional.h>
#include <gtsam/linear/GaussianFactor.h>
#include <gtsam/linear/HessianFactor.h>
#include <gtsam/linear/JacobianFactor.h>
#include <gtsam/linear/GaussianFactorGraph.h>

#include <boost/foreach.hpp>
#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 {

/* ************************************************************************* */
string SlotEntry::toString() const {
  ostringstream oss;
  oss << "SlotEntry: slot=" << slot << ", dim=" << dimension;
  return oss.str();
}

/* ************************************************************************* */
Scatter::Scatter(const GaussianFactorGraph& gfg,
    boost::optional<const Ordering&> ordering) {
  gttic(Scatter_Constructor);
  static const size_t none = std::numeric_limits<size_t>::max();

  // First do the set union.
  BOOST_FOREACH(const GaussianFactor::shared_ptr& factor, gfg) {
    if (factor) {
      for (GaussianFactor::const_iterator variable = factor->begin();
          variable != factor->end(); ++variable) {
        // TODO: Fix this hack to cope with zero-row Jacobians that come from BayesTreeOrphanWrappers
        const JacobianFactor* asJacobian =
            dynamic_cast<const JacobianFactor*>(factor.get());
        if (!asJacobian || asJacobian->cols() > 1)
          this->insert(
              make_pair(*variable, SlotEntry(none, factor->getDim(variable))));
      }
    }
  }

  // If we have an ordering, pre-fill the ordered variables first
  size_t slot = 0;
  if (ordering) {
    BOOST_FOREACH(Key key, *ordering) {
      const_iterator entry = find(key);
      if (entry == end())
        throw std::invalid_argument(
            "The ordering provided to the HessianFactor Scatter constructor\n"
                "contained extra variables that did not appear in the factors to combine.");
      at(key).slot = (slot++);
    }
  }

  // Next fill in the slot indices (we can only get these after doing the set
  // union.
  BOOST_FOREACH(value_type& var_slot, *this) {
    if (var_slot.second.slot == none)
      var_slot.second.slot = (slot++);
  }
}

/* ************************************************************************* */
HessianFactor::HessianFactor() :
                          info_(cref_list_of<1>(1))
{
  linearTerm().setZero();
  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_(0,0) = G;
  info_(0,1) = g;
  info_(1,1)(0,0) = 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_(0,0) = Sigma.inverse(); // G
  info_(0,1) = info_(0,0).selfadjointView() * mu; // g
  info_(1,1)(0,0) = mu.dot(info_(0,1).knownOffDiagonal().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_(0,0) = G11;
  info_(0,1) = G12;
  info_(0,2) = g1;
  info_(1,1) = G22;
  info_(1,2) = g2;
  info_(2,2)(0,0) = 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_(0,0) = G11;
  info_(0,1) = G12;
  info_(0,2) = G13;
  info_(0,3) = g1;
  info_(1,1) = G22;
  info_(1,2) = G23;
  info_(1,3) = g2;
  info_(2,2) = G33;
  info_(2,3) = g3;
  info_(3,3)(0,0) = 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){
      info_(i, j) = Gs[index++];
    }
    info_(i, variable_count) = gs[i];
  }
  info_(variable_count, variable_count)(0,0) = f;
}

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

/* ************************************************************************* */
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);
  boost::optional<Scatter> computedScatter;
  if(!scatter) {
    computedScatter = Scatter(factors);
    scatter = computedScatter;
  }

  // Allocate and copy keys
  gttic(allocate);
  // Allocate with dimensions for each variable plus 1 at the end for the information vector
  keys_.resize(scatter->size());
  FastVector<DenseIndex> dims(scatter->size() + 1);
  BOOST_FOREACH(const Scatter::value_type& key_slotentry, *scatter) {
    keys_[key_slotentry.second.slot] = key_slotentry.first;
    dims[key_slotentry.second.slot] = key_slotentry.second.dimension;
  }
  dims.back() = 1;
  info_ = SymmetricBlockMatrix(dims);
  info_.full().triangularView().setZero();
  gttoc(allocate);

  // Form A' * A
  gttic(update);
  BOOST_FOREACH(const GaussianFactor::shared_ptr& factor, factors)
  {
    if(factor) {
      if(const HessianFactor* hessian = dynamic_cast<const HessianFactor*>(factor.get()))
        updateATA(*hessian, *scatter);
      else if(const JacobianFactor* jacobian = dynamic_cast<const JacobianFactor*>(factor.get()))
        updateATA(*jacobian, *scatter);
      else
        throw invalid_argument("GaussianFactor is neither Hessian nor Jacobian");
    }
  }
  gttoc(update);
}

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

/* ************************************************************************* */
bool HessianFactor::equals(const GaussianFactor& lf, double tol) const {
  if(!dynamic_cast<const HessianFactor*>(&lf))
    return false;
  else {
    if(!Factor::equals(lf, tol))
      return false;
    Matrix thisMatrix = this->info_.full().selfadjointView();
    thisMatrix(thisMatrix.rows()-1, thisMatrix.cols()-1) = 0.0;
    Matrix rhsMatrix = static_cast<const HessianFactor&>(lf).info_.full().selfadjointView();
    rhsMatrix(rhsMatrix.rows()-1, rhsMatrix.cols()-1) = 0.0;
    return equal_with_abs_tol(thisMatrix, rhsMatrix, tol);
  }
}

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

/* ************************************************************************* */
Matrix HessianFactor::information() const
{
  return info_.range(0, this->size(), 0, this->size()).selfadjointView();
}

/* ************************************************************************* */
VectorValues HessianFactor::hessianDiagonal() const {
  VectorValues d;
  // Loop over all variables
  for (DenseIndex j = 0; j < (DenseIndex)size(); ++j) {
    // Get the diagonal block, and insert its diagonal
    Matrix B = info_(j, j).selfadjointView();
    d.insert(keys_[j],B.diagonal());
  }
  return d;
}

/* ************************************************************************* */
// TODO: currently assumes all variables of the same size 9 and keys arranged from 0 to n
void HessianFactor::hessianDiagonal(double* d) const {

  // Use eigen magic to access raw memory
  typedef Eigen::Matrix<double, 9, 1> DVector;
  typedef Eigen::Map<DVector> DMap;

  // Loop over all variables in the factor
  for (DenseIndex pos = 0; pos < (DenseIndex)size(); ++pos) {
    Key j = keys_[pos];
    // Get the diagonal block, and insert its diagonal
    const Matrix& B = info_(pos, pos).selfadjointView();
    DMap(d + 9 * j) += B.diagonal();
  }
}

/* ************************************************************************* */
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
    Matrix B = info_(j, j).selfadjointView();
    blocks.insert(make_pair(keys_[j],B));
  }
  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();
  double xtg = 0, xGx = 0;
  // extract the relevant subset of the VectorValues
  // NOTE may not be as efficient
  const Vector x = c.vector(this->keys());
  xtg = x.dot(linearTerm());
  xGx = x.transpose() * info_.range(0, this->size(), 0, this->size()).selfadjointView() *  x;
  return 0.5 * (f - 2.0 * xtg +  xGx);
}

/* ************************************************************************* */
void HessianFactor::updateATA(const HessianFactor& update, const Scatter& scatter)
{
  gttic(updateATA);
  // This function updates 'combined' with the information in 'update'. 'scatter' maps variables in
  // the update factor to slots in the combined factor.

  // First build an array of slots
  gttic(slots);
  FastVector<DenseIndex> slots(update.size());
  DenseIndex slot = 0;
  BOOST_FOREACH(Key j, update) {
    slots[slot] = scatter.at(j).slot;
    ++ slot;
  }
  gttoc(slots);

  // Apply updates to the upper triangle
  gttic(update);
  size_t nrInfoBlocks = this->info_.nBlocks();
  for(DenseIndex j2=0; j2<update.info_.nBlocks(); ++j2) {
    DenseIndex slot2 = (j2 == (DenseIndex)update.size()) ? nrInfoBlocks-1 : slots[j2];
    for(DenseIndex j1=0; j1<=j2; ++j1) {
      DenseIndex slot1 = (j1 == (DenseIndex)update.size()) ? nrInfoBlocks-1 : slots[j1];
      info_(slot1, slot2) += update.info_(j1, j2);
    }
  }
  gttoc(update);
}

/* ************************************************************************* */
void HessianFactor::updateATA(const JacobianFactor& update, const Scatter& scatter) {

  // This function updates 'combined' with the information in 'update'.
  // 'scatter' maps variables in the update factor to slots in the combined
  // factor.

  gttic(updateATA);

  if(update.rows() > 0)
  {
    // First build an array of slots
    gttic(slots);
    FastVector<DenseIndex> slots(update.size());
    DenseIndex slot = 0;
    BOOST_FOREACH(Key j, update) {
      slots[slot] = scatter.at(j).slot;
      ++ slot;
    }
    gttoc(slots);

    gttic(whiten);
    // Whiten the factor if it has a noise model
    boost::optional<JacobianFactor> _whitenedFactor;
    const JacobianFactor* whitenedFactor;
    if(update.get_model())
    {
      if(update.get_model()->isConstrained())
        throw invalid_argument("Cannot update HessianFactor from JacobianFactor with constrained noise model");

      _whitenedFactor = update.whiten();
      whitenedFactor = &(*_whitenedFactor);
    }
    else
    {
      whitenedFactor = &update;
    }
    gttoc(whiten);

    const VerticalBlockMatrix& updateBlocks = whitenedFactor->matrixObject();

    // Apply updates to the upper triangle
    gttic(update);
    DenseIndex nrInfoBlocks = this->info_.nBlocks(), nrUpdateBlocks = updateBlocks.nBlocks();
    for(DenseIndex j2 = 0; j2 < nrUpdateBlocks; ++j2) { // Horizontal block of Hessian
      DenseIndex slot2 = (j2 == (DenseIndex)update.size()) ? nrInfoBlocks-1 : slots[j2];
      assert(slot2 >= 0 && slot2 <= nrInfoBlocks);
      for(DenseIndex j1 = 0; j1 <= j2; ++j1) { // Vertical block of Hessian
        DenseIndex slot1 = (j1 == (DenseIndex)update.size()) ? nrInfoBlocks-1 : slots[j1];
        assert(slot1 >= 0 && slot1 < nrInfoBlocks);
        if(slot1 != slot2)
          info_(slot1, slot2).knownOffDiagonal() += updateBlocks(j1).transpose() * updateBlocks(j2);
        else
          info_(slot1, slot2).selfadjointView().rankUpdate(updateBlocks(j1).transpose());
      }
    }
    gttoc(update);
  }
}

/* ************************************************************************* */
GaussianFactor::shared_ptr HessianFactor::negate() const
{
  shared_ptr result = boost::make_shared<This>(*this);
  result->info_.full().triangularView() = -result->info_.full().triangularView().nestedExpression(); // Negate the information matrix of the result
  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(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_(i, j).knownOffDiagonal() * xj;
    // blocks on the diagonal are only half
    y[i] += info_(j, j).selfadjointView() * xj;
    // for below diagonal, we take transpose block from upper triangular part
    for (i = j + 1; i < (DenseIndex)size(); ++i)
      y[i] += info_(i, j).knownOffDiagonal() * 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
  }
}

/* ************************************************************************* */
void HessianFactor::multiplyHessianAdd(double alpha, const double* x,
    double* yvalues, vector<size_t> offsets) const {

  // Use eigen magic to access raw memory
  typedef Eigen::Matrix<double, Eigen::Dynamic, 1> DVector;
  typedef Eigen::Map<DVector> DMap;
  typedef Eigen::Map<const DVector> ConstDMap;

  // 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(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) {
    DenseIndex i = 0;
    for (; i < j; ++i)
      y[i] += info_(i, j).knownOffDiagonal()
          * ConstDMap(x + offsets[keys_[j]],
              offsets[keys_[j] + 1] - offsets[keys_[j]]);
    // blocks on the diagonal are only half
    y[i] += info_(j, j).selfadjointView()
        * ConstDMap(x + offsets[keys_[j]],
            offsets[keys_[j] + 1] - offsets[keys_[j]]);
    // for below diagonal, we take transpose block from upper triangular part
    for (i = j + 1; i < (DenseIndex) size(); ++i)
      y[i] += info_(i, j).knownOffDiagonal()
          * ConstDMap(x + offsets[keys_[j]],
              offsets[keys_[j] + 1] - offsets[keys_[j]]);
  }

  // copy to yvalues
  for (DenseIndex i = 0; i < (DenseIndex) size(); ++i)
    DMap(yvalues + offsets[keys_[i]], offsets[keys_[i] + 1] - offsets[keys_[i]]) +=
        alpha * y[i];
}

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

/* ************************************************************************* */
// TODO: currently assumes all variables of the same size 9 and keys arranged from 0 to n
void HessianFactor::gradientAtZero(double* d) const {

  // Use eigen magic to access raw memory
  typedef Eigen::Matrix<double, 9, 1> DVector;
  typedef Eigen::Map<DVector> DMap;

  // Loop over all variables in the factor
    for (DenseIndex pos = 0; pos < (DenseIndex)size(); ++pos) {
      Key j = keys_[pos];
      // Get the diagonal block, and insert its diagonal
      DVector dj =  -info_(pos,size()).knownOffDiagonal();
      DMap(d + 9 * j) += dj;
    }
}

/* ************************************************************************* */
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 {
    jointFactor = boost::make_shared<HessianFactor>(factors, Scatter(factors, keys));
  } 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
  GaussianConditional::shared_ptr conditional;
  try {
    VerticalBlockMatrix Ab = jointFactor->info_.choleskyPartial(keys.size());
    conditional = boost::make_shared<GaussianConditional>(jointFactor->keys(), keys.size(), Ab);
    // Erase the eliminated keys in the remaining factor
    jointFactor->keys_.erase(jointFactor->begin(), jointFactor->begin() + keys.size());
  } catch(CholeskyFailed&) {
    throw IndeterminantLinearSystemException(keys.front());
  }

  // 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


// Copyright 2020 Josh Pieper, jjp@pobox.com.
//
// 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 "base/bezier.h"

#include <Eigen/Core>

#include <boost/test/auto_unit_test.hpp>

template <typename T>
using DUT = mjmech::base::Bezier<T>;

BOOST_AUTO_TEST_CASE(BezierTest) {
  {
    DUT<double> dut(0.0, 1.0);
    BOOST_TEST(dut.position(0.0) == 0.0);
    BOOST_TEST(dut.velocity(0.0) == 0.0);
    BOOST_TEST(dut.acceleration(0.0) == 6.0);

    BOOST_TEST(dut.position(0.5) == 0.5);
    BOOST_TEST(dut.velocity(0.5) == 1.5);
    BOOST_TEST(dut.acceleration(0.5) == 0.0);

    BOOST_TEST(dut.position(1.0) == 1.0);
    BOOST_TEST(dut.velocity(1.0) == 0.0);
    BOOST_TEST(dut.acceleration(1.0) == -6.0);
  }

  {
    DUT<double> dut(-2, 8);
    BOOST_TEST(dut.position(0.0) == -2);
    BOOST_TEST(dut.velocity(0.0) == 0.0);
    BOOST_TEST(dut.acceleration(0.0) == 60.0);

    BOOST_TEST(dut.position(0.5) == 3.0);
    BOOST_TEST(dut.velocity(0.5) == 15);
    BOOST_TEST(dut.acceleration(0.5) == 0.0);

    BOOST_TEST(dut.position(1.0) == 8.0);
    BOOST_TEST(dut.velocity(1.0) == 0.0);
    BOOST_TEST(dut.acceleration(1.0) == -60.0);
  }
}

BOOST_AUTO_TEST_CASE(BezierEigen) {
  DUT<Eigen::Vector3d> dut(Eigen::Vector3d(0, 5, 10),
                           Eigen::Vector3d(1, 10, 20));
  BOOST_TEST(dut.position(0.0) == Eigen::Vector3d(0, 5, 10));
  BOOST_TEST(dut.velocity(0.0) == Eigen::Vector3d(0, 0, 0));
  BOOST_TEST(dut.acceleration(0.0) == Eigen::Vector3d(6, 30, 60));

  BOOST_TEST(dut.position(0.5) == Eigen::Vector3d(0.5, 7.5, 15));
  BOOST_TEST(dut.velocity(0.5) == Eigen::Vector3d(1.5, 7.5, 15));
  BOOST_TEST(dut.acceleration(0.5) == Eigen::Vector3d(0., 0., 0.));

  BOOST_TEST(dut.position(1.0) == Eigen::Vector3d(1, 10, 20));
  BOOST_TEST(dut.velocity(1.0) == Eigen::Vector3d(0, 0, 0));
  BOOST_TEST(dut.acceleration(1.0) == Eigen::Vector3d(-6, -30, -60));
}


#ifndef LIDAR_UNDISTORTION_HPP_
#define LIDAR_UNDISTORTION_HPP_

#include <pcl/point_types.h>
#include <pcl/io/pcd_io.h>
#include <pcl/common/common.h>
#include <pcl/common/eigen.h>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <iostream>

class LidarUndistortion
{
public:
  LidarUndistortion() {}

  // Ref:LeGO-LOAM(BSD-3 LICENSE)
  // https://github.com/RobustFieldAutonomyLab/LeGO-LOAM/blob/master/LeGO-LOAM/src/featureAssociation.cpp#L431-L459
  void getImu(
    Eigen::Vector3f angular_velo, Eigen::Vector3f acc, const Eigen::Quaternionf quat,
    const double imu_time /*[sec]*/)
  {
    float roll, pitch, yaw;
    Eigen::Affine3f affine(quat);
    pcl::getEulerAngles(affine, roll, pitch, yaw);

    imu_ptr_last_ = (imu_ptr_last_ + 1) % imu_que_length_;

    if ((imu_ptr_last_ + 1) % imu_que_length_ == imu_ptr_front_) {
      imu_ptr_front_ = (imu_ptr_front_ + 1) % imu_que_length_;
    }

    imu_time_[imu_ptr_last_] = imu_time;
    imu_roll_[imu_ptr_last_] = roll;
    imu_pitch_[imu_ptr_last_] = pitch;
    imu_yaw_[imu_ptr_last_] = yaw;
    imu_acc_x_[imu_ptr_last_] = acc.x();
    imu_acc_y_[imu_ptr_last_] = acc.y();
    imu_acc_z_[imu_ptr_last_] = acc.z();
    imu_angular_velo_x_[imu_ptr_last_] = angular_velo.x();
    imu_angular_velo_y_[imu_ptr_last_] = angular_velo.y();
    imu_angular_velo_z_[imu_ptr_last_] = angular_velo.z();

    Eigen::Matrix3f rot = quat.toRotationMatrix();
    acc = rot * acc;
    // angular_velo = rot * angular_velo;

    int imu_ptr_back = (imu_ptr_last_ - 1 + imu_que_length_) % imu_que_length_;
    double time_diff = imu_time_[imu_ptr_last_] - imu_time_[imu_ptr_back];
    if (time_diff < scan_period_) {
      imu_shift_x_[imu_ptr_last_] =
        imu_shift_x_[imu_ptr_back] + imu_velo_x_[imu_ptr_back] * time_diff + acc(0) * time_diff *
        time_diff * 0.5;
      imu_shift_y_[imu_ptr_last_] =
        imu_shift_y_[imu_ptr_back] + imu_velo_y_[imu_ptr_back] * time_diff + acc(1) * time_diff *
        time_diff * 0.5;
      imu_shift_z_[imu_ptr_last_] =
        imu_shift_z_[imu_ptr_back] + imu_velo_z_[imu_ptr_back] * time_diff + acc(2) * time_diff *
        time_diff * 0.5;

      imu_velo_x_[imu_ptr_last_] = imu_velo_x_[imu_ptr_back] + acc(0) * time_diff;
      imu_velo_y_[imu_ptr_last_] = imu_velo_y_[imu_ptr_back] + acc(1) * time_diff;
      imu_velo_z_[imu_ptr_last_] = imu_velo_z_[imu_ptr_back] + acc(2) * time_diff;

      imu_angular_rot_x_[imu_ptr_last_] = imu_angular_rot_x_[imu_ptr_back] + angular_velo(0) *
        time_diff;
      imu_angular_rot_y_[imu_ptr_last_] = imu_angular_rot_y_[imu_ptr_back] + angular_velo(1) *
        time_diff;
      imu_angular_rot_z_[imu_ptr_last_] = imu_angular_rot_z_[imu_ptr_back] + angular_velo(2) *
        time_diff;
    }
  }

  // Ref:LeGO-LOAM(BSD-3 LICENSE)
  // https://github.com/RobustFieldAutonomyLab/LeGO-LOAM/blob/master/LeGO-LOAM/src/featureAssociation.cpp#L491-L619
  void adjustDistortion(
    pcl::PointCloud<pcl::PointXYZI>::Ptr & cloud,
    const double scan_time /*[sec]*/)
  {
    bool half_passed = false;
    int cloud_size = cloud->points.size();

    float start_ori = -std::atan2(cloud->points[0].y, cloud->points[0].x);
    float end_ori = -std::atan2(cloud->points[cloud_size - 1].y, cloud->points[cloud_size - 1].x);
    if (end_ori - start_ori > 3 * M_PI) {
      end_ori -= 2 * M_PI;
    } else if (end_ori - start_ori < M_PI) {
      end_ori += 2 * M_PI;
    }
    float ori_diff = end_ori - start_ori;

    Eigen::Vector3f rpy_start, shift_start, velo_start, rpy_cur, shift_cur, velo_cur;
    Eigen::Vector3f shift_from_start;
    Eigen::Matrix3f r_s_i, r_c;
    Eigen::Vector3f adjusted_p;
    float ori_h;
    for (int i = 0; i < cloud_size; ++i) {
      pcl::PointXYZI & p = cloud->points[i];
      ori_h = -std::atan2(p.y, p.x);
      if (!half_passed) {
        if (ori_h < start_ori - M_PI * 0.5) {
          ori_h += 2 * M_PI;
        } else if (ori_h > start_ori + M_PI * 1.5) {
          ori_h -= 2 * M_PI;
        }

        if (ori_h - start_ori > M_PI) {
          half_passed = true;
        }
      } else {
        ori_h += 2 * M_PI;
        if (ori_h < end_ori - 1.5 * M_PI) {
          ori_h += 2 * M_PI;
        } else if (ori_h > end_ori + 0.5 * M_PI) {
          ori_h -= 2 * M_PI;
        }
      }

      float rel_time = (ori_h - start_ori) / ori_diff * scan_period_;

      if (imu_ptr_last_ > 0) {
        imu_ptr_front_ = imu_ptr_last_iter_;
        while (imu_ptr_front_ != imu_ptr_last_) {
          if (scan_time + rel_time > imu_time_[imu_ptr_front_]) {
            break;
          }
          imu_ptr_front_ = (imu_ptr_front_ + 1) % imu_que_length_;
        }

        if (scan_time + rel_time > imu_time_[imu_ptr_front_]) {
          rpy_cur(0) = imu_roll_[imu_ptr_front_];
          rpy_cur(1) = imu_pitch_[imu_ptr_front_];
          rpy_cur(2) = imu_yaw_[imu_ptr_front_];
          shift_cur(0) = imu_shift_x_[imu_ptr_front_];
          shift_cur(1) = imu_shift_y_[imu_ptr_front_];
          shift_cur(2) = imu_shift_z_[imu_ptr_front_];
          velo_cur(0) = imu_velo_x_[imu_ptr_front_];
          velo_cur(1) = imu_velo_y_[imu_ptr_front_];
          velo_cur(2) = imu_velo_z_[imu_ptr_front_];
        } else {
          int imu_ptr_back = (imu_ptr_front_ - 1 + imu_que_length_) % imu_que_length_;
          float ratio_front = (scan_time + rel_time - imu_time_[imu_ptr_back]) /
            (imu_time_[imu_ptr_front_] - imu_time_[imu_ptr_back]);
          float ratio_back = 1.0 - ratio_front;
          rpy_cur(0) = imu_roll_[imu_ptr_front_] * ratio_front + imu_roll_[imu_ptr_back] *
            ratio_back;
          rpy_cur(1) = imu_pitch_[imu_ptr_front_] * ratio_front + imu_pitch_[imu_ptr_back] *
            ratio_back;
          rpy_cur(2) = imu_yaw_[imu_ptr_front_] * ratio_front + imu_yaw_[imu_ptr_back] * ratio_back;
          shift_cur(0) = imu_shift_x_[imu_ptr_front_] * ratio_front + imu_shift_x_[imu_ptr_back] *
            ratio_back;
          shift_cur(1) = imu_shift_y_[imu_ptr_front_] * ratio_front + imu_shift_y_[imu_ptr_back] *
            ratio_back;
          shift_cur(2) = imu_shift_z_[imu_ptr_front_] * ratio_front + imu_shift_z_[imu_ptr_back] *
            ratio_back;
          velo_cur(0) = imu_velo_x_[imu_ptr_front_] * ratio_front + imu_velo_x_[imu_ptr_back] *
            ratio_back;
          velo_cur(1) = imu_velo_y_[imu_ptr_front_] * ratio_front + imu_velo_y_[imu_ptr_back] *
            ratio_back;
          velo_cur(2) = imu_velo_z_[imu_ptr_front_] * ratio_front + imu_velo_z_[imu_ptr_back] *
            ratio_back;
        }

        r_c = (
          Eigen::AngleAxisf(rpy_cur(2), Eigen::Vector3f::UnitZ()) *
          Eigen::AngleAxisf(rpy_cur(1), Eigen::Vector3f::UnitY()) *
          Eigen::AngleAxisf(rpy_cur(0), Eigen::Vector3f::UnitX())
          ).toRotationMatrix();

        if (i == 0) {
          rpy_start = rpy_cur;
          shift_start = shift_cur;
          velo_start = velo_cur;
          r_s_i = r_c.inverse();
        } else {
          shift_from_start = shift_cur - shift_start - velo_start * rel_time;
          adjusted_p = r_s_i * (r_c * Eigen::Vector3f(p.x, p.y, p.z) + shift_from_start);
          p.x = adjusted_p.x();
          p.y = adjusted_p.y();
          p.z = adjusted_p.z();
        }
      }
      imu_ptr_last_iter_ = imu_ptr_front_;
    }
  }

  void setScanPeriod(const double scan_period /*[sec]*/)
  {
    scan_period_ = scan_period;
  }

private:
  double scan_period_{0.1};
  static const int imu_que_length_{200};
  int imu_ptr_front_{0}, imu_ptr_last_{-1}, imu_ptr_last_iter_{0};

  std::array<double, imu_que_length_> imu_time_;
  std::array<float, imu_que_length_> imu_roll_;
  std::array<float, imu_que_length_> imu_pitch_;
  std::array<float, imu_que_length_> imu_yaw_;

  std::array<float, imu_que_length_> imu_acc_x_;
  std::array<float, imu_que_length_> imu_acc_y_;
  std::array<float, imu_que_length_> imu_acc_z_;
  std::array<float, imu_que_length_> imu_velo_x_;
  std::array<float, imu_que_length_> imu_velo_y_;
  std::array<float, imu_que_length_> imu_velo_z_;
  std::array<float, imu_que_length_> imu_shift_x_;
  std::array<float, imu_que_length_> imu_shift_y_;
  std::array<float, imu_que_length_> imu_shift_z_;

  std::array<float, imu_que_length_> imu_angular_velo_x_;
  std::array<float, imu_que_length_> imu_angular_velo_y_;
  std::array<float, imu_que_length_> imu_angular_velo_z_;
  std::array<float, imu_que_length_> imu_angular_rot_x_;
  std::array<float, imu_que_length_> imu_angular_rot_y_;
  std::array<float, imu_que_length_> imu_angular_rot_z_;
};

#endif  // LIDAR_UNDISTORTION_HPP_


/* ----------------------------------------------------------------------- *//**
 *
 * @file Ordinal_impl.hpp
 *
 *//* ----------------------------------------------------------------------- */

#ifndef MADLIB_MODULES_GLM_ORDINAL_IMPL_HPP
#define MADLIB_MODULES_GLM_ORDINAL_IMPL_HPP

#include <dbconnector/dbconnector.hpp>
#include <boost/math/distributions.hpp>
#include <limits>

namespace madlib {

namespace modules {

namespace glm {

// Use Eigen
using namespace dbal::eigen_integration;

template <class Container, class Family, class Link>
inline
OrdinalAccumulator<Container,Family,Link>::OrdinalAccumulator(
        Init_type& inInitialization)
  : Base(inInitialization), optimizer(inInitialization) {

    this->initialize();
}

/**
 * @brief Reset the accumulator before accumulating the first tuple
 */
template <class Container, class Family, class Link>
inline
void
OrdinalAccumulator<Container,Family,Link>::reset() {
    num_rows = 0;
    terminated = false;
    loglik = 0.;
    optimizer.reset();
}

/**
 * @brief Bind all elements of the state to the data in the stream
 *
 * The bind() is special in that even after running operator>>() on an element,
 * there is no guarantee yet that the element can indeed be accessed. It is
 * cruicial to first check this.
 *
 * Provided that this methods correctly lists all member variables, all other
 * methods can, however, rely on that fact that all variables are correctly
 * initialized and accessible.
 */
template <class Container, class Family, class Link>
inline
void
OrdinalAccumulator<Container,Family,Link>::bind(
        ByteStream_type& inStream) {

    inStream
        >> num_features
        >> num_categories
        >> num_rows
        >> terminated
        >> loglik
        >> optimizer;

    // bind vcov and optimizer.hessian onto the same memory as we don't need them
    // at the same time
    vcov.rebind(optimizer.hessian.memoryHandle(),
                optimizer.hessian.rows(),
                optimizer.hessian.cols());
}

/**
 * @brief Update the accumulation state by feeding a tuple
 */
template <class Container, class Family, class Link>
inline
OrdinalAccumulator<Container,Family,Link>&
OrdinalAccumulator<Container,Family,Link>::operator<<(
        const tuple_type& inTuple) {

    const MappedColumnVector& x = std::get<0>(inTuple);
    const int& y = static_cast<int>(std::get<1>(inTuple));
    // convert y to an indicator vector
    ColumnVector vecY(num_categories-1);
    vecY.fill(0);
    if (y != (num_categories-1)) { vecY(y) = 1; }

    // The following checks were introduced with MADLIB-138. It still seems
    // useful to have clear error messages in case of infinite input values.

    if (!dbal::eigen_integration::isfinite(x)) {
        warning("Design matrix is not finite.");
    } else if (x.size() > std::numeric_limits<uint16_t>::max()) {
        warning("Number of independent variables cannot be "
            "larger than 65535.");
    } else if (num_features != static_cast<uint16_t>(x.size())) {
        warning("Inconsistent numbers of independent variables.");
    } else {
        // normal case
        uint16_t N = num_features;
        uint16_t C = static_cast<uint16_t>(num_categories - 1);

        uint32_t state_size = 6 + (N + C)*(N + C) + 2 * (N + C);
        // GPDB limits the single array size to be 1GB, which means that the size
        // of a double array cannot be large than 134217727 because
        // (134217727 * 8) / (1024 * 1024) = 1023. And solve
        // state_size = x^2 + 2^x + 6 <= 134217727 will give x <= 11584.
        if(state_size > 134217727)
            throw std::domain_error(
                "The sum of number of independent variables and number of "
                 "categories cannot be larger than 11584.");
        ColumnVector mu(C);
        ColumnVector ita(C);
        Matrix G_prime(C, C);
        Matrix V(C, C);

        if (optimizer.beta.norm() == 0.) {
            // beta is 0 in the first iteration
            Link::init(mu);
            Link::link_func(mu,ita);
        } else {
            //beta has non-zero value
            double Xtbeta = trans(x)*optimizer.beta.segment(C,N);
            for (int i=0; i<ita.size(); i++)
                ita(i) = optimizer.beta(i) - Xtbeta;
            Link::mean_func(ita, mu);
       }

        Link::mean_derivative(ita, G_prime);
        Family::variance(mu, V);
        loglik += Family::loglik(vecY, mu, 1);
        Matrix Vinv(V.inverse());

        // calcualte X_trans_W_X
        Matrix GtVinvG(trans(G_prime)*Vinv*G_prime);
        ColumnVector h_vec(C);
        h_vec.fill(0);
        for (int i=0;i<C;i++)
            for (int j=0;j<C;j++)
                 h_vec(i) += G_prime(i,j);

        Matrix GtVinvhXt(trans(G_prime)*Vinv*h_vec*trans(x));
        GtVinvhXt = -1*GtVinvhXt;
        Matrix XhtVinvG(x*trans(h_vec)*Vinv*G_prime);
        XhtVinvG = -1*XhtVinvG;
        Matrix htVinvhXXt(x*trans(x));
        double temp  = trans(h_vec)*Vinv*h_vec;
        htVinvhXXt = temp * htVinvhXXt;

        for (int i=0;i<C;i++)
            for (int j=0;j<C;j++)
                optimizer.hessian(i,j) += GtVinvG(i,j);
        for (int i=0;i<C;i++)
            for (int j=C;j<(C+N);j++)
                optimizer.hessian(i,j) += GtVinvhXt(i,j-C);

        for (int i=C;i<(C+N);i++)
            for (int j=0;j<C;j++)
                optimizer.hessian(i,j) += XhtVinvG(i-C,j);

        for (int i=C;i<(C+N);i++)
            for (int j=C;j<(C+N);j++)
                optimizer.hessian(i,j) += htVinvhXXt(i-C,j-C);

        // calculate X_trans_W_Y
        ColumnVector YVinvG(trans(vecY - mu)*Vinv*G_prime);
        ColumnVector YVinvhX(x);
        temp = trans(vecY - mu)*Vinv*h_vec;
        YVinvhX = -1*temp * YVinvhX;

        optimizer.grad.segment(0, C) += YVinvG;
        optimizer.grad.segment(C, N) += YVinvhX;

        num_rows ++;
        return *this;
    } // all tests passed and we were in the else

    // error case
    terminated = true;
    return *this;
}

/**
 * @brief Merge with another accumulation state
 */
template <class Container, class Family, class Link>
template <class C, class F, class L>
inline
OrdinalAccumulator<Container,Family,Link>&
OrdinalAccumulator<Container,Family,Link>::operator<<(
        const OrdinalAccumulator<C,F,L>& inOther) {
    if (this->empty()) {
        *this = inOther;
    } else if (inOther.empty()) {
    } else if (num_features != inOther.num_features) {
        warning("Inconsistent numbers of independent variables.");
        terminated = true;
    } else {
        num_rows += inOther.num_rows;
        loglik += inOther.loglik;
        optimizer.grad += inOther.optimizer.grad;
        optimizer.hessian += inOther.optimizer.hessian;
    }

    return *this;
}

/**
 * @brief Copy from a previous state
 */
template <class Container, class Family, class Link>
template <class C, class F, class L>
inline
OrdinalAccumulator<Container,Family,Link>&
OrdinalAccumulator<Container,Family,Link>::operator=(
        const OrdinalAccumulator<C,F,L>& inOther) {
    this->copy(inOther);
    return *this;
}

/**
 * @brief Apply the accumulated intra-state values to inter-state member
 */
template <class Container, class Family, class Link>
inline
void
OrdinalAccumulator<Container,Family,Link>::apply() {
    if (!dbal::eigen_integration::isfinite(optimizer.hessian) ||
            !dbal::eigen_integration::isfinite(optimizer.grad)) {
        warning("Ordinal:Hessian or gradient is not finite. One possibility is that intercept is included in the independent variables. If that is the case, please drop the intercept and rerun the function.");
        terminated = true;
    } else {
                vcov = optimizer.hessian.inverse();
                optimizer.beta += vcov * optimizer.grad;
    }

}
// -----------------------------------------------------------------------

template <class Container>
OrdinalResult::OrdinalResult(const OrdinalAccumulator<Container>& state) {
    compute(state);
}

/**
 * @brief Transform an accumulation state into a result
 */
template <class Container>
inline
OrdinalResult&
OrdinalResult::compute(const OrdinalAccumulator<Container>& state) {
    // Vector of coefficients: For efficiency reasons, we want to return this
    // by reference, so we need to bind to db memory
    Allocator& allocator = defaultAllocator();
    Index N = static_cast<Index>(state.num_features);
    Index C = static_cast<Index>(state.num_categories - 1);

    coef_alpha.rebind(allocator.allocateArray<double>(C));
    std_err_alpha.rebind(allocator.allocateArray<double>(C));
    z_stats_alpha.rebind(allocator.allocateArray<double>(C));
    p_values_alpha.rebind(allocator.allocateArray<double>(C));

    coef_beta.rebind(allocator.allocateArray<double>(N));
    std_err_beta.rebind(allocator.allocateArray<double>(N));
    z_stats_beta.rebind(allocator.allocateArray<double>(N));
    p_values_beta.rebind(allocator.allocateArray<double>(N));

    // computation
    loglik = state.loglik;
    coef_alpha = state.optimizer.beta.segment(0,C);
    ColumnVector temp(state.vcov.diagonal().cwiseSqrt());
    std_err_alpha = temp.segment(0,C);
    z_stats_alpha = coef_alpha.cwiseQuotient(std_err_alpha);
    for (Index i = 0; i < C; i ++) {
        p_values_alpha(i) = 2. * prob::cdf( prob::normal(), -std::abs(z_stats_alpha(i)));
    }

    coef_beta = state.optimizer.beta.segment(C,N);
    std_err_beta = temp.segment(C,N);
    z_stats_beta = coef_beta.cwiseQuotient(std_err_beta);
    for (Index i = 0; i < N; i ++) {
        p_values_beta(i) = 2. * prob::cdf( prob::normal(), -std::abs(z_stats_beta(i)));
    }
    num_rows_processed = state.num_rows;

    return *this;
}

} // namespace glm

} // namespace modules

} // namespace madlib

#endif // defined(MADLIB_MODULES_GLM_ORDINAL_IMPL_HPP)


#include <mutex>
#include <memory>
#include <iostream>
#include <boost/circular_buffer.hpp>

#include <ros/ros.h>
#include <pcl_ros/point_cloud.h>
#include <tf_conversions/tf_eigen.h>
#include <tf/transform_broadcaster.h>

#include <nav_msgs/Odometry.h>
#include <sensor_msgs/Imu.h>
#include <sensor_msgs/PointCloud2.h>
#include <geometry_msgs/PoseWithCovarianceStamped.h>
#include <geometry_msgs/PoseStamped.h>

#include <nodelet/nodelet.h>
#include <pluginlib/class_list_macros.h>

#include <pcl/filters/voxel_grid.h>

#include <pclomp/ndt_omp.h>

#include <amsi/pose_estimator.hpp>
#include <amsi/gnss_tools.hpp>

#include <nmea_msgs/Sentence.h> // message from DGPS of University of California, Berkeley.
#include <sensor_msgs/NavSatFix.h>

#include <Eigen/Dense>

namespace amsi {

class amsiNodelet : public nodelet::Nodelet {
public:
  using PointT = pcl::PointXYZI;

  amsiNodelet() {
  }
  virtual ~amsiNodelet() {
  }

  void onInit() override {
    nh = getNodeHandle();
    mt_nh = getMTNodeHandle();
    private_nh = getPrivateNodeHandle();

    processing_time.resize(16);
    initialize_params();

    use_imu = private_nh.param<bool>("use_imu", true);
    invert_imu = private_nh.param<bool>("invert_imu", false);
    if(use_imu) {
      NODELET_INFO("enable imu-based prediction");
      imu_sub = mt_nh.subscribe("/imu_raw", 256, &amsiNodelet::imu_callback, this);
    }

    use_dgps = private_nh.param<bool>("use_dgps", true);
    if(use_dgps) {
      NODELET_INFO("enable dgps-based correction");
      nmea_sub = nh.subscribe("/nmea_sentence", 32, &amsiNodelet::nmea_callback, this);
    }

    points_sub = mt_nh.subscribe("/points_raw", 5, &amsiNodelet::points_callback, this);
    globalmap_sub = nh.subscribe("/points_map", 1, &amsiNodelet::globalmap_callback, this);
    initialpose_sub = nh.subscribe("/initialpose", 8, &amsiNodelet::initialpose_callback, this);
    ndtpose_sub = nh.subscribe("/ndt_pose", 8, &amsiNodelet::ndtpose_callback, this);

    pose_pub = nh.advertise<nav_msgs::Odometry>("/gnss_ins", 5, false); // UKF-based GNSS/INS fusion
    enu_pose_pub = nh.advertise<geometry_msgs::PoseStamped>("/gnss_enu", 5, false);
    gnss_navsat_pose_pub = nh.advertise<sensor_msgs::NavSatFix>("/gnss_navsat", 5, false);
    gnss_ins_navsat_pose_pub = nh.advertise<sensor_msgs::NavSatFix>("/gnss_ins_navsat", 5, false);
    aligned_pub = nh.advertise<sensor_msgs::PointCloud2>("/aligned_points", 5, false);
  }

private:
  void initialize_params() {
    // intialize scan matching method
    double downsample_resolution = private_nh.param<double>("downsample_resolution", 0.1);
    std::string ndt_neighbor_search_method = private_nh.param<std::string>("ndt_neighbor_search_method", "DIRECT7");

    double ndt_resolution = private_nh.param<double>("ndt_resolution", 1.0);
    boost::shared_ptr<pcl::VoxelGrid<PointT>> voxelgrid(new pcl::VoxelGrid<PointT>());
    voxelgrid->setLeafSize(downsample_resolution, downsample_resolution, downsample_resolution);
    downsample_filter = voxelgrid;

    pclomp::NormalDistributionsTransform<PointT, PointT>::Ptr ndt(new pclomp::NormalDistributionsTransform<PointT, PointT>());
    ndt->setTransformationEpsilon(0.01);
    ndt->setResolution(ndt_resolution);
    if(ndt_neighbor_search_method == "DIRECT1") {
      NODELET_INFO("search_method DIRECT1 is selected");
      ndt->setNeighborhoodSearchMethod(pclomp::DIRECT1);
    } else if(ndt_neighbor_search_method == "DIRECT7") {
      NODELET_INFO("search_method DIRECT7 is selected");
      ndt->setNeighborhoodSearchMethod(pclomp::DIRECT7);
    } else {
      if(ndt_neighbor_search_method == "KDTREE") {
        NODELET_INFO("search_method KDTREE is selected");
      } else {
        NODELET_WARN("invalid search method was given");
        NODELET_WARN("default method is selected (KDTREE)");
      }
      ndt->setNeighborhoodSearchMethod(pclomp::KDTREE);
    }
    // ndt->setStepSize(0.01);
    // ndt->setMaximumIterations(30);
    registration = ndt;

    // initialize pose estimator
    if(private_nh.param<bool>("specify_init_pose", true)) {
      NODELET_INFO("initialize pose estimator with specified parameters!!");
      pose_estimator.reset(new amsi::PoseEstimator(registration,
        ros::Time::now(),
        Eigen::Vector3f(private_nh.param<double>("init_pos_x", 0.0), private_nh.param<double>("init_pos_y", 0.0), private_nh.param<double>("init_pos_z", 0.0)),
        Eigen::Quaternionf(private_nh.param<double>("init_ori_w", 1.0), private_nh.param<double>("init_ori_x", 0.0), private_nh.param<double>("init_ori_y", 0.0), private_nh.param<double>("init_ori_z", 0.0)),
        private_nh.param<double>("cool_time_duration", 0.5)
      ));
    }
  }

private:
  /**
   * @brief callback for imu data
   * @param imu_msg
   */
  void imu_callback(const sensor_msgs::ImuConstPtr& imu_msg) {
    std::lock_guard<std::mutex> lock(imu_data_mutex);
    imu_data.push_back(imu_msg);
  }

  /**
   * @brief callback for DGPS in nmea format
   * @param DGPS message
   */
  void nmea_callback(const nmea_msgs::SentenceConstPtr& nmea_msg) {
    std::lock_guard<std::mutex> lock(dgps_data_mutex);
    std::vector<std::string> str_vec_ptr;
    std::string token;
    std::stringstream ss(nmea_msg->sentence);
    bool find_SOL_COMPUTED =0;
    while (getline(ss, token, ' '))
    {
      if(token == "SOL_COMPUTED") // solutions are computed
      {
        // std::cout<<"message obtained"<<std::endl;
        find_SOL_COMPUTED = true;
      }
      if( find_SOL_COMPUTED ) // find flag SOL_COMPUTED
      {
        str_vec_ptr.push_back(token);
      }
    }

    if(find_SOL_COMPUTED)
    {
      sensor_msgs::NavSatFix navfix_ ;
      navfix_.header = nmea_msg->header;
      std::cout << std::setprecision(17);
      double lat = strtod((str_vec_ptr[2]).c_str(), NULL);
      double lon = strtod((str_vec_ptr[3]).c_str(), NULL);
      double alt = strtod((str_vec_ptr[4]).c_str(), NULL);
      std::cout << std::setprecision(17);

      navfix_.latitude = lat;
      navfix_.longitude = lon;
      navfix_.altitude = alt;
      gnss_navsat_pose_pub.publish(navfix_); // gnss standalone
      if(ini_navf.latitude == NULL)
      {
        ini_navf = navfix_;
        std::cout<<"ini_navf.header  -> "<<ini_navf.header<<std::endl;
        originllh.resize(3, 1);
        originllh(0) = navfix_.longitude;
        originllh(1) = navfix_.latitude;
        originllh(2) = navfix_.altitude;
      }
      dgps_data.push_back(navfix_);
      Eigen::MatrixXd curLLh; //
      curLLh.resize(3, 1);
      curLLh(0) = navfix_.longitude;
      curLLh(1) = navfix_.latitude;
      curLLh(2) = navfix_.altitude;

      Eigen::MatrixXd ecef; //
      ecef.resize(3, 1);
      ecef = gnss_tools_.llh2ecef(curLLh);
      ENU_.header = navfix_.header;
      Eigen::MatrixXd eigenENU;; //
      eigenENU.resize(3, 1);
      eigenENU = gnss_tools_.ecef2enu(originllh,ecef);
      ENU_.pose.position.x = eigenENU(0);
      ENU_.pose.position.y = eigenENU(1);
      ENU_.pose.position.z = eigenENU(2);
      ENU_.pose.orientation.x = 0.23;
      ENU_.pose.orientation.y = 0.25;
      ENU_.pose.orientation.z = 0.12;
      ENU_.pose.orientation.w = 0.5;
      enu_pose_pub.publish(ENU_); // publish dgps in enu
      dgpsENU_data.push_back(ENU_);

      std::cout << std::setprecision(3);
      auto cur_tim = ros::WallTime::now();

      ecef = gnss_tools_.enu2ecef(originllh,eigenENU);
      curLLh = gnss_tools_.ecef2llh(ecef);
      navfix_.longitude = curLLh(0);
      navfix_.latitude  = curLLh(1);
      navfix_.altitude  = curLLh(2);
      gnss_ins_navsat_pose_pub.publish(navfix_); // ukf-based

      std::cout<<"  E: "<<float(eigenENU(0))<<"  N: "<<float(eigenENU(1))<<"  U: "<<float(eigenENU(2))<<std::endl;
    }
  }

  /**
   * @brief callback for point cloud data
   * @param points_msg
   */
  void points_callback(const sensor_msgs::PointCloud2ConstPtr& points_msg) {
    std::lock_guard<std::mutex> estimator_lock(pose_estimator_mutex);
    if(!pose_estimator) {
      NODELET_ERROR("waiting for initial pose input!!");
      return;
    }

    // if(!globalmap) {
    //   NODELET_ERROR("globalmap has not been received!!");
    //   return;
    // }

    const auto& stamp = points_msg->header.stamp;
    pcl::PointCloud<PointT>::Ptr cloud(new pcl::PointCloud<PointT>());
    pcl::fromROSMsg(*points_msg, *cloud);

    if(cloud->empty()) {
      NODELET_ERROR("cloud is empty!!");
      return;
    }

    auto filtered = downsample(cloud);

    // predict
    if(!use_imu) {
      pose_estimator->predict(stamp, Eigen::Vector3f::Zero(), Eigen::Vector3f::Zero()); // no imu, constant vel model
      std::cout<<"const velocity prediction"<<std::endl;
    } else {
      std::lock_guard<std::mutex> lock(imu_data_mutex);
      auto imu_iter = imu_data.begin();
      for(imu_iter; imu_iter != imu_data.end(); imu_iter++) {
        if(stamp < (*imu_iter)->header.stamp) {
          break;
        }
        const auto& acc = (*imu_iter)->linear_acceleration;
        const auto& gyro = (*imu_iter)->angular_velocity;
        double gyro_sign = invert_imu ? -1.0 : 1.0;
        pose_estimator->predict((*imu_iter)->header.stamp, Eigen::Vector3f(acc.x, acc.y, acc.z), gyro_sign * Eigen::Vector3f(gyro.x, gyro.y, gyro.z));
      }
      imu_data.erase(imu_data.begin(), imu_iter);
    }

    // correct
    auto t1 = ros::WallTime::now();
    auto aligned = pose_estimator->correct(filtered,ENU_);
    auto t2 = ros::WallTime::now();

    processing_time.push_back((t2 - t1).toSec());
    double avg_processing_time = std::accumulate(processing_time.begin(), processing_time.end(), 0.0) / processing_time.size();
    // NODELET_INFO_STREAM("processing_time: " << avg_processing_time * 1000.0 << "[msec]");

    if(aligned_pub.getNumSubscribers()) {
      aligned->header.frame_id = "map";
      aligned->header.stamp = cloud->header.stamp;
      aligned_pub.publish(aligned);
    }

    publish_odometry(points_msg->header.stamp, pose_estimator->matrix());
  }

  /**
   * @brief callback for globalmap input
   * @param points_msg
   */
  void globalmap_callback(const sensor_msgs::PointCloud2ConstPtr& points_msg) {
    NODELET_INFO("globalmap received!");
    pcl::PointCloud<PointT>::Ptr cloud(new pcl::PointCloud<PointT>());
    pcl::fromROSMsg(*points_msg, *cloud);
    globalmap = cloud;

    registration->setInputTarget(globalmap);
  }

  /**
   * @brief callback for initial pose input ("2D Pose Estimate" on rviz)
   * @param pose_msg
   */
  void initialpose_callback(const geometry_msgs::PoseWithCovarianceStampedConstPtr& pose_msg) {
    NODELET_INFO("initial pose received!!");
    std::lock_guard<std::mutex> lock(pose_estimator_mutex);
    const auto& p = pose_msg->pose.pose.position;
    const auto& q = pose_msg->pose.pose.orientation;
    pose_estimator.reset(
          new amsi::PoseEstimator(
            registration,
            ros::Time::now(),
            Eigen::Vector3f(p.x, p.y, p.z),
            Eigen::Quaternionf(q.w, q.x, q.y, q.z),
            private_nh.param<double>("cool_time_duration", 0.5))
    );
  }

  /**
   * @brief callback for ndt pose input ("3D Pose Estimate" from Autoware)
   * @param pose_msg
   */
  void ndtpose_callback(const geometry_msgs::PoseStampedConstPtr& pose_msg) {
    NODELET_INFO("ndt pose received!!");
    ndt_pose = *pose_msg;
  }

  /**
   * @brief downsampling
   * @param cloud   input cloud
   * @return downsampled cloud
   */
  pcl::PointCloud<PointT>::ConstPtr downsample(const pcl::PointCloud<PointT>::ConstPtr& cloud) const {
    if(!downsample_filter) {
      return cloud;
    }

    pcl::PointCloud<PointT>::Ptr filtered(new pcl::PointCloud<PointT>());
    downsample_filter->setInputCloud(cloud);
    downsample_filter->filter(*filtered);
    filtered->header = cloud->header;

    return filtered;
  }

  /**
   * @brief publish odometry
   * @param stamp  timestamp
   * @param pose   odometry pose to be published
   */
  void publish_odometry(const ros::Time& stamp, const Eigen::Matrix4f& pose) {
    // broadcast the transform over tf
    geometry_msgs::TransformStamped odom_trans = matrix2transform(stamp, pose, "map", "velodyne");
    // pose_broadcaster.sendTransform(odom_trans);

    // publish the transform
    nav_msgs::Odometry odom;
    odom.header.stamp = stamp;
    odom.header.frame_id = "map";

    odom.pose.pose.position.x = pose(0, 3);
    odom.pose.pose.position.y = pose(1, 3);
    odom.pose.pose.position.z = pose(2, 3);
    odom.pose.pose.orientation = odom_trans.transform.rotation;

    odom.child_frame_id = "velodyne";
    odom.twist.twist.linear.x = 0.0;
    odom.twist.twist.linear.y = 0.0;
    odom.twist.twist.angular.z = 0.0;

    pose_pub.publish(odom);
  }

  /**
   * @brief convert a Eigen::Matrix to TransformedStamped
   * @param stamp           timestamp
   * @param pose            pose matrix
   * @param frame_id        frame_id
   * @param child_frame_id  child_frame_id
   * @return transform
   */
  geometry_msgs::TransformStamped matrix2transform(const ros::Time& stamp, const Eigen::Matrix4f& pose, const std::string& frame_id, const std::string& child_frame_id) {
    Eigen::Quaternionf quat(pose.block<3, 3>(0, 0));
    quat.normalize();
    geometry_msgs::Quaternion odom_quat;
    odom_quat.w = quat.w();
    odom_quat.x = quat.x();
    odom_quat.y = quat.y();
    odom_quat.z = quat.z();

    geometry_msgs::TransformStamped odom_trans;
    odom_trans.header.stamp = stamp;
    odom_trans.header.frame_id = frame_id;
    odom_trans.child_frame_id = child_frame_id;

    odom_trans.transform.translation.x = pose(0, 3);
    odom_trans.transform.translation.y = pose(1, 3);
    odom_trans.transform.translation.z = pose(2, 3);
    odom_trans.transform.rotation = odom_quat;

    return odom_trans;
  }

private:
  // ROS
  ros::NodeHandle nh;
  ros::NodeHandle mt_nh;
  ros::NodeHandle private_nh;

  bool use_imu;
  bool invert_imu;
  bool use_dgps;

  ros::Subscriber imu_sub;
  ros::Subscriber points_sub;
  ros::Subscriber globalmap_sub;
  ros::Subscriber initialpose_sub;
  ros::Subscriber ndtpose_sub; // subscribe /ndt_pose from Autoware
  ros::Subscriber nmea_sub; // nmea sentence from DGPS

  ros::Publisher pose_pub;
  ros::Publisher enu_pose_pub; // enu publisher
  ros::Publisher gnss_navsat_pose_pub; // navsat
  ros::Publisher gnss_ins_navsat_pose_pub; // navsat
  ros::Publisher aligned_pub;
  tf::TransformBroadcaster pose_broadcaster;

  // imu input buffer
  std::mutex imu_data_mutex;
  std::mutex dgps_data_mutex;
  std::vector<sensor_msgs::ImuConstPtr> imu_data;
  std::vector<sensor_msgs::NavSatFix> dgps_data;
  geometry_msgs::PoseStamped ENU_; // dgps pose in ENU
  std::vector<geometry_msgs::PoseStamped> dgpsENU_data;

  sensor_msgs::NavSatFix ini_navf ; // initial sensor msg
  Eigen::MatrixXd originllh; // origin llh

  // ndt pose from Autoware
  geometry_msgs::PoseStamped ndt_pose;

  // globalmap and registration method
  pcl::PointCloud<PointT>::Ptr globalmap;
  pcl::Filter<PointT>::Ptr downsample_filter;
  pcl::Registration<PointT, PointT>::Ptr registration;

  // pose estimator
  std::mutex pose_estimator_mutex;
  std::unique_ptr<amsi::PoseEstimator> pose_estimator;

  // gnss_tools
  GNSS_Tools gnss_tools_;

  // processing time buffer
  boost::circular_buffer<double> processing_time;
};

}

PLUGINLIB_EXPORT_CLASS(amsi::amsiNodelet, nodelet::Nodelet)


/*!
@file
Defines `boost::hana::detail::tuple_cartesian_product`.

@copyright Louis Dionne 2015
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_HANA_DETAIL_TUPLE_CARTESIAN_PRODUCT_HPP
#define BOOST_HANA_DETAIL_TUPLE_CARTESIAN_PRODUCT_HPP

#include <boost/hana/detail/array.hpp>
#include <boost/hana/detail/generate_integer_sequence.hpp>
#include <boost/hana/detail/std/forward.hpp>
#include <boost/hana/detail/std/integer_sequence.hpp>
#include <boost/hana/detail/std/size_t.hpp>
#include <boost/hana/foldable.hpp>
#include <boost/hana/tuple.hpp>

namespace boost { namespace hana { namespace detail {
    struct _tuple_cartesian_product {
        using Size = detail::std::size_t;

        template <unsigned Which, Size ...>
        struct indexN;

        template <unsigned Which, Size I, Size J>
        struct indexN<Which, I, J> {
            constexpr auto operator()(...) const {
                detail::array<detail::array<Size, I * J>, 2> arrays{};
                for (Size i = 0, index = 0; i < I; ++i) {
                    for (Size j = 0; j < J; ++j, ++index) {
                        arrays[0][index] = i;
                        arrays[1][index] = j;
                    }
                }
                return arrays[Which];
            }
        };

        template <unsigned Which, Size I, Size J, Size K>
        struct indexN<Which, I, J, K> {
            constexpr auto operator()(...) const {
                detail::array<detail::array<Size, I * J * K>, 3> arrays{};
                for (Size i = 0, index = 0; i < I; ++i) {
                    for (Size j = 0; j < J; ++j) {
                        for (Size k = 0; k < K; ++k, ++index) {
                            arrays[0][index] = i;
                            arrays[1][index] = j;
                            arrays[2][index] = k;
                        }
                    }
                }
                return arrays[Which];
            }
        };

        template <typename ...T, typename ...U,
                  Size ...i,     Size ...j>
        static constexpr auto
        cartesian_prod_impl(_tuple<T...> const& a,
                            _tuple<U...> const& b,
                            detail::std::index_sequence<i...>,
                            detail::std::index_sequence<j...>)
        {
            return hana::make_tuple(hana::make_tuple(
                detail::get<i>(a), detail::get<j>(b)
            )...);
        }

        template <typename ...T, typename ...U, typename ...V,
                 Size ...i,      Size ...j,     Size ...k>
        static constexpr auto
        cartesian_prod_impl(_tuple<T...> const& a,
                            _tuple<U...> const& b,
                            _tuple<V...> const& c,
                            detail::std::index_sequence<i...>,
                            detail::std::index_sequence<j...>,
                            detail::std::index_sequence<k...>)
        {
            return hana::make_tuple(hana::make_tuple(
                detail::get<i>(a), detail::get<j>(b), detail::get<k>(c)
            )...);
        }

        template <typename ...T, typename ...U,
                  Size ...i,     Size ...j>
        static constexpr auto
        cartesian_prod_impl_cat2(_tuple<T...> const& a,
                            _tuple<U...> const& b,
                            detail::std::index_sequence<i...>,
                            detail::std::index_sequence<j...>)
        {
            return hana::make_tuple(hana::concat(
                detail::get<i>(a), detail::get<j>(b)
            )...);
        }

        template <typename ...Tuples>
        static constexpr decltype(auto)
        cartesian_prod_helper_cat2(Tuples&& ...tuples) {
            constexpr long long lengths[] = {tuple_detail::size<Tuples>{}...};
            constexpr Size total_length = hana::product(lengths);
            return cartesian_prod_impl_cat2(
                detail::std::forward<Tuples>(tuples)...,
                detail::generate_index_sequence<total_length,
                    indexN<0, tuple_detail::size<Tuples>{}...>
                >{},
                detail::generate_index_sequence<total_length,
                    indexN<1, tuple_detail::size<Tuples>{}...>
                >{}
            );
        }

        template <Size ...i, typename ...Tuples>
        static constexpr decltype(auto)
        cartesian_prod_helper(detail::std::index_sequence<i...>, Tuples&& ...tuples) {
            constexpr long long lengths[] = {tuple_detail::size<Tuples>{}...};
            constexpr Size total_length = hana::product(lengths);
            return cartesian_prod_impl(
                detail::std::forward<Tuples>(tuples)...,
                detail::generate_index_sequence<total_length,
                    indexN<i, tuple_detail::size<Tuples>{}...>
                >{}...
            );
        }

        template <typename Tuple>
        constexpr decltype(auto) operator()(Tuple&& tuple) const
        { return hana::transform(detail::std::forward<Tuple>(tuple), hana::make_tuple); }

        template <typename ...Tuples>
        constexpr decltype(auto) operator()(Tuples&& ...tuples) const {
            static_assert(sizeof...(Tuples) == 2 || sizeof...(Tuples) == 3,
            "this overload should only pick up 2 and 3 tuples");

            return cartesian_prod_helper(
                detail::std::make_index_sequence<sizeof...(Tuples)>{},
                detail::std::forward<Tuples>(tuples)...
            );
        }

        template <typename A, typename B, typename C, typename D1, typename ...Dn>
        constexpr decltype(auto)
        operator()(A&& a, B&& b, C&& c, D1&& d1, Dn&& ...dn) const {
            return hana::transform( // Flatten to get A x B x C x D1 x ... x Dn
                cartesian_prod_helper_cat2( // <- Compute (A x B x C) x (D1 x ... x Dn)
                    (*this)( // <- Compute A x B x C
                        detail::std::forward<A>(a),
                        detail::std::forward<B>(b),
                        detail::std::forward<C>(c)
                    ),
                    (*this)( // <- Compute D1 x ... x Dn
                        detail::std::forward<D1>(d1),
                        detail::std::forward<Dn>(dn)...
                    )
                )
            , id);
        }
    };
    constexpr _tuple_cartesian_product tuple_cartesian_product{};
}}} // end namespace boost::hana::detail

#endif // !BOOST_HANA_DETAIL_TUPLE_CARTESIAN_PRODUCT_HPP


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef NT2_ELLIPTIC_FUNCTIONS_GENERIC_AM_HPP_INCLUDED
#define NT2_ELLIPTIC_FUNCTIONS_GENERIC_AM_HPP_INCLUDED

#include <nt2/elliptic/functions/am.hpp>
#include <boost/dispatch/attributes.hpp>
#include <nt2/include/constants/eps.hpp>
#include <nt2/include/functions/simd/abs.hpp>
#include <nt2/include/functions/simd/sin.hpp>
#include <nt2/include/functions/simd/sqrt.hpp>
#include <nt2/sdk/meta/scalar_of.hpp>

namespace nt2 { namespace ext
{
  BOOST_DISPATCH_IMPLEMENT  (am_, tag::cpu_,
                             (A0)(A1)(A2)(A3),
                             (generic_<floating_<A0> >)
                             (generic_<floating_<A1> >)
                             (scalar_<floating_<A2> >)
                             (scalar_<integer_<A3> >)
                            )
  {
    typedef A0 result_type;

    BOOST_FORCEINLINE result_type operator()(const A0 & u, const A1 & x,
                                             const A2 & tol, const A3 & choice) const
    {
      switch (choice)
      {
      case 'a': return am(u,nt2::sin(x), tol);
      case 'm': return am(u,nt2::sqrt(nt2::abs(x)), tol);
      default : return am(u,x,tol);
      }
    }
  };

  BOOST_DISPATCH_IMPLEMENT  (am_, tag::cpu_,
                             (A0)(A1)(A2),
                             (generic_<floating_<A0> >)
                             (generic_<floating_<A1> >)
                             (scalar_<integer_<A2> >)
                            )
  {
    typedef A0 result_type;

    BOOST_FORCEINLINE result_type operator()(const A0 & u, const A1 & x,
                                             const A2 & choice) const
    {
      typedef typename meta::scalar_of<result_type>::type sA0;
      return am(u,x,Eps<sA0>(),choice);
    }
  };

  BOOST_DISPATCH_IMPLEMENT  (am_, tag::cpu_,
                              (A0)(A1),
                              (generic_<floating_<A0> >)
                              (generic_<floating_<A1> >)
                             )
  {
    typedef A0 result_type;
    BOOST_FORCEINLINE result_type operator()(const A0 & u, const A1 & x) const
    {
      typedef typename meta::scalar_of<result_type>::type sA0;
      return am(u,x,Eps<sA0>());
    }
  };
} }
#endif


/*
 * Marginalizer.cpp
 *
 *  Created on: May 24, 2013
 *      Author: mbuerki
 */

#include "aslam/backend/Marginalizer.hpp"

#include <aslam/backend/DenseQrLinearSystemSolver.hpp>
#include <Eigen/QR>
#include <Eigen/Dense>
#include <aslam/backend/DenseMatrix.hpp>

#include <iostream>

#include <sm/logging.hpp>
#include <sm/timing/Timer.hpp>

namespace aslam {
namespace backend {

void marginalize(
			std::vector<aslam::backend::DesignVariable*>& inDesignVariables,
			std::vector<aslam::backend::ErrorTerm*>& inErrorTerms,
			int numberOfInputDesignVariablesToRemove,
			bool useMEstimator,
			boost::shared_ptr<aslam::backend::MarginalizationPriorErrorTerm>& outPriorErrorTermPtr,
			Eigen::MatrixXd& outCov,
			std::vector<aslam::backend::DesignVariable*>& outDesignVariablesInRTop,
			size_t numTopRowsInCov,
			size_t /* numThreads */)
{
		  SM_WARN_STREAM_COND(inDesignVariables.size() == 0, "Zero input design variables in the marginalizer!");

		  // check for duplicates!
		  std::unordered_set<aslam::backend::DesignVariable*> inDvSetHT;
		  for(auto it = inDesignVariables.begin(); it != inDesignVariables.end(); ++it)
		  {
			  auto ret = inDvSetHT.insert(*it);
			  SM_ASSERT_TRUE(aslam::Exception, ret.second, "Error! Duplicate design variables in input list!");
		  }
		  std::unordered_set<aslam::backend::ErrorTerm*> inEtSetHT;
		  for(auto it = inErrorTerms.begin(); it != inErrorTerms.end(); ++it)
		  {
			  auto ret = inEtSetHT.insert(*it);
			  SM_ASSERT_TRUE(aslam::Exception, ret.second, "Error! Duplicate error term in input list!");
		  }
		  SM_DEBUG_STREAM("NO duplicates in input design variables or input error terms found.");

          // Partition the design varibles into removed/remaining.
		  int dimOfDesignVariablesToRemove = 0;
		  std::vector<aslam::backend::DesignVariable*> remainingDesignVariables;
		  int k = 0;
		  size_t dimOfDvsInTopBlock = 0;
		  for(std::vector<aslam::backend::DesignVariable*>::const_iterator it = inDesignVariables.begin(); it != inDesignVariables.end(); ++it)
		  {
			  if (k < numberOfInputDesignVariablesToRemove)
			  {
				  dimOfDesignVariablesToRemove += (*it)->minimalDimensions();
			  } else
			  {
				  remainingDesignVariables.push_back(*it);
			  }

			  if(dimOfDvsInTopBlock < numTopRowsInCov)
			  {
				  outDesignVariablesInRTop.push_back(*it);
			  }
			  dimOfDvsInTopBlock += (*it)->minimalDimensions();
			  k++;
		  }

		  // store original block indices to prevent side effects
		  std::vector<int> originalBlockIndices;
		  std::vector<int> originalColumnBase;
			// assign block indices
			int columnBase = 0;
			for (size_t i = 0; i < inDesignVariables.size(); ++i) {
				originalBlockIndices.push_back(inDesignVariables[i]->blockIndex());
				originalColumnBase.push_back(inDesignVariables[i]->columnBase());

				inDesignVariables[i]->setBlockIndex(i);
				inDesignVariables[i]->setColumnBase(columnBase);
			  columnBase += inDesignVariables[i]->minimalDimensions();
			}

			int dim = 0;
			std::vector<size_t> originalRowBase;
			for(std::vector<aslam::backend::ErrorTerm*>::iterator it = inErrorTerms.begin(); it != inErrorTerms.end(); ++it)
			{
				originalRowBase.push_back((*it)->rowBase());
				(*it)->setRowBase(dim);
				dim += (*it)->dimension();
			}

		  aslam::backend::DenseQrLinearSystemSolver qrSolver;
          qrSolver.initMatrixStructure(inDesignVariables, inErrorTerms, false);

		  SM_INFO_STREAM("Marginalization optimization problem initialized with " << inDesignVariables.size() << " design variables and " << inErrorTerms.size() << " error terrms");
		  SM_INFO_STREAM("The Jacobian matrix is " << dim << " x " << columnBase);

		  qrSolver.evaluateError(1, useMEstimator);
		  qrSolver.buildSystem(1, useMEstimator);

		  const Eigen::MatrixXd& jacobian = qrSolver.getJacobian();
		  const Eigen::VectorXd& b = qrSolver.e();

		  // check dimension of jacobian
		  int jrows = jacobian.rows();
		  int jcols = jacobian.cols();

		  int dimOfRemainingDesignVariables = jcols - dimOfDesignVariablesToRemove;
		  //int dimOfPriorErrorTerm = jrows;

		  // check the rank
		  Eigen::FullPivLU<Eigen::MatrixXd> lu_decomp(jacobian);
		  //lu_decomp.setThreshold(1e-20);
		  double threshold = lu_decomp.threshold();
		  int rank = lu_decomp.rank();
		  int fullRank = std::min(jacobian.rows(), jacobian.cols());
		  SM_DEBUG_STREAM("Rank of jacobian: " << rank << " (full rank: " << fullRank << ", threshold: " << threshold << ")");
		  bool rankDeficient = rank < fullRank;
		  if(rankDeficient)
		  {
			  SM_WARN("Marginalization jacobian is rank deficient!");
		  }
		  //SM_ASSERT_FALSE(aslam::Exception, rankDeficient, "Right now, we don't want the jacobian to be rank deficient - ever...");

		  Eigen::MatrixXd R_reduced;
		  Eigen::VectorXd d_reduced;
		  if (jrows < jcols)
		  {
			  SM_THROW(aslam::Exception, "underdetermined LSE!");
			  // underdetermined LSE, don't do QR
			  R_reduced = jacobian.block(0, dimOfDesignVariablesToRemove, jrows, jcols - dimOfDesignVariablesToRemove);
			  d_reduced = b;

		  } else {
			  // PTF: Do we know what will happen when the jacobian matrix is rank deficient?
			  // MB: yes, bad things!

              // do QR decomposition
			  sm::timing::Timer myTimer("QR Decomposition");
              Eigen::HouseholderQR<Eigen::MatrixXd> qr(jacobian);
			  Eigen::MatrixXd Q = qr.householderQ();
			  Eigen::MatrixXd R = qr.matrixQR().triangularView<Eigen::Upper>();
			  Eigen::VectorXd d = Q.transpose()*b;
			  myTimer.stop();

			  if(numTopRowsInCov > 0)
			  {
				sm::timing::Timer myTimer("Covariance computation");
				Eigen::FullPivLU<Eigen::MatrixXd> lu_decomp(R);
				Eigen::MatrixXd Rinv = lu_decomp.inverse();
				Eigen::MatrixXd covariance = Rinv * Rinv.transpose();
				outCov = covariance.block(0, 0, numTopRowsInCov, numTopRowsInCov);
				myTimer.stop();
			  }

//			  size_t numRowsToKeep = rank - dimOfDesignVariablesToRemove;
//			  SM_ASSERT_TRUE_DBG(aslam::Exception, rankDeficient || (numRowsToKeep == dimOfRemainingDesignVariables), "must be the same if full rank!");

			  // get the top left block
			  SM_ASSERT_GE(aslam::Exception, R.rows(), numTopRowsInCov, "Cannot extract " << numTopRowsInCov << " rows of R because it only has " << R.rows() << " rows.");
			  SM_ASSERT_GE(aslam::Exception, R.cols(), numTopRowsInCov, "Cannot extract " << numTopRowsInCov << " cols of R because it only has " << R.cols() << " cols.");
			  //outRtop = R.block(0, 0, numTopRowsInRtop, numTopRowsInRtop);

              // cut off the zero rows at the bottom
              R_reduced = R.block(dimOfDesignVariablesToRemove, dimOfDesignVariablesToRemove, dimOfRemainingDesignVariables, dimOfRemainingDesignVariables);
			  //R_reduced = R.block(dimOfDesignVariablesToRemove, dimOfDesignVariablesToRemove, numRowsToKeep, dimOfRemainingDesignVariables);

              d_reduced = d.segment(dimOfDesignVariablesToRemove, dimOfRemainingDesignVariables);
              //d_reduced = d.segment(dimOfDesignVariablesToRemove, numRowsToKeep);
              //dimOfPriorErrorTerm = dimOfRemainingDesignVariables;
		  }

		  // now create the new error term
		  boost::shared_ptr<aslam::backend::MarginalizationPriorErrorTerm> err(new aslam::backend::MarginalizationPriorErrorTerm(remainingDesignVariables, d_reduced, R_reduced));

		  outPriorErrorTermPtr.swap(err);

		  // restore initial block indices to prevent side effects
          for (size_t i = 0; i < inDesignVariables.size(); ++i) {
              inDesignVariables[i]->setBlockIndex(originalBlockIndices[i]);
              inDesignVariables[i]->setColumnBase(originalColumnBase[i]);
          }
          int index = 0;
          for(std::vector<aslam::backend::ErrorTerm*>::iterator it = inErrorTerms.begin(); it != inErrorTerms.end(); ++it)
          {
              (*it)->setRowBase(originalRowBase[index++]);
          }
}

} /* namespace backend */
} /* namespace aslam */


#ifndef CORE_H
#define CORE_H

#include <Eigen/Core>
#include <memory>

namespace sequential_line_search
{
    class PreferentialBayesianOptimizer;
} // namespace sequential_line_search

namespace mathtoolbox
{
    class DataNormalizer;
}

class Core
{
public:
    Core();

    void reset();

    static Core& getInstance()
    {
        static Core core;
        return core;
    }

    std::shared_ptr<sequential_line_search::PreferentialBayesianOptimizer> m_optimizer;
    std::shared_ptr<mathtoolbox::DataNormalizer>                           m_normalizer;

    Eigen::VectorXd m_y;

    double evaluateObjectiveFunction(const Eigen::VectorXd& x) const;

    // For optimization
    void            proceedOptimization();
    Eigen::VectorXd m_x_max;
    double          m_y_max;
};

#endif // CORE_H


/*!
 * @file     mppi_methods.cpp
 * @author   Giuseppe Rizzi
 * @date     11.06.2020
 * @version  1.0
 * @brief    description
 */

#include <gtest/gtest.h>
#include <Eigen/Dense>
#include <array>
#include <chrono>
#include "double_integrator.cpp"
#include "mppi/controller/mppi.h"
#include "mppi/sampler/gaussian_sampler.h"
#include "mppi/utils/logging.h"

using namespace Eigen;
using namespace mppi;

class PathIntegralTest : public ::testing::Test {
 protected:
  void SetUp() override {
    config.caching_factor = 0.2;
    config.rollouts = 4;
    config.lambda = 1;
    config.step_size = 0.1;
    config.horizon = 0.5;
    dynamics = std::make_shared<DoubleIntegratorDynamics>();
    cost = std::make_shared<DoubleIntegratorCost>();
    sampler = std::make_shared<mppi::GaussianSampler>(1);
    solver_ = std::make_unique<PathIntegral>(dynamics, cost, sampler, config);
  }
  SolverConfig config;
  std::unique_ptr<PathIntegral> solver_;
  std::shared_ptr<DoubleIntegratorDynamics> dynamics;
  std::shared_ptr<DoubleIntegratorCost> cost;
  mppi::GaussianSampler::sampler_ptr sampler;

  void print_rollouts() {
    for (size_t k = 0; k < solver_->config_.rollouts; k++) {
      std::cout << "Rollout " << k << ": " << std::endl;
      std::cout << solver_->rollouts_[k] << std::endl;
    }
  }

  void verbose_run(Eigen::VectorXd x0, const double t0) {
    log_info("Resetting observation.");
    solver_->set_observation(x0, t0);

    log_info("Preparing rollouts.");
    solver_->prepare_rollouts();

    log_info("Printing rollouts info.");
    print_rollouts();

    log_info("Optimal rollout");
    std::cout << solver_->get_optimal_rollout();

    log_info("Optimal rollout cached");
    std::cout << solver_->get_optimal_rollout_cache();

    log_info("Sampling trajectories");
    solver_->sample_trajectories();

    log_info("Printing rollouts info.");
    print_rollouts();

    log_info("Optimizing.");
    solver_->optimize();

    log_info("Optimal rollout");
    std::cout << solver_->get_optimal_rollout();

    log_info("Optimal rollout cached");
    std::cout << solver_->get_optimal_rollout_cache();

    log_info("Swapping policies");
    solver_->swap_policies();

    log_info("Optimal rollout");
    std::cout << solver_->get_optimal_rollout();

    log_info("Optimal rollout cached");
    std::cout << solver_->get_optimal_rollout_cache();
  }
};

TEST_F(PathIntegralTest, VerboseRun) {
  Eigen::Vector2d x0{0.0, 0.0};
  double t0 = 0.0;

  std::cout << std::string(100, '=') << std::endl;
  std::cout << " Iteraion 0 " << std::endl;
  std::cout << std::string(100, '=') << std::endl;
  verbose_run(x0, t0);

  t0 = 0.25;
  std::cout << std::string(100, '=') << std::endl;
  std::cout << " Iteraion 1 " << std::endl;
  std::cout << std::string(100, '=') << std::endl;
  verbose_run(x0, t0);

  double query_t = 0.3;
  PathIntegral::input_t u;
  log_info("Querying control input at t=" + std::to_string(query_t));
  solver_->get_input(x0, u, query_t);
  std::cout << u.transpose() << std::endl;
}

TEST_F(PathIntegralTest, InputQueries) {
  Eigen::Vector2d x0{0.0, 0.0};
  Eigen::VectorXd u = Eigen::VectorXd::Zero(1);
  double t0 = 0.0;

  config.horizon = 2.0;
  solver_ = std::make_unique<PathIntegral>(dynamics, cost, sampler, config);

  solver_->set_observation(x0, t0);
  solver_->update_policy();
  std::cout << "Optimal trajectory: " << std::endl
            << solver_->get_optimal_rollout_cache() << std::endl;

  double t;
  std::cout << "Aligned time." << std::endl;
  for (int i = 0; i < 10; i++) {
    t = t0 + i * config.step_size;
    solver_->get_input(x0, u, t);
    std::cout << "u*[" << t << "]: " << u.transpose() << std::endl;
  }

  std::cout << "Unaligned time." << std::endl;
  for (int i = 0; i < 10; i++) {
    t = t0 + 0.001 + i * config.step_size;
    solver_->get_input(x0, u, t);
    std::cout << "u*[" << t << "]: " << u.transpose() << std::endl;
  }
}

TEST_F(PathIntegralTest, AsyncRun) {
  Eigen::Vector2d x0{0.0, 0.0};
  Eigen::VectorXd u = Eigen::VectorXd::Zero(1);

  config.horizon = 2.0;
  solver_ = std::make_unique<PathIntegral>(dynamics, cost, sampler, config);

  double t0 = 0.0;
  std::atomic_bool run_once;
  run_once = false;

  solver_->set_observation(x0, t0);

  // optimization thread
  std::thread t1([&, this] {
    size_t counter = 0;
    while (counter < 5) {
      std::cout << "Update #" << counter << std::endl;
      solver_->update_policy();
      counter++;
      std::this_thread::sleep_for(std::chrono::milliseconds(50));
      std::cout << "Optimal rollout: " << std::endl
                << solver_->get_optimal_rollout_cache() << std::endl;
      run_once = true;
    }
  });

  // set observation and get control thread
  std::thread t2([&, this] {
    size_t counter = 0;
    double t;
    while (counter < 5) {
      if (run_once) {
        std::cout << "Reset #" << counter << std::endl;
        t = t0 + counter * config.step_size;
        solver_->set_observation(x0, t);
        solver_->get_input(x0, u, t);
        counter++;
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
        std::cout << "u*[" << t << "]: " << u.transpose() << std::endl;
      }
    }
  });

  t1.join();
  t2.join();
}

int main(int argc, char **argv) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}


/*
 * DistributionRightHandSide.cc
 *
 *  Created on: 30.06.2017
 *      Author: thies
 */

#include <deal.II/base/exceptions.h>
#include <deal.II/base/types.h>
#include <deal.II/base/work_stream.h>
#include <deal.II/fe/fe.h>
#include <deal.II/fe/fe_update_flags.h>
#include <deal.II/fe/fe_values.h>

#include <base/DiscretizedFunction.h>
#include <forward/DistributionRightHandSide.h>

#include <vector>

namespace wavepi {
namespace forward {
using namespace dealii;
using namespace wavepi::base;

template <int dim>
DistributionRightHandSide<dim>::AssemblyScratchData::AssemblyScratchData(const FiniteElement<dim> &fe,
                                                                         const Quadrature<dim> &quad)
    : fe_values(fe, quad, update_values | update_gradients | update_quadrature_points | update_JxW_values) {}

template <int dim>
DistributionRightHandSide<dim>::AssemblyScratchData::AssemblyScratchData(const AssemblyScratchData &scratch_data)
    : fe_values(scratch_data.fe_values.get_fe(), scratch_data.fe_values.get_quadrature(),
                update_values | update_gradients | update_quadrature_points | update_JxW_values) {}

template <int dim>
DistributionRightHandSide<dim>::DistributionRightHandSide(Function<dim> *f1, Function<dim> *f2) {
  this->f1 = f1;
  this->f2 = f2;
}

template <int dim>
void DistributionRightHandSide<dim>::copy_local_to_global(Vector<double> &result, const AssemblyCopyData &copy_data) {
  for (unsigned int i = 0; i < copy_data.local_dof_indices.size(); ++i)
    result(copy_data.local_dof_indices[i]) += copy_data.cell_rhs(i);
}

template <int dim>
void DistributionRightHandSide<dim>::local_assemble_cc(const Function<dim> *const f1, const Function<dim> *const f2,
                                                       const typename DoFHandler<dim>::active_cell_iterator &cell,
                                                       AssemblyScratchData &scratch_data, AssemblyCopyData &copy_data) {
  const unsigned int dofs_per_cell = scratch_data.fe_values.get_fe().dofs_per_cell;
  const unsigned int n_q_points    = scratch_data.fe_values.get_quadrature().size();

  copy_data.cell_rhs.reinit(dofs_per_cell);
  copy_data.local_dof_indices.resize(dofs_per_cell);
  scratch_data.fe_values.reinit(cell);

  for (unsigned int q_point = 0; q_point < n_q_points; ++q_point) {
    if (f1 != nullptr) {
      const double val1 = f1->value(scratch_data.fe_values.quadrature_point(q_point));

      for (unsigned int i = 0; i < dofs_per_cell; ++i)
        copy_data.cell_rhs(i) +=
            val1 * scratch_data.fe_values.shape_value(i, q_point) * scratch_data.fe_values.JxW(q_point);
    }

    if (f2 != nullptr) {
      auto val2 = f2->gradient(scratch_data.fe_values.quadrature_point(q_point));

      for (unsigned int i = 0; i < dofs_per_cell; ++i)
        copy_data.cell_rhs(i) +=
            val2 * scratch_data.fe_values.shape_grad(i, q_point) * scratch_data.fe_values.JxW(q_point);
    }
  }

  cell->get_dof_indices(copy_data.local_dof_indices);
}

template <int dim>
void DistributionRightHandSide<dim>::local_assemble_dd(const Vector<double> &f1, const Vector<double> &f2,
                                                       const typename DoFHandler<dim>::active_cell_iterator &cell,
                                                       AssemblyScratchData &scratch_data, AssemblyCopyData &copy_data) {
  const unsigned int dofs_per_cell = scratch_data.fe_values.get_fe().dofs_per_cell;
  const unsigned int n_q_points    = scratch_data.fe_values.get_quadrature().size();

  copy_data.cell_rhs.reinit(dofs_per_cell);
  copy_data.local_dof_indices.resize(dofs_per_cell);
  scratch_data.fe_values.reinit(cell);

  cell->get_dof_indices(copy_data.local_dof_indices);

  for (unsigned int q_point = 0; q_point < n_q_points; ++q_point)
    for (unsigned int i = 0; i < dofs_per_cell; ++i)
      for (unsigned int k = 0; k < dofs_per_cell; ++k)
        copy_data.cell_rhs(i) += (f1[copy_data.local_dof_indices[k]] * scratch_data.fe_values.shape_value(k, q_point) *
                                      scratch_data.fe_values.shape_value(i, q_point) +
                                  f2[copy_data.local_dof_indices[k]] * scratch_data.fe_values.shape_grad(k, q_point) *
                                      scratch_data.fe_values.shape_grad(i, q_point)) *
                                 scratch_data.fe_values.JxW(q_point);
}

template <int dim>
void DistributionRightHandSide<dim>::local_assemble_cd(const Function<dim> *const f1, const Vector<double> &f2,
                                                       const typename DoFHandler<dim>::active_cell_iterator &cell,
                                                       AssemblyScratchData &scratch_data, AssemblyCopyData &copy_data) {
  const unsigned int dofs_per_cell = scratch_data.fe_values.get_fe().dofs_per_cell;
  const unsigned int n_q_points    = scratch_data.fe_values.get_quadrature().size();

  copy_data.cell_rhs.reinit(dofs_per_cell);
  copy_data.local_dof_indices.resize(dofs_per_cell);
  scratch_data.fe_values.reinit(cell);

  cell->get_dof_indices(copy_data.local_dof_indices);

  for (unsigned int q_point = 0; q_point < n_q_points; ++q_point) {
    if (f1 != nullptr) {
      const double val1 = f1->value(scratch_data.fe_values.quadrature_point(q_point));

      for (unsigned int i = 0; i < dofs_per_cell; ++i)
        copy_data.cell_rhs(i) +=
            val1 * scratch_data.fe_values.shape_value(i, q_point) * scratch_data.fe_values.JxW(q_point);
    }

    for (unsigned int i = 0; i < dofs_per_cell; ++i) {
      for (unsigned int k = 0; k < dofs_per_cell; ++k)
        copy_data.cell_rhs(i) += f2[copy_data.local_dof_indices[k]] * scratch_data.fe_values.shape_grad(k, q_point) *
                                 scratch_data.fe_values.shape_grad(i, q_point) * scratch_data.fe_values.JxW(q_point);
    }
  }
}

template <int dim>
void DistributionRightHandSide<dim>::local_assemble_dc(const Vector<double> &f1, const Function<dim> *const f2,
                                                       const typename DoFHandler<dim>::active_cell_iterator &cell,
                                                       AssemblyScratchData &scratch_data, AssemblyCopyData &copy_data) {
  const unsigned int dofs_per_cell = scratch_data.fe_values.get_fe().dofs_per_cell;
  const unsigned int n_q_points    = scratch_data.fe_values.get_quadrature().size();

  copy_data.cell_rhs.reinit(dofs_per_cell);
  copy_data.local_dof_indices.resize(dofs_per_cell);
  scratch_data.fe_values.reinit(cell);

  cell->get_dof_indices(copy_data.local_dof_indices);

  for (unsigned int q_point = 0; q_point < n_q_points; ++q_point)
    for (unsigned int i = 0; i < dofs_per_cell; ++i) {
      for (unsigned int k = 0; k < dofs_per_cell; ++k)
        copy_data.cell_rhs(i) += f1[copy_data.local_dof_indices[k]] * scratch_data.fe_values.shape_value(k, q_point) *
                                 scratch_data.fe_values.shape_value(i, q_point) * scratch_data.fe_values.JxW(q_point);

      if (f2 != nullptr) {
        auto val2 = f2->gradient(scratch_data.fe_values.quadrature_point(q_point));

        for (unsigned int i = 0; i < dofs_per_cell; ++i)
          copy_data.cell_rhs(i) +=
              val2 * scratch_data.fe_values.shape_grad(i, q_point) * scratch_data.fe_values.JxW(q_point);
      }
    }
}

template <int dim>
void DistributionRightHandSide<dim>::create_right_hand_side(const DoFHandler<dim> &dof, const Quadrature<dim> &quad,
                                                            Vector<double> &rhs) const {
  f1->set_time(this->get_time());
  f2->set_time(this->get_time());

  auto f1_d = dynamic_cast<DiscretizedFunction<dim> *>(f1);
  auto f2_d = dynamic_cast<DiscretizedFunction<dim> *>(f2);

  if (f1_d != nullptr)
    Assert(f1_d->get_function_coefficients(f1_d->get_time_index()).size() == dof.n_dofs(),
           ExcDimensionMismatch(f1_d->get_function_coefficients(f1_d->get_time_index()).size(), dof.n_dofs()));

  if (f2_d != nullptr)
    Assert(f2_d->get_function_coefficients(f2_d->get_time_index()).size() == dof.n_dofs(),
           ExcDimensionMismatch(f2_d->get_function_coefficients(f2_d->get_time_index()).size(), dof.n_dofs()));

  if (f1_d != nullptr && f2_d != nullptr)
    WorkStream::run(
        dof.begin_active(), dof.end(),
        std::bind(&DistributionRightHandSide<dim>::local_assemble_cc, f1, f2, std::placeholders::_1,
                  std::placeholders::_2, std::placeholders::_3),
        std::bind(&DistributionRightHandSide<dim>::copy_local_to_global, std::ref(rhs), std::placeholders::_1),
        AssemblyScratchData(dof.get_fe(), quad), AssemblyCopyData());
  else if (f1_d == nullptr && f2_d != nullptr)
    WorkStream::run(
        dof.begin_active(), dof.end(),
        std::bind(&DistributionRightHandSide<dim>::local_assemble_cd, f1,
                  std::ref(f2_d->get_function_coefficients(f2_d->get_time_index())), std::placeholders::_1,
                  std::placeholders::_2, std::placeholders::_3),
        std::bind(&DistributionRightHandSide<dim>::copy_local_to_global, std::ref(rhs), std::placeholders::_1),
        AssemblyScratchData(dof.get_fe(), quad), AssemblyCopyData());
  else if (f1_d != nullptr && f2_d == nullptr)
    WorkStream::run(
        dof.begin_active(), dof.end(),
        std::bind(&DistributionRightHandSide<dim>::local_assemble_dc,
                  std::ref(f1_d->get_function_coefficients(f1_d->get_time_index())), f2, std::placeholders::_1,
                  std::placeholders::_2, std::placeholders::_3),
        std::bind(&DistributionRightHandSide<dim>::copy_local_to_global, std::ref(rhs), std::placeholders::_1),
        AssemblyScratchData(dof.get_fe(), quad), AssemblyCopyData());
  else
    WorkStream::run(
        dof.begin_active(), dof.end(),
        std::bind(&DistributionRightHandSide<dim>::local_assemble_cc, f1, f2, std::placeholders::_1,
                  std::placeholders::_2, std::placeholders::_3),
        std::bind(&DistributionRightHandSide<dim>::copy_local_to_global, std::ref(rhs), std::placeholders::_1),
        AssemblyScratchData(dof.get_fe(), quad), AssemblyCopyData());
}

template class DistributionRightHandSide<1>;
template class DistributionRightHandSide<2>;
template class DistributionRightHandSide<3>;

} /* namespace forward */
} /* namespace wavepi */


/*
* TS Elements
* Copyright 2015-2018 M. Newhouse
* Released under the MIT license.
*/

#pragma once

#include "utility/rect.hpp"
#include "utility/vector2.hpp"

#include <boost/container/small_vector.hpp>
#include <boost/variant.hpp>

#include <cstdint>
#include <array>

namespace ts
{
  namespace resources
  {
    namespace collision_shapes
    {
      struct Circle
      {
        Vector2f center;
        float radius;
        std::uint32_t height = 1;
      };

      struct Polygon
      {
        struct Point
        {
          Vector2f position;
        };

        std::array<Point, 8> points;
        std::uint32_t num_points;
        std::uint32_t height = 1;
      };

      struct SubShape
      {
        boost::variant<Circle, Polygon> data;
        float bounciness;
      };
    }

    struct CollisionShape
    {
      boost::container::small_vector<collision_shapes::SubShape, 4> sub_shapes;
    };
  }
}