#include <algorithm>
#include <boost/range/irange.hpp>
#include <cmath>
#include <map>
#include <memory>
#include <set>

#include "bounding_edge.hpp"
#include "kd_tree_node.hpp"
#include "rtc_log.hpp"
#include "scoped_timer.hpp"

namespace rtc
{
kd_tree::kd_tree(const rtc::scene_model& ss) : bbox{ss.points}
{
  SCOPE_TIME_COUNTER;
  const auto& p = ss.points;
  const auto& t = ss.triangles;
  const std::uint32_t max_depth = 1.3F * std::log2(t.size()) + 8;

  auto make_bbox = [](auto& p, auto& tr) noexcept {
    return rtc::bounding_box{p[tr.vertex_a()], p[tr.vertex_b()], p[tr.vertex_c()]};
  };

  rtc::bounding_box node_bbox{p};
  std::vector<std::uint32_t> v(t.size());
  std::vector<bounding_box> primitive_bboxes;
  edge_buffer_array_t edge_buffer{
      edge_buffer_t(2 * t.size()), edge_buffer_t(2 * t.size()), edge_buffer_t(2 * t.size())};

  primitive_bboxes.reserve(t.size());

  std::for_each(t.begin(), t.end(), [&](auto& t) { primitive_bboxes.push_back(make_bbox(p, t)); });
  std::generate(v.begin(), v.end(), [i = 0]() mutable { return i++; });

  build_tree(root, node_bbox, std::move(v), primitive_bboxes, edge_buffer, max_depth);
}

auto kd_tree::cbegin(const rtc::math_ray& ray) const noexcept -> kd_tree::const_iterator
{
  if (auto range = bbox.intersection_values_for(ray))
  {
    const auto [tmin, tmax] = range.value();
    return const_iterator{ray, {root.get(), tmin, tmax}};
  }

  return const_iterator{};
}

auto kd_tree::cend(const rtc::math_ray&) const noexcept -> kd_tree::const_iterator { return const_iterator{}; }

kd_tree::~kd_tree() = default;
kd_tree::kd_tree(kd_tree&&) noexcept = default;
auto kd_tree::operator=(kd_tree&&) noexcept -> kd_tree& = default;

auto kd_tree::compute_node_split_paramters(edge_buffer_array_t& edge_buffer,
                                           const std::vector<std::uint32_t>& tr,
                                           const rtc::bounding_box& node_bbox,
                                           const std::vector<rtc::bounding_box>& primitive_bboxes)
{
  constexpr int isect_cost{80};
  constexpr int traversal_cost{1};
  constexpr rtc_float empty_bonus{0.5F};
  const auto node_diagonal{node_bbox.diagonal()};
  const rtc_float invTotalSA{1.0F / node_bbox.surface_area()};
  const rtc_float old_cost{isect_cost * static_cast<rtc_float>(tr.size())};

  int retries{0};
  int best_axis{-1};
  int best_offset{-1};
  rtc::axis axis{node_bbox.maximum_extent()};
  rtc_float best_cost{std::numeric_limits<rtc_float>::max()};

  do
  {
    const rtc::axis axis2{next(axis)}, axis3{next(axis2)};
    int below{}, above(tr.size());

    for (const auto i : boost::irange<std::size_t>(0, tr.size()))
    {
      const rtc::bounding_box& box{primitive_bboxes[tr[i]]};
      edge_buffer[int(axis)][2 * i] = {box.min_boundary().axis(axis), tr[i], bounding_edge_point::start};
      edge_buffer[int(axis)][2 * i + 1] = {box.max_boundary().axis(axis), tr[i], bounding_edge_point::end};
    }

    std::sort(&edge_buffer[int(axis)][0], &edge_buffer[int(axis)][0] + 2 * tr.size());

    for (const auto i : boost::irange<std::size_t>(0, 2 * tr.size()))
    {
      const rtc::bounding_edge_point& bp{edge_buffer[int(axis)][i]};
      const auto edge_t = bp.value;

      if (bp.type == bounding_edge_point::end)
        --above;
      if (bp.type == bounding_edge_point::start)
        ++below;

      if (edge_t > node_bbox.min_boundary().axis(axis) && edge_t < node_bbox.max_boundary().axis(axis))
      {
        const rtc_float belowSA{2 * (node_diagonal.axis(axis2) * node_diagonal.axis(axis3) +
                                     (edge_t - node_bbox.min_boundary().axis(axis)) *
                                         (node_diagonal.axis(axis2) + node_diagonal.axis(axis3)))};

        const rtc_float aboveSA{2 * (node_diagonal.axis(axis2) * node_diagonal.axis(axis3) +
                                     (node_bbox.max_boundary().axis(axis) - edge_t) *
                                         (node_diagonal.axis(axis2) + node_diagonal.axis(axis3)))};

        const rtc_float pBelow{belowSA * invTotalSA};
        const rtc_float pAbove{aboveSA * invTotalSA};
        const rtc_float eb{(above == 0 || below == 0) ? empty_bonus : 0};
        const rtc_float cost{traversal_cost + isect_cost * (1 - eb) * (pBelow * below + pAbove * above)};

        if (cost < best_cost)
        {
          best_cost = cost;
          best_axis = int(axis);
          best_offset = i;
        }
      }
    }
  } while ((best_axis == -1) && ((axis = next(axis), retries++) < 2));

  return std::make_tuple(best_axis, best_offset, best_cost, old_cost);
}

auto kd_tree::split_triangles(std::vector<std::uint32_t>&& tr_init,
                              const edge_buffer_array_t& edge_buffer,
                              const std::uint32_t best_axis,
                              const std::uint32_t best_offset)
{
  std::vector<std::uint32_t> tr{std::move(tr_init)}, left_set{}, right_set{};

  left_set.reserve(best_offset);
  right_set.reserve(2 * tr.size() - best_offset);

  for (auto i = 0; i < best_offset; ++i)
  {
    if (edge_buffer[best_axis][i].type == bounding_edge_point::start)
    {
      left_set.emplace_back(edge_buffer[best_axis][i].tr_index);
    }
  }

  for (auto i = best_offset + 1; i < 2 * tr.size(); ++i)
  {
    if (edge_buffer[best_axis][i].type == bounding_edge_point::end)
    {
      right_set.emplace_back(edge_buffer[best_axis][i].tr_index);
    }
  }
  return std::make_tuple(std::move(left_set), std::move(right_set));
}

void kd_tree::build_tree(std::unique_ptr<tree_node>& node,
                         rtc::bounding_box node_bbox,
                         std::vector<std::uint32_t> tr,
                         const std::vector<rtc::bounding_box>& primitive_bboxes,
                         edge_buffer_array_t& edge_buffer,
                         const std::uint32_t depth,
                         std::uint32_t bad_refines)
{
  const auto create_leaf_node = [&node, &tr] {
    node = std::make_unique<tree_node>();
    node->triangles = std::make_unique<kd_tree::value_type>();
    node->triangles->insert(node->triangles->end(), tr.begin(), tr.end());
  };

  if (tr.size() <= 1 || !depth)
    return create_leaf_node();

  const auto [axis, offset, cost, old_cost] =
      compute_node_split_paramters(edge_buffer, tr, node_bbox, primitive_bboxes);

  bad_refines += (cost > old_cost) ? 1 : 0;

  auto no_use_of_spliting_node = [old_cost = old_cost, cost = cost, axis = axis, &bad_refines, &tr] {
    return ((cost > 4 * old_cost) && (tr.size() < 16)) || (axis == -1) || (bad_refines == 3);
  };

  if (no_use_of_spliting_node())
    return create_leaf_node();

  auto [left_set, right_set] = split_triangles(std::move(tr), edge_buffer, axis, offset);
  bounding_box left_bbox{node_bbox}, right_bbox{node_bbox};

  node = std::make_unique<tree_node>();
  node->axis = {edge_buffer[axis][offset].value, static_cast<rtc::axis>(axis)};

  left_bbox.max_boundary().axis(node->axis.split) = right_bbox.min_boundary().axis(node->axis.split) = node->axis.value;

  build_tree(node->left, left_bbox, std::move(left_set), primitive_bboxes, edge_buffer, depth - 1, bad_refines);
  build_tree(node->right, right_bbox, std::move(right_set), primitive_bboxes, edge_buffer, depth - 1, bad_refines);
}

}  // namespace rtc


#ifndef STAN_MATH_TORSTEN_PKMODEL_MODELPARAMETERS_HPP
#define STAN_MATH_TORSTEN_PKMODEL_MODELPARAMETERS_HPP

#include <stan/math/prim/fun/to_array_1d.hpp>
#include <stan/math/torsten/ev_history.hpp>
#include <stan/math/torsten/PKModel/ExtractVector.hpp>
#include <stan/math/torsten/PKModel/SearchReal.hpp>
#include <Eigen/Dense>
#include <algorithm>
#include <vector>

namespace torsten {

/**
 * The ModelParameters class defines objects that contain the parameters of
 * a model at a given time.
 */
template<typename T_time, typename T_parameters, typename T_biovar, typename T_tlag>
struct ModelParameters {
  double time_;
  std::vector<T_parameters> theta_;
  int nrow, ncol;
  std::vector<T_biovar> biovar_;
  std::vector<T_tlag> tlag_;

  ModelParameters() {}

  ModelParameters(const T_time& time,
                  const Eigen::Matrix<T_parameters, Eigen::Dynamic, Eigen::Dynamic>& K,
                  const std::vector<T_biovar>& biovar,
                  const std::vector<T_tlag>& tlag)
    : time_(stan::math::value_of(time)), theta_(K.size()), nrow(K.rows()), ncol(K.cols()), biovar_(biovar), tlag_(tlag)
  {
    Eigen::Matrix<T_parameters, -1, -1>::Map(theta_.data(), nrow, ncol) = K;
  }

  ModelParameters(const T_time& time,
                  const std::vector<T_parameters>& theta,
                  const std::vector<T_biovar>& biovar,
                  const std::vector<T_tlag>& tlag)
    : time_(stan::math::value_of(time)), theta_(theta), biovar_(biovar), tlag_(tlag) {}

  /**
   * Adds parameters. Useful for the mixed solver, where
   * we want to augment the parameters with the intial PK
   * states when calling the numerical integrator.
   */
  template <typename T>
  ModelParameters<T_time, T, T_biovar, T_tlag>
  augment(const std::vector<T>& thetaAdd) const {
    std::vector<T> theta(theta_.size());
    for (size_t i = 0; i < theta.size(); i++) theta[i] = theta_[i];
    for (size_t i = 0; i < thetaAdd.size(); i++) theta.push_back(thetaAdd[i]);
    return
      ModelParameters<T_time, T, T_biovar, T_tlag>
        (time_, theta, biovar_, tlag_);
  }

  template <typename T>
  ModelParameters<T_time, T, T_biovar, T_tlag>
  augment(const Eigen::Matrix<T, Eigen::Dynamic, 1>& thetaAdd)
  const {
    return augment(stan::math::to_array_1d(thetaAdd));
  }

  /**
   * Edit time stored in parameter object.
   */
  void time(double time) {
    time_ = time;
  }

  int CountParameters() const {
    return theta_.size();
  }

