#include <arti_base_control/vehicle.h>
#include <arti_base_control/axle.h>
#include <arti_base_control/utils.h>
#include <boost/range/irange.hpp>
#include <Eigen/Core>
#include <Eigen/QR>
#include <functional>
#include <memory>

namespace arti_base_control
{
VehicleVelocityConstraint::VehicleVelocityConstraint(double a_v_x_, double a_v_y_, double a_v_theta_, double b_)
  : a_v_x(a_v_x_), a_v_y(a_v_y_), a_v_theta(a_v_theta_), b(b_)
{
}

Vehicle::Vehicle(const ros::NodeHandle& nh, const JointActuatorFactoryPtr& motor_factory, bool process_ackermann)
  : nh_(nh), motor_factory_(motor_factory), reconfigure_server_(nh), process_ackermann_(process_ackermann)
{
  reconfigure_server_.setCallback(std::bind(&Vehicle::reconfigure, this, std::placeholders::_1));
}

void Vehicle::reconfigure(VehicleConfig& config)
{
  config_ = config;

  if (config.max_velocity_linear == 0.0)
  {
    ROS_ERROR("Parameter max_velocity_linear is not set");
  }

  if (config.allowed_brake_velocity == 0.0)
  {
    ROS_WARN("Parameter allowed_brake_velocity is not set");
  }

  if (config.brake_velocity == 0.0)
  {
    ROS_WARN("Parameter brake_velocity is not set");
  }

  if (config.brake_current == 0.0)
  {
    ROS_WARN("Parameter brake_current is not set");
  }

  if (axles_.empty())
  {
    XmlRpc::XmlRpcValue axles_param;
    if (nh_.getParam("axles", axles_param))
    {
      if (axles_param.getType() == XmlRpc::XmlRpcValue::TypeStruct)
      {
        const ros::NodeHandle axles_nh(nh_, "axles");
        for (const XmlRpc::XmlRpcValue::ValueStruct::value_type& axle_param : axles_param)
        {
          axles_.emplace_back(
            std::make_shared<Axle>(ros::NodeHandle(axles_nh, axle_param.first), config_, motor_factory_));
        }
      }
      else
      {
        ROS_ERROR_STREAM("axles parameter has invalid type, must be map");
      }
    }
    else
    {
      ROS_ERROR_STREAM("axles parameter is missing");
    }
  }

  wheelbase_ = config_.wheelbase;
  if (process_ackermann_)
  {
    if (wheelbase_ == 0.0)
    {
      for (const AxlePtr& axle : axles_)
      {
        const AxleConfig& axle_config = axle->getConfig();

        if (axle_config.is_steered)
        {
          wheelbase_ = std::fabs(axle_config.position_x - config_.icr_x);
          if (wheelbase_ != 0.0)
          {
            break;
          }
        }
      }

      if (wheelbase_ == 0.0)
      {
        ROS_WARN_STREAM("Wheelbase is not set and could not be determined automatically, this prevents control via"
                        " Ackermann messages");
      }
    }
  }

  for (const AxlePtr& axle : axles_)
  {
    axle->setVehicleConfig(config_);
  }
}

void Vehicle::setVelocity(const ackermann_msgs::AckermannDrive& velocity, const ros::Time& time)
{
  if (!process_ackermann_)
  {
    ROS_ERROR("got ackerman command but should not process ackerman commands");
    return;
  }

  const double steering_angle = limit(normalizeSteeringAngle(velocity.steering_angle), config_.max_steering_angle,
                                      "steering angle");
  const double sin_steering_angle = std::sin(steering_angle);
  const double cos_steering_angle = std::cos(steering_angle);

  // Limit linear velocity to stay below angular velocity limit
  // (angular_velocity = tan(steering_angle) * linear_velocity / wheelbase; this can become infinite when
  // steering_angle approaches 90 degrees):
  double angular_velocity = 0.0;
  double linear_velocity = limit(velocity.speed, config_.max_velocity_linear, "linear velocity");
  if (linear_velocity != 0.0 && wheelbase_ != 0.0)
  {
    const double a = std::fabs(linear_velocity * sin_steering_angle);
    const double b = std::fabs(config_.max_velocity_angular * wheelbase_ * cos_steering_angle);
    if (a > b)
    {
      const double limited_linear_velocity = linear_velocity * b / a;
      ROS_WARN_NAMED("limit", "linear velocity (%f) exceeded maximum (%f) due to angular velocity constraint and was"
                              " limited", linear_velocity, limited_linear_velocity);
      linear_velocity = limited_linear_velocity;
      angular_velocity = config_.max_velocity_angular * (linear_velocity < 0.0 ? -1.0 : 1.0)
                         * (steering_angle < 0.0 ? -1.0 : 1.0);
    }
    else
    {
      angular_velocity = linear_velocity * sin_steering_angle / (wheelbase_ * cos_steering_angle);
    }
  }

  for (const AxlePtr& axle : axles_)
  {
    const AxleConfig& axle_config = axle->getConfig();

    double axle_steering_angle = 0.0;
    if (axle_config.is_steered)
    {
      axle_steering_angle = normalizeSteeringAngle(
        std::atan2(sin_steering_angle * (axle_config.position_x - config_.icr_x), cos_steering_angle * wheelbase_));
    }

    axle->setVelocity(linear_velocity, angular_velocity, axle_steering_angle, time);
  }
}

void Vehicle::setVelocity(const geometry_msgs::Twist& velocity, const ros::Time& time)
{
  const double linear_velocity = limit(velocity.linear.x, config_.max_velocity_linear, "linear velocity");
  double angular_velocity = limit(velocity.angular.z, config_.max_velocity_angular, "angular velocity");

  if (angular_velocity != 0)
  {
    const double a = std::fabs(std::cos(config_.max_steering_angle) * wheelbase_ * angular_velocity);
    const double b = std::fabs(std::sin(config_.max_steering_angle) * linear_velocity);
    if (a > b)
    {
      const double limited_angular_velocity = angular_velocity * b / a;
      ROS_WARN_NAMED("limit", "angular velocity (%f) exceeded maximum (%f) due to steering angle constraint and was"
                              " limited", angular_velocity, limited_angular_velocity);
      angular_velocity = limited_angular_velocity;
    }
  }

  for (const AxlePtr& axle : axles_)
  {
    const AxleConfig& axle_config = axle->getConfig();

    double axle_steering_angle = 0.0;
    if (axle_config.is_steered)
    {
      axle_steering_angle = normalizeSteeringAngle(
        std::atan2((axle_config.position_x - config_.icr_x) * angular_velocity, linear_velocity));
    }

    axle->setVelocity(linear_velocity, angular_velocity, axle_steering_angle, time);
  }
}

VehicleState Vehicle::getState(const ros::Time& time) const
{
  VehicleState state;
  for (const AxlePtr& axle : axles_)
  {
    state.axle_states.emplace_back(axle->getState(time));
  }
  return state;
}

void Vehicle::getVelocity(const VehicleState& state, geometry_msgs::Twist& velocity) const
{
  if (state.axle_states.size() != axles_.size())
  {
    throw std::runtime_error("number of axle states in vehicle state differs from number of vehicle axles");
  }

  // Compute vehicle velocity using pseudo inverse of velocity constraints:
  VehicleVelocityConstraints constraints;
  for (const size_t i : boost::irange<size_t>(0, axles_.size()))
  {
    axles_.at(i)->getVelocityConstraints(state.axle_states.at(i), constraints);
  }

  Eigen::MatrixXd a(Eigen::MatrixXd::Zero(constraints.size(), 3));
  Eigen::VectorXd b(Eigen::VectorXd::Zero(constraints.size()));
  for (size_t i = 0; i < constraints.size(); ++i)
  {
    a(i, 0) = constraints[i].a_v_x;
    a(i, 1) = constraints[i].a_v_y;
    a(i, 2) = constraints[i].a_v_theta;
    b(i) = constraints[i].b;
  }

  const Eigen::Vector3d x = a.fullPivHouseholderQr().solve(b);

  velocity.linear.x = x(0);
  velocity.linear.y = x(1);
  velocity.linear.z = 0.0;
  velocity.angular.x = 0.0;
  velocity.angular.y = 0.0;
  velocity.angular.z = x(2);
}

void Vehicle::getVelocity(const VehicleState& state, ackermann_msgs::AckermannDrive& velocity) const
{
  if (state.axle_states.size() != axles_.size())
  {
    throw std::runtime_error("number of axle states in vehicle state differs from number of vehicle axles");
  }

  geometry_msgs::Twist twist;
  getVelocity(state, twist);
  velocity.speed = twist.linear.x;

  // Compute steering angle as average of axle's steering angles:
  double accumulated_steering_angle = 0;
  size_t steered_axles_count = 0;

  for (const size_t i : boost::irange<size_t>(0, axles_.size()))
  {
    const AxleState& axle_state = state.axle_states.at(i);
    if (axle_state.steering_motor_state)
    {
      const double axle_steering_angle = axle_state.steering_motor_state->position;
      const AxleConfig& axle_config = axles_.at(i)->getConfig();

      accumulated_steering_angle += normalizeSteeringAngle(std::atan2(
        std::sin(axle_steering_angle) * wheelbase_,
        std::cos(axle_steering_angle) * (axle_config.position_x - config_.icr_x)));
      ++steered_axles_count;
    }
  }

  if (steered_axles_count != 0)
  {
    velocity.steering_angle = accumulated_steering_angle / static_cast<double>(steered_axles_count);
  }
  else
  {
    velocity.steering_angle = 0.0;
  }
}

void Vehicle::getJointStates(const VehicleState& state, JointStates& joint_states) const
{
  if (state.axle_states.size() != axles_.size())
  {
    throw std::runtime_error("number of axle states in vehicle state differs from number of vehicle axles");
  }

  for (const size_t i : boost::irange<size_t>(0, axles_.size()))
  {
    axles_.at(i)->getJointStates(state.axle_states.at(i), joint_states);
  }
}

boost::optional<double> Vehicle::getSupplyVoltage()
{
  double supply_voltage = 0.0;
  int num_measurements = 0;

  for (const AxlePtr& axle : axles_)
  {
    const boost::optional<double> axle_supply_voltage = axle->getSupplyVoltage();
    if (axle_supply_voltage)
    {
      supply_voltage += *axle_supply_voltage;
      ++num_measurements;
    }
  }

  if (num_measurements != 0)
  {
    return supply_voltage / num_measurements;
  }

  return boost::none;
}

double Vehicle::limit(const double value, const double max, const char* name)
{
  const double limited_value = std::min(std::max(-max, value), max);
  if (limited_value != value)
  {
    ROS_WARN_NAMED("limit", "%s (%f) exceeded maximum (%f) and was limited", name, value, max);
  }
  return limited_value;
}
}


/**
 * @file methods/kde/kde_model.hpp
 * @author Roberto Hueso
 *
 * Model for KDE. It abstracts different types of tree, kernels, etc.
 *
 * mlpack is free software; you may redistribute it and/or modify it under the
 * terms of the 3-clause BSD license.  You should have received a copy of the
 * 3-clause BSD license along with mlpack.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */
#ifndef MLPACK_METHODS_KDE_MODEL_HPP
#define MLPACK_METHODS_KDE_MODEL_HPP

// Include trees.
#include <mlpack/core/tree/binary_space_tree.hpp>
#include <mlpack/core/tree/octree.hpp>
#include <mlpack/core/tree/cover_tree.hpp>
#include <mlpack/core/tree/rectangle_tree.hpp>

// Include core.
#include <mlpack/core.hpp>

// Remaining includes.
#include <boost/variant.hpp>
#include "kde.hpp"

namespace mlpack {
namespace kde {

//! Alias template.
template<typename KernelType,
         template<typename TreeMetricType,
                  typename TreeStatType,
                  typename TreeMatType> class TreeType>
using KDEType = KDE<KernelType,
                    metric::EuclideanDistance,
                    arma::mat,
                    TreeType,
                    TreeType<metric::EuclideanDistance,
                             kde::KDEStat,
                             arma::mat>::template DualTreeTraverser,
                    TreeType<metric::EuclideanDistance,
                             kde::KDEStat,
                             arma::mat>::template SingleTreeTraverser>;

/**
 * KernelNormalizer holds a set of methods to normalize estimations applying
 * in each case the appropiate kernel normalizer function.
 */
class KernelNormalizer
{
 private:
  // SFINAE check if Normalizer function is present.
  HAS_MEM_FUNC(Normalizer, HasNormalizer);

 public:
  //! Normalization not needed.
  template<typename KernelType>
  static void ApplyNormalizer(
      KernelType& /* kernel */,
      const size_t /* dimension */,
      arma::vec& /* estimations */,
      const typename std::enable_if<
          !HasNormalizer<KernelType, double(KernelType::*)(size_t)>::value>::
          type* = 0)
  { return; }

  //! Normalize kernels that have normalizer.
  template<typename KernelType>
  static void ApplyNormalizer(
      KernelType& kernel,
      const size_t dimension,
      arma::vec& estimations,
      const typename std::enable_if<
          HasNormalizer<KernelType, double(KernelType::*)(size_t)>::value>::
          type* = 0)
  {
    estimations /= kernel.Normalizer(dimension);
  }
};

/**
 * DualMonoKDE computes a Kernel Density Estimation on the given KDEType.
 * It performs a monochromatic KDE.
 */
class DualMonoKDE : public boost::static_visitor<void>
{
 private:
  //! Vector to store the KDE results.
  arma::vec& estimations;

 public:
  //! Alias template necessary for Visual C++ compiler.
  template<typename KernelType,
           template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  using KDETypeT = KDEType<KernelType, TreeType>;

  //! Default DualMonoKDE on some KDEType.
  template<typename KernelType,
           template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  void operator()(KDETypeT<KernelType, TreeType>* kde) const;

  // TODO Implement specific cases where a leaf size can be selected.

  //! DualMonoKDE constructor.
  DualMonoKDE(arma::vec& estimations);
};

/**
 * DualBiKDE computes a Kernel Density Estimation on the given KDEType.
 * It performs a bichromatic KDE.
 */
class DualBiKDE : public boost::static_visitor<void>
{
 private:
  //! Query set dimensionality.
  const size_t dimension;

  //! The query set for the KDE.
  const arma::mat& querySet;

  //! Vector to store the KDE results.
  arma::vec& estimations;

 public:
  //! Alias template necessary for Visual C++ compiler.
  template<typename KernelType,
           template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  using KDETypeT = KDEType<KernelType, TreeType>;

  //! Default DualBiKDE on some KDEType.
  template<typename KernelType,
           template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  void operator()(KDETypeT<KernelType, TreeType>* kde) const;

  // TODO Implement specific cases where a leaf size can be selected.

  //! DualBiKDE constructor. Takes ownership of the given querySet.
  DualBiKDE(arma::mat&& querySet, arma::vec& estimations);
};

/**
 * TrainVisitor trains a given KDEType using a reference set.
 */
class TrainVisitor : public boost::static_visitor<void>
{
 private:
  //! The reference set used for training.
  arma::mat&& referenceSet;

 public:
  //! Default TrainVisitor on some KDEType.
  template<typename KernelType,
           template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  void operator()(KDEType<KernelType, TreeType>* kde) const;

  // TODO Implement specific cases where a leaf size can be selected.

  //! TrainVisitor constructor. Takes ownership of the given referenceSet.
  TrainVisitor(arma::mat&& referenceSet);
};

/**
 * BandwidthVisitor modifies the bandwidth of a KDEType kernel.
 */
class BandwidthVisitor : public boost::static_visitor<void>
{
 private:
  //! Relative error tolerance.
  const double bandwidth;

 public:
  //! Default BandwidthVisitor on some KDEType.
  template<typename KernelType,
           template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  void operator()(KDEType<KernelType, TreeType>* kde) const;

  //! BandwidthVisitor constructor.
  BandwidthVisitor(const double bandwidth);
};

/**
 * RelErrorVisitor modifies relative error tolerance for a KDEType.
 */
class RelErrorVisitor : public boost::static_visitor<void>
{
 private:
  //! Relative error tolerance.
  const double relError;

 public:
  //! Default RelErrorVisitor on some KDEType.
  template<typename KernelType,
           template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  void operator()(KDEType<KernelType, TreeType>* kde) const;

  //! RelErrorVisitor constructor.
  RelErrorVisitor(const double relError);
};

/**
 * AbsErrorVisitor modifies absolute error tolerance for a KDEType.
 */
class AbsErrorVisitor : public boost::static_visitor<void>
{
 private:
  //! Absolute error tolerance.
  const double absError;

 public:
  //! Default AbsErrorVisitor on some KDEType.
  template<typename KernelType,
           template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  void operator()(KDEType<KernelType, TreeType>* kde) const;

  //! AbsErrorVisitor constructor.
  AbsErrorVisitor(const double absError);
};

/**
 * MonteCarloVisitor activates or deactivates Monte Carlo for a given KDEType.
 */
class MonteCarloVisitor : public boost::static_visitor<void>
{
 private:
  //! Whether to use Monte Carlo estimations or not.
  const bool monteCarlo;

 public:
  //! Default MonteCarloVisitor on some KDEType.
  template<typename KernelType,
           template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  void operator()(KDEType<KernelType, TreeType>* kde) const;

  //! MonteCarloVisitor constructor.
  MonteCarloVisitor(const bool monteCarlo);
};

/**
 * MCProbabilityVisitor sets the Monte Carlo probability for a given KDEType.
 */
class MCProbabilityVisitor : public boost::static_visitor<void>
{
 private:
  //! Monte Carlo probability.
  const double probability;

 public:
  //! Default MCProbabilityVisitor on some KDEType.
  template<typename KernelType,
           template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  void operator()(KDEType<KernelType, TreeType>* kde) const;

  //! MCProbabilityVisitor constructor.
  MCProbabilityVisitor(const double probability);
};

/**
 * MCSampleSizeVisitor sets the Monte Carlo intial sample size for a given
 * KDEType.
 */
class MCSampleSizeVisitor : public boost::static_visitor<void>
{
 private:
  //! Monte Carlo sample size.
  const size_t sampleSize;

 public:
  //! Default MCSampleSizeVisitor on some KDEType.
  template<typename KernelType,
           template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  void operator()(KDEType<KernelType, TreeType>* kde) const;

  //! MCSampleSizeVisitor constructor.
  MCSampleSizeVisitor(const size_t sampleSize);
};

/**
 * MCEntryCoefVisitor sets the Monte Carlo entry coefficient.
 */
class MCEntryCoefVisitor : public boost::static_visitor<void>
{
 private:
  //! Monte Carlo entry coefficient.
  const double entryCoef;

 public:
  //! Default MCEntryCoefVisitor on some KDEType.
  template<typename KernelType,
           template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  void operator()(KDEType<KernelType, TreeType>* kde) const;

  //! MCEntryCoefVisitor constructor.
  MCEntryCoefVisitor(const double entryCoef);
};

/**
 * MCBreakCoefVisitor sets the Monte Carlo break coefficient.
 */
class MCBreakCoefVisitor : public boost::static_visitor<void>
{
 private:
  //! Monte Carlo break coefficient.
  const double breakCoef;

 public:
  //! Default MCBreakCoefVisitor on some KDEType.
  template<typename KernelType,
           template<typename TreeMetricType,
                    typename TreeStatType,
                    typename TreeMatType> class TreeType>
  void operator()(KDEType<KernelType, TreeType>* kde) const;

  //! MCBreakCoefVisitor constructor.
  MCBreakCoefVisitor(const double breakCoef);
};

/**
 * ModeVisitor exposes the Mode() method of the KDEType.
 */
class ModeVisitor : public boost::static_visitor<KDEMode&>
{
 public:
  //! Return mode of KDEType instance.
  template<typename KDEType>
  KDEMode& operator()(KDEType* kde) const;
};

class DeleteVisitor : public boost::static_visitor<void>
{
 public:
  //! Delete KDEType instance.
  template<typename KDEType>
  void operator()(KDEType* kde) const;
};

class KDEModel
{
 public:
  enum TreeTypes
  {
    KD_TREE,
    BALL_TREE,
    COVER_TREE,
    OCTREE,
    R_TREE
  };

  enum KernelTypes
  {
    GAUSSIAN_KERNEL,
    EPANECHNIKOV_KERNEL,
    LAPLACIAN_KERNEL,
    SPHERICAL_KERNEL,
    TRIANGULAR_KERNEL
  };

 private:
  //! Bandwidth of the kernel.
  double bandwidth;

  //! Relative error tolerance.
  double relError;

  //! Absolute error tolerance.
  double absError;

  //! Type of kernel.
  KernelTypes kernelType;

  //! Type of tree.
  TreeTypes treeType;

  //! Whether Monte Carlo estimations will be used.
  bool monteCarlo;

  //! Probability of estimation being bounded by relative error when using
  //! Monte Carlo estimations.
  double mcProb;

  //! Size of the initial sample for Monte Carlo estimations.
  size_t initialSampleSize;

  //! Entry coefficient for Monte Carlo estimations.
  double mcEntryCoef;

  //! Break coefficient for Monte Carlo estimations.
  double mcBreakCoef;

  /**
   * kdeModel holds an instance of each possible combination of KernelType and
   * TreeType. It is initialized using BuildModel.
   */
  boost::variant<KDEType<kernel::GaussianKernel, tree::KDTree>*,
                 KDEType<kernel::GaussianKernel, tree::BallTree>*,
                 KDEType<kernel::GaussianKernel, tree::StandardCoverTree>*,
                 KDEType<kernel::GaussianKernel, tree::Octree>*,
                 KDEType<kernel::GaussianKernel, tree::RTree>*,
                 KDEType<kernel::EpanechnikovKernel, tree::KDTree>*,
                 KDEType<kernel::EpanechnikovKernel, tree::BallTree>*,
                 KDEType<kernel::EpanechnikovKernel, tree::StandardCoverTree>*,
                 KDEType<kernel::EpanechnikovKernel, tree::Octree>*,
                 KDEType<kernel::EpanechnikovKernel, tree::RTree>*,
                 KDEType<kernel::LaplacianKernel, tree::KDTree>*,
                 KDEType<kernel::LaplacianKernel, tree::BallTree>*,
                 KDEType<kernel::LaplacianKernel, tree::StandardCoverTree>*,
                 KDEType<kernel::LaplacianKernel, tree::Octree>*,
                 KDEType<kernel::LaplacianKernel, tree::RTree>*,
                 KDEType<kernel::SphericalKernel, tree::KDTree>*,
                 KDEType<kernel::SphericalKernel, tree::BallTree>*,
                 KDEType<kernel::SphericalKernel, tree::StandardCoverTree>*,
                 KDEType<kernel::SphericalKernel, tree::Octree>*,
                 KDEType<kernel::SphericalKernel, tree::RTree>*,
                 KDEType<kernel::TriangularKernel, tree::KDTree>*,
                 KDEType<kernel::TriangularKernel, tree::BallTree>*,
                 KDEType<kernel::TriangularKernel, tree::StandardCoverTree>*,
                 KDEType<kernel::TriangularKernel, tree::Octree>*,
                 KDEType<kernel::TriangularKernel, tree::RTree>*> kdeModel;