  // access functions   // FIX ME - name should be get_theta.
  double get_time() const { return time_; }
  std::vector<T_parameters> get_RealParameters(bool return_matrix) const {
    if (return_matrix) {
      auto k = get_K();
      std::vector<T_parameters> par(k.size());
      for (size_t j = 0; j < par.size(); ++j) par[j] = k(j);
      return par;
    } else {
      return theta_;
    }
  }
  std::vector<T_biovar> get_biovar() const {
    return biovar_;
  }
  std::vector<T_tlag> get_tlag() const {
    return tlag_;
  }
  Eigen::Matrix<T_parameters, Eigen::Dynamic, Eigen::Dynamic> get_K() const {
    Eigen::Matrix<T_parameters, -1, -1> res(nrow, ncol);
    res = Eigen::Matrix<T_parameters, -1, -1>::Map(theta_.data(), nrow, ncol);
    return res;
  }

  inline ModelParameters<double, double, double, double> to_value() const {
    return ModelParameters<double, double, double, double>
      (time_,
       stan::math::value_of(theta_),
       stan::math::value_of(biovar_),
       stan::math::value_of(tlag_));
  }
};

  template<typename T>
  struct is_matrix : std::false_type {};

  template<typename T, int R, int C>
  struct is_matrix<Eigen::Matrix<T, R, C> > : std::true_type {};

/**
 * The ModelParameterHistory class defines objects that contain a vector
 * of ModelParameters, along with a series of functions that operate on
 * them.
 */
template<typename T_time, typename T4_container, typename T5, typename T6>
struct ModelParameterHistory {
  using T4 = typename stan::value_type<T4_container>::type;
  using Param = std::pair<double, std::array<int, 3> >;

  static const bool has_matrix_param;

  std::vector<Param> index;
  const std::vector<T4_container>& theta_;
  const std::vector<std::vector<T5> >& biovar_;
  const std::vector<std::vector<T6> >& tlag_;

  template<typename T0>
   ModelParameterHistory(const std::vector<T0>& time,
                        const std::vector<T4_container>& theta,
                        const std::vector<std::vector<T5> >& biovar,
                        const std::vector<std::vector<T6> >& tlag) :
     ModelParameterHistory(0, time.size(), time,
                           0, theta.size(), theta,
                           0, biovar.size(), biovar,
                           0, tlag.size(), tlag)
  {}

  /*
   * For population data in form of ragged array, we need to
   * generate individual parameter history given the entire
   * population data and the location of the
   * inidividual. However, @c theta, @c biovar and @c tlag
   * could have different lengths, so for each variable we
   * need an index that points to the range that belongs to the individual.
   * Note that if all three variables are of size 1, their
   * time is set to be the first entry of the @c time vector
   */
  template<typename T0>
  ModelParameterHistory(int ibegin, int isize,
                        const std::vector<T0>& time,
                        int ibegin_theta, int isize_theta,
                        const std::vector<T4_container>& theta,
                        int ibegin_biovar, int isize_biovar,
                        const std::vector<std::vector<T5> >& biovar,
                        int ibegin_tlag, int isize_tlag,
                        const std::vector<std::vector<T6> >& tlag) :
    index(isize),
    theta_(theta),
    biovar_(biovar),
    tlag_(tlag)
  {
    for (int i = 0; i < isize; ++i) {
      int j = isize_theta   > 1 ? ibegin_theta  + i : ibegin_theta;
      int k = isize_biovar  > 1 ? ibegin_biovar + i : ibegin_biovar;
      int l = isize_tlag    > 1 ? ibegin_tlag   + i : ibegin_tlag;
      index[i] = std::make_pair<double, std::array<int, 3> >(stan::math::value_of(time[ibegin + i]), {j, k, l });
    }
    Sort();
  }

  const T4_container& model_param(int i) const {
    return theta_[index[i].second[0]];
  }

  ModelParameters<T_time, T4, T5, T6> GetModelParameters(int i) const {
    return ModelParameters<T_time, T4, T5, T6>(std::get<0>(index[i]), theta_[std::get<1>(index[i])[0]], biovar_[std::get<1>(index[i])[1]], tlag_[std::get<1>(index[i])[2]]);
  }

  /**
   * MPV.size gives us the number of events.
   * MPV[i].RealParameters.size gives us the number of
   * ODE parameters for the ith event.
   *
   * FIX ME - rename this GetValueTheta
   */
  inline const T4& GetValue(int iEvent, int iParameter) const {
    return theta_[std::get<1>(index[iEvent])[0]][iParameter];
  }

  inline const T5& GetValueBio(int iEvent, int iParameter) const {
    return biovar_[std::get<1>(index[iEvent])[1]][iParameter];
  }

  inline const T6& GetValueTlag(int iEvent, int iParameter) const {
    return tlag_[std::get<1>(index[iEvent])[2]][iParameter];
  }

  inline int get_size() const {
    return index.size();
  }

  void Sort() {
    std::sort(index.begin(), index.end(),
              [](const Param& a, const Param& b)
              { return std::get<0>(a) < std::get<0>(b); });
  }

  bool Check() {
  // check that elements are in chronological order.
    int i = index.size() - 1;
    bool ordered = true;

    while (i > 0 && ordered) {
      ordered = (std::get<0>(index[i]) >= std::get<0>(index[i-1]));
      i--;
    }
    return ordered;
  }

  /**
   * COMPLETE MODEL PARAMETERS
   *
   * Completes parameters so that it contains model parameters for each event
   * in events. If parameters contains only one set of parameters (case where
   * the parameters are constant), this set is replicated for each event in
   * events. Otherwise a new parameter vector is added for each new event
   * (isnew = true). This parameter vector is identical to the parameter vector
   * at the subsequent event. If the new event occurs at a time posterior to
   * the time of the last event, than the new vector parameter equals the
   * parameter vector of the last event. This amounts to doing an LOCF
   * (Last Observation Carried Forward):
   * Three cases:
   * (a) The time of the new event is higher than the time of the last
   *     parameter vector in parameters (k = len_parameters).
   *     Create a parameter vector at the the time of the new event,
   *     with the parameters of the last parameter vector.
   *     (Last Observation Carried Forward)
   * (b) The time of the new event matches the time of a parameter vector
   *     in parameters. This parameter vector gets replicated.
   * (c) (a) is not verified and no parameter vector occurs at the time
   *     of the new event. A new parameter vector is created at the time
   *     of the new event, and its parameters are equal to the parameters
   *     of the subsequent parameter vector in parameters.
   *
   * Since both @c index and events @c index are sorted
   * in time, we always move paramtter pointer to events
   *
   * Events and Parameters are sorted at the end of the procedure.
   *
   * @param[in] parameters at each event
   * @param[in] events elements (following NONMEM convention) at each event
   * @return - modified parameters and events.
   */
  template<typename T0, typename T_p1, typename T_p2, typename T_p3>
  void CompleteParameterHistory(torsten::EventHistory<T0, T_p1, T_p2, T_p3, T6>& events) {
    int nEvent = events.size();
    assert(nEvent > 0);
    int len_Parameters = index.size();  // numbers of events for which parameters are determined
    assert(len_Parameters > 0);

    if (!Check()) Sort();
    if (!events.Check()) events.Sort();
    index.resize(nEvent);

    int iEvent = 0;
    for (int i = 0; i < len_Parameters - 1; i++) {
      while (events.isnew(iEvent)) iEvent++;  // skip new events
      assert(std::get<0>(index[i]) == events.time(iEvent));  // compare time of "old' events to time of parameters.
      iEvent++;
    }

    if (len_Parameters == 1)  {
      for (int i = 0; i < nEvent; i++) {
        index[i] = std::make_pair<double, std::array<int, 3> >(stan::math::value_of(events.time(i)) , std::array<int,3>(std::get<1>(index[0])));
        events.index[i][3] = 0;
      }
    } else {  // parameters are event dependent.
      std::vector<double> times(nEvent, 0);
      for (int i = 0; i < nEvent; i++) times[i] = index[i].first;
      iEvent = 0;

      Param newParameter;
      int j = 0;
      std::vector<Param>::const_iterator lower = index.begin();
      std::vector<Param>::const_iterator it_param_end = index.begin() + len_Parameters;
      for (int iEvent = 0; iEvent < nEvent; ++iEvent) {
        if (events.isnew(iEvent)) {
          // Find the index corresponding to the time of the new event in the
          // times vector.
          const double t = stan::math::value_of(events.time(iEvent));
          lower = std::lower_bound(lower, it_param_end, t,
                                   [](const Param& t1, const double& t2) {return t1.first < t2;});
          newParameter = lower == (it_param_end) ? index[len_Parameters-1] : *lower;
          newParameter.first = t;
          index[len_Parameters + j] = newParameter;
          events.index[iEvent][3] = 0;
          j++;
        }
      }
    }
    Sort();
  }
};

template<typename T_time, typename T4_container, typename T5, typename T6>
const bool ModelParameterHistory<T_time, T4_container, T5, T6>::has_matrix_param = is_matrix<T4_container>::value;

}

#endif


#ifndef DART_NEURAL_SNAPSHOT_HPP_
#define DART_NEURAL_SNAPSHOT_HPP_

#include <unordered_map>
#include <vector>

#include <Eigen/Dense>

#include "dart/neural/DifferentiableContactConstraint.hpp"
#include "dart/neural/NeuralConstants.hpp"
#include "dart/neural/NeuralUtils.hpp"
#include "dart/neural/WithRespectTo.hpp"
#include "dart/performance/PerformanceLog.hpp"
#include "dart/simulation/World.hpp"

namespace dart {

using namespace performance;

namespace neural {

class BackpropSnapshot
{
  friend class MappedBackpropSnapshot;

public:
  /// This saves a snapshot from a forward pass, with all the info we need in
  /// order to efficiently compute a backwards pass. Crucially, the positions
  /// must all be snapshots from before the timestep, yet this constructor must
  /// be called after the timestep.
  BackpropSnapshot(
      simulation::WorldPtr world,
      Eigen::VectorXs preStepPosition,
      Eigen::VectorXs preStepVelocity,
      Eigen::VectorXs preStepTorques,
      Eigen::VectorXs preConstraintVelocities,
      Eigen::VectorXs preStepLCPCache);

  /// This computes the implicit backprop without forming intermediate
  /// Jacobians. It takes a LossGradient with the position and velocity vectors
  /// filled it, though the loss with respect to torque is ignored and can be
  /// null. It returns a LossGradient with all three values filled in, position,
  /// velocity, and torque.
  void backprop(
      simulation::WorldPtr world,
      LossGradient& thisTimestepLoss,
      const LossGradient& nextTimestepLoss,
      PerformanceLog* perfLog = nullptr,
      bool exploreAlternateStrategies = false);

  /// This computes backprop in the high-level RL API's space, use `state` and
  /// `action` as the primitives we're taking gradients wrt to.
  LossGradientHighLevelAPI backpropState(
      simulation::WorldPtr world,
      const Eigen::VectorXs& nextTimestepStateLossGrad,
      PerformanceLog* perfLog = nullptr,
      bool exploreAlternateStrategies = false);

  /// This zeros out any components of the gradient that would want to push us
  /// out of the box-bounds encoded in the world for pos, vel, or force.
  void clipLossGradientsToBounds(
      simulation::WorldPtr world,
      Eigen::VectorXs& lossWrtPos,
      Eigen::VectorXs& lossWrtVel,
      Eigen::VectorXs& lossWrtForce);

  /// This computes and returns the whole vel-vel jacobian. For backprop, you
  /// don't actually need this matrix, you can compute backprop directly. This
  /// is here if you want access to the full Jacobian for some reason.
  const Eigen::MatrixXs& getVelVelJacobian(
      simulation::WorldPtr world, PerformanceLog* perfLog = nullptr);

  /// This computes and returns the whole pos-vel jacobian. For backprop, you
  /// don't actually need this matrix, you can compute backprop directly. This
  /// is here if you want access to the full Jacobian for some reason.
  const Eigen::MatrixXs& getPosVelJacobian(
      simulation::WorldPtr world, PerformanceLog* perfLog = nullptr);

  /// This computes and returns the whole force-vel jacobian. For backprop, you
  /// don't actually need this matrix, you can compute backprop directly. This
  /// is here if you want access to the full Jacobian for some reason.
  const Eigen::MatrixXs& getControlForceVelJacobian(
      simulation::WorldPtr world, PerformanceLog* perfLog = nullptr);

  /// This computes and returns the whole mass-vel jacobian. For backprop, you
  /// don't actually need this matrix, you can compute backprop directly. This
  /// is here if you want access to the full Jacobian for some reason.
  const Eigen::MatrixXs& getMassVelJacobian(
      simulation::WorldPtr world, PerformanceLog* perfLog = nullptr);

  /// This computes and returns the whole pos-pos jacobian. For backprop, you
  /// don't actually need this matrix, you can compute backprop directly. This
  /// is here if you want access to the full Jacobian for some reason.
  const Eigen::MatrixXs& getPosPosJacobian(
      simulation::WorldPtr world, PerformanceLog* perfLog = nullptr);

  /// This computes and returns the whole vel-pos jacobian. For backprop, you
  /// don't actually need this matrix, you can compute backprop directly. This
  /// is here if you want access to the full Jacobian for some reason.
  const Eigen::MatrixXs& getVelPosJacobian(
      simulation::WorldPtr world, PerformanceLog* perfLog = nullptr);

  /// This computes and returns the component of the pos-pos and pos-vel
  /// jacobians due to bounce approximation. For backprop, you don't actually
  /// need this matrix, you can compute backprop directly. This is here if you
  /// want access to the full Jacobian for some reason.
  const Eigen::MatrixXs& getBounceApproximationJacobian(
      simulation::WorldPtr world, PerformanceLog* perfLog = nullptr);

  /// This returns the Jacobian for state_t -> state_{t+1}.
  Eigen::MatrixXs getStateJacobian(simulation::WorldPtr world);

  /// This returns the Jacobian for action_t -> state_{t+1}.
  Eigen::MatrixXs getActionJacobian(simulation::WorldPtr world);

  /// Returns a concatenated vector of all the Skeletons' position()'s in the
  /// World, in order in which the Skeletons appear in the World's
  /// getSkeleton(i) returns them, BEFORE the timestep.
  Eigen::VectorXs getPreStepPosition();

  /// Returns a concatenated vector of all the Skeletons' velocity()'s in the
  /// World, in order in which the Skeletons appear in the World's
  /// getSkeleton(i) returns them, BEFORE the timestep.
  Eigen::VectorXs getPreStepVelocity();

  /// Returns a concatenated vector of all the joint torques that were applied
  /// during the forward pass, BEFORE the timestep.
  Eigen::VectorXs getPreStepTorques();

  /// Returns a concatenated vector of all the Skeletons' velocity()'s in the
  /// World, in order in which the Skeletons appear in the World's
  /// getSkeleton(i) returns them, AFTER integrating forward dynamics but BEFORE
  /// running the LCP.
  Eigen::VectorXs getPreConstraintVelocity();

  /// Returns a concatenated vector of all the Skeletons' position()'s in the
  /// World, in order in which the Skeletons appear in the World's
  /// getSkeleton(i) returns them, AFTER the timestep.
  Eigen::VectorXs getPostStepPosition();

  /// Returns a concatenated vector of all the Skeletons' velocity()'s in the
  /// World, in order in which the Skeletons appear in the World's
  /// getSkeleton(i) returns them, AFTER the timestep.
  Eigen::VectorXs getPostStepVelocity();

  /// Returns a concatenated vector of all the joint torques that were applied
  /// during the forward pass, AFTER the timestep.
  Eigen::VectorXs getPostStepTorques();

  /// Returns the LCP's cached solution from before the step
  const Eigen::VectorXs& getPreStepLCPCache();

  /////////////////////////////////////////////////////////////////////////////
  /// Just public for testing
  /////////////////////////////////////////////////////////////////////////////

  /// This returns the A_c matrix. You shouldn't ever need this matrix, it's
  /// just here to enable testing.
  Eigen::MatrixXs getClampingConstraintMatrix(simulation::WorldPtr world);

  /// This returns the V_c matrix. You shouldn't ever need this matrix, it's
  /// just here to enable testing.
  Eigen::MatrixXs getMassedClampingConstraintMatrix(simulation::WorldPtr world);

  /// This returns the A_ub matrix. You shouldn't ever need this matrix, it's
  /// just here to enable testing.
  Eigen::MatrixXs getUpperBoundConstraintMatrix(simulation::WorldPtr world);

  /// This returns the V_c matrix. You shouldn't ever need this matrix, it's
  /// just here to enable testing.
  Eigen::MatrixXs getMassedUpperBoundConstraintMatrix(
      simulation::WorldPtr world);

  /// This returns the E matrix. You shouldn't ever need this matrix, it's
  /// just here to enable testing.
  Eigen::MatrixXs getUpperBoundMappingMatrix();

  /// This returns the B matrix. You shouldn't ever need this matrix, it's
  /// just here to enable testing.
  Eigen::MatrixXs getBouncingConstraintMatrix(simulation::WorldPtr world);

  /// This returns the mass matrix for the whole world, a block diagonal
  /// concatenation of the skeleton mass matrices.
  Eigen::MatrixXs getMassMatrix(
      simulation::WorldPtr world, bool forFiniteDifferencing = false);

  /// This returns the inverse mass matrix for the whole world, a block diagonal
  /// concatenation of the skeleton inverse mass matrices.
  Eigen::MatrixXs getInvMassMatrix(
      simulation::WorldPtr world, bool forFiniteDifferencing = false);

  /// This is the subset of the A matrix from the original LCP that corresponds
  /// to clamping indices.
  Eigen::MatrixXs getClampingAMatrix();

  /// This returns the pos-C(pos,vel) Jacobian for the whole world, a block
  /// diagonal concatenation of the skeleton pos-C(pos,vel) Jacobians.
  Eigen::MatrixXs getPosCJacobian(simulation::WorldPtr world);

  /// This returns the vel-C(pos,vel) Jacobian for the whole world, a block
  /// diagonal concatenation of the skeleton vel-C(pos,vel) Jacobians.
  Eigen::MatrixXs getVelCJacobian(simulation::WorldPtr world);

  /// This computes and returns the whole vel-vel jacobian by finite
  /// differences. This is SUPER SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferenceVelVelJacobian(
      simulation::WorldPtr world, bool useRidders = true);

  /// This computes and returns the whole vel-vel jacobian by Ridders
  /// extrapolated finite differences. This is SUPER DUPER SLOW,
  /// and is only here for testing.
  Eigen::MatrixXs finiteDifferenceRiddersVelVelJacobian(
      simulation::WorldPtr world);

  /// This computes and returns the whole pos-C(pos,vel) jacobian by finite
  /// differences. This is SUPER SUPER SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferencePosVelJacobian(
      simulation::WorldPtr world, bool useRidders = true);

  /// This computes and returns the whole pos-C(pos,vel) jacobian by Ridders
  /// extrapolated finite differences. This is SUPER DUPER SLOW,
  /// and is only here for testing.
  Eigen::MatrixXs finiteDifferenceRiddersPosVelJacobian(
      simulation::WorldPtr world);