 public:
  /**
   * Initialize KDEModel.
   *
   * @param bandwidth Bandwidth to use for the kernel.
   * @param relError Maximum relative error tolerance for each point in the
   *                 model. For example, 0.05 means that each value must be
   *                 within 5% of the true KDE value.
   * @param absError Maximum absolute error tolerance for each point in the
   *                 model. For example, 0.1 means that for each point the
   *                 value can have a maximum error of 0.1 units.
   * @param kernelType Type of kernel to use.
   * @param treeType Type of tree to use.
   * @param monteCarlo Whether to use Monte Carlo estimations when possible.
   * @param mcProb of a Monte Carlo estimation to be bounded by relative
   *        error tolerance.
   * @param initialSampleSize Initial sample size for Monte Carlo estimations.
   * @param mcEntryCoef Coefficient to control how much larger does the amount
   *                    of node descendants has to be compared to the initial
   *                    sample size in order for it to be a candidate for Monte
   *                    Carlo estimations.
   * @param mcBreakCoef Coefficient to control what fraction of the node's
   *                    descendants evaluated is the limit before Monte Carlo
   *                    estimation recurses.
   */
  KDEModel(const double bandwidth = 1.0,
           const double relError = KDEDefaultParams::relError,
           const double absError = KDEDefaultParams::absError,
           const KernelTypes kernelType = KernelTypes::GAUSSIAN_KERNEL,
           const TreeTypes treeType = TreeTypes::KD_TREE,
           const bool monteCarlo = KDEDefaultParams::mode,
           const double mcProb = KDEDefaultParams::mcProb,
           const size_t initialSampleSize = KDEDefaultParams::initialSampleSize,
           const double mcEntryCoef = KDEDefaultParams::mcEntryCoef,
           const double mcBreakCoef = KDEDefaultParams::mcBreakCoef);

  //! Copy constructor of the given model.
  KDEModel(const KDEModel& other);

  //! Move constructor of the given model. Takes ownership of the model.
  KDEModel(KDEModel&& other);

  /**
   * Copy the given model.
   *
   * Use std::move if the object to copy is no longer needed.
   *
   * @param other KDEModel to copy.
   */
  KDEModel& operator=(KDEModel other);

  //! Destroy the KDEModel object.
  ~KDEModel();

  //! Serialize the KDE model.
  template<typename Archive>
  void serialize(Archive& ar, const unsigned int version);

  //! Get the bandwidth of the kernel.
  double Bandwidth() const { return bandwidth; }

  //! Modify the bandwidth of the kernel.
  void Bandwidth(const double newBandwidth);

  //! Get the relative error tolerance.
  double RelativeError() const { return relError; }

  //! Modify the relative error tolerance.
  void RelativeError(const double newRelError);

  //! Get the absolute error tolerance.
  double AbsoluteError() const { return absError; }

  //! Modify the absolute error tolerance.
  void AbsoluteError(const double newAbsError);

  //! Get the tree type of the model.
  TreeTypes TreeType() const { return treeType; }

  //! Modify the tree type of the model.
  TreeTypes& TreeType() { return treeType; }

  //! Get the kernel type of the model.
  KernelTypes KernelType() const { return kernelType; }

  //! Modify the kernel type of the model.
  KernelTypes& KernelType() { return kernelType; }

  //! Get whether the model is using Monte Carlo estimations or not.
  bool MonteCarlo() const { return monteCarlo; }

  //! Modify whether the model is using Monte Carlo estimations or not.
  void MonteCarlo(const bool newMonteCarlo);

  //! Get Monte Carlo probability of error being bounded by relative error.
  double MCProbability() const { return mcProb; }

  //! Modify Monte Carlo probability of error being bounded by relative error.
  void MCProbability(const double newMCProb);

  //! Get the initial sample size for Monte Carlo estimations.
  size_t MCInitialSampleSize() const { return initialSampleSize; }

  //! Modify the initial sample size for Monte Carlo estimations.
  void MCInitialSampleSize(const size_t newSampleSize);

  //! Get Monte Carlo entry coefficient.
  double MCEntryCoefficient() const { return mcEntryCoef; }

  //! Modify Monte Carlo entry coefficient.
  void MCEntryCoefficient(const double newEntryCoef);

  //! Get Monte Carlo break coefficient.
  double MCBreakCoefficient() const { return mcBreakCoef; }

  //! Modify Monte Carlo break coefficient.
  void MCBreakCoefficient(const double newBreakCoef);

  //! Get the mode of the model.
  KDEMode Mode() const;

  //! Modify the mode of the model.
  KDEMode& Mode();

  /**
   * Build the KDE model with the given parameters and then trains it with the
   * given reference data.
   * Takes possession of the reference set to avoid a copy, so the reference set
   * will not be usable after this.
   *
   * @param referenceSet Set of reference points.
   */
  void BuildModel(arma::mat&& referenceSet);

  /**
   * Perform kernel density estimation on the given query set.
   * Takes possession of the query set to avoid a copy, so the query set
   * will not be usable after this. If possible, it returns normalized
   * estimations.
   *
   * @pre The model has to be previously created with BuildModel.
   * @param querySet Set of query points.
   * @param estimations Vector where the results will be stored in the same
   *                    order as the query points.
   */
  void Evaluate(arma::mat&& querySet, arma::vec& estimations);

  /**
   * Perform kernel density estimation on the reference set.
   * If possible, it returns normalized estimations.
   *
   * @pre The model has to be previously created with BuildModel.
   * @param estimations Vector where the results will be stored in the same
   *                    order as the query points.
   */
  void Evaluate(arma::vec& estimations);

 private:
  //! Clean memory.
  void CleanMemory();
};

} // namespace kde
} // namespace mlpack

//! Set the serialization version of the KDEModel class.
BOOST_TEMPLATE_CLASS_VERSION(template<>, mlpack::kde::KDEModel, 1);

#include "kde_model_impl.hpp"

#endif


// Author: Mihai-Dusmanu (mihaidusmanu)

#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>

namespace py = pybind11;

#include <Eigen/Core>

extern "C" {
#include <lib/VLFeat/sift.h>
}

#include <colmap/feature/sift.h>
#include <colmap/feature/utils.h>
#include <colmap/lib/SiftGPU/SiftGPU.h>
#ifdef CUDA_ENABLED
#include "GL/glew.h"
#endif

using namespace colmap;

#include "helpers.h"
#include "utils.h"

#include <iostream>

#define kdim 4
#define ddim 128

template <typename dtype>
using pyimage_t =
    Eigen::Matrix<dtype, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;

typedef Eigen::Matrix<float, Eigen::Dynamic, ddim, Eigen::RowMajor>
    descriptors_t;
typedef Eigen::Matrix<float, Eigen::Dynamic, kdim, Eigen::RowMajor> keypoints_t;
typedef Eigen::VectorXf scores_t;
typedef std::tuple<keypoints_t, scores_t, descriptors_t> sift_output_t;

static std::map<int, std::unique_ptr<std::mutex>> sift_gpu_mutexes;

class Sift {
   public:
    Sift(SiftExtractionOptions options, Device device)
        : options_(options), use_gpu_(IsGPU(device)) {
        VerifySiftGPUParams(use_gpu_);
        if (use_gpu_) {
#ifdef CUDA_ENABLED
            sift_gpu.reset(new SiftGPU);
            CreateSiftGPUExtractor(options_, sift_gpu.get());
            if (sift_gpu_mutexes.count(sift_gpu->gpu_index) == 0) {
                sift_gpu_mutexes.emplace(
                    sift_gpu->gpu_index,
                    std::unique_ptr<std::mutex>(new std::mutex()));
            }
#endif
        }
    }

    template <typename dtype>
    sift_output_t Extract(Eigen::Ref<const pyimage_t<dtype>> image,
                          bool do_normalize) {
        THROW_CHECK_LE(image.rows(), options_.max_image_size);
        THROW_CHECK_LE(image.cols(), options_.max_image_size);

        descriptors_t descriptors;
        keypoints_t keypoints;
        scores_t scores;

        if (use_gpu_) {
#ifdef CUDA_ENABLED
            // TODO: SiftGPU changes the image inplace --> we copy it for now
            pyimage_t<dtype> image_copy = image;
            std::tie(keypoints, scores, descriptors) = ExtractGPU(image_copy);
#endif
        } else {
            std::tie(keypoints, scores, descriptors) = ExtractCPU(image);
        }
        if (do_normalize) {
            // Save and normalize the descriptors.
            if (options_.normalization ==
                SiftExtractionOptions::Normalization::L2) {
                descriptors = L2NormalizeFeatureDescriptors(descriptors);
            } else if (options_.normalization ==
                       SiftExtractionOptions::Normalization::L1_ROOT) {
                descriptors = L1RootNormalizeFeatureDescriptors(descriptors);
            }
        }

        return std::make_tuple(keypoints, scores, descriptors);
    }
    const SiftExtractionOptions& Options() const { return options_; };
    Device GetDevice() const {
        return (use_gpu_) ? Device::CUDA : Device::CPU;
    };

   private:
#ifdef CUDA_ENABLED
    template <typename dtype>
    sift_output_t ExtractGPU(const pyimage_t<dtype>& image /* [h, w] */) {
        THROW_CHECK_EQ(options_.max_image_size, sift_gpu->GetMaxDimension());
        THROW_CHECK(!options_.estimate_affine_shape);
        THROW_CHECK(!options_.domain_size_pooling);

        std::unique_lock<std::mutex> lock(
            *sift_gpu_mutexes[sift_gpu->gpu_index]);

        // Note, that this produces slightly different results than using
        // SiftGPU directly for RGB->GRAY conversion, since it uses different
        // weights.
        int code;
        if (std::is_same<dtype, float>::value) {
            code = sift_gpu->RunSIFT(image.cols(),
                                     image.rows(),
                                     image.data(),
                                     GL_LUMINANCE,
                                     GL_FLOAT);
        } else if (std::is_same<dtype, uint8_t>::value) {
            code = sift_gpu->RunSIFT(image.cols(),
                                     image.rows(),
                                     image.data(),
                                     GL_LUMINANCE,
                                     GL_UNSIGNED_BYTE);
        } else {
            THROW_EXCEPTION(std::runtime_error,
                            "SIFT GPU only support float/byte images.");
        }
        const int kSuccessCode = 1;

        THROW_CHECK_EQ(code, kSuccessCode);

        const size_t num_features =
            static_cast<size_t>(sift_gpu->GetFeatureNum());

        std::vector<SiftKeypoint> keypoints_data(num_features);

        descriptors_t descriptors(num_features, ddim);
        keypoints_t keypoints(num_features, kdim);
        scores_t scores = Eigen::VectorXf::Ones(num_features);

        // Download the extracted keypoints and descriptors.
        sift_gpu->GetFeatureVector(keypoints_data.data(), descriptors.data());

        for (size_t i = 0; i < num_features; ++i) {
            keypoints(i, 0) = keypoints_data[i].x;
            keypoints(i, 1) = keypoints_data[i].y;
            keypoints(i, 2) = keypoints_data[i].s;
            keypoints(i, 3) = keypoints_data[i].o;
        }

        return std::make_tuple(keypoints, scores, descriptors);
    }
#endif

    sift_output_t ExtractCPU(Eigen::Ref<const pyimage_t<float>> image) {
        // Create a new instance of SIFT detector & descriptor.
        VlSiftFilt* sift = vl_sift_new(image.cols(),
                                       image.rows(),
                                       options_.num_octaves,
                                       options_.octave_resolution,
                                       options_.first_octave);
        vl_sift_set_edge_thresh(sift, options_.edge_threshold);
        vl_sift_set_peak_thresh(sift, options_.peak_threshold);

        // Build image pyramid.
        bool is_first_octave = true;
        descriptors_t descriptors;
        keypoints_t keypoints;
        scores_t scores;

        int num_feats = 0;

        while (true) {
            if (is_first_octave) {
                if (vl_sift_process_first_octave(sift, image.data())) {
                    break;
                }
                is_first_octave = false;
            } else {
                if (vl_sift_process_next_octave(sift)) {
                    break;
                }
            }

            // Detect keypoints.
            vl_sift_detect(sift);

            // Octave.
            const float* dog = sift->dog;
            const int octave_width = sift->octave_width;
            const int octave_height = sift->octave_height;

            // Extract detected keypoints.
            const VlSiftKeypoint* vl_keypoints = vl_sift_get_keypoints(sift);
            const int num_keypoints = vl_sift_get_nkeypoints(sift);
            if (num_keypoints == 0) {
                continue;
            }

            // Process keypoints.
            for (int i = 0; i < num_keypoints; ++i) {
                // Extract feature orientations.
                double angles[4];
                int num_orientations;
                if (options_.upright) {
                    num_orientations = 1;
                    angles[0] = 0.0;
                } else {
                    num_orientations = vl_sift_calc_keypoint_orientations(
                        sift, angles, &vl_keypoints[i]);
                }
                for (int o = 0; o < num_orientations; ++o) {
                    keypoints.conservativeResize(num_feats + 1, kdim);
                    scores.conservativeResize(num_feats + 1);
                    descriptors.conservativeResize(num_feats + 1, ddim);
                    // Construct keypoint.
                    keypoints(num_feats, 0) = vl_keypoints[i].x;
                    keypoints(num_feats, 1) = vl_keypoints[i].y;
                    keypoints(num_feats, 2) = vl_keypoints[i].sigma;
                    keypoints(num_feats, 3) = angles[o];
                    // Construct descriptor.
                    vl_sift_calc_keypoint_descriptor(
                        sift,
                        descriptors.row(num_feats).data(),
                        &vl_keypoints[i],
                        angles[o]);
                    // Recover score.
                    scores(num_feats) =
                        dog[vl_keypoints[i].is * octave_width * octave_height +
                            vl_keypoints[i].iy * octave_width +
                            vl_keypoints[i].ix];
                    num_feats++;
                }
            }
        }

        // Delete the SIFT object.
        vl_sift_delete(sift);

        return std::make_tuple(keypoints, scores, descriptors);
    }

    sift_output_t ExtractCPU(Eigen::Ref<const pyimage_t<uint8_t>> image) {
        pyimage_t<float> image_f = image.cast<float>();
        image_f.array() /= 255.0f;
        return ExtractCPU(image_f);
    }

    std::unique_ptr<SiftGPU> sift_gpu;
    SiftExtractionOptions options_;
    bool use_gpu_ = false;
};

// Backward compatibility
sift_output_t extract_sift(const py::array_t<float> image,
                           const int num_octaves,
                           const int octave_resolution,
                           const int first_octave,
                           const float edge_thresh,
                           const float peak_thresh,
                           const bool upright) {
    SiftExtractionOptions options;
    options.num_octaves = num_octaves;
    options.octave_resolution = octave_resolution;
    options.first_octave = first_octave;
    options.edge_threshold = edge_thresh;
    options.peak_threshold = peak_thresh;
    options.upright = upright;
    Sift sift(options, Device::CPU);
    return sift.Extract(image.cast<Eigen::Ref<const pyimage_t<float>>>(),
                        false);
}

void init_sift(py::module& m) {
    m.def("extract_sift",
          &extract_sift,
          py::arg("image"),
          py::arg("num_octaves") = 4,
          py::arg("octave_resolution") = 3,
          py::arg("first_octave") = 0,
          py::arg("edge_thresh") = 10.0,
          py::arg("peak_thresh") = 0.01,
          py::arg("upright") = false,
          "Extract SIFT features.");

    // For backwards consistency
    py::dict sift_options;
    sift_options["peak_threshold"] = 0.01;
    sift_options["first_octave"] = 0;
    sift_options["max_image_size"] = 7000;

    py::class_<Sift>(m, "Sift")
        .def(py::init<SiftExtractionOptions, Device>(),
             py::arg("options") = sift_options,
             py::arg("device") = Device::AUTO)
        .def("extract",
             &Sift::Extract<float>,
             py::arg("image"),
             py::arg("do_normalize") = false)
        .def("extract",
             &Sift::Extract<uint8_t>,
             py::arg("image").noconvert(),
             py::arg("do_normalize") = false)
        .def_property_readonly("options", &Sift::Options)
        .def_property_readonly("device", &Sift::GetDevice);
}


// Dylan Losey, July 1, 2019.
// Print measure force at EE

#include <iostream>
#include <cmath>
#include <Eigen/Core>

#include <franka/exception.h>
#include <franka/robot.h>
#include <franka/model.h>

class CartesianVelocityController {
  public:
    CartesianVelocityController(franka::Model& model);
    franka::CartesianVelocities operator()(const franka::RobotState& robot_state, franka::Duration period);
  private:
    franka::Model* modelPtr;
    double time = 0.0;
    double radius = 0.1;
    double frequency = 2*M_PI*0.05;
    double threshold = 0.0;
    std::array<double, 3> pdot_est;
};

CartesianVelocityController::CartesianVelocityController(franka::Model& model) {
  modelPtr = &model;
}

franka::CartesianVelocities CartesianVelocityController::operator()(const franka::RobotState& robot_state,
                                                   franka::Duration period) {

  time += period.toSec();

  std::array<double, 6> What = robot_state.O_F_ext_hat_K;
  std::array<double, 3> Fhat = {{What[0], What[1], What[2]}};

  for (size_t i = 2; i < 3; i++) {
    if (std::abs(Fhat[i]) > threshold) {
      std::cout << i << " axis = " << Fhat[i] << " at time " << time << std::endl;
    }
  }

  double v_x = tanh(10.0 * time) * radius * std::cos(frequency*time) * frequency;
  double v_z = tanh(10.0 * time) * radius * -std::sin(frequency*time) * frequency;
  franka::CartesianVelocities xidot_c = {{v_x, v_x, v_z, 0.0, 0.0, 0.0}};
  return xidot_c;

}

int main(int argc, char** argv) {
  if (argc != 2) {
    std::cerr << "Usage: " << argv[0] << " <robot-hostname>" << std::endl;
    return -1;
  }

  try {
    franka::Robot robot(argv[1]);
    franka::Model model = robot.loadModel();
    robot.setLoad(0.2, {{0,0,0}}, {{1,0,0,0,1,0,0,0,1}});
    robot.automaticErrorRecovery();

    robot.setCollisionBehavior(
        {{40.0, 40.0, 40.0, 40.0, 40.0, 40.0, 40.0}}, {{40.0, 40.0, 40.0, 40.0, 40.0, 40.0, 40.0}},
        {{40.0, 40.0, 40.0, 40.0, 40.0, 40.0}}, {{40.0, 40.0, 40.0, 40.0, 40.0, 40.0}});
    robot.setCartesianImpedance(
	{{100.0, 100.0, 100.0, 1.0, 1.0, 1.0}});

    CartesianVelocityController motion_generator(model);
    robot.control(motion_generator, franka::ControllerMode::kCartesianImpedance);

  } catch (franka::Exception const& e) {
    std::cout << e.what() << std::endl;
    return -1;
  }

  return 0;
}


// smooth_feedback: Control theory on Lie groups
// https://github.com/pettni/smooth_feedback
//
// Licensed under the MIT License <http://opensource.org/licenses/MIT>.
//
// Copyright (c) 2021 Petter Nilsson, John B. Mains
//
// 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 SMOOTH__FEEDBACK__QP_SOLVER_HPP_
#define SMOOTH__FEEDBACK__QP_SOLVER_HPP_

/**
 * @file
 * @brief Quadratic Program solver.
 */

#include <Eigen/Cholesky>
#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <Eigen/SparseCholesky>

#include <chrono>
#include <iomanip>
#include <iostream>
#include <limits>
#include <optional>

#include "qp.hpp"
#include "utils/sparse.hpp"

namespace smooth::feedback {

/**
 * @brief Options for solve_qp
 */
struct QPSolverParams
{
  /// print solver info to stdout
  bool verbose = false;

  /// relaxation parameter
  float alpha = 1.6;
  /// first dual step size
  float rho = 0.1;
  /// second dual step length
  float sigma = 1e-6;

  /// scale problem
  bool scaling = true;

  /// absolute threshold for convergence
  float eps_abs = 1e-3;
  /// relative threshold for convergence
  float eps_rel = 1e-3;
  /// threshold for primal infeasibility
  float eps_primal_inf = 1e-4;
  /// threshold for dual infeasibility
  float eps_dual_inf = 1e-4;

  /// max number of iterations (default no limit)
  std::optional<uint32_t> max_iter = {};

  /// max solution time (default no limit)
  std::optional<std::chrono::nanoseconds> max_time = {};

  /// iterations between checking stopping criterion
  uint32_t stop_check_iter = 25;