  /// This computes and returns the whole force-vel jacobian by finite
  /// differences. This is SUPER SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferenceForceVelJacobian(
      simulation::WorldPtr world, bool useRidders = true);

  /// This computes and returns the whole force-vel jacobian by Ridders
  /// extrapolated finite differences. This is SUPER DUPER SLOW,
  /// and is only here for testing.
  Eigen::MatrixXs finiteDifferenceRiddersForceVelJacobian(
      simulation::WorldPtr world);

  /// This computes and returns the whole mass-vel jacobian by finite
  /// differences. This is SUPER SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferenceMassVelJacobian(
      simulation::WorldPtr world, bool useRidders = true);

  /// This computes and returns the whole mass-vel jacobian by Ridders
  /// extrapolated finite differences. This is SUPER DUPER SLOW,
  /// and is only here for testing.
  Eigen::MatrixXs finiteDifferenceRiddersMassVelJacobian(
      simulation::WorldPtr world);

  /// This computes and returns the whole pos-pos jacobian by finite
  /// differences. This is SUPER SUPER SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferencePosPosJacobian(
      simulation::WorldPtr world,
      std::size_t subdivisions = 20,
      bool useRidders = true);

  /// This computes and returns the whole pos-pos jacobian by Ridders
  /// extrapolated finite differences. This is SUPER DUPER SUPER SLOW,
  /// and is only here for testing.
  Eigen::MatrixXs finiteDifferenceRiddersPosPosJacobian(
      simulation::WorldPtr world, std::size_t subdivisions = 20);

  /// This computes and returns the whole vel-pos jacobian by finite
  /// differences. This is SUPER SUPER SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferenceVelPosJacobian(
      simulation::WorldPtr world,
      std::size_t subdivisions = 20,
      bool useRidders = true);

  /// This computes and returns the whole vel-pos jacobian by Ridders
  /// extrapolated finite differences. This is SUPER DUPER SUPER SLOW,
  /// and is only here for testing.
  Eigen::MatrixXs finiteDifferenceRiddersVelPosJacobian(
      simulation::WorldPtr world, std::size_t subdivisions = 20);

  /// This computes and returns the whole wrt-vel jacobian by finite
  /// differences. This is SUPER SUPER SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferenceVelJacobianWrt(
      simulation::WorldPtr world, WithRespectTo* wrt, bool useRidders = true);

  /// This computes and returns the whole wrt-vel jacobian by Ridders
  /// extrapolated finite differences. This is SUPER SUPER SLOW,
  /// and is only here for testing.
  Eigen::MatrixXs finiteDifferenceRiddersVelJacobianWrt(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This computes and returns the whole wrt-pos jacobian by finite
  /// differences. This is SUPER SUPER SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferencePosJacobianWrt(
      simulation::WorldPtr world, WithRespectTo* wrt, bool useRidders = true);

  /// This computes and returns the whole wrt-pos jacobian by Ridders
  /// extrapolated finite differences. This is SUPER SUPER SLOW,
  /// and is only here for testing.
  Eigen::MatrixXs finiteDifferenceRiddersPosJacobianWrt(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This returns the P_c matrix. You shouldn't ever need this matrix, it's
  /// just here to enable testing.
  Eigen::MatrixXs getProjectionIntoClampsMatrix(
      simulation::WorldPtr world, bool forFiniteDifferencing = false);

  /// This replaces x with the result of M*x in place, without explicitly
  /// forming M
  Eigen::VectorXs implicitMultiplyByMassMatrix(
      simulation::WorldPtr world, const Eigen::VectorXs& x);

  /// This replaces x with the result of Minv*x in place, without explicitly
  /// forming Minv
  Eigen::VectorXs implicitMultiplyByInvMassMatrix(
      simulation::WorldPtr world, const Eigen::VectorXs& x);

  /// TODO(keenon): Remove me
  Eigen::MatrixXs getScratchAnalytical(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// TODO(keenon): Remove me
  Eigen::MatrixXs getScratchFiniteDifference(
      simulation::WorldPtr world, WithRespectTo* wrt, bool useRidders = true);

  /// TODO(keenon): Remove me
  Eigen::MatrixXs getScratchFiniteDifferenceRidders(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This predicts what the next velocity will be using our linear algebra
  /// formula. This is only here for testing, to compare it against the actual
  /// result of a timestep.
  ///
  /// The `morePreciseButSlower` flag tells this function to do brute force
  /// steps to get constraint matrices A_c and A_ub, rather than use 1st order
  /// approximations. This is important because when we're doing
  /// finite-differencing over tiny EPS (1e-9) then tiny errors in the 1st order
  /// approximations blow up to become huge errors in gradients.
  Eigen::VectorXs getAnalyticalNextV(
      simulation::WorldPtr world, bool morePreciseButSlower = false);

  /// This computes and returns the whole wrt-vel jacobian. For backprop, you
  /// don't actually need this matrix, you can compute backprop directly. This
  /// is here if you want access to the full Jacobian for some reason.
  Eigen::MatrixXs getVelJacobianWrt(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This computes and returns the whole wrt-pos jacobian. For backprop, you
  /// don't actually need this matrix, you can compute backprop directly. This
  /// is here if you want access to the full Jacobian for some reason.
  Eigen::MatrixXs getPosJacobianWrt(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This returns the jacobian of constraint force, holding everyhing constant
  /// except the value of WithRespectTo
  Eigen::MatrixXs getJacobianOfConstraintForce(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This returns the analytical expression for the Jacobian of Q*b, holding b
  /// constant, if there are some upper-bound indices
  Eigen::MatrixXs dQ_WithUB(
      simulation::WorldPtr world,
      Eigen::MatrixXs& Minv,
      Eigen::MatrixXs& A_c,
      Eigen::MatrixXs& E,
      Eigen::MatrixXs& A_c_ub_E,
      Eigen::VectorXs rhs,
      WithRespectTo* wrt);
  /// This returns the analytical expression for the Jacobian of Q^T*b, holding
  /// b constant, if there are some upper-bound indices
  Eigen::MatrixXs dQT_WithUB(
      simulation::WorldPtr world,
      Eigen::MatrixXs& Minv,
      Eigen::MatrixXs& A_c,
      Eigen::MatrixXs& E,
      Eigen::MatrixXs& A_ub,
      Eigen::VectorXs rhs,
      WithRespectTo* wrt);
  /// This returns the analytical expression for the Jacobian of Q*b, holding b
  /// constant, if there are no upper-bound indices
  Eigen::MatrixXs dQ_WithoutUB(
      simulation::WorldPtr world,
      Eigen::MatrixXs& Minv,
      Eigen::MatrixXs& A_c,
      Eigen::VectorXs rhs,
      WithRespectTo* wrt);

  /// This returns the jacobian of Qb, holding b constant, with respect to
  /// wrt, by finite differencing
  Eigen::MatrixXs finiteDifferenceJacobianOfQb(
      simulation::WorldPtr world,
      Eigen::VectorXs b,
      WithRespectTo* wrt,
      bool useRidders = true);

  /// This returns the jacobian of Qb, holding b constant, with respect to
  /// wrt, by finite differencing
  Eigen::MatrixXs finiteDifferenceRiddersJacobianOfQb(
      simulation::WorldPtr world, Eigen::VectorXs b, WithRespectTo* wrt);

  /// This returns the vector of constants that get added to the diagonal of Q
  /// to guarantee that Q is full-rank
  Eigen::VectorXs getConstraintForceMixingDiagonal();

  /// This returns the jacobian of Q^{-1}b, holding b constant, with respect to
  /// wrt
  Eigen::MatrixXs getJacobianOfLCPConstraintMatrixClampingSubset(
      simulation::WorldPtr world, Eigen::VectorXs b, WithRespectTo* wrt);

  /// This returns the jacobian of Q^{-1}b, holding b constant, with respect to
  /// wrt, by finite differencing
  Eigen::MatrixXs finiteDifferenceJacobianOfLCPConstraintMatrixClampingSubset(
      simulation::WorldPtr world,
      Eigen::VectorXs b,
      WithRespectTo* wrt,
      bool useRidders = true);

  /// This returns the jacobian of Q^{-1}b, holding b constant, with respect to
  /// wrt, by Ridders extrapolated finite differencing
  Eigen::MatrixXs
  finiteDifferenceRiddersJacobianOfLCPConstraintMatrixClampingSubset(
      simulation::WorldPtr world, Eigen::VectorXs b, WithRespectTo* wrt);

  /// This returns the jacobian of b (from Q^{-1}b) with respect to wrt
  Eigen::MatrixXs getJacobianOfLCPOffsetClampingSubset(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This returns the jacobian of b (from Q^{-1}b) with respect to wrt, by
  /// finite differencing
  Eigen::MatrixXs finiteDifferenceJacobianOfLCPOffsetClampingSubset(
      simulation::WorldPtr world, WithRespectTo* wrt, bool useRidders = true);

  /// This returns the jacobian of b (from Q^{-1}b) with respect to wrt, by
  /// Ridders extrapolated finite differencing
  Eigen::MatrixXs finiteDifferenceRiddersJacobianOfLCPOffsetClampingSubset(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This returns the jacobian of b (from Q^{-1}b) with respect to wrt, by
  /// finite differencing
  Eigen::MatrixXs finiteDifferenceJacobianOfLCPEstimatedOffsetClampingSubset(
      simulation::WorldPtr world, WithRespectTo* wrt, bool useRidders = true);

  /// This returns the jacobian of b (from Q^{-1}b) with respect to wrt, by
  /// Ridders extrapolated finite differencing
  Eigen::MatrixXs
  finiteDifferenceRiddersJacobianOfLCPEstimatedOffsetClampingSubset(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This returns the subset of the A matrix used by the original LCP for just
  /// the clamping constraints. It relates constraint force to constraint
  /// acceleration. It's a mass matrix, just in a weird frame.
  void computeLCPConstraintMatrixClampingSubset(
      simulation::WorldPtr world,
      Eigen::MatrixXs& Q,
      const Eigen::MatrixXs& A_c,
      const Eigen::MatrixXs& A_ub,
      const Eigen::MatrixXs& E);

  /// This returns the subset of the b vector used by the original LCP for just
  /// the clamping constraints. It's just the relative velocity at the clamping
  /// contact points.
  void computeLCPOffsetClampingSubset(
      simulation::WorldPtr world,
      Eigen::VectorXs& b,
      const Eigen::MatrixXs& A_c);

  /// This computes and returns an estimate of the constraint impulses for the
  /// clamping constraints. This is based on a linear approximation of the
  /// constraint impulses.
  Eigen::VectorXs estimateClampingConstraintImpulses(
      simulation::WorldPtr world,
      const Eigen::MatrixXs& A_c,
      const Eigen::MatrixXs& A_ub,
      const Eigen::MatrixXs& E);

  /// This returns the jacobian of P_c * v, holding everyhing constant except
  /// the value of WithRespectTo
  Eigen::MatrixXs getJacobianOfProjectionIntoClampsMatrix(
      simulation::WorldPtr world, Eigen::VectorXs v, WithRespectTo* wrt);

  /// This computes and returns the jacobian of P_c * v by finite
  /// differences. This is SUPER SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferenceJacobianOfProjectionIntoClampsMatrix(
      simulation::WorldPtr world,
      Eigen::VectorXs v,
      WithRespectTo* wrt,
      bool useRidders = true);

  /// This computes and returns the jacobian of P_c * v by Ridders extrapolated
  /// finite differences. This is SUPER SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferenceRiddersJacobianOfProjectionIntoClampsMatrix(
      simulation::WorldPtr world, Eigen::VectorXs v, WithRespectTo* wrt);

  /// This returns the jacobian of M^{-1}(pos, inertia) * tau, holding
  /// everything constant except the value of WithRespectTo
  Eigen::MatrixXs getJacobianOfMinv(
      simulation::WorldPtr world, Eigen::VectorXs tau, WithRespectTo* wrt);

  /// This computes and returns the jacobian of M^{-1}(pos, inertia) * tau by
  /// finite differences. This is SUPER SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferenceJacobianOfMinv(
      simulation::WorldPtr world,
      Eigen::VectorXs tau,
      WithRespectTo* wrt,
      bool useRidders = true);

  /// This computes and returns the jacobian of M^{-1}(pos, inertia) * tau by
  /// Ridders extrapolated finite differences. This is SUPER SLOW, and is
  /// only here for testing.
  Eigen::MatrixXs finiteDifferenceRiddersJacobianOfMinv(
      simulation::WorldPtr world, Eigen::VectorXs tau, WithRespectTo* wrt);

  /// This returns the jacobian of M(pos, inertia) * v, holding
  /// everything constant except the value of WithRespectTo
  Eigen::MatrixXs getJacobianOfM(
      simulation::WorldPtr world, Eigen::VectorXs v, WithRespectTo* wrt);

  /// This computes and returns the jacobian of M(pos, inertia) * v by
  /// finite differences. This is SUPER SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferenceJacobianOfM(
      simulation::WorldPtr world,
      Eigen::VectorXs v,
      WithRespectTo* wrt,
      bool useRidders = true);

  /// This computes and returns the jacobian of M(pos, inertia) * v by
  /// Ridders extrapolated finite differences. This is SUPER SLOW, and is
  /// only here for testing.
  Eigen::MatrixXs finiteDifferenceRiddersJacobianOfM(
      simulation::WorldPtr world, Eigen::VectorXs v, WithRespectTo* wrt);

  /// This returns the jacobian of C(pos, inertia, vel), holding everything
  /// constant except the value of WithRespectTo
  Eigen::MatrixXs getJacobianOfC(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This returns the jacobian of C(pos, inertia, vel), holding everything
  /// constant except the value of WithRespectTo
  Eigen::MatrixXs computeJacobianOfC(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This computes and returns the jacobian of C(pos, inertia, vel) by finite
  /// differences. This is SUPER SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferenceJacobianOfC(
      simulation::WorldPtr world, WithRespectTo* wrt, bool useRidders = true);

  /// This computes and returns the jacobian of C(pos, inertia, vel) by Ridders
  /// extrapolated finite differences. This is SUPER SLOW, and is only here
  /// for testing.
  Eigen::MatrixXs finiteDifferenceRiddersJacobianOfC(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This returns the jacobian of M^{-1}(pos, inertia) * (C(pos, inertia, vel)
  /// + mPreStepTorques), holding everything constant except the value of
  /// WithRespectTo
  Eigen::MatrixXs getJacobianOfMinvC(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This computes and returns the jacobian of M^{-1}(pos, inertia) * C(pos,
  /// inertia, vel) by finite differences. This is SUPER SLOW, and is only here
  /// for testing.
  Eigen::MatrixXs finiteDifferenceJacobianOfMinvC(
      simulation::WorldPtr world, WithRespectTo* wrt, bool useRidders = false);

  /// This computes and returns the jacobian of M^{-1}(pos, inertia) * C(pos,
  /// inertia, vel) by Ridders extrapolated finite differences. This is SUPER
  /// SLOW, and is only here for testing.
  Eigen::MatrixXs finiteDifferenceRiddersJacobianOfMinvC(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This returns a fast approximation to A_c in the neighborhood of the
  /// original
  Eigen::MatrixXs estimateClampingConstraintMatrixAt(
      simulation::WorldPtr world, Eigen::VectorXs pos);

  /// This returns a fast approximation to A_ub in the neighborhood of the
  /// original
  Eigen::MatrixXs estimateUpperBoundConstraintMatrixAt(
      simulation::WorldPtr world, Eigen::VectorXs pos);

  /// Only for testing: VERY SLOW. This returns the actual value of A_c at the
  /// desired position.
  Eigen::MatrixXs getClampingConstraintMatrixAt(
      simulation::WorldPtr world, Eigen::VectorXs pos);

  /// Only for testing: VERY SLOW. This returns the actual value of A_ub at the
  /// desired position.
  Eigen::MatrixXs getUpperBoundConstraintMatrixAt(
      simulation::WorldPtr world, Eigen::VectorXs pos);

  /// Only for testing: VERY SLOW. This returns the actual value of E at the
  /// desired position.
  Eigen::MatrixXs getUpperBoundMappingMatrixAt(
      simulation::WorldPtr world, Eigen::VectorXs pos);

  /// Only for testing: VERY SLOW. This returns the actual value of the bounce
  /// diagonals at the desired position.
  Eigen::VectorXs getBounceDiagonalsAt(
      simulation::WorldPtr world, Eigen::VectorXs pos);

  /// This computes the Jacobian of A_c*f0 with respect to position using
  /// impulse tests.
  Eigen::MatrixXs getJacobianOfClampingConstraints(
      simulation::WorldPtr world, Eigen::VectorXs f0);

  /// This computes the finite difference Jacobian of A_c*f0 with respect to
  /// position. This is AS SLOW AS FINITE DIFFERENCING THE WHOLE ENGINE, which
  /// is way too slow to use in practice.
  Eigen::MatrixXs finiteDifferenceJacobianOfClampingConstraints(
      simulation::WorldPtr world, Eigen::VectorXs f0, bool useRidders = true);

  /// This computes the finite difference Jacobian of A_c*f0 with respect to
  /// position. This is AS SLOW AS FINITE DIFFERENCING THE WHOLE ENGINE, which
  /// is way too slow to use in practice. Using Ridders, so even slower.
  Eigen::MatrixXs finiteDifferenceRiddersJacobianOfClampingConstraints(
      simulation::WorldPtr world, Eigen::VectorXs f0);

  /// This computes the Jacobian of A_c^T*v0 with respect to position using
  /// impulse tests.
  Eigen::MatrixXs getJacobianOfClampingConstraintsTranspose(
      simulation::WorldPtr world, Eigen::VectorXs v0);

  /// This computes the finite difference Jacobian of A_c^T*v0 with respect to
  /// position. This is AS SLOW AS FINITE DIFFERENCING THE WHOLE ENGINE, which
  /// is way too slow to use in practice.
  Eigen::MatrixXs finiteDifferenceJacobianOfClampingConstraintsTranspose(
      simulation::WorldPtr world, Eigen::VectorXs v0, bool useRidders = true);

  /// This computes the finite difference Jacobian of A_c^T*v0 with respect to
  /// position. This is AS SLOW AS FINITE DIFFERENCING THE WHOLE ENGINE, which
  /// is way too slow to use in practice. Using Ridders, so even slower.
  Eigen::MatrixXs finiteDifferenceRiddersJacobianOfClampingConstraintsTranspose(
      simulation::WorldPtr world, Eigen::VectorXs v0);

  /// This computes the Jacobian of A_ub*(E*f0) with respect to position using
  /// impulse tests.
  Eigen::MatrixXs getJacobianOfUpperBoundConstraints(
      simulation::WorldPtr world, Eigen::VectorXs E_f0);

  /// This computes the finite difference Jacobian of A_ub*E*f0 with respect to
  /// position. This is AS SLOW AS FINITE DIFFERENCING THE WHOLE ENGINE, which
  /// is way too slow to use in practice.
  Eigen::MatrixXs finiteDifferenceJacobianOfUpperBoundConstraints(
      simulation::WorldPtr world, Eigen::VectorXs f0, bool useRidders = true);

  /// This computes the finite difference Jacobian of A_ub*E*f0 with respect to
  /// position. This is AS SLOW AS FINITE DIFFERENCING THE WHOLE ENGINE, which
  /// is way too slow to use in practice. Uses Ridders method.
  Eigen::MatrixXs finiteDifferenceRiddersJacobianOfUpperBoundConstraints(
      simulation::WorldPtr world, Eigen::VectorXs f0);

  /// This computes the Jacobian of A_ub^T*v0 with respect to position using
  /// impulse tests.
  Eigen::MatrixXs getJacobianOfUpperBoundConstraintsTranspose(
      simulation::WorldPtr world, Eigen::VectorXs v0);

  /// This returns the jacobian of constraint force, holding everything constant
  /// except the value of WithRespectTo
  Eigen::MatrixXs finiteDifferenceJacobianOfConstraintForce(
      simulation::WorldPtr world, WithRespectTo* wrt, bool useRidders = true);

  /// This returns the jacobian of constraint force, holding everything constant
  /// except the value of WithRespectTo. Uses Ridders method.
  Eigen::MatrixXs finiteDifferenceRiddersJacobianOfConstraintForce(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// This returns the jacobian of estimated constraint force, without actually
  /// running forward passes, holding everyhing constant except the value of
  /// WithRespectTo
  Eigen::MatrixXs finiteDifferenceJacobianOfEstimatedConstraintForce(
      simulation::WorldPtr world, WithRespectTo* wrt, bool useRidders = true);

  /// This returns the jacobian of estimated constraint force, without actually
  /// running forward passes, holding everyhing constant except the value of
  /// WithRespectTo. Uses Ridders method.
  Eigen::MatrixXs finiteDifferenceRiddersJacobianOfEstimatedConstraintForce(
      simulation::WorldPtr world, WithRespectTo* wrt);

  /// These was the mX() vector used to construct this. Pretty much only here
  /// for testing.
  Eigen::VectorXs getContactConstraintImpulses();

  /// These was the fIndex() vector used to construct this. Pretty much only
  /// here for testing.
  Eigen::VectorXi getContactConstraintMappings();

  /// Returns the vector of the coefficients on the diagonal of the bounce
  /// matrix. These are 1+restitutionCoeff[i].
  Eigen::VectorXs getBounceDiagonals();

  /// Returns the vector of the restitution coeffs, sized for the number of
  /// bouncing collisions.
  Eigen::VectorXs getRestitutionDiagonals();

  /// Returns the penetration correction hack "bounce" (or 0 if the contact is
  /// not inter-penetrating or is actively bouncing) at each contact point.
  Eigen::VectorXs getPenetrationCorrectionVelocities();

  /// Returns the constraint impulses along the clamping constraints
  Eigen::VectorXs getClampingConstraintImpulses();

  /// Returns the relative velocities along the clamping constraints
  Eigen::VectorXs getClampingConstraintRelativeVels();

  /// Returns the velocity change caused by illegal impulses in the LCP this
  /// timestep
  Eigen::VectorXs getVelocityDueToIllegalImpulses();

  /// Returns the velocity pre-LCP
  Eigen::VectorXs getPreLCPVelocity();

  /// Returns true if there were any bounces in this snapshot.
  bool hasBounces();

  /// Returns the number of contacts (regardless of state) in this snapshot.
  std::size_t getNumContacts();

  /// Returns the number of clamping contacts in this snapshot.
  std::size_t getNumClamping();

  /// Returns the number of upper bound contacts in this snapshot.
  std::size_t getNumUpperBound();

  /// These are the gradient constraint matrices from the LCP solver
  std::vector<std::shared_ptr<ConstrainedGroupGradientMatrices>>
      mGradientMatrices;

  /// This is the clamping constraints from all the constrained
  /// groups, concatenated together
  std::vector<std::shared_ptr<DifferentiableContactConstraint>>
  getDifferentiableConstraints();

  /// This is the clamping constraints from all the constrained
  /// groups, concatenated together
  std::vector<std::shared_ptr<DifferentiableContactConstraint>>
  getClampingConstraints();

  /// This is the upper bound constraints from all the constrained
  /// groups, concatenated together
  std::vector<std::shared_ptr<DifferentiableContactConstraint>>
  getUpperBoundConstraints();

  /// This verifies that the two matrices are equal to some tolerance, and if
  /// they're not it prints the information needed to replicated this scenario
  /// and it exits the program.
  void equalsOrCrash(
      std::shared_ptr<simulation::World> world,
      Eigen::MatrixXs analytical,
      Eigen::MatrixXs bruteForce,
      std::string name);

  /// This compares our analytical sub-Jacobians (like dMinv), to attempt to
  /// diagnose where there are differences creeping in between our finite
  /// differencing and our analytical results.
  void diagnoseSubJacobianErrors(
      std::shared_ptr<simulation::World> world, WithRespectTo* wrt);

  /// This prints code to the console to replicate a scenario for testing.
  void printReplicationInstructions(std::shared_ptr<simulation::World> world);

  /// Returns true if we were able to standardize our LCP results, false if we
  /// weren't
  bool areResultsStandardized() const;

  /// If this is true, we use finite-differencing to compute all of the
  /// requested Jacobians. This override can be useful to verify if there's a
  /// bug in the analytical Jacobians that's causing learning to not converge.
  void setUseFDOverride(bool override);

  /// If this is true, we check all Jacobians against their finite-differencing
  /// counterparts at runtime. If they aren't sufficiently close, we immediately
  /// crash the program and print what went wrong and some simple replication
  /// instructions.
  void setSlowDebugResultsAgainstFD(bool slowDebug);

  /// This does a battery of tests comparing the speeds to compute all the
  /// different Jacobians, both with finite differencing and analytically, and
  /// prints the results to std out.
  void benchmarkJacobians(
      std::shared_ptr<simulation::World> world, int numSamples);

protected:
  /// If this is true, we use finite-differencing to compute all of the
  /// requested Jacobians. This override can be useful to verify if there's a
  /// bug in the analytical Jacobians that's causing learning to not converge.
  bool mUseFDOverride;

  /// If this is true, we check all Jacobians against their finite-differencing
  /// counterparts at runtime. If they aren't sufficiently close, we immediately
  /// crash the program and print what went wrong and some simple replication
  /// instructions.
  bool mSlowDebugResultsAgainstFD;

  /// This is the global timestep length. This is included here because it shows
  /// up as a constant in some of the matrices.
  s_t mTimeStep;

  /// This is the total DOFs for this World
  std::size_t mNumDOFs;

  /// This is the number of total dimensions on all the constraints active in
  /// the world
  std::size_t mNumConstraintDim;

  /// This is the number of total constraint dimensions that are clamping
  std::size_t mNumClamping;

  /// This is the number of total constraint dimensions that are upper bounded
  std::size_t mNumUpperBound;

  /// This is the number of total constraint dimensions that are upper bounded
  std::size_t mNumBouncing;

  /// These are the offsets into the total degrees of freedom for each skeleton
  std::unordered_map<std::string, std::size_t> mSkeletonOffset;

  /// These are the number of degrees of freedom for each skeleton
  std::unordered_map<std::string, std::size_t> mSkeletonDofs;

  /// The position of all the DOFs of the world BEFORE the timestep
  Eigen::VectorXs mPreStepPosition;

  /// The velocities of all the DOFs of the world BEFORE the timestep
  Eigen::VectorXs mPreStepVelocity;

  /// The torques on all the DOFs of the world BEFORE the timestep
  Eigen::VectorXs mPreStepTorques;

  /// The LCP's initial cached value BEFORE the timestep
  Eigen::VectorXs mPreStepLCPCache;

  /// The velocities of all the DOFs of the world AFTER an unconstrained forward
  /// step, but BEFORE the LCP runs
  Eigen::VectorXs mPreConstraintVelocities;

  /// The position of all the DOFs of the world AFTER the timestep
  Eigen::VectorXs mPostStepPosition;

  /// The velocities of all the DOFs of the world AFTER the timestep
  /// created
  Eigen::VectorXs mPostStepVelocity;

  /// The torques on all the DOFs of the world AFTER the timestep
  Eigen::VectorXs mPostStepTorques;

private:
  /// These are mCached versions of the various Jacobians
  bool mCachedPosPosDirty;
  Eigen::MatrixXs mCachedPosPos;
  bool mCachedPosVelDirty;
  Eigen::MatrixXs mCachedPosVel;
  bool mCachedBounceApproximationDirty;
  Eigen::MatrixXs mCachedBounceApproximation;
  bool mCachedVelPosDirty;
  Eigen::MatrixXs mCachedVelPos;
  bool mCachedVelVelDirty;
  Eigen::MatrixXs mCachedVelVel;
  bool mCachedForcePosDirty;
  Eigen::MatrixXs mCachedForcePos;
  bool mCachedForceVelDirty;
  Eigen::MatrixXs mCachedForceVel;
  bool mCachedMassVelDirty;
  Eigen::MatrixXs mCachedMassVel;
  bool mCachedPosCDirty;
  Eigen::MatrixXs mCachedPosC;
  bool mCachedVelCDirty;
  Eigen::MatrixXs mCachedVelC;

  Eigen::VectorXs scratch(simulation::WorldPtr world);

  enum MatrixToAssemble
  {
    CLAMPING,
    MASSED_CLAMPING,
    UPPER_BOUND,
    MASSED_UPPER_BOUND,
    BOUNCING
  };

  Eigen::MatrixXs assembleMatrix(
      simulation::WorldPtr world, MatrixToAssemble whichMatrix);

  enum BlockDiagonalMatrixToAssemble
  {
    MASS,
    INV_MASS,
    POS_C,
    VEL_C
  };

  Eigen::MatrixXs assembleBlockDiagonalMatrix(
      simulation::WorldPtr world,
      BlockDiagonalMatrixToAssemble whichMatrix,
      bool forFiniteDifferencing = false);

  enum VectorToAssemble
  {
    CONTACT_CONSTRAINT_IMPULSES,
    CONTACT_CONSTRAINT_MAPPINGS,
    BOUNCE_DIAGONALS,
    RESTITUTION_DIAGONALS,
    PENETRATION_VELOCITY_HACK,
    CLAMPING_CONSTRAINT_IMPULSES,
    CLAMPING_CONSTRAINT_RELATIVE_VELS,
    VEL_DUE_TO_ILLEGAL,
    PRE_STEP_VEL,
    PRE_STEP_TAU,
    PRE_LCP_VEL,
    CFM_CONSTANTS
  };
  template <typename Vec>
  Vec assembleVector(VectorToAssemble whichVector);

  template <typename Vec>
  const Vec& getVectorToAssemble(
      std::shared_ptr<ConstrainedGroupGradientMatrices> matrices,
      VectorToAssemble whichVector);
};

using BackpropSnapshotPtr = std::shared_ptr<BackpropSnapshot>;

} // namespace neural
} // namespace dart

#endif


// This file is part of the dune-gdt project:
//   http://users.dune-project.org/projects/dune-gdt
// Copyright holders: Felix Schindler
// License: BSD 2-Clause License (http://opensource.org/licenses/BSD-2-Clause)
//
// Contributors: Kirsten Weber

#ifndef DUNE_GDT_EVALUATION_ELLIPTIC_HH
#define DUNE_GDT_EVALUATION_ELLIPTIC_HH

#include <tuple>

#include <boost/numeric/conversion/cast.hpp>

#include <dune/common/dynmatrix.hh>
#include <dune/common/typetraits.hh>

#include <dune/stuff/functions/interfaces.hh>

#include "interface.hh"

namespace Dune {
namespace GDT {
namespace LocalEvaluation {

// forward
template< class DiffusionFactorImp, class DiffusionTensorImp = void >
class Elliptic;

namespace internal {

/**
 * \brief Traits for the Elliptic evaluation (variant for given diffusion factor and tensor).
 * \sa    EllipticTraits (below) for a variant if only a diffusion is given.
 */
template< class DiffusionFactorType, class DiffusionTensorType >
class EllipticTraits
{
  static_assert(Stuff::is_localizable_function< DiffusionFactorType >::value,
                "DiffusionFactorType has to be a localizable function!");
  static_assert(Stuff::is_localizable_function< DiffusionTensorType >::value,
                "DiffusionTensorType has to be a localizable function!");
  static_assert(std::is_same< typename DiffusionFactorType::EntityType,
                              typename DiffusionTensorType::EntityType >::value,
                "EntityTypes have to agree!");
  static_assert(std::is_same< typename DiffusionFactorType::DomainFieldType,
                              typename DiffusionTensorType::DomainFieldType >::value,
                "DomainFieldTypes have to agree!");
  static_assert(DiffusionFactorType::dimDomain == DiffusionTensorType::dimDomain,
                "Dimensions have to agree!");
public:
  typedef Elliptic< DiffusionFactorType, DiffusionTensorType > derived_type;
  typedef std::tuple< std::shared_ptr< typename DiffusionFactorType::LocalfunctionType >,
                      std::shared_ptr< typename DiffusionTensorType::LocalfunctionType > > LocalfunctionTupleType;
  typedef typename DiffusionFactorType::EntityType      EntityType;
  typedef typename DiffusionFactorType::DomainFieldType DomainFieldType;
  static const size_t                                   dimDomain = DiffusionFactorType::dimDomain;
}; // class EllipticTraits

/**
 * \brief Traits for the Elliptic evaluation (variant for a given diffusion).
 * \note  It does not matter if that function plays the role of the diffusion factor (scalar) or the diffusion
 *        tensor (matrix).
 * \sa    EllipticTraits (above) for a variant if a diffusion factor and a diffusion tensor is given.
 */
template< class DiffusionType >
class EllipticTraits< DiffusionType, void >
{
  static_assert(Stuff::is_localizable_function< DiffusionType >::value,
                "DiffusionType has to be a localizable function!");
public:
  typedef Elliptic< DiffusionType, void >         derived_type;
  typedef typename DiffusionType::EntityType      EntityType;
  typedef typename DiffusionType::DomainFieldType DomainFieldType;
  typedef std::tuple< std::shared_ptr< typename DiffusionType::LocalfunctionType > > LocalfunctionTupleType;
  static const size_t                             dimDomain = DiffusionType::dimDomain;
}; // class EllipticTraits< ..., void >

} // namespace internal

/**
 * \brief Computes an elliptic evaluation (variant for given diffusion factor and tensor).
 * \sa    Elliptic (below) for a variant if only a diffusion is given.
 */
template< class DiffusionFactorImp, class DiffusionTensorImp >
class Elliptic
  : public LocalEvaluation::Codim0Interface< internal::EllipticTraits< DiffusionFactorImp, DiffusionTensorImp >, 2 >
{
public:
  typedef DiffusionFactorImp                                                   DiffusionFactorType;
  typedef DiffusionTensorImp                                                   DiffusionTensorType;
  typedef internal::EllipticTraits< DiffusionFactorType, DiffusionTensorType > Traits;
  typedef typename Traits::LocalfunctionTupleType                              LocalfunctionTupleType;
  typedef typename Traits::EntityType                                          EntityType;
  typedef typename Traits::DomainFieldType                                     DomainFieldType;
  static const size_t                                                          dimDomain = Traits::dimDomain;

  Elliptic(const DiffusionFactorType& diffusion_factor,
           const DiffusionTensorType& diffusion_tensor)
    : diffusion_factor_(diffusion_factor)
    , diffusion_tensor_(diffusion_tensor)
  {}

  LocalfunctionTupleType localFunctions(const EntityType& entity) const
  {
    return std::make_tuple(diffusion_factor_.local_function(entity),
                           diffusion_tensor_.local_function(entity));
  }

  /**
   * \brief extracts the local functions and calls the correct order() method
   */
  template< class R, size_t rT, size_t rCT, size_t rA, size_t rCA >
  size_t order(const LocalfunctionTupleType& local_functions_tuple,
               const Stuff::LocalfunctionSetInterface
                   < EntityType, DomainFieldType, dimDomain, R, rT, rCT >& testBase,
               const Stuff::LocalfunctionSetInterface
                   < EntityType, DomainFieldType, dimDomain, R, rA, rCA >& ansatzBase) const
  {
    const auto local_diffusion_factor = std::get< 0 >(local_functions_tuple);
    const auto local_diffusion_tensor = std::get< 1 >(local_functions_tuple);
    return order(*local_diffusion_factor, *local_diffusion_tensor, testBase, ansatzBase);
  } // ... order(...)

  /**
   * \brief extracts the local functions and calls the correct evaluate() method
   */
  template< class R, size_t rT, size_t rCT, size_t rA, size_t rCA >
  void evaluate(const LocalfunctionTupleType& local_functions_tuple,
                const Stuff::LocalfunctionSetInterface
                    < EntityType, DomainFieldType, dimDomain, R, rT, rCT >& testBase,
                const Stuff::LocalfunctionSetInterface
                    < EntityType, DomainFieldType, dimDomain, R, rA, rCA >& ansatzBase,
                const Dune::FieldVector< DomainFieldType, dimDomain >& localPoint,
                Dune::DynamicMatrix< R >& ret) const
  {
    const auto local_diffusion_factor = std::get< 0 >(local_functions_tuple);
    const auto local_diffusion_tensor = std::get< 1 >(local_functions_tuple);
    evaluate(*local_diffusion_factor, *local_diffusion_tensor, testBase, ansatzBase, localPoint, ret);
  }

private:
  template< class R, size_t rDF, size_t rCDF, size_t rDT, size_t rCDT, size_t rT, size_t rCT, size_t rA, size_t rCA >
  size_t order(const Stuff::LocalfunctionInterface
                   < EntityType, DomainFieldType, dimDomain, R, rDF, rCDF >& local_diffusion_factor,
               const Stuff::LocalfunctionInterface
                   < EntityType, DomainFieldType, dimDomain, R, rDT, rCDT >& local_diffusion_tensor,
               const Stuff::LocalfunctionSetInterface
                   < EntityType, DomainFieldType, dimDomain, R, rT, rCT >& testBase,
               const Stuff::LocalfunctionSetInterface
                   < EntityType, DomainFieldType, dimDomain, R, rA, rCA >& ansatzBase) const
  {
    return local_diffusion_factor.order()
        + local_diffusion_tensor.order()
        + std::max(ssize_t(testBase.order()) - 1, ssize_t(0))
        + std::max(ssize_t(ansatzBase.order()) - 1, ssize_t(0));
  } // ... order(...)

  template< class R, size_t r >
  void evaluate(const Stuff::LocalfunctionInterface
                    < EntityType, DomainFieldType, dimDomain, R, 1, 1 >& local_diffusion_factor,
                const Stuff::LocalfunctionInterface
                    < EntityType, DomainFieldType, dimDomain, R, 1, 1 >& local_diffusion_tensor,
                const Stuff::LocalfunctionSetInterface
                    < EntityType, DomainFieldType, dimDomain, R, r, 1 >& testBase,
                const Stuff::LocalfunctionSetInterface
                    < EntityType, DomainFieldType, dimDomain, R, r, 1 >& ansatzBase,
                const Dune::FieldVector< DomainFieldType, dimDomain >& localPoint,
                Dune::DynamicMatrix< R >& ret) const
  {
    // evaluate local functions
    const auto local_diffusion_factor_value = local_diffusion_factor.evaluate(localPoint);
    const auto local_diffusion_tensor_value = local_diffusion_tensor.evaluate(localPoint);
    // evaluate test gradient
    const auto rows = testBase.size();
    const auto testGradients = testBase.jacobian(localPoint);
    // evaluate ansatz gradient
    const auto cols = ansatzBase.size();
    const auto ansatzGradients = ansatzBase.jacobian(localPoint);
    // compute products
    assert(ret.rows() >= rows);
    assert(ret.cols() >= cols);
    for (size_t ii = 0; ii < rows; ++ii) {
      auto& retRow = ret[ii];
      for (size_t jj = 0; jj < cols; ++jj) {
        retRow[jj] = local_diffusion_factor_value * local_diffusion_tensor_value
                     * (ansatzGradients[jj][0] * testGradients[ii][0]);
      }
    }
  } // ... evaluate(...)

  template< class R >
  void evaluate(const Stuff::LocalfunctionInterface
                    < EntityType, DomainFieldType, dimDomain, R, 1, 1 >& local_diffusion_factor,
                const Stuff::LocalfunctionInterface
                    < EntityType, DomainFieldType, dimDomain, R, 2, 2 >& local_diffusion_tensor,
                const Stuff::LocalfunctionSetInterface
                    < EntityType, DomainFieldType, dimDomain, R, 1, 1 >& testBase,
                const Stuff::LocalfunctionSetInterface
                    < EntityType, DomainFieldType, dimDomain, R, 1, 1 >& ansatzBase,
                const Dune::FieldVector< DomainFieldType, dimDomain >& localPoint,
                Dune::DynamicMatrix< R >& ret) const
  {
    evaluate_matrix_valued_(local_diffusion_factor, local_diffusion_tensor, testBase, ansatzBase, localPoint, ret);
  }

  template< class R >
  void evaluate(const Stuff::LocalfunctionInterface
                    < EntityType, DomainFieldType, dimDomain, R, 1, 1 >& local_diffusion_factor,
                const Stuff::LocalfunctionInterface
                    < EntityType, DomainFieldType, dimDomain, R, 3, 3 >& local_diffusion_tensor,
                const Stuff::LocalfunctionSetInterface
                    < EntityType, DomainFieldType, dimDomain, R, 1, 1 >& testBase,
                const Stuff::LocalfunctionSetInterface
                    < EntityType, DomainFieldType, dimDomain, R, 1, 1 >& ansatzBase,
                const Dune::FieldVector< DomainFieldType, dimDomain >& localPoint,
                Dune::DynamicMatrix< R >& ret) const
  {
    evaluate_matrix_valued_(local_diffusion_factor, local_diffusion_tensor, testBase, ansatzBase, localPoint, ret);
  }

  template< class R >
  void evaluate_matrix_valued_(const Stuff::LocalfunctionInterface
                                   < EntityType, DomainFieldType, dimDomain, R, 1, 1 >& local_diffusion_factor,
                               const Stuff::LocalfunctionInterface
                                   < EntityType, DomainFieldType, dimDomain, R, dimDomain, dimDomain >& local_diffusion_tensor,
                               const Stuff::LocalfunctionSetInterface
                                   < EntityType, DomainFieldType, dimDomain, R, 1, 1 >& testBase,
                               const Stuff::LocalfunctionSetInterface
                                   < EntityType, DomainFieldType, dimDomain, R, 1, 1 >& ansatzBase,
                               const Dune::FieldVector< DomainFieldType, dimDomain >& localPoint,
                               Dune::DynamicMatrix< R >& ret) const
  {
    typedef typename Stuff::LocalfunctionSetInterface
        < EntityType, DomainFieldType, dimDomain, R, 1, 1 >::JacobianRangeType JacobianRangeType;
    // evaluate local functions
    const auto local_diffusion_factor_value = local_diffusion_factor.evaluate(localPoint);
    auto local_diffusion_tensor_value = local_diffusion_tensor.evaluate(localPoint);
    local_diffusion_tensor_value *= local_diffusion_factor_value[0];
    // evaluate test gradient
    const size_t rows = testBase.size();
    std::vector< JacobianRangeType > testGradients(rows, JacobianRangeType(0));
    testBase.jacobian(localPoint, testGradients);
    // evaluate ansatz gradient
    const size_t cols = ansatzBase.size();
    std::vector< JacobianRangeType > ansatzGradients(cols, JacobianRangeType(0));
    ansatzBase.jacobian(localPoint, ansatzGradients);
    // compute products
    assert(ret.rows() >= rows);
    assert(ret.cols() >= cols);
    FieldVector< DomainFieldType, dimDomain > product(0.0);
    for (size_t ii = 0; ii < rows; ++ii) {
      auto& retRow = ret[ii];
      for (size_t jj = 0; jj < cols; ++jj) {
        local_diffusion_tensor_value.mv(ansatzGradients[jj][0], product);
        retRow[jj] = product * testGradients[ii][0];
      }
    }
  } // ... evaluate_matrix_valued_(...)

  const DiffusionFactorType& diffusion_factor_;
  const DiffusionTensorType& diffusion_tensor_;
}; // class Elliptic

/**
 * \brief Computes an elliptic evaluation (variant for a given diffusion).
 * \note  It does not matter if that function plays the role of the diffusion factor (scalar) or the diffusion
 *        tensor (matrix).
 * \sa    Elliptic (above) for a variant if a diffusion factor and a diffusion tensor is given.
 */
template< class DiffusionImp >
class Elliptic< DiffusionImp, void >
  : public LocalEvaluation::Codim0Interface< internal::EllipticTraits< DiffusionImp, void >, 2 >
{
public:
  typedef DiffusionImp                                    DiffusionType;
  typedef internal::EllipticTraits< DiffusionType, void > Traits;
  typedef typename Traits::LocalfunctionTupleType         LocalfunctionTupleType;
  typedef typename Traits::EntityType                     EntityType;
  typedef typename Traits::DomainFieldType                DomainFieldType;
  static const size_t                                     dimDomain = Traits::dimDomain;

  explicit Elliptic(const DiffusionType& inducingFunction)
    : diffusion_(inducingFunction)
  {}

  /// \name Required by LocalEvaluation::Codim0Interface< ..., 2 >
  /// \{

  LocalfunctionTupleType localFunctions(const EntityType& entity) const
  {
    return std::make_tuple(diffusion_.local_function(entity));
  }

  /**
   * \brief extracts the local functions and calls the correct order() method
   */
  template< class R, size_t rT, size_t rCT, size_t rA, size_t rCA >
  size_t order(const LocalfunctionTupleType& localFuncs,
               const Stuff::LocalfunctionSetInterface
                   < EntityType, DomainFieldType, dimDomain, R, rT, rCT >& testBase,
               const Stuff::LocalfunctionSetInterface
                   < EntityType, DomainFieldType, dimDomain, R, rA, rCA >& ansatzBase) const
  {
    return order(*std::get< 0 >(localFuncs), testBase, ansatzBase);
  }

  /**
   * \brief extracts the local functions and calls the correct evaluate() method
   */
  template< class R, size_t rT, size_t rCT, size_t rA, size_t rCA >
  void evaluate(const LocalfunctionTupleType& localFuncs,
                const Stuff::LocalfunctionSetInterface
                    < EntityType, DomainFieldType, dimDomain, R, rT, rCT >& testBase,
                const Stuff::LocalfunctionSetInterface
                    < EntityType, DomainFieldType, dimDomain, R, rA, rCA >& ansatzBase,
                const Dune::FieldVector< DomainFieldType, dimDomain >& localPoint,
                Dune::DynamicMatrix< R >& ret) const
  {
    evaluate(*std::get< 0 >(localFuncs), testBase, ansatzBase, localPoint, ret);
  }

  /// \}
  /// \name Actual implementations of order
  /// \{

  /**
   *  \return localFunction.order() + (testBase.order() - 1) + (ansatzBase.order() - 1)
   */
  template< class R, size_t rL, size_t rCL, size_t rT, size_t rCT, size_t rA, size_t rCA >
  size_t order(const Stuff::LocalfunctionInterface< EntityType, DomainFieldType, dimDomain, R, rL, rCL >& localFunction,
               const Stuff::LocalfunctionSetInterface< EntityType, DomainFieldType, dimDomain, R, rT, rCT >& testBase,
               const Stuff::LocalfunctionSetInterface< EntityType, DomainFieldType, dimDomain, R, rA, rCA >& ansatzBase)
  const
  {
    return localFunction.order()
        + boost::numeric_cast< size_t >(std::max(ssize_t(testBase.order())   - 1, ssize_t(0)))
        + boost::numeric_cast< size_t >(std::max(ssize_t(ansatzBase.order()) - 1, ssize_t(0)));
  } // ... order( ... )

  /// \}
  /// \name Actual implementations of evaluate
  /// \{

  /**
   *  \brief  Computes an elliptic evaluation for a scalar local function and scalar or vector valued basefunctionsets.
   *  \tparam R RangeFieldType
   */
  template< class R, size_t r >
  void evaluate(const Stuff::LocalfunctionInterface< EntityType, DomainFieldType, dimDomain, R, 1, 1 >& localFunction,
                const Stuff::LocalfunctionSetInterface< EntityType, DomainFieldType, dimDomain, R, r, 1 >& testBase,
                const Stuff::LocalfunctionSetInterface< EntityType, DomainFieldType, dimDomain, R, r, 1 >& ansatzBase,
                const Dune::FieldVector< DomainFieldType, dimDomain >& localPoint,
                Dune::DynamicMatrix< R >& ret) const
  {
    typedef typename Stuff::LocalfunctionSetInterface
        < EntityType, DomainFieldType, dimDomain, R, r, 1 >::JacobianRangeType JacobianRangeType;
    // evaluate local function
    const auto functionValue = localFunction.evaluate(localPoint);
    // evaluate test gradient
    const size_t rows = testBase.size();
    std::vector< JacobianRangeType > testGradients(rows, JacobianRangeType(0));
    testBase.jacobian(localPoint, testGradients);
    // evaluate ansatz gradient
    const size_t cols = ansatzBase.size();
    std::vector< JacobianRangeType > ansatzGradients(cols, JacobianRangeType(0));
    ansatzBase.jacobian(localPoint, ansatzGradients);
    // compute products
    assert(ret.rows() >= rows);
    assert(ret.cols() >= cols);
    for (size_t ii = 0; ii < rows; ++ii) {
      auto& retRow = ret[ii];
      for (size_t jj = 0; jj < cols; ++jj) {
        retRow[jj] = functionValue * (ansatzGradients[jj][0] * testGradients[ii][0]);
      }
    }
  } // ... evaluate< ..., 1, ... >(...)

  /**
   *  \brief  Computes an elliptic evaluation for a 2x2 matrix-valued local function and matrix-valued basefunctionsets.
   *  \tparam R RangeFieldType
   *  \note   Unfortunately we need this explicit specialization, otherwise the compiler will complain for 1d grids.
   */
  template< class R >
  void evaluate(const Stuff::LocalfunctionInterface< EntityType, DomainFieldType, dimDomain, R, 2, 2 >& localFunction,
                const Stuff::LocalfunctionSetInterface< EntityType, DomainFieldType, dimDomain, R, 1, 1 >& testBase,
                const Stuff::LocalfunctionSetInterface< EntityType, DomainFieldType, dimDomain, R, 1, 1 >& ansatzBase,
                const Dune::FieldVector< DomainFieldType, dimDomain >& localPoint,
                Dune::DynamicMatrix< R >& ret) const
  {
    evaluate_matrix_valued_(localFunction, testBase, ansatzBase, localPoint, ret);
  }

  /**
   *  \brief  Computes an elliptic evaluation for a 3x3 matrix-valued local function and matrix-valued basefunctionsets.
   *  \tparam R RangeFieldType
   *  \note   Unfortunately we need this explicit specialization, otherwise the compiler will complain for 1d grids.
   */
  template< class R >
  void evaluate(const Stuff::LocalfunctionInterface< EntityType, DomainFieldType, dimDomain, R, 3, 3 >& localFunction,
                const Stuff::LocalfunctionSetInterface< EntityType, DomainFieldType, dimDomain, R, 1, 1 >& testBase,
                const Stuff::LocalfunctionSetInterface< EntityType, DomainFieldType, dimDomain, R, 1, 1 >& ansatzBase,
                const Dune::FieldVector< DomainFieldType, dimDomain >& localPoint,
                Dune::DynamicMatrix< R >& ret) const
  {
    evaluate_matrix_valued_(localFunction, testBase, ansatzBase, localPoint, ret);
  }

  /// \}

private:
  template< class R >
  void evaluate_matrix_valued_(const Stuff::LocalfunctionInterface
                                   < EntityType, DomainFieldType, dimDomain, R, dimDomain, dimDomain >& localFunction,
                               const Stuff::LocalfunctionSetInterface
                                   < EntityType, DomainFieldType, dimDomain, R, 1, 1 >& testBase,
                               const Stuff::LocalfunctionSetInterface
                                   < EntityType, DomainFieldType, dimDomain, R, 1, 1 >& ansatzBase,
                               const Dune::FieldVector< DomainFieldType, dimDomain >& localPoint,
                               Dune::DynamicMatrix< R >& ret) const
  {
    // evaluate local function
    const auto functionValue = localFunction.evaluate(localPoint);
    // evaluate test gradient
    const size_t rows = testBase.size();
    const auto testGradients = testBase.jacobian(localPoint);
    // evaluate ansatz gradient
    const size_t cols = ansatzBase.size();
    const auto ansatzGradients = ansatzBase.jacobian(localPoint);
    // compute products
    assert(ret.rows() >= rows);
    assert(ret.cols() >= cols);
    FieldVector< DomainFieldType, dimDomain > product(0.0);
    for (size_t ii = 0; ii < rows; ++ii) {
      auto& retRow = ret[ii];
      for (size_t jj = 0; jj < cols; ++jj) {
        functionValue.mv(ansatzGradients[jj][0], product);
        retRow[jj] = product * testGradients[ii][0];
      }
    }
  } // ... evaluate_matrix_valued_(...)

  const DiffusionType& diffusion_;
}; // class Elliptic< ...., void >

} // namespace LocalEvaluation
} // namespace GDT
} // namespace Dune

#endif // DUNE_GDT_EVALUATION_ELLIPTIC_HH


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

#ifndef CRYPTO3_BLOCK_KASUMI_HPP
#define CRYPTO3_BLOCK_KASUMI_HPP

#include <boost/endian/arithmetic.hpp>
#include <boost/endian/conversion.hpp>

#include <nil/crypto3/block/detail/kasumi/kasumi_functions.hpp>

#include <nil/crypto3/block/detail/block_stream_processor.hpp>
#include <nil/crypto3/block/detail/cipher_modes.hpp>

namespace nil {
    namespace crypto3 {
        namespace block {
            /*!
             * @brief Kasumi. A 64-bit cipher used in 3GPP mobile phone protocols.
             * There is no reason to use it outside of this context.
             *
             * @ingroup block
             */
            class kasumi {
            protected:
                typedef detail::kasumi_functions policy_type;

                constexpr static const std::size_t key_schedule_size = policy_type::key_schedule_size;
                typedef typename policy_type::key_schedule_type key_schedule_type;

            public:
                constexpr static const std::size_t rounds = policy_type::rounds;

                constexpr static const std::size_t word_bits = policy_type::word_bits;
                typedef typename policy_type::word_type word_type;

                constexpr static const std::size_t block_bits = policy_type::block_bits;
                constexpr static const std::size_t block_words = policy_type::block_words;
                typedef typename policy_type::block_type block_type;

                constexpr static const std::size_t key_bits = policy_type::key_bits;
                constexpr static const std::size_t key_words = policy_type::key_words;
                typedef typename policy_type::key_type key_type;

                template<class Mode, typename StateAccumulator, std::size_t ValueBits>
                struct stream_processor {
                    struct params_type {

                        constexpr static const std::size_t value_bits = ValueBits;
                        constexpr static const std::size_t length_bits = policy_type::word_bits * 2;
                    };

                    typedef block_stream_processor<Mode, StateAccumulator, params_type> type;
                };

                typedef typename stream_endian::little_octet_big_bit endian_type;

                kasumi(const key_type &key) {
                    schedule_key(key_schedule, key);
                }

                ~kasumi() {
                    key_schedule.fill(0);
                }

                inline block_type encrypt(const block_type &plaintext) const {
                    return encrypt_block(plaintext, key_schedule);
                }

                inline block_type decrypt(const block_type &ciphertext) const {
                    return decrypt_block(ciphertext, key_schedule);
                }

            protected:
                inline block_type encrypt_block(const block_type &plaintext,
                                                const key_schedule_type &key_schedule) const {
                    word_type B0 = boost::endian::native_to_big(plaintext[0]);
                    word_type B1 = boost::endian::native_to_big(plaintext[1]);
                    word_type B2 = boost::endian::native_to_big(plaintext[2]);
                    word_type B3 = boost::endian::native_to_big(plaintext[3]);

                    for (size_t j = 0; j != rounds; j += 2) {
                        const word_type *K = &key_schedule[8 * j];

                        word_type R = B1 ^ (policy_type::template rotl<1>(B0) & K[0]);
                        word_type L = B0 ^ (policy_type::template rotl<1>(R) | K[1]);

                        L = policy_type::FI(L ^ K[2], K[3]) ^ R;
                        R = policy_type::FI(R ^ K[4], K[5]) ^ L;
                        L = policy_type::FI(L ^ K[6], K[7]) ^ R;

                        R = B2 ^= R;
                        L = B3 ^= L;

                        R = policy_type::FI(R ^ K[10], K[11]) ^ L;
                        L = policy_type::FI(L ^ K[12], K[13]) ^ R;
                        R = policy_type::FI(R ^ K[14], K[15]) ^ L;

                        R ^= (policy_type::template rotl<1>(L) & K[8]);
                        L ^= (policy_type::template rotl<1>(R) | K[9]);

                        B0 ^= L;
                        B1 ^= R;
                    }

                    return {boost::endian::big_to_native(B0), boost::endian::big_to_native(B1),
                            boost::endian::big_to_native(B2), boost::endian::big_to_native(B3)};
                }

                inline block_type decrypt_block(const block_type &ciphertext,
                                                const key_schedule_type &key_schedule) const {
                    word_type B0 = boost::endian::native_to_big(ciphertext[0]);
                    word_type B1 = boost::endian::native_to_big(ciphertext[1]);
                    word_type B2 = boost::endian::native_to_big(ciphertext[2]);
                    word_type B3 = boost::endian::native_to_big(ciphertext[3]);

                    for (size_t j = 0; j != rounds; j += 2) {
                        const word_type *K = &key_schedule[8 * (6 - j)];

                        word_type L = B2, R = B3;

                        L = policy_type::FI(L ^ K[10], K[11]) ^ R;
                        R = policy_type::FI(R ^ K[12], K[13]) ^ L;
                        L = policy_type::FI(L ^ K[14], K[15]) ^ R;

                        L ^= (policy_type::template rotl<1>(R) & K[8]);
                        R ^= (policy_type::template rotl<1>(L) | K[9]);

                        R = B0 ^= R;
                        L = B1 ^= L;

                        L ^= (policy_type::template rotl<1>(R) & K[0]);
                        R ^= (policy_type::template rotl<1>(L) | K[1]);

                        R = policy_type::FI(R ^ K[2], K[3]) ^ L;
                        L = policy_type::FI(L ^ K[4], K[5]) ^ R;
                        R = policy_type::FI(R ^ K[6], K[7]) ^ L;

                        B2 ^= L;
                        B3 ^= R;
                    }

                    return {boost::endian::big_to_native(B0), boost::endian::big_to_native(B1),
                            boost::endian::big_to_native(B2), boost::endian::big_to_native(B3)};
                }

                key_schedule_type key_schedule;

                void schedule_key(key_schedule_type &key_schedule, const key_type &key) {
                    std::array<word_type, 16> K = {0};
                    for (size_t i = 0; i != rounds; ++i) {
                        K[i] = boost::endian::native_to_big(key[i]);
                        K[i + 8] = K[i] ^ policy_type::round_constants[i];
                    }

                    for (size_t i = 0; i != rounds; ++i) {
                        key_schedule[8 * i] = policy_type::template rotl<2>(K[(i + 0) % 8]);
                        key_schedule[8 * i + 1] = policy_type::template rotl<1>(K[(i + 2) % 8 + 8]);
                        key_schedule[8 * i + 2] = policy_type::template rotl<5>(K[(i + 1) % 8]);
                        key_schedule[8 * i + 3] = K[(i + 4) % 8 + 8];
                        key_schedule[8 * i + 4] = policy_type::template rotl<8>(K[(i + 5) % 8]);
                        key_schedule[8 * i + 5] = K[(i + 3) % 8 + 8];
                        key_schedule[8 * i + 6] = policy_type::template rotl<13>(K[(i + 6) % 8]);
                        key_schedule[8 * i + 7] = K[(i + 7) % 8 + 8];
                    }

                    K.fill(0);
                }
            };
        }    // namespace block
    }        // namespace crypto3
}    // namespace nil
#endif