  /// run solution polishing (uses dynamic memory)
  bool polish = true;
  /// number of iterations to refine polish
  uint32_t polish_iter = 5;
  /// regularization parameter for polishing
  float delta = 1e-6;
};

namespace detail {

template<typename Pbm>
using qp_solution_t = QPSolution<
  decltype(Pbm::A)::RowsAtCompileTime,
  decltype(Pbm::A)::ColsAtCompileTime,
  typename decltype(Pbm::A)::Scalar>;

/**
 * @brief Polish solution of quadratic program
 *
 * @tparam Pbm problem type
 *
 * @param[in] pbm problem formulation
 * @param[in, out] sol solution to polish
 * @param[in] prm solver options
 * @param[in] c cost scaling
 * @param[in] sx variable scaling
 * @param[in] sy constraint scaling
 *
 * @warning This function allocates dynamic memory even for static-sized problems.
 */
template<typename Pbm, typename D1, typename D2>
bool polish_qp(
  const Pbm & pbm,
  qp_solution_t<Pbm> & sol,
  const QPSolverParams & prm,
  const typename decltype(Pbm::A)::Scalar c,
  const Eigen::MatrixBase<D1> & sx,
  const Eigen::MatrixBase<D2> & sy)
{
  using AmatT                  = decltype(Pbm::A);
  using Scalar                 = typename AmatT::Scalar;
  static constexpr bool sparse = std::is_base_of_v<Eigen::SparseMatrixBase<AmatT>, AmatT>;

  static constexpr Scalar inf = std::numeric_limits<Scalar>::infinity();
  static constexpr Scalar eps = std::numeric_limits<Scalar>::epsilon();

  static constexpr Eigen::Index N = AmatT::ColsAtCompileTime;
  const Eigen::Index n = pbm.A.cols(), m = pbm.A.rows();

  // FIND ACTIVE CONSTRAINT SETS

  Eigen::Index nl = 0, nu = 0;
  for (Eigen::Index idx = 0; idx < m; ++idx) {
    if (sol.dual[idx] < -100 * eps && pbm.l[idx] != -inf) { nl++; }
    if (sol.dual[idx] > 100 * eps && pbm.u[idx] != inf) { nu++; }
  }

  Eigen::VectorXi LU_idx(nl + nu);
  for (Eigen::Index idx = 0, lcntr = 0, ucntr = 0; idx < m; ++idx) {
    if (sol.dual[idx] < -100 * eps && pbm.l[idx] != -inf) { LU_idx(lcntr++) = idx; }
    if (sol.dual[idx] > 100 * eps && pbm.u[idx] != inf) { LU_idx(nl + ucntr++) = idx; }
  }

  // FORM REDUCED SYSTEMS (27) AND (30)

  // square symmetric system matrix
  using HT = std::conditional_t<sparse, Eigen::SparseMatrix<Scalar>, Eigen::Matrix<Scalar, -1, -1>>;
  HT H(n + nl + nu, n + nl + nu), Hp(n + nl + nu, n + nl + nu);

  // fill up H
  if constexpr (sparse) {
    // preallocate nonzeros
    Eigen::VectorXi nnz(n + nl + nu);
    for (auto i = 0u; i != n; ++i) {
      nnz(i) = pbm.P.outerIndexPtr()[i + 1] - pbm.P.outerIndexPtr()[i];
    }
    for (auto i = 0u; i != nl + nu; ++i) {
      nnz(n + i) = pbm.A.outerIndexPtr()[LU_idx(i) + 1] - pbm.A.outerIndexPtr()[LU_idx(i)];
    }
    H.reserve(nnz);
    Hp.reserve(nnz + Eigen::VectorXi::Ones(n + nl + nu));

    // fill P in top left block
    for (auto k = 0u; k < n; ++k) {
      for (Eigen::InnerIterator it(pbm.P, k); it; ++it) {
        const Scalar pij              = c * sx(it.col()) * sx(it.row()) * it.value();
        H.insert(it.row(), it.col())  = pij;
        Hp.insert(it.row(), it.col()) = pij;
      }
    }

    // fill selected rows of A in top right block
    for (auto a_row = 0u; a_row != nl + nu; ++a_row) {
      for (Eigen::InnerIterator it(pbm.A, LU_idx(a_row)); it; ++it) {
        const Scalar Aij               = sy(it.row()) * sx(it.col()) * it.value();
        H.insert(it.col(), n + a_row)  = Aij;
        Hp.insert(it.col(), n + a_row) = Aij;
      }
    }
  } else {
    H.setZero();
    H.topLeftCorner(n, n) = c * sx.asDiagonal() * pbm.P * sx.asDiagonal();
    for (auto i = 0u; i != nl + nu; ++i) {
      H.col(n + i).template head<N>(n) = sy(LU_idx(i)) * pbm.A.row(LU_idx(i)) * sx.asDiagonal();
    }
    Hp = H;
  }

  // add perturbing diagonal elements to Hp
  if constexpr (sparse) {
    for (auto i = 0u; i != n; ++i) { Hp.coeffRef(i, i) += prm.delta; }
    for (auto i = 0u; i != nl + nu; ++i) { Hp.coeffRef(n + i, n + i) -= prm.delta; }
    H.makeCompressed();
    Hp.makeCompressed();
  } else {
    Hp.topLeftCorner(n, n) += Eigen::VectorX<Scalar>::Constant(n, prm.delta).asDiagonal();
    Hp.bottomRightCorner(nl + nu, nl + nu) -=
      Eigen::VectorX<Scalar>::Constant(nl + nu, prm.delta).asDiagonal();
  }

  Eigen::VectorX<Scalar> h(n + nl + nu);
  h.head(n) = -c * sx.cwiseProduct(pbm.q);
  for (auto i = 0u; i != nl; ++i) { h(n + i) = sy(LU_idx(i)) * pbm.l(LU_idx(i)); }
  for (auto i = 0u; i != nu; ++i) { h(n + nl + i) = sy(LU_idx(nl + i)) * pbm.u(LU_idx(nl + i)); }

  // ITERATIVE REFINEMENT

  // factorize Hp
  std::conditional_t<
    sparse,
    Eigen::SimplicialLDLT<decltype(H), Eigen::Upper>,
    Eigen::LDLT<decltype(H), Eigen::Upper>>
    ldlt(Hp);

  if (ldlt.info()) { return false; }

  Eigen::VectorX<Scalar> t_hat = Eigen::VectorX<Scalar>::Zero(n + nl + nu);
  for (auto i = 0u; i != prm.polish_iter; ++i) {
    t_hat += ldlt.solve(h - H.template selfadjointView<Eigen::Upper>() * t_hat);
  }

  // UPDATE SOLUTION

  sol.primal = t_hat.template head<N>(n);
  for (auto i = 0u; i < nl; ++i) { sol.dual(LU_idx(i)) = t_hat(n + i); }
  for (auto i = 0u; i < nu; ++i) { sol.dual(LU_idx(nl + i)) = t_hat(n + nl + i); }

  return true;
}

/**
 * @brief Copy-able wrapper around Eigen LDLT types.
 */
template<typename LDLTt>
struct LDLTWrapper
{
  bool first{true};  // true if ldlt object needs factorization
  LDLTt ldlt{};

  LDLTWrapper() = default;
  LDLTWrapper(const LDLTWrapper &) : first{true}, ldlt{} {}
  LDLTWrapper(LDLTWrapper &&) : first{true}, ldlt{} {}
  LDLTWrapper & operator=(const LDLTWrapper &)
  {
    first = true;
    return *this;
  }
  LDLTWrapper & operator=(LDLTWrapper &)
  {
    first = true;
    return *this;
  }
  ~LDLTWrapper() = default;
};

}  // namespace detail

/**
 * @brief Solver for quadratic programs
 *
 * Use this class to efficiently solve many QPs with the same problem structure.
 *
 * For one-off QPs, see solve_qp().
 */
template<typename Pbm>
class QPSolver
{
  using AmatT                  = decltype(Pbm::A);
  using Scalar                 = typename AmatT::Scalar;
  static constexpr bool sparse = std::is_base_of_v<Eigen::SparseMatrixBase<AmatT>, AmatT>;

  // static sizes
  static constexpr Eigen::Index M = AmatT::RowsAtCompileTime;
  static constexpr Eigen::Index N = AmatT::ColsAtCompileTime;
  static constexpr Eigen::Index K = (N == -1 || M == -1) ? Eigen::Index(-1) : N + M;

  // typedefs
  using Rn = Eigen::Vector<Scalar, N>;
  using Rm = Eigen::Vector<Scalar, M>;
  using Rk = Eigen::Vector<Scalar, K>;
  using Ht = std::conditional_t<sparse, Eigen::SparseMatrix<Scalar>, Eigen::Matrix<Scalar, K, K>>;
  using LDLTt = std::
    conditional_t<sparse, Eigen::SimplicialLDLT<Ht, Eigen::Upper>, Eigen::LDLT<Ht, Eigen::Upper>>;

  static inline const Scalar inf = std::numeric_limits<Scalar>::infinity();

public:
  /**
   * @brief Default constructor.
   */
  QPSolver(const QPSolverParams & prm = {}) : prm_(prm) {}

  /**
   * @brief Construct and allocate working memory.
   *
   * @param pbm template problem.
   *
   * Memory is allocated for solving problems with same structure as pbm.
   */
  QPSolver(const Pbm & pbm, const QPSolverParams & prm = {}) : prm_(prm) { analyze(pbm); }

  /// @brief Default copy constructor
  QPSolver(const QPSolver &) = default;
  /// @brief Default move constructor
  QPSolver(QPSolver &&) = default;
  /// @brief Default copy assignment
  QPSolver & operator=(const QPSolver &) = default;
  /// @brief Default move assignment
  QPSolver & operator=(QPSolver &&) = default;

  /**
   * @brief Access most recent QP solution.
   */
  const QPSolution<M, N, Scalar> & sol() const { return sol_; }

  /**
   * @brief Prepare for solving problems.
   */
  void analyze(const Pbm & pbm)
  {
    const Eigen::Index n = pbm.A.cols(), m = pbm.A.rows(), k = n + m;

    // solution
    sol_.primal.setZero(n);
    sol_.dual.setZero(m);

    // scaling variables and working memory
    c_ = 1;
    sx_.setOnes(n);
    sy_.setOnes(m);
    sx_inc_.setZero(n);
    sy_inc_.setZero(m);

    // solve working memory
    z_.resize(m);
    z_next_.resize(m);
    rho_.resize(m);
    p_.resize(k);
    pt_.resize(k);

    // stopping working memory
    x_us_.resize(n);
    dx_us_.resize(n);
    y_us_.resize(m);
    dy_us_.resize(m);
    z_us_.resize(m);
    Px_.resize(n);
    Aty_.resize(n);
    Ax_.resize(m);

    // preallocate nonzeros in H
    H_.setZero();
    H_.resize(k, k);
    if constexpr (sparse) {
      Eigen::VectorXi nnz(k);
      for (auto i = 0u; i < n; ++i) {
        nnz(i) = pbm.P.outerIndexPtr()[i + 1] - pbm.P.outerIndexPtr()[i] + 1;
      }
      for (auto i = 0u; i < m; ++i) {
        nnz(n + i) = pbm.A.outerIndexPtr()[i + 1] - pbm.A.outerIndexPtr()[i] + 1;
      }
      H_.reserve(nnz);
    }
  }

  /**
   * @brief Solve quadratic program.
   */
  const QPSolution<M, N, Scalar> & solve(
    const Pbm & pbm,
    std::optional<std::reference_wrapper<const QPSolution<M, N, Scalar>>> warmstart = {})
  {
    // update problem scaling
    if (prm_.scaling) { scale(pbm); }

    // dynamic sizes
    const Eigen::Index n = pbm.A.cols(), m = pbm.A.rows();

    // cast parameters to scalar type
    const Scalar rho_bar    = static_cast<Scalar>(prm_.rho);
    const Scalar alpha      = static_cast<Scalar>(prm_.alpha);
    const Scalar alpha_comp = Scalar(1) - alpha;
    const Scalar sigma      = static_cast<Scalar>(prm_.sigma);

    // return code: when set algorithm is finished
    std::optional<QPSolutionStatus> ret_code = std::nullopt;

    for (auto i = 0u; i != m; ++i) {
      if (pbm.l(i) == inf || pbm.u(i) == -inf || pbm.u(i) - pbm.l(i) < Scalar(0.)) {
        ret_code = QPSolutionStatus::PrimalInfeasible;  // feasible set trivially empty
      }

      // set rho depending on constraint type
      if (pbm.l(i) == -inf && pbm.u(i) == inf) {
        rho_(i) = Scalar(1e-6);  // unbounded
      } else if (sy_(i) * abs(pbm.l(i) - pbm.u(i)) < 1e-5) {
        rho_(i) = Scalar(1e3) * rho_bar;  // equality
      } else {
        rho_(i) = rho_bar;  // inequality
      }
    }

    const auto t0 = std::chrono::high_resolution_clock::now();

    // fill square symmetric system matrix H = [P A'; A 0]
    if constexpr (sparse) {
      if (H_.isCompressed()) { H_.coeffs().setZero(); }

      for (auto i = 0u; i < pbm.P.outerSize(); ++i) {
        for (Eigen::InnerIterator it(pbm.P, i); it; ++it) {
          if (it.col() >= it.row()) {
            H_.coeffRef(it.row(), it.col()) = c_ * sx_(it.row()) * sx_(it.col()) * it.value();
          }
        }
      }
      block_add_identity(H_, 0, 0, n, sigma);
      for (auto i = 0u; i < pbm.A.outerSize(); ++i) {
        for (Eigen::InnerIterator it(pbm.A, i); it; ++it) {
          H_.coeffRef(it.col(), n + it.row()) = sy_(it.row()) * sx_(it.col()) * it.value();
        }
      }
      for (auto row = 0u; row < m; ++row) {
        H_.coeffRef(n + row, n + row) = Scalar(-1) / rho_(row);
      }

      if (!H_.isCompressed()) { H_.makeCompressed(); }
    } else {
      H_.setZero();

      H_.template topLeftCorner<N, N>(n, n) = c_ * sx_.asDiagonal() * pbm.P * sx_.asDiagonal();
      H_.template topLeftCorner<N, N>(n, n) +=
        Eigen::Vector<Scalar, N>::Constant(n, sigma).asDiagonal();
      H_.template topRightCorner<N, M>(n, m) =
        (sy_.asDiagonal() * pbm.A * sx_.asDiagonal()).transpose();
      H_.template bottomRightCorner<M, M>(m, m) = (-rho_).cwiseInverse().asDiagonal();
    }

    const auto t_fill = std::chrono::high_resolution_clock::now();

    if (prm_.verbose) {
      using std::cout, std::left, std::setw, std::right;
      // clang-format off
      cout << "========================= QP Solver =========================" << '\n';
      cout << "Solving " << (sparse ? "sparse" : "dense") << " QP with n=" << n << ", m=" << m << '\n';
      cout << setw(8)  << right << "ITER"
           << setw(14) << right << "OBJ"
           << setw(14) << right << "PRI_RES"
           << setw(14) << right << "DUA_RES"
           << setw(10) << right << "TIME" << '\n';
      // clang-format on
    }

    // factorize H
    if constexpr (sparse) {
      if (ldlt_.first) { ldlt_.ldlt.analyzePattern(H_); }
      ldlt_.ldlt.factorize(H_);  // makes copy for permuted matrix..
      ldlt_.first = false;
    } else {
      ldlt_.ldlt.compute(H_);
    }

    const auto t_factor = std::chrono::high_resolution_clock::now();

    if (ldlt_.ldlt.info()) { ret_code = QPSolutionStatus::Unknown; }

    // initialize solver variables
    if (warmstart.has_value()) {
      // warmstart variables must be scaled
      sol_.primal  = sx_.cwiseInverse().cwiseProduct(warmstart.value().get().primal);
      sol_.dual    = c_ * sy_.cwiseInverse().cwiseProduct(warmstart.value().get().dual);
      z_.noalias() = sy_.asDiagonal() * pbm.A * warmstart.value().get().primal;
    } else {
      sol_.primal.setZero();
      sol_.dual.setZero();
      z_.setZero();
    }

    // main optimization loop
    auto iter = 0u;
    for (; (!prm_.max_iter || iter != prm_.max_iter.value()) && !ret_code; ++iter) {
      p_.template segment<N>(0, n) = sigma * sol_.primal - c_ * sx_.asDiagonal() * pbm.q;
      p_.template segment<M>(n, m) = z_ - rho_.cwiseInverse().cwiseProduct(sol_.dual);
      if constexpr (sparse) {
        // p_ = ldlt_.solve(p_);
        // manual solve because ldlt_.solve() uses temporaries...
        // Pinv L D L' P x = b  <==> x = Pinv * (L' \ Dinv(L \ (P * b)))
        pt_.noalias() = ldlt_.ldlt.permutationP() * p_;
        ldlt_.ldlt.matrixL().solveInPlace(pt_);
        pt_.applyOnTheLeft(ldlt_.ldlt.vectorD().cwiseInverse().asDiagonal());
        ldlt_.ldlt.matrixU().solveInPlace(pt_);
        p_.noalias() = ldlt_.ldlt.permutationPinv() * pt_;
      } else {
        ldlt_.ldlt.solveInPlace(p_);
      }

      if (iter % prm_.stop_check_iter == 1) {
        // termination checking requires difference, store old scaled values
        dx_us_ = sol_.primal, dy_us_ = sol_.dual;
      }

      sol_.primal = alpha * p_.template segment<N>(0, n) + alpha_comp * sol_.primal;
      z_next_     = (alpha * rho_.cwiseInverse().cwiseProduct(p_.template segment<M>(n, m))
                 + alpha_comp * rho_.cwiseInverse().cwiseProduct(sol_.dual) + z_)
                  .cwiseMax(sy_.cwiseProduct(pbm.l))
                  .cwiseMin(sy_.cwiseProduct(pbm.u));
      sol_.dual = alpha_comp * sol_.dual + alpha * p_.template segment<M>(n, m)
                + rho_.cwiseProduct(z_) - rho_.cwiseProduct(z_next_);
      std::swap(z_, z_next_);

      if (iter % prm_.stop_check_iter == 1) {
        // unscale solution
        x_us_  = sx_.cwiseProduct(sol_.primal);
        y_us_  = sy_.cwiseProduct(sol_.dual) / c_;
        z_us_  = sy_.cwiseInverse().cwiseProduct(z_);
        dx_us_ = sx_.cwiseProduct(sol_.primal - dx_us_);
        dy_us_ = sy_.cwiseProduct(sol_.dual - dy_us_) / c_;

        // check stopping criteria for unscaled problem and unscaled variables
        ret_code = check_stopping(pbm);

        if (prm_.verbose) {
          using std::cout, std::setw, std::right, std::chrono::microseconds;
          // clang-format off
          cout << setw(7) << right << iter << ":"
            << std::scientific
            << setw(14) << right << (0.5 * pbm.P * x_us_ + pbm.q).dot(x_us_)
            << setw(14) << right << (pbm.A * x_us_ - z_us_).template lpNorm<Eigen::Infinity>()
            << setw(14) << right << (pbm.P * x_us_ + pbm.q + pbm.A.transpose() * y_us_).template lpNorm<Eigen::Infinity>()
            << setw(10) << right << duration_cast<microseconds>(std::chrono::high_resolution_clock::now() - t0).count()
            << '\n';
          // clang-format on
        }

        // check for timeout
        if (!ret_code) {
          if (
            prm_.max_time
            && std::chrono::high_resolution_clock::now() > t0 + prm_.max_time.value()) {
            ret_code = QPSolutionStatus::MaxTime;
          }
        }
      }
    }

    const auto t_iter = std::chrono::high_resolution_clock::now();

    // polish solution if optimal
    if (ret_code.has_value() && ret_code.value() == QPSolutionStatus::Optimal && prm_.polish) {
      if (detail::polish_qp(pbm, sol_, prm_, c_, sx_, sy_)) {
        if (prm_.verbose) {
          // unscale solution
          x_us_ = sx_.cwiseProduct(sol_.primal);
          y_us_ = sy_.cwiseProduct(sol_.dual) / c_;
          z_us_ = sy_.cwiseInverse().cwiseProduct(z_);

          using std::cout, std::setw, std::right, std::chrono::microseconds;
          // clang-format off
          cout << setw(8) << right << "polish:"
            << std::scientific
            << setw(14) << right << (0.5 * pbm.P * x_us_ + pbm.q).dot(x_us_)
            << setw(14) << right << (pbm.A * x_us_ - z_us_).template lpNorm<Eigen::Infinity>()
            << setw(14) << right << (pbm.P * x_us_ + pbm.q + pbm.A.transpose() * y_us_).template lpNorm<Eigen::Infinity>()
            << setw(10) << right << duration_cast<microseconds>(std::chrono::high_resolution_clock::now() - t0).count()
            << '\n';
          // clang-format on
        }

      } else {
        if (prm_.verbose) { std::cout << "Polish failed" << '\n'; }
        sol_.code = QPSolutionStatus::PolishFailed;
      }
    }

    const auto t_polish = std::chrono::high_resolution_clock::now();

    // unscale solution
    sol_.code      = ret_code.value_or(QPSolutionStatus::MaxIterations);
    sol_.primal    = sx_.cwiseProduct(sol_.primal);
    sol_.dual      = sy_.cwiseProduct(sol_.dual) / c_;
    sol_.objective = sol_.primal.dot(0.5 * pbm.P * sol_.primal + pbm.q);
    sol_.iter      = iter;

    if (prm_.verbose) {
      using std::cout, std::left, std::right, std::setw, std::chrono::microseconds;

      // clang-format off
      cout << "QP solver summary:" << '\n';
      cout << "Result " << static_cast<int>(sol_.code) << '\n';

      cout << setw(25) << left << "Iterations"        << setw(10) << right << iter - 1                                               << '\n';
      cout << setw(26) << left << "Total time (µs)"   << setw(10) << right << duration_cast<microseconds>(t_polish - t0).count()     << '\n';
      cout << setw(25) << left << "  Matrix filling"  << setw(10) << right << duration_cast<microseconds>(t_fill - t0).count()       << '\n';
      cout << setw(25) << left << "  Factorization"   << setw(10) << right << duration_cast<microseconds>(t_factor - t_fill).count() << '\n';
      cout << setw(25) << left << "  Iteration"       << setw(10) << right << duration_cast<microseconds>(t_iter - t_factor).count() << '\n';
      cout << setw(25) << left << "  Polish"          << setw(10) << right << duration_cast<microseconds>(t_polish - t_iter).count() << '\n';
      cout << "=============================================================" << '\n';
      // clang-format on
    }

    return sol_;
  }

protected:
  /**
   * @brief Check stopping criteria for solver.
   */
  std::optional<QPSolutionStatus> check_stopping(const Pbm & pbm)
  {
    const Eigen::Index m = pbm.A.rows();

    // norm function
    static const auto norm = [](auto && t) -> Scalar {
      return t.template lpNorm<Eigen::Infinity>();
    };

    // OPTIMALITY

    // check primal
    Ax_.noalias()        = pbm.A * x_us_;
    const Scalar Ax_norm = norm(Ax_);
    Ax_ -= z_us_;
    if (norm(Ax_) <= prm_.eps_abs + prm_.eps_rel * std::max<Scalar>(Ax_norm, norm(z_us_))) {
      // primal succeeded, check dual
      Px_.noalias()           = pbm.P * x_us_;
      Aty_.noalias()          = pbm.A.transpose() * y_us_;
      const Scalar dual_scale = std::max<Scalar>({norm(Px_), norm(pbm.q), norm(Aty_)});
      Px_ += pbm.q + Aty_;
      if (norm(Px_) <= prm_.eps_abs + prm_.eps_rel * dual_scale) {
        return QPSolutionStatus::Optimal;
      }
    }

    // PRIMAL INFEASIBILITY

    Aty_.noalias()        = pbm.A.transpose() * dy_us_;  // NOTE new value A' * dy
    const Scalar Edy_norm = norm(dy_us_);

    Scalar u_dyp_plus_l_dyn = Scalar(0);
    for (auto i = 0u; i != m; ++i) {
      if (pbm.u(i) != inf) {
        u_dyp_plus_l_dyn += pbm.u(i) * std::max<Scalar>(Scalar(0), dy_us_(i));
      } else if (dy_us_(i) > prm_.eps_primal_inf * Edy_norm) {
        // contributes +inf to sum --> no certificate
        u_dyp_plus_l_dyn = inf;
        break;
      }
      if (pbm.l(i) != -inf) {
        u_dyp_plus_l_dyn += pbm.l(i) * std::min<Scalar>(Scalar(0), dy_us_(i));
      } else if (dy_us_(i) < -prm_.eps_primal_inf * Edy_norm) {
        // contributes +inf to sum --> no certificate
        u_dyp_plus_l_dyn = inf;
        break;
      }
    }

    if (std::max<Scalar>(norm(Aty_), u_dyp_plus_l_dyn) < prm_.eps_primal_inf * Edy_norm) {
      return QPSolutionStatus::PrimalInfeasible;
    }

    // DUAL INFEASIBILITY

    Ax_.noalias()        = pbm.A * dx_us_;  // note new value A * dx
    const Scalar dx_norm = norm(dx_us_);
    Px_.noalias()        = pbm.P * dx_us_;

    bool dual_infeasible = (norm(Px_) <= prm_.eps_dual_inf * dx_norm)
                        && (pbm.q.dot(dx_us_) <= prm_.eps_dual_inf * dx_norm);
    for (auto i = 0u; i != m && dual_infeasible; ++i) {
      if (pbm.u(i) == inf) {
        dual_infeasible &= (Ax_(i) >= -prm_.eps_dual_inf * dx_norm);
      } else if (pbm.l(i) == -inf) {
        dual_infeasible &= (Ax_(i) <= prm_.eps_dual_inf * dx_norm);
      } else {
        dual_infeasible &= std::abs(Ax_(i)) < prm_.eps_dual_inf * dx_norm;
      }
    }

    if (dual_infeasible) { return QPSolutionStatus::DualInfeasible; }

    return std::nullopt;
  }

  /**
   * @brief Re-scale QP.
   *
   * The scaled problem is defined as
   *
   * * \f$ P_s = c S_x P S_x \f$,
   * * \f$ q_s = c q S_x \f$,
   * * \f$ A_s = S_y A S_x \f$,
   * * \f$ l_s = S_y l \f$,
   * * \f$ u_s = S_y u \f$,
   *
   * where Sx = diag(sx), Sy = diag(sy).
   *
   * The relation between scaled variables and original variables are
   *
   * * Primal: \f$ x_s = S_x^{-1} x \f$,
   * * Dual: \f$ y_s = c S_y^{-1} y \f$.
   *
   * The objective of the rescaling is the make the columns of
   * \f[
   *   \begin{bmatrix} \bar P & \bar A^T \\ \bar A & 0 \end{bmatrix}
   * \f]
   * have similar \f$ l_\infty \f$ norm, and similarly for the columns of
   * \f[
   *  \begin{bmatrix} \bar P & \bar q \end{bmatrix}.
   * \f]
   */
  void scale(const Pbm & pbm)
  {
    sx_.setOnes();
    sy_.setOnes();

    sx_inc_.setZero();

    // find "norm" of cost function
    for (auto i = 0u; i < pbm.P.outerSize(); ++i) {
      for (Eigen::InnerIterator it(pbm.P, i); it; ++it) {
        sx_inc_(it.col()) = std::max(sx_inc_(it.col()), std::abs(it.value()));
      }
    }

    // if there are "zero cols"
    for (auto i = 0u; i != sx_.size(); ++i) {
      if (sx_inc_(i) == 0) { sx_inc_(i) = 1; }
    }

    // scale cost function
    c_ = Scalar(1) / std::max({1e-6, sx_inc_.mean(), pbm.q.template lpNorm<Eigen::Infinity>()});

    int iter = 0;

    // calculate inf-norm for every column of [Ps As' ; As 0]
    do {
      sx_inc_.setZero();
      sy_inc_.setZero();
      for (auto k = 0u; k < pbm.P.outerSize(); ++k) {
        for (Eigen::InnerIterator it(pbm.P, k); it; ++it) {
          // upper left block of H
          sx_inc_(it.col()) = std::max({
            sx_inc_(it.col()),
            std::abs(c_ * sx_(it.row()) * sx_(it.col()) * it.value()),
          });
        }
      }
      for (auto k = 0u; k < pbm.A.outerSize(); ++k) {
        for (Eigen::InnerIterator it(pbm.A, k); it; ++it) {
          const Scalar Aij  = std::abs(sy_(it.row()) * sx_(it.col()) * it.value());
          sx_inc_(it.col()) = std::max(sx_inc_(it.col()), Aij);  // bottom left block of H
          sy_inc_(it.row()) = std::max(sy_inc_(it.row()), Aij);  // upper right block of H
        }
      }

      // if there are "zero cols" we don't scale
      for (auto i = 0u; i < sx_.size(); ++i) {
        if (sx_inc_(i) == 0) { sx_inc_(i) = 1; }
      }
      for (auto i = 0u; i < sy_.size(); ++i) {
        if (sy_inc_(i) == 0) { sy_inc_(i) = 1; }
      }

      sx_.applyOnTheLeft(sx_inc_.cwiseMax(1e-8).cwiseInverse().cwiseSqrt().asDiagonal());
      sy_.applyOnTheLeft(sy_inc_.cwiseMax(1e-8).cwiseInverse().cwiseSqrt().asDiagonal());
    } while (
      iter++ < 10
      && std::max((sx_inc_.array() - 1).abs().maxCoeff(), (sy_inc_.array() - 1).abs().maxCoeff())
           > 0.1);
  }

private:
  // solver parameters
  QPSolverParams prm_{};

  // solution
  QPSolution<M, N, Scalar> sol_{};

  // scaling variables and working memory
  Scalar c_{0};
  Rn sx_{}, sx_inc_{};
  Rm sy_{}, sy_inc_{};

  // solve working memory
  Rm z_{}, z_next_{}, rho_{};
  Rk p_{}, pt_{};

  // stopping working memory
  Rn x_us_{}, dx_us_{};
  Rn Px_{}, Aty_{};
  Rm Ax_{};
  Rm y_us_{}, z_us_{}, dy_us_{};

  // system matrix and decomposition
  Ht H_{};
  detail::LDLTWrapper<LDLTt> ldlt_{};
};

/**
 * @brief Solve a quadratic program using the operator splitting approach.
 *
 * @tparam Pbm problem type (QuadraticProgram or QuadraticProgramSparse)
 *
 * @param pbm problem formulation
 * @param prm solver options
 * @param warmstart provide initial guess for primal and dual variables
 * @return solution as QuasraticProgramSolution<M, N>
 *
 * @note dynamic problem sizes (`M == -1 || N == -1`) are supported
 *
 * This is a third-party implementation of the algorithm described in the following paper:
 * * Stellato, B., Banjac, G., Goulart, P. et al.
 * **OSQP: an operator splitting solver for quadratic programs.**
 * *Math. Prog. Comp.* 12, 637–672 (2020).
 * https://doi.org/10.1007/s12532-020-00179-2
 *
 * For the official C implementation, see https://osqp.org/.
 */
template<typename Pbm>
detail::qp_solution_t<Pbm> solve_qp(
  const Pbm & pbm,
  const QPSolverParams & prm,
  std::optional<std::reference_wrapper<const detail::qp_solution_t<Pbm>>> warmstart = {})
{
  QPSolver solver(pbm, prm);
  return solver.solve(pbm, warmstart);
}

}  // namespace smooth::feedback

#endif  // SMOOTH__FEEDBACK__QP_SOLVER_HPP_


/* Siconos is a program dedicated to modeling, simulation and control
 * of non smooth dynamical systems.
 *
 * Copyright 2016 INRIA.
 *
 * 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 "EventDriven.hpp"
#include "LsodarOSI.hpp"
#include "Model.hpp"
#include "EventsManager.hpp"
#include "Event.hpp"
#include "NonSmoothDynamicalSystem.hpp"
#include "SimulationGraphs.hpp"

#include "ControlLinearAdditionalTermsED.hpp"
#include "ControlManager.hpp"
#include "ControlLsodarSimulation.hpp"
#include "ControlSimulation_impl.hpp"
#include "Observer.hpp"
#include "Actuator.hpp"

#include <boost/progress.hpp>
#include <boost/timer.hpp>

ControlLsodarSimulation::ControlLsodarSimulation(double t0, double T, double h):
  ControlSimulation(t0, T, h)
{
  _processIntegrator.reset(new LsodarOSI());
  _processSimulation.reset(new EventDriven(_processTD, 0));
  _processSimulation->setName("plant simulation");
  _processSimulation->insertIntegrator(_processIntegrator);
  std11::static_pointer_cast<LsodarOSI>(_processIntegrator)->setExtraAdditionalTerms(
      std11::shared_ptr<ControlLinearAdditionalTermsED>(new ControlLinearAdditionalTermsED()));

  _DSG0 = _model->nonSmoothDynamicalSystem()->topology()->dSG(0);
  _IG0 = _model->nonSmoothDynamicalSystem()->topology()->indexSet0();

  // Control part
  _CM.reset(new ControlManager(_processSimulation));
}

void ControlLsodarSimulation::run()
{
  EventsManager& eventsManager = *_processSimulation->eventsManager();
  unsigned k = 0;
  boost::progress_display show_progress(_N);
  boost::timer time;
  time.restart();
  EventDriven& sim = static_cast<EventDriven&>(*_processSimulation);

  while (sim.hasNextEvent())
  {
    if (eventsManager.needsIntegration())
    {
      sim.advanceToEvent();
    }
    sim.processEvents();
    Event& currentEvent = *eventsManager.currentEvent();
    if (currentEvent.getType() == ACTUATOR_EVENT)
    {
      // this is necessary since we changed the control input, hence the RHS
      sim.setIstate(1);
    }
    if (sim.hasNextEvent() && eventsManager.nextEvent()->getType() == TD_EVENT) // We store only on TD_EVENT, this should be settable
    {
      (*_dataM)(k, 0) = sim.startingTime();
      storeData(k);
      ++k;
      if (!_silent)
      {
       ++show_progress;
      }
    }
  }

  /* saves last status */
  (*_dataM)(k, 0) = sim.startingTime();
  storeData(k);
  ++k;

  _elapsedTime = time.elapsed();
  _dataM->resize(k, _nDim + 1);
}


/*
Copyright 2017 InitialDLab

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.
*/
#include <vector>
#include <iostream>
#include <string>
#include <cstdlib>

#include <boost/timer/timer.hpp>

#include <rtree/rtree.h>
#include "random_shuffle.h"

#include "mongo_types.h"

//#define USE_OSM

std::string IN_FN = "/tmp/in.txt";
//std::string OUT_FN = "/tmp/out.bin";
#ifdef USE_OSM
std::string OUT_FN = "./built/osm_shuffle_1.shf";
#else
std::string OUT_FN = "./built/geo_shuffle.shf";
#endif

using point = mongo_types::point3d;
using box = mongo_types::box3d;
using entry = mongo_types::entry;

template <class T>
void do_query_for_rs(T & rs)
{
    boost::timer::cpu_timer timer;
    auto cursor = rs.sample_query(
#ifdef USE_OSM
            box(
                point(40.64967346, 16.55130386, 1249275136),
                point(46.23781204, 27.72757912, 1335750784)
                )
#else
            box(
                point(39.83063889,115.985733,287327584),
                point(39.96952438,116.7704773,288192224)
                )
#endif
            );
    std::cout << "after cursor " << (timer.elapsed().wall / 1000000.0) << std::endl;

#ifdef USE_OSM
    size_t K_piece = 20000;
    size_t K = 200000;
#else
    size_t K_piece = 50000;
    size_t K = 50000;
#endif
    std::vector<entry> buffer;
    {
        boost::timer::auto_cpu_timer _timer;
        _timer.stop();
        size_t samples_got = 0;
        while(samples_got < K)
        {
            buffer.clear();
            _timer.resume();
            cursor.get_samples(K_piece, std::back_inserter(buffer));
            _timer.stop();
            if(buffer.empty()) break;
            samples_got += buffer.size();
            //std::cout << samples_got << ' ' << (timer.elapsed().wall / 1000000.0) << std::endl;
        }
        std::cout << samples_got << ' ' << (timer.elapsed().wall / 1000000.0) << std::endl;
    }
    auto stats = cursor.get_stats();
    std::cerr << "cache_read " << stats.cached_entries_read << std::endl;
    std::cerr << "disk_read " << stats.disk_entries_read << std::endl;
    std::cerr << "disk_block_read " << stats.blocks_read << std::endl;
}

void clear_cache()
{
    int i = system("/uusoc/scratch/magpie18/robertc/sample/clearCache");
    if(i != 0)
        std::cerr << "WARNING: clear_cache() failed!" << std::endl;
}

void do_query()
{
#ifdef USE_OSM
    rtree::IOLayersParameters parameters;
    parameters.max_top_layer_io_node_count = 1024;
    auto memory_limit = rtree::rtree<entry>::estimate_memory_usage(2210518076ULL, parameters);
#else
    auto memory_limit = rtree::rtree<entry>::estimate_memory_usage(24876977);
#endif
    std::cerr << "memory_limit(G) " << ((double)memory_limit)/1024/1024/1024 << std::endl;
    //debug
    memory_limit = 24876977*30;
    //RandomShuffle<entry> rs(OUT_FN, 1024*16*sizeof(entry));
    RandomShuffle<entry> rs(OUT_FN, memory_limit);
    std::cerr << "create rs " << std::endl;
    std::cerr << "entry_count " << rs.entry_count << std::endl;
    std::cerr << "block_size " << rs.block_size << std::endl;

    //clear_cache();
    do_query_for_rs(rs);

    std::cerr << "destroy rs" << std::endl;
}

int main(int argc, char **argv)
{
    if(argc != 2)
    {
        std::cerr << "Need exactly 1 argument" << std::endl;
        return -1;
    }

    std::string cmd = argv[1];

    if(false) { }
#if 0
    else if(cmd == "build")
    {
        {
            std::ofstream fout(IN_FN);
            for(int i = 0; i < 16*4096; ++i)
                fout << i << std::endl;
        }

        RandomShuffle<entry>::build(IN_FN, OUT_FN,
                [](std::string const& line)->entry{ return entry(std::stoi(line)); },
                RandomShuffle<entry>::Parameters(),
                1024*1024*1024
                );
    }
#endif
    else if (cmd == "query")
    {
        do_query();
    }
    else
    {
        std::cerr << "Invalid argument" << std::endl;
        return -1;
    }

    return 0;
}


#include <algorithm>                         // copy, min
#include <cstring>
#include <ios>
#include <iosfwd>                            // streamsize
#include <boost/iostreams/categories.hpp>    // source_tag, sinkt_tag
#include <ostream>
#include <boost/iostreams/device/file.hpp>
#include <boost/iostreams/stream.hpp>

namespace io = boost::iostreams;

namespace boost { namespace iostreams { namespace example {

template <typename Container>
class container_device {
   public:
      typedef typename Container::value_type char_type;
      typedef seekable_device_tag            category;
      container_device(Container& container) : container_(container), pos_(0) { }

      std::streamsize read(char_type* s, std::streamsize n) {
         using namespace std;
         streamsize amt = static_cast<streamsize>(container_.size() - pos_);
         streamsize result = (min) (n, amt);
         if (result != 0) {
            std::copy( container_.begin() + pos_, container_.begin() + pos_ + result, s);
            pos_ += result;
            return result;
         } else {
            return -1; //EOF
         }
      }

      std::streamsize write(const char_type* s, std::streamsize n){
         using namespace std;
         streamsize result = 0;
         if (pos_ != container_.size()) {
            streamsize amt = static_cast<streamsize>(container_.size() - pos_);
            result = (min)(n, amt);
            std::copy(s, s + result, container_.begin() + pos_);
            pos_ += result;
         }
         if (result < n){
            container_.insert(container_.end(), s + result, s + n);
            pos_ = container_.size();
         }
         return n;
      }

      stream_offset seek(stream_offset off, std::ios_base::seekdir way) {
         using namespace std;

         // Determine new value of pos_
         stream_offset next;
         if (way == ios_base::beg) {
            next = off;
         } else if (way == ios_base::cur){
            next = pos_ + off;
         } else if (way == ios_base::end){
            next = container_.size() + off - 1;
         } else {
            throw ios_base::failure("bad seek direction");
         }

         // Check for errors
         if (next < 0 || next >= static_cast<stream_offset>(container_.size()))
            throw ios_base::failure("bad seek offset");

         pos_ = next;
         return pos_;
      }
   private:
      typedef typename Container::size_type           size_type;
      Container& container_;
      size_type pos_;
};

template <typename Container>
class container_source {
   private:
   typedef typename Container::size_type size_type;
   Container& container_;
   size_type  pos_;

   public:
      typedef typename Container::value_type char_type;
      typedef boost::iostreams::source_tag                    category;
      container_source(Container& container) : container_(container), pos_(0) { }
      std::streamsize read(char_type* s, std::streamsize n) {
            using namespace std;
            streamsize amt = static_cast<streamsize>(container_.size() - pos_);
            streamsize result = (min) (n, amt);
            if (result != 0) {
               std::copy( container_.begin() + pos_, container_.begin() + pos_ + result, s);
               pos_ += result;
               return result;
            } else {
               return -1; // EOF
            }
      }
      Container& container() { return container_; }
};

template <typename Container>
class container_sink {
   private:
      Container& container_;
   public:
      typedef typename Container::value_type char_type;
      typedef sink_tag                       category;
      container_sink(Container& container) : container_(container) {}
      std::streamsize write(const char_type* s, std::streamsize n){
         container_.insert(container_.end(), s, s + n);
         return n;
      }
      Container& container() {return container_;}
};

} } } // End namespace boost::iostream::example


//==============================================================================
//         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_EULER_CONSTANTS_STIRLINGSPLITLIM_HPP_INCLUDED
#define NT2_EULER_CONSTANTS_STIRLINGSPLITLIM_HPP_INCLUDED

#include <boost/simd/constant/hierarchy.hpp>
#include <boost/simd/constant/register.hpp>

namespace nt2
{
  namespace tag
  {
    BOOST_SIMD_CONSTANT_REGISTER( Stirlingsplitlim, double
                                , 0, 0x41d628f6         // 26.77f
                                , 0x4061e083ba3443d4ll  // 143.01608
                                )
  }
  namespace ext
  {
    template<class Site, class... Ts>
    BOOST_FORCEINLINE generic_dispatcher<tag::Stirlingsplitlim, Site> dispatching_Stirlingsplitlim(adl_helper, boost::dispatch::meta::unknown_<Site>, boost::dispatch::meta::unknown_<Ts>...)
    {
      return generic_dispatcher<tag::Stirlingsplitlim, Site>();
    }
    template<class... Args>
    struct impl_Stirlingsplitlim;
  }

  BOOST_SIMD_CONSTANT_IMPLEMENTATION(tag::Stirlingsplitlim, Stirlingsplitlim);
}

#endif


//==============================================================================
//         Copyright 2003 - 2012   LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2012   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 BOOST_SIMD_SDK_SIMD_META_EXTENSION_OF_HPP_INCLUDED
#define BOOST_SIMD_SDK_SIMD_META_EXTENSION_OF_HPP_INCLUDED

#include <boost/simd/sdk/simd/extensions/meta/tags.hpp>
#include <boost/simd/sdk/memory/aligned_array_fwd.hpp>
#include <boost/dispatch/meta/na.hpp>
#include <boost/simd/sdk/simd/native_fwd.hpp>

namespace boost { namespace simd { namespace meta
{
  template<class T,class Dummy=void, class Dummy1 = void>
  struct extension_of : dispatch::meta::na_ {};

  template<class T, class X>
  struct extension_of< boost::simd::native<T,X> > { typedef X type; };

  template<class T, class N, std::size_t Align>
  struct extension_of< boost::simd::memory::aligned_array<T, N::value, Align>, T, N>
  {
    typedef tag::simd_emulation_< N::value * sizeof(T) > type;
  };
} } }

#endif


#define BOOST_MPL11_NO_ASSERTIONS
#include <boost/mpl11/detail/right_folds/until_naive.hpp>
#include <boost/mpl11/list.hpp>

<%= render('_main.erb') do |f, state, xs|
  <<-EOS
  boost::mpl11::detail::right_folds::until_naive<
      boost::mpl11::is_empty,
      #{f}::type::apply,
      #{state},
      boost::mpl11::list<#{xs.join(', ')}>
  >
  EOS
end %>


#include "PrismaticJoint.h"
#include <Eigen/Geometry>

#include "RigidBodyManipulator.h" // todo: remove this when I remove setupOldKinematicTree

using namespace Eigen;

PrismaticJoint::PrismaticJoint(const std::string& name, const Isometry3d& transform_to_parent_body, const Vector3d& translation_axis) :
    FixedAxisOneDoFJoint(name, transform_to_parent_body, spatialJointAxis(translation_axis)), translation_axis(translation_axis)
{
  assert(abs(translation_axis.norm()-1)<1e-10);
}

PrismaticJoint::~PrismaticJoint()
{
  // empty
}

Isometry3d PrismaticJoint::jointTransform(const Eigen::Ref<const VectorXd>& q) const
{
  Isometry3d ret;
  ret.linear().setIdentity();
  ret.translation() = q[0] * translation_axis;
  ret.makeAffine();
  return ret;
}

Matrix<double, TWIST_SIZE, 1> PrismaticJoint::spatialJointAxis(const Vector3d& translation_axis)
{
  Matrix<double, TWIST_SIZE, 1> ret;
  ret.topRows<3>() = Vector3d::Zero();
  ret.bottomRows<3>() = translation_axis;
  return ret;
}

void PrismaticJoint::setupOldKinematicTree(RigidBodyManipulator* model, int body_ind, int position_num_start, int velocity_num_start) const
{
  FixedAxisOneDoFJoint::setupOldKinematicTree(model,body_ind,position_num_start,velocity_num_start);
  model->bodies[body_ind]->pitch = INF;

  Vector3d z_axis(0.0,0.0,1.0);
  if (translation_axis.dot(z_axis)<1-1e-4) {
    Vector4d a;
    a << translation_axis.cross(z_axis), acos(translation_axis.dot(z_axis));
    if ((std::abs(a(0))<1e-4) && (std::abs(a(1))<1e-4) && (std::abs(a(2))<1e-4))
      a.head(3) << 0.0, 1.0, 0.0;
    model->bodies[body_ind]->T_body_to_joint.topLeftCorner(3,3) = axis2rotmat(a);
  }

}


#pragma once
#include <pcl/PCLPointCloud2.h>
#include <pcl/common/time.h>
#include <pcl/console/print.h>
#include <pcl/console/time.h>  // TicToc
#include <pcl/conversions.h>
#include <pcl/filters/statistical_outlier_removal.h>
#include <pcl/visualization/pcl_visualizer.h>
#include <pcl_conversions/pcl_conversions.h>
#include <iostream>
#include <map>
#include <string>
#include <vector>

#include <cv_bridge/cv_bridge.h>
#include <image_geometry/pinhole_camera_model.h>
#include <image_transport/image_transport.h>
#include <algorithm>
#include <opencv2/opencv.hpp>

#include <sensor_msgs/image_encodings.h>
#include <tf/transform_listener.h>

#include <eigen_conversions/eigen_msg.h>
#include <Eigen/StdVector>
#include <mil_vision_lib/cv_tools.hpp>
#include <mil_vision_lib/pcl_tools.hpp>

#include "ros/ros.h"

#include "sub8_msgs/VisionRequest.h"
#include "sub8_msgs/VisionRequest2D.h"

// For stack-tracing on seg-fault
#define BACKWARD_HAS_BFD 1
#include <sub8_build_tools/backward.hpp>

// ROS_NAMESPACE=/camera/front/left rosrun image_proc image_proc
// rosbag play ./holding_buoy_mil.bag -r 0.1
// [1]
// http://docs.ros.org/hydro/api/image_geometry/html/c++/classimage__geometry_1_1PinholeCameraModel.html
// [2] http://docs.pointclouds.org/trunk/classpcl_1_1visualization_1_1_p_c_l_visualizer.html
/* TODO
  - Smarter buoy tracking
  - Better cacheing (instead of immediate assignment)
*/

class Sub8BuoyDetector
{
public:
  Sub8BuoyDetector();
  ~Sub8BuoyDetector();
  void compute_loop(const ros::TimerEvent &);
  void cloud_callback(const sensor_msgs::PointCloud2::ConstPtr &);
  void image_callback(const sensor_msgs::ImageConstPtr &msg, const sensor_msgs::CameraInfoConstPtr &info_msg);
  bool get_last_image(cv::Mat &last_image);
  bool determine_buoy_position(const image_geometry::PinholeCameraModel &cam_model, const std::string &target_color,
                               const cv::Mat &image_raw, const sub::PointCloudT::Ptr &point_cloud_raw,
                               Eigen::Vector3f &center);
  bool segment_buoy(cv::Mat &input_image, cv::Point &center, std::vector<sub::Contour> &output_contours,
                    std::string &target_name);
  bool request_buoy_position_2d(sub8_msgs::VisionRequest2D::Request &req, sub8_msgs::VisionRequest2D::Response &resp);
  bool request_buoy_position(sub8_msgs::VisionRequest::Request &req, sub8_msgs::VisionRequest::Response &resp);
  // Visualize
  int vp1;
  int vp2;
  boost::shared_ptr<pcl::visualization::PCLVisualizer> viewer;
  sub::RvizVisualizer rviz;

  ros::Timer compute_timer;
  ros::Subscriber data_sub;
  ros::NodeHandle nh;
  ros::ServiceServer service_2d;
  ros::ServiceServer service_3d;
  double buoy_radius;

  std::map<std::string, sub::Range> color_ranges;
  Eigen::Vector3f last_bump_target;

  image_transport::CameraSubscriber image_sub;
  image_transport::ImageTransport image_transport;
  image_transport::Publisher image_pub;
  cv::Mat last_draw_image;

  image_geometry::PinholeCameraModel cam_model;

  ros::Time image_time;
  ros::Time last_cloud_time;
  sub::PointCloudT::Ptr current_cloud;
  sensor_msgs::ImageConstPtr last_image_msg;

  bool got_cloud, got_image, line_added, computing, need_new_cloud;
};