/*
 *  utils.hpp
 *
 *	Author(s): Tamas D. Nagy
 *	Created on: 2016-11-08
 *
 *  Useful functions to provide compatibility
 *  of basic datatypes, mainly in matematical
 *  calculations.
 *
 */

#ifndef DVRK_UTILS_HPP_
#define DVRK_UTILS_HPP_

#include <iostream>
#include <sstream>
#include <vector>
#include <cmath>
#include <Eigen/Dense>
#include <Eigen/Geometry>
#include <limits>
#include "irob_utils/pose.hpp"

#include <std_msgs/Float32.h>
#include <irob_msgs/FloatArray.h>
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/Pose.h>
#include <geometry_msgs/Point.h>
#include <geometry_msgs/PointStamped.h>
#include <irob_msgs/ToolPose.h>
#include <irob_msgs/ToolPoseStamped.h>
#include <irob_msgs/Environment.h>

namespace saf {

typedef enum InterpolationMethod
{LINEAR, BEZIER} InterpolationMethod;

inline double degToRad(double deg) {
  return (deg / 180.0) * M_PI;
}

inline double radToDeg(double rad) {
  return (rad * 180.0) / M_PI;
}

template<typename T>
std::ostream& operator<<(std::ostream& out, const std::vector<T>& v) {
  out << "[";
  size_t last = v.size() - 1;
  for(size_t i = 0; i < v.size(); ++i) {
    out << v[i];
    if (i != last)
      out << ", ";
  }
  out << "]";
  return out;
}

// Interpolation
template<typename T>
inline T interpolate(double a, T const& x1, T const& x2) {
  return ((1.0-a) * x1) + ((a) * x2);
}

template <>
inline Pose interpolate(double a, const Pose& x1, const Pose& x2) {
  return x1.interpolate(a, x2);
}

template <>
inline Eigen::Quaternion<double> interpolate(double a,
                                             const Eigen::Quaternion<double>& x1,
                                             const Eigen::Quaternion<double>& x2) {
  return x1.slerp(a, x2);
}

// Distance
template<typename T>
inline double distanceEuler(T const& x1, T const& x2) {
  return std::abs(x2 - x1);
}

template <>
inline double distanceEuler(const Pose& x1, const Pose& x2) {
  Pose::Distance d = x1.dist(x2);
  double weighted_cartesian = std::abs(d.cartesian) * 10000.0;
  double weighted_angle = std::abs(d.angle);
  double weighted_jaw = radToDeg(std::abs(d.jaw));
  if (weighted_cartesian >= weighted_angle
      && weighted_cartesian >= weighted_jaw)
    return std::abs(d.cartesian);
  if (weighted_angle >= weighted_jaw)
    // in degrees, should be converted to rad?
    return std::abs(d.angle);
  return std::abs(d.jaw);
}

template <>
inline double distanceEuler(const Eigen::Vector3d& x1,
                            const Eigen::Vector3d& x2) {
  return std::abs((x2-x1).norm());
}

// Conversion from ROS msg
template<typename MsgT, typename DataT>
inline DataT unwrapMsg(const MsgT& msg);

template <>
inline irob_msgs::Environment unwrapMsg(const irob_msgs::Environment& msg){
  return msg;
}

template <>
inline double unwrapMsg(const std_msgs::Float32& msg){
  return msg.data;
}

template <>
inline Pose unwrapMsg(const geometry_msgs::Pose& msg){
  Pose ret(msg, 0);
  return ret;
}

template <>
inline Pose unwrapMsg(const irob_msgs::ToolPose& msg){
  Pose ret(msg);
  return ret;
}

template <>
inline Pose unwrapMsg(const geometry_msgs::PoseStamped& msg){
  Pose ret(msg, 0);
  return ret;
}

template <>
inline Pose unwrapMsg(const irob_msgs::ToolPoseStamped& msg){
  Pose ret(msg);
  return ret;
}

template <>
inline Eigen::Vector3d unwrapMsg(const geometry_msgs::Point& msg){
  Eigen::Vector3d ret(msg.x, msg.y, msg.z);
  return ret;
}

template <>
inline Eigen::Vector3d unwrapMsg(const geometry_msgs::PointStamped& msg){
  Eigen::Vector3d ret(msg.point.x, msg.point.y, msg.point.z);
  return ret;
}

template <>
inline Eigen::Quaternion<double> unwrapMsg(const geometry_msgs::Quaternion& msg){
  Eigen::Quaternion<double> ret(msg.w, msg.x, msg.y, msg.z);
  return ret;
}

template <>
inline std::vector<double> unwrapMsg(const irob_msgs::FloatArray& msg){
  return msg.data;
}

// Conversion to ROS msg
template<typename MsgT, typename DataT>
inline MsgT wrapToMsg(const DataT& data);

template <>
inline std_msgs::Float32 wrapToMsg(const double& data){
  std_msgs::Float32 msg;
  msg.data = data;
  return msg;
}

template <>
inline geometry_msgs::Pose wrapToMsg(const Pose& data){
  return data.toRosPose();
}

template <>
inline irob_msgs::ToolPose wrapToMsg(const Pose& data){
  return data.toRosToolPose();
}

template <>
inline geometry_msgs::Point wrapToMsg(const Eigen::Vector3d& data){
  geometry_msgs::Point msg;
  msg.x = data.x();
  msg.y = data.y();
  msg.z = data.z();
  return msg;
}

template <>
inline geometry_msgs::PointStamped wrapToMsg(const Eigen::Vector3d& data){
  geometry_msgs::PointStamped msg;
  msg.point.x = data.x();
  msg.point.y = data.y();
  msg.point.z = data.z();
  return msg;
}

template <>
inline geometry_msgs::Quaternion wrapToMsg(
    const Eigen::Quaternion<double>& data){
  geometry_msgs::Quaternion msg;
  msg.w = data.w();
  msg.x = data.x();
  msg.y = data.y();
  msg.z = data.z();
  return msg;
}

// NaN
template<typename DataT>
inline DataT makeNaN();

template <>
inline double makeNaN(){
  return std::numeric_limits<double>::quiet_NaN();
}

template <>
inline Pose makeNaN(){
  Pose ret(std::numeric_limits<double>::quiet_NaN(),
           std::numeric_limits<double>::quiet_NaN(),
           std::numeric_limits<double>::quiet_NaN(),
           std::numeric_limits<double>::quiet_NaN(),
           std::numeric_limits<double>::quiet_NaN(),
           std::numeric_limits<double>::quiet_NaN(),
           std::numeric_limits<double>::quiet_NaN(),
           std::numeric_limits<double>::quiet_NaN());
  return ret;
}

template <>
inline Eigen::Vector3d makeNaN(){
  Eigen::Vector3d ret(std::numeric_limits<double>::quiet_NaN(),
                      std::numeric_limits<double>::quiet_NaN(),
                      std::numeric_limits<double>::quiet_NaN());
  return ret;
}

template <>
inline Eigen::Quaternion<double> makeNaN(){
  Eigen::Quaternion<double> ret(std::numeric_limits<double>::quiet_NaN(),
                                std::numeric_limits<double>::quiet_NaN(),
                                std::numeric_limits<double>::quiet_NaN(),
                                std::numeric_limits<double>::quiet_NaN());
  return ret;
}

template <>
inline std_msgs::Float32 makeNaN(){
  std_msgs::Float32 msg;
  msg.data = std::numeric_limits<double>::quiet_NaN();
  return msg;
}

template <>
inline irob_msgs::FloatArray makeNaN(){
  irob_msgs::FloatArray msg;
  msg.data.push_back(std::numeric_limits<double>::quiet_NaN());
  return msg;
}

template <>
inline geometry_msgs::Pose makeNaN(){
  Pose nanp = makeNaN<Pose>();
  return nanp.toRosPose();
}

template <>
inline irob_msgs::ToolPose makeNaN(){
  Pose nanp = makeNaN<Pose>();
  return nanp.toRosToolPose();
}

template <>
inline irob_msgs::Environment makeNaN(){
  irob_msgs::Environment nanp;
  nanp.valid = irob_msgs::Environment::INVALID;
  return nanp;
}

template <>
inline geometry_msgs::Point makeNaN(){
  geometry_msgs::Point msg;
  msg.x = std::numeric_limits<double>::quiet_NaN();
  msg.y = std::numeric_limits<double>::quiet_NaN();
  msg.z = std::numeric_limits<double>::quiet_NaN();
  return msg;
}

template <>
inline geometry_msgs::PointStamped makeNaN(){
  geometry_msgs::PointStamped msg;
  msg.point.x = std::numeric_limits<double>::quiet_NaN();
  msg.point.y = std::numeric_limits<double>::quiet_NaN();
  msg.point.z = std::numeric_limits<double>::quiet_NaN();
  return msg;
}

template <>
inline geometry_msgs::Quaternion makeNaN(){
  geometry_msgs::Quaternion msg;
  msg.w = std::numeric_limits<double>::quiet_NaN();
  msg.x = std::numeric_limits<double>::quiet_NaN();
  msg.y = std::numeric_limits<double>::quiet_NaN();
  msg.z = std::numeric_limits<double>::quiet_NaN();
  return msg;
}

// isnan
template<typename DataT>
inline bool isnan(const DataT& d);

template <>
inline bool isnan(const double& d)
{
  return std::isnan(d);
}

template <>
inline bool isnan(const Pose& d)
{
  return (std::isnan(d.position.x())
          || std::isnan(d.position.y())
          || std::isnan(d.position.z())
          || std::isnan(d.orientation.x())
          || std::isnan(d.orientation.y())
          || std::isnan(d.orientation.z())
          || std::isnan(d.orientation.w())
          || std::isnan(d.jaw));
}

template <>
inline bool isnan(const Eigen::Vector3d& d)
{
  return (std::isnan(d.x())
          || std::isnan(d.y())
          || std::isnan(d.z()));
}

template <>
inline bool isnan(const Eigen::Quaternion<double>& d)
{
  return (std::isnan(d.x())
          || std::isnan(d.y())
          || std::isnan(d.z())
          || std::isnan(d.w()));
}

template <>
inline bool isnan(const std_msgs::Float32& d)
{

  return (std::isnan(d.data));
}

template <>
inline bool isnan(const geometry_msgs::Pose& d)
{
  return (std::isnan(d.position.x)
          || std::isnan(d.position.y)
          || std::isnan(d.position.z)
          || std::isnan(d.orientation.x)
          || std::isnan(d.orientation.y)
          || std::isnan(d.orientation.z)
          || std::isnan(d.orientation.w));
}

template <>
inline bool isnan(const irob_msgs::ToolPose& d)
{
  return (std::isnan(d.position.x)
          || std::isnan(d.position.y)
          || std::isnan(d.position.z)
          || std::isnan(d.orientation.x)
          || std::isnan(d.orientation.y)
          || std::isnan(d.orientation.z)
          || std::isnan(d.orientation.w)
          || std::isnan(d.jaw));
}

template <>
inline bool isnan(const geometry_msgs::Point& d)
{
  return (std::isnan(d.x)
          || std::isnan(d.y)
          || std::isnan(d.z) );
}

template <>
inline bool isnan(const geometry_msgs::PointStamped& d)
{
  return isnan<geometry_msgs::Point>(d.point);
}

template <>
inline bool isnan(const geometry_msgs::Quaternion& d)
{
  return (std::isnan(d.x)
          || std::isnan(d.y)
          || std::isnan(d.z)
          || std::isnan(d.w));
}

// Unit vector + rotation to quat
template<typename QuatT, typename VecT>
inline QuatT vecToQuat(const VecT& vec, double angle);

template <>
inline Eigen::Quaternion<double> vecToQuat(const Eigen::Vector3d& vec,
                                           double angle){
  Eigen::Quaternion<double> quat_start(0.0, 0.707107, 0.707106, 0.0);
  double angle_rad = (angle / 180.0) * M_PI;
  Eigen::Matrix3d R1m;
  R1m = Eigen::AngleAxisd(0.0, Eigen::Vector3d::UnitX())
      * Eigen::AngleAxisd(0.0,  Eigen::Vector3d::UnitY())
      * Eigen::AngleAxisd(angle_rad, Eigen::Vector3d::UnitZ());
  Eigen::Quaternion<double> R1(R1m);

  Eigen::Quaternion<double> ret = R1 * quat_start;

  Eigen::Vector3d vec_start(0.0, 0.0, -1.0);
  Eigen::Quaternion<double> R2 =
      Eigen::Quaternion<double>::FromTwoVectors(vec_start, vec);
  ret = R2 * ret;
  return ret;
}

// Quat to unit vector
template<typename QuatT, typename VecT>
inline VecT quatToVec(const QuatT& quat);

template <>
inline Eigen::Vector3d quatToVec(const Eigen::Quaternion<double>& quat){

  Eigen::Quaternion<double> quat_start(0.0, 0.707107, 0.707106, 0.0);

  Eigen::Quaternion<double> R = quat * quat_start.inverse();

  Eigen::Vector3d vec_start(0.0, 0.0, -1.0);
  Eigen::Vector3d ret = R * vec_start;
  return ret;
}

}

#endif /* DVRK_UTILS_HPP_ */


// [[Rcpp::depends(BH)]]

#include <cmath>
#include <boost/math/special_functions/beta.hpp>

#include "include/lm.h"

double pval_t(double t, double df) {
    if (df > 50) {
        return 1 + boost::math::erf(-std::abs(t) / sqrt(2));
    }
    double x = df / (t * t + df);
    return boost::math::ibeta(df / 2, 0.5, x);
}

lm::lm(unsigned _n) :lt(_n * 2), b(_n), v(_n), vM(_n), R(4), Q(_n * 2), n(_n) {}

double lm::norm(const std::vector<double>& v) {
    double value = 0.0;
    for (double x: v) {
        value += x * x;
    }
    return std::sqrt(value);
}

void lm::set(int i, double x, double y, double w, bool bias) {
    double sq = std::sqrt(w);
    if (bias) {
        lt[i] = sq;
    }
    lt[i + n] = x * sq;
    b[i] = y * sq;
}

void lm::solve() {
    preprocess(v, lt, 0, n, 0);

    // Q_1 * A
    std::vector<double> RM(n);
    double cross = 0.0;
    for (int j = 0; j < n; j++) {
        cross += v[j] * lt[n + j];
    }
    for (int i = 1; i < n; i++) {
        RM[i] = lt[n + i] - 2 * v[i] * cross;
    }

    vM[0] = 0;
    preprocess(vM, RM, 0, n, 1);
    compute_Q();
    compute_R();
    solve_system();
}

void lm::preprocess(std::vector<double>& ret, const std::vector<double>& lt, int l, int r, int s) {
    int n = r - l;
    for (int i = s; i < n; i++) {
        ret[i] = lt[i + l];
    }
    ret[s] -= norm(ret);
    double l2 = norm(ret);
    if (l2 > 0) {
        for (int i = 0; i < n; i++) {
            ret[i] /= l2;
        }
    }
}

void lm::compute_Q() {
    double cross = 0.0;
    for (int k = 0; k < n; k++) {
        cross += v[k] * vM[k];
    }
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < 2; j++) {
            double sum = 0;
            sum += 4 * v[i] * vM[j] * cross;
            Q[i + j * n] = sum - 2 * (v[i] * v[j] + vM[i] * vM[j]);
        }
    }
    Q[0] += 1;
    Q[n + 1] += 1;
}

void lm::compute_R() {
    for (int i = 0; i < 2; i++) {
        for (int j = 0; j < 2; j++) {
            if (i == 1 && j == 0) {
                continue;
            }
            double sum = 0;
            for (int k = 0; k < n; k++) {
                sum += Q[k + i * n] * lt[k + j * n];
            }
            R[i + j * 2] = sum;
        }
    }
    R[1] = 0;

}

void lm::solve_system() {
    double b1 = 0.0;
    double b2 = 0.0;
    for (int i = 0; i < n; i++) {
        b1 += Q[i] * b[i];
        b2 += Q[i + n] * b[i];
    }
    k = b2 / R[3];
    bias = (b1 - k * R[2]) / R[0];
}

double lm::compute_t(double df) {
    double a = R[0] * R[0];
    double b = R[0] * R[2];
    double c = R[2] * R[2] + R[3] * R[3];
    double inv =  a / (a * c - b * b);
    return k / std::sqrt((compute_rss() / df) * inv);

}

double lm::compute_rss() {
    double rss = 0;
    for (int i = 0; i < n; i++) {
        rss += pow(bias * lt[i] + k * lt[n + i] - b[i], 2.0);
    }
    return rss;
}

double lm::get_lambda() {
    return k;
}

double lm::get_bias() {
    return bias;
}


// Petter Strandmark 2012.

#include <algorithm>
#include <cstdio>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <vector>

#include <Eigen/Dense>

#include <spii/spii.h>
#include <spii/solver.h>

namespace spii {

// Holds a point in the Nelder-Mead simplex.
// Equipped with a comparison operator for sorting.
struct SimplexPoint
{
	Eigen::VectorXd x;
	double value;

	bool operator<(const SimplexPoint& rhs) const
	{
		return this->value < rhs.value;
	}
};

// If required for debugging.
std::ostream& operator<<(std::ostream& out, const SimplexPoint& point)
{
	out << point.x.transpose() << " : " << point.value;
	return out;
}

}  // namespace spii

namespace std
{
	template<>
	void swap<spii::SimplexPoint>(spii::SimplexPoint& lhs, spii::SimplexPoint& rhs)
	{
		lhs.x.swap(rhs.x);
		swap(lhs.value, rhs.value);
	}
}

namespace spii {

void initialize_simplex(const Function& function,
                        const Eigen::VectorXd& x0,
                        std::vector<SimplexPoint>* simplex)
{
	size_t n = function.get_number_of_scalars();
	Eigen::VectorXd absx0 = x0;
	for (size_t i = 0; i < n; ++i) {
		absx0[i] = std::abs(x0[i]);
	}
	double scale = std::max(absx0.maxCoeff(), 1.0);
	const double nd = static_cast<double>(n);
	double alpha1 = scale / (nd * std::sqrt(2.0)) * (std::sqrt(nd+1)- 1 + nd);
	double alpha2 = scale / (nd * std::sqrt(2.0)) * (std::sqrt(nd+1) - 1);
	Eigen::VectorXd alpha2_vec(x0.size());
	alpha2_vec.setConstant(alpha2);

	simplex->at(0).x = x0;
	for (size_t i = 1; i < n + 1; ++i) {
		simplex->at(i).x  = x0 + alpha2_vec;
		simplex->at(i).x[i-1] = x0[i-1] + alpha1;
	}

	for (size_t i = 0; i < n + 1; ++i) {
		simplex->at(i).value = function.evaluate(simplex->at(i).x);
	}

	std::sort(simplex->begin(), simplex->end());
}

void NelderMeadSolver::solve(const Function& function,
                             SolverResults* results) const
{
	double global_start_time = wall_time();

	// Dimension of problem.
	size_t n = function.get_number_of_scalars();

	if (n == 0) {
		results->exit_condition = SolverResults::FUNCTION_TOLERANCE;
		return;
	}

	// The Nelder-Mead simplex.
	std::vector<SimplexPoint> simplex(n + 1);

	// Copy the user state to the current point.
	Eigen::VectorXd x;
	function.copy_user_to_global(&x);

	initialize_simplex(function, x, &simplex);

	SimplexPoint mean_point;
	SimplexPoint reflection_point;
	SimplexPoint expansion_point;
	mean_point.x.resize(n);
	reflection_point.x.resize(n);
	expansion_point.x.resize(n);

	double fmin  = std::numeric_limits<double>::quiet_NaN();
	double fmax  = std::numeric_limits<double>::quiet_NaN();
	double fval  = std::numeric_limits<double>::quiet_NaN();
	double area  = std::numeric_limits<double>::quiet_NaN();
	double area0 = std::numeric_limits<double>::quiet_NaN();
	double length  = std::numeric_limits<double>::quiet_NaN();
	double length0 = std::numeric_limits<double>::quiet_NaN();

	Eigen::MatrixXd area_mat(n, n);

	//
	// START MAIN ITERATION
	//
	results->startup_time   += wall_time() - global_start_time;
	results->exit_condition = SolverResults::INTERNAL_ERROR;
	int iter = 0;
	int n_shrink_in_a_row = 0;
	while (true) {

		//
		// In each iteration, the worst point in the simplex
		// is replaced with a new one.
		//
		double start_time = wall_time();

		mean_point.x.setZero();
		fval = 0;
		// Compute the mean of the best n points.
		for (size_t i = 0; i < n; ++i) {
			mean_point.x += simplex[i].x;
			fval         += simplex[i].value;
		}
		fval         /= double(n);
		mean_point.x /= double(n);
		fmin = simplex[0].value;
		fmax = simplex[n].value;

		const char* iteration_type = "n/a";

		// Compute the reflexion point and evaluate it.
		reflection_point.x = 2.0 * mean_point.x - simplex[n].x;
		reflection_point.value = function.evaluate(reflection_point.x);

		bool is_shrink = false;
		if (simplex[0].value <= reflection_point.value &&
			reflection_point.value < simplex[n - 1].value) {
			// Reflected point is neither better nor worst in the
			// new simplex.
			std::swap(reflection_point, simplex[n]);
			iteration_type = "Reflect 1";
		}
		else if (reflection_point.value < simplex[0].value) {
			// Reflected point is better than the current best; try
			// to go farther along this direction.

			// Compute expansion point.
			expansion_point.x = 3.0 * mean_point.x - 2.0 * simplex[n].x;
			expansion_point.value = function.evaluate(expansion_point.x);

			if (expansion_point.value < reflection_point.value) {
				std::swap(expansion_point, simplex[n]);
				iteration_type = "Expansion";
			}
			else {
				std::swap(reflection_point, simplex[n]);
				iteration_type = "Reflect 2";
			}
		}
		else {
			// Reflected point is still worse than x[n]; contract.
			bool success = false;

			if (simplex[n - 1].value <= reflection_point.value &&
			    reflection_point.value < simplex[n].value) {
				// Try to perform "outside" contraction.
				expansion_point.x = 1.5 * mean_point.x - 0.5 * simplex[n].x;
				expansion_point.value = function.evaluate(expansion_point.x);

				if (expansion_point.value <= reflection_point.value) {
					std::swap(expansion_point, simplex[n]);
					success = true;
					iteration_type = "Outside contraction";
				}
			}
			else {
				// Try to perform "inside" contraction.
				expansion_point.x = 0.5 * mean_point.x + 0.5 * simplex[n].x;
				expansion_point.value = function.evaluate(expansion_point.x);

				if (expansion_point.value < simplex[n].value) {
					std::swap(expansion_point, simplex[n]);
					success = true;
					iteration_type = "Inside contraction";
				}
			}

			if (! success) {
				// Neither outside nor inside contraction was acceptable;
				// shrink the simplex toward the best point.
				for (size_t i = 1; i < n + 1; ++i) {
					simplex[i].x = 0.5 * (simplex[0].x + simplex[i].x);
					simplex[i].value = function.evaluate(simplex[i].x);
					iteration_type = "Shrink";
					is_shrink = true;
				}
			}
		}

		std::sort(simplex.begin(), simplex.end());

		results->function_evaluation_time += wall_time() - start_time;

		//
		// Test stopping criteriea
		//
		start_time = wall_time();

		// Compute the area of the simplex.
		length = 0;
		for (size_t i = 0; i < n; ++i) {
			area_mat.col(i) = simplex[i].x - simplex[n].x;
			length = std::max(length, area_mat.col(i).norm());
		}
		area = std::abs(area_mat.determinant());
		if (iter == 0) {
			area0 = area;
			length0 = length;
		}

		if (area / area0 < this->area_tolerance) {
			results->exit_condition = SolverResults::GRADIENT_TOLERANCE;
			break;
		}

		if (area == 0) {
			results->exit_condition = SolverResults::GRADIENT_TOLERANCE;
			break;
		}

		if (length / length0 < this->length_tolerance) {
			results->exit_condition = SolverResults::GRADIENT_TOLERANCE;
			break;
		}

		if (is_shrink) {
			n_shrink_in_a_row++;
		}
		else {
			n_shrink_in_a_row = 0;
		}
		if (n_shrink_in_a_row > 50) {
			results->exit_condition = SolverResults::GRADIENT_TOLERANCE;
			break;
		}

		if (iter >= this->maximum_iterations) {
			results->exit_condition = SolverResults::NO_CONVERGENCE;
			break;
		}

		if (this->callback_function) {
			CallbackInformation information;
			information.objective_value = simplex[0].value;
			information.x = &simplex[0].x;

			if (!callback_function(information)) {
				results->exit_condition = SolverResults::USER_ABORT;
				break;
			}
		}
		results->stopping_criteria_time += wall_time() - start_time;

		//
		// Restarting
		//
		//if (area / area1 < 1e-10) {
		//	x = simplex[0].x;
		//	initialize_simplex(function, x, &simplex);
		//	area1 = area;
		//	if (this->log_function) {
		//		this->log_function("Restarted.");
		//	}
		//}

		//
		// Log the results of this iteration.
		//
		start_time = wall_time();

		int log_interval = 1;
		if (iter > 30) {
			log_interval = 10;
		}
		if (iter > 200) {
			log_interval = 100;
		}
		if (iter > 2000) {
			log_interval = 1000;
		}
		if (this->log_function && iter % log_interval == 0) {
			char str[1024];
				if (iter == 0) {
					this->log_function("Itr     min(f)     avg(f)     max(f)    area    length   type");
				}
				std::sprintf(str, "%6d %+.3e %+.3e %+.3e %.3e %.3e %s",
					iter, fmin, fval, fmax, area, length, iteration_type);
			this->log_function(str);
		}
		results->log_time += wall_time() - start_time;

		iter++;
	}

	// Return the best point as solution.
	function.copy_global_to_user(simplex[0].x);
	results->total_time += wall_time() - global_start_time;

	if (this->log_function) {
		char str[1024];
		std::sprintf(str, " end   %+.3e                       %.3e %.3e", fval, area, length);
		this->log_function(str);
	}
}

}  // namespace spii


/**
 * @file particle.hpp
 * @author Vahid Bastani
 *
 * Generic particle filter.
 */
#ifndef SSMPACK_FILTER_PARTICLE_HPP
#define SSMPACK_FILTER_PARTICLE_HPP

#include "ssmkit/distribution/conditional.hpp"
#include "ssmkit/filter/recursive_bayesian_base.hpp"
#include "ssmkit/process/hierarchical.hpp"
#include "ssmkit/process/markov.hpp"
#include "ssmkit/process/memoryless.hpp"
#include <armadillo>

#include <tuple>

namespace ssmkit {
namespace filter {

using process::Hierarchical;
using process::Markov;
using process::Memoryless;
using distribution::Conditional;

/** Particle Filter.
 */
template <class Process, class Resampler>
class Particle : public RecursiveBayesianBase<Particle<Process, Resampler>> {
 public:
  /** Type of the state posterior
   *
   * \f$ \{\mathbf{x}^{(i)}_t,\omega^{(i)}\}_{i=1}^{M}\f$
   */
  using CompeleteState = std::tuple<arma::mat, arma::vec>;

 private:
  //! Particle weights \f$ \{\omega^{(i)}\}_{i=1}^{M}\f$.
  arma::vec w_;
  //! State particles \f$ \{\mathbf{x}^{(i)}_t\}_{i=1}^{M}\f$.
  arma::mat state_par_;
  //! The process model
  Process process_;
  //! Resampling algorithm
  Resampler resampler_;
  //! Number of particles \f$M\f$.
  unsigned long num_;

 private:
  //! Normalizes the sum of weights to one
  void normalizeWeights(void) { w_ = w_ / arma::sum(w_); }

 public:
  /** Constructor
   *
   * returns a Particle filter object.
   *
   * @param process The process model object that the PF is defined for
   * @param resampler The resampling algorithm
   * @param particles_num Number of particles \f$M\f$
   */
  Particle(Process process, Resampler resampler, unsigned long particles_num)
      : process_{process}, resampler_{resampler}, num_{particles_num} {
    // initialized w_ and state_par_
    w_.resize(num_);
    // take one sample to find out dimension
    auto tmp = process_.template getProcess<0>().getInitialPDF().random();
    state_par_.resize(tmp.size(), num_);
  }
  /** Prediction
   *
   * Performs prediction step.
   * \f{equation}{\mathbf{x}^{(i)}_t \sim p(\mathbf{x}_t| \tilde{\mathbf{x}}^{(i)}_{t-1}, y^d_1, \cdots, y^d_{N_d})
   * \quad \text{for} \quad i=1,\cdots,M\f}
   *
   * @param args... Control variables \f$y^d_1, \cdots, y^d_{N_d}\f$ of the dynamic process, if any.
   */
  template <class... Args>
  void predict(const Args &... args) {
    state_par_.each_col([this, &args...](arma::vec &v, const Args &... args) {
      v = process_.template getProcess<0>().getCPDF().random(v, args...);
    });
  }
  /** Correction
   *
   * Performs correction step.
   *
   * \f{equation}{ \omega^{(i)} = \tilde{\omega}^{(i)} p(\mathbf{z}_t| \mathbf{x}^{(i)}_t, y^m_1, \cdots, y^m_{N_m}) \f}
   * \f{equation}{\{\mathbf{x}^{(i)}_t,\omega^{(i)}\}_{i=1}^{M}
   * \overset{\mbox{resample}}{\longrightarrow}
   * \{\tilde{\mathbf{x}}^{(i)}_t,\tilde{\omega}^{(i)}\}_{i=1}^{M}\f}
   *
   * @param measurement Measurement vector \f$\mathbf{z}_t\f$.
   * @param args... Control variables \f$y^m_1, \cdots, y^m_{N_m}\f$ of the measurement process, if any.
   * @return Estimated state \f$\{\tilde{\mathbf{x}}^{(i)}_t,\tilde{\omega}^{(i)}\}_{i=1}^{M}\f$
   */
  template <class Measurement, class... TArgs>
  CompeleteState correct(const Measurement &measurement,
                         const TArgs &... args) {
    unsigned long cnt = 0;
    w_.for_each([this, &cnt, &measurement, &args...](double &e) {
      e *= process_.template getProcess<1>().getCPDF().likelihood(
          measurement, state_par_.col(cnt++), args...);
    });

    normalizeWeights();

    resampler_(state_par_, w_);

    return std::make_tuple(state_par_, w_);
  }
  /** Initialization
   *
   * @return Estimated state \f$\{\tilde{\mathbf{x}}^{(i)}_0,\tilde{\omega}^{(i)}\}_{i=1}^{M}\f$
   */
  CompeleteState initialize() {
    state_par_.each_col([this](arma::vec &v) {
      v = process_.template getProcess<0>().getInitialPDF().random();
    });

    unsigned long cnt = 0;
    w_.for_each([this, &cnt](double &e) {
      e = process_.template getProcess<0>().getInitialPDF().likelihood(
          state_par_.col(cnt++));
    });

    normalizeWeights();

    return std::make_tuple(state_par_, w_);
  }
  //! @return Estimated state \f$\{\tilde{\omega}^{(i)}\}_{i=1}^{M}\f$
  const arma::vec &getWeights(void) const { return w_; }
  //! @return Estimated state \f$\{\tilde{\mathbf{x}}^{(i)}_t\}_{i=1}^{M}\f$
  const arma::mat &getStateParticles(void) const { return state_par_; }
};

/**
 */
template <class StatePDF, class StateParamMap, class InitialPDF,
          class MeasurementPDF, class MeasurementParamMap, class Resampler>
auto makeParticle(
    Hierarchical<Markov<StatePDF, StateParamMap, InitialPDF>,
                 Memoryless<MeasurementPDF, MeasurementParamMap>> process,
    Resampler resampler, unsigned long particle_num) {
  return Particle<Hierarchical<Markov<StatePDF, StateParamMap, InitialPDF>,
                               Memoryless<MeasurementPDF, MeasurementParamMap>>,
                  Resampler>(process, resampler, particle_num);
}

} // namespace filter
} // namespace ssmkit

#endif // SSMPACK_FILTER_PARTICLE_HPP


// This ray-tracing code is adapted from Hugues Thomas' project.
#include <math.h>
#include <iostream>
#include <unordered_map>
#include <algorithm>
#include <set>
#include <cstdint>
#include <cstdio>
#include <boost/algorithm/string.hpp>
#include "ray_tracing.hpp"
#include "estimation.hpp"
#include "utils.hpp"
#include "pointmatcher/PointMatcher.h"

static Eigen::Vector3d max_point(Eigen::MatrixXd & pc) {
    Eigen::Vector3d maxP = Eigen::Vector3d::Zero();
    for (uint i = 0; i < pc.cols(); ++i) {
        if (pc(0, i) > maxP(0))
            maxP(0) = pc(0, i);
        if (pc(1, i) > maxP(1))
            maxP(1) = pc(1, i);
        if (pc(2, i) > maxP(2))
            maxP(2) = pc(2, i);
    }
    return maxP;
}

static Eigen::Vector3d min_point(Eigen::MatrixXd & pc) {
    Eigen::Vector3d minP = Eigen::Vector3d::Zero();
    for (uint i = 0; i < pc.cols(); ++i) {
        if (pc(0, i) < minP(0))
            minP(0) = pc(0, i);
        if (pc(1, i) < minP(1))
            minP(1) = pc(1, i);
        if (pc(2, i) < minP(2))
            minP(2) = pc(2, i);
    }
    return minP;
}

static void cart2pol_(Eigen::MatrixXd &cart_frame, Eigen::MatrixXd &polar_frame) {
    polar_frame = Eigen::MatrixXd::Ones(cart_frame.rows(), cart_frame.cols());
    for (uint i = 0; i < cart_frame.cols(); ++i) {
        Eigen::Vector3d p = cart_frame.block(0, i, 3, 1);
        float rho = p.norm();
        float phi = atan2f(p(1), p(0));
        float theta = atan2f(sqrt(p(0) * p(0) + p(1) * p(1)), p(2));
        polar_frame(0, i) = rho;
        polar_frame(1, i) = theta;
        polar_frame(2, i) = phi + M_PI / 2;
    }
}

static Eigen::Vector3d cart2pol(Eigen::Vector3d &p) {
    float rho = p.norm();
    float phi = atan2f(p(1), p(0));
    float theta = atan2f(sqrt(p(0) * p(0) + p(1) * p(1)), p(2));
    return Eigen::Vector3d(rho, theta, phi + M_PI / 2);
}

static void compare_map_to_frame(Eigen::MatrixXd &aligned_frame, Eigen::MatrixXd &map_points,
    Eigen::MatrixXd &map_normals, Eigen::Matrix4d T_sensor_map, std::unordered_map<VoxKey, size_t> &map_samples,
    float theta_dl, float phi_dl, float map_dl, std::vector<float> &movable_probs, std::vector<int> &movable_counts) {

    float inv_theta_dl = 1.0 / theta_dl;
    float inv_phi_dl = 1.0 / phi_dl;
    float inv_map_dl = 1.0 / map_dl;
    float max_angle = 5 * M_PI / 12;
    float min_vert_cos = cos(M_PI / 3);

    // Mask of the map points not updated yet
    std::vector<bool> not_updated(map_points.cols(), true);

    // Get limits
    Eigen::Vector3d min_P = min_point(aligned_frame) - Eigen::Vector3d(map_dl, map_dl, map_dl);
    Eigen::Vector3d max_P = max_point(aligned_frame) + Eigen::Vector3d(map_dl, map_dl, map_dl);

    // Update full voxels
    // Loop over aligned_frame
    VoxKey k0, k;
    for (uint i = 0; i < aligned_frame.cols(); ++i) {
        // Corresponding key
        k0.x = (int)floor(aligned_frame(0, i) * inv_map_dl);
        k0.y = (int)floor(aligned_frame(1, i) * inv_map_dl);
        k0.z = (int)floor(aligned_frame(2, i) * inv_map_dl);
        // Update the adjacent cells
        for (k.x = k0.x - 1; k.x < k0.x + 2; k.x++) {
            for (k.y = k0.y - 1; k.y < k0.y + 2; k.y++) {
                for (k.z = k0.z - 1; k.z < k0.z + 2; k.z++) {
                    // Update count and movable at this point
                    if (map_samples.count(k) > 0) {
                        // Only update once
                        size_t i0 = map_samples[k];
                        if (not_updated[i0]) {
                            not_updated[i0] = false;
                            movable_counts[i0] += 1;
                        }
                    }
                }
            }
        }
    }

    // Create the free frustum grid

    // get frame in polar coordinates
    Eigen::MatrixXd frame = T_sensor_map * aligned_frame;
    Eigen::MatrixXd polar_frame;
    cart2pol_(frame, polar_frame);

    // Get grid limits
    Eigen::Vector3d minCorner = min_point(polar_frame);
    Eigen::Vector3d maxCorner = max_point(polar_frame);
    Eigen::Vector3d originCorner = minCorner - Eigen::Vector3d(0, 0.5 * theta_dl, 0.5 * phi_dl);

    // Dimensions of the grid
    size_t grid_n_theta = (size_t)floor((maxCorner(1) - originCorner(1)) / theta_dl) + 1;
    size_t grid_n_phi = (size_t)floor((maxCorner(2) - originCorner(2)) / phi_dl) + 1;

    // Initialize variables
    std::vector<float> frustrum_radiuses(grid_n_theta * grid_n_phi, -1.0);

    // Fill the frustrum radiuses
    for (uint i = 0; i < polar_frame.cols(); ++i) {
        // Corresponding key
        Eigen::Vector3d p = polar_frame.block(0, i, 3, 1);
        // Position of point in grid
        size_t i_theta = (size_t)floor((p(1) - originCorner(1)) * inv_theta_dl);
        size_t i_phi = (size_t)floor((p(2) - originCorner(2)) * inv_phi_dl);
        size_t gridIdx = i_theta + grid_n_theta * i_phi;

        // Update the radius in cell
        if (frustrum_radiuses[gridIdx] < 0)
            frustrum_radiuses[gridIdx] = p(0);
        else if (p(0) < frustrum_radiuses[gridIdx])
            frustrum_radiuses[gridIdx] = p(0);
    }

    // Apply margin to free ranges
    float margin = map_dl;
    float frustrum_alpha = theta_dl / 2;
    for (auto &r : frustrum_radiuses) {
        float adapt_margin = r * frustrum_alpha;
        if (margin < adapt_margin)
            r -= adapt_margin;
        else
            r -= margin;
    }

    // Apply frustum casting

    // update free pixels
    float min_r = 2 * map_dl;
    Eigen::Matrix3d C_sensor_map = T_sensor_map.block(0, 0, 3, 3);

    for (uint i = 0; i < map_points.cols(); ++i) {
        // Ignore points updated just now
        if (!not_updated[i])
            continue;
        Eigen::Vector4d p = map_points.block(0, i, 4, 1);
        // Ignore points outside area of the frame
        if (p(0) > max_P(0) || p(1) > max_P(1) || p(2) > max_P(2) ||
            p(0) < min_P(0) || p(1) < min_P(1) || p(2) < min_P(2)) {
            continue;
        }

        // Align point in frame coordinates (and normal)
        p = T_sensor_map * p;
        Eigen::Vector3d xyz = p.block(0, 0, 3, 1);
        Eigen::Vector3d nxyz = C_sensor_map * map_normals.block(0, i, 3, 1);

        // Project in polar coordinates
        Eigen::Vector3d rtp = cart2pol(xyz);

        // Position of point in grid
        size_t i_theta = (size_t)floor((rtp(1) - originCorner(1)) * inv_theta_dl);
        size_t i_phi = (size_t)floor((rtp(2) - originCorner(2)) * inv_phi_dl);
        size_t gridIdx = i_theta + grid_n_theta * i_phi;

        // Update movable prob
        if (rtp(0) > min_r && rtp(0) < frustrum_radiuses[gridIdx]) {
            // Do not update if normal is horizontal and perpendicular to ray (to avoid removing walls)
            if (abs(nxyz(2)) > min_vert_cos) {
                movable_counts[i] += 1;
                movable_probs[i] += 1.0;
            } else {
                float angle = acos(std::min(fabs(xyz.dot(nxyz) / rtp(0)), 1.0));
                if (angle < max_angle) {
                    movable_counts[i] += 1;
                    movable_probs[i] += 1.0;
                }
            }
        }
    }
}

static void getNameFromPath(std::string path, std::string &name) {
    std::vector<std::string> parts;
    boost::split(parts, path, boost::is_any_of("/"));
    std::string endy_bit = parts[parts.size() - 1];
    boost::split(parts, endy_bit, boost::is_any_of("."));
    name = parts[0];
}

int main(int argc, const char *argv[]) {
    std::string root, config;
    std::string new_map_name = "map_no_movable.ply";
    validateArgs(argc, argv, root, config);
    // Params:
    float map_dl = 0.15;
    float theta_dl = 1.29 * M_PI / 180;
    float phi_dl = 0.1 * M_PI / 180;

    // Init point map
    std::cout << "Loading map..." << std::endl;
    DP map = DP::load(root + "map/map.ply");
    std::cout << "Finished loading map" << std::endl;

    // Create the pointmap voxels
    std::unordered_map<VoxKey, size_t> map_samples;
    uint N = map.features.cols();
    map_samples.reserve(N);
    float inv_map_dl = 1.0 / map_dl;
    VoxKey k0;

    for (uint i = 0; i < N; ++i) {
        k0.x = (int)floor(map.features(0, i) * inv_map_dl);
        k0.y = (int)floor(map.features(1, i) * inv_map_dl);
        k0.z = (int)floor(map.features(2, i) * inv_map_dl);
        if (map_samples.count(k0) < 1) {
            map_samples.emplace(k0, i);
        }
    }

    // Init map movable probabilities and counts
    std::vector<float> movable_probs(N, 0);
    std::vector<int> movable_counts(N, 0);

    // Calculate normal vectors for the map
    std::shared_ptr<PM::DataPointsFilter> normalFilter = PM::get().DataPointsFilterRegistrar.create(
        "SurfaceNormalDataPointsFilter", {{"knn", toParam(10)}, {"epsilon", toParam(5)}, {"keepNormals", toParam(1)},
        {"keepDensities", toParam(0)}});

    std::cout << "Calculating normals for the entire map..." << std::endl;
    map = normalFilter->filter(map);
    std::cout << "Finished calculating normals" << std::endl;
    uint norm_row = map.getDescriptorStartingRow("normals");
    Eigen::MatrixXd map_normals = map.descriptors.block(norm_row, 0, 3, N);

    // Start movable detection

    std::vector<std::string> frame_names;
    getMapFrames(root, frame_names);

    for (uint i = 0; i < frame_names.size(); ++i) {
        std::cout << "Ray tracing frame " << i << " / " << frame_names.size() - 1 << std::endl;
        // Load frame / ply file
        DP frame = DP::load(frame_names[i]);
        // Load pose for this frame
        Eigen::Matrix4d T_map_sensor = Eigen::Matrix4d::Identity();
        std::string name;
        getNameFromPath(frame_names[i], name);
        load_transform(root + "map/frame_poses/" + name + ".txt", T_map_sensor);
        Eigen::Matrix4d T_sensor_map = get_inverse_tf(T_map_sensor);
        // Perform ray-tracing to identity movable points
        compare_map_to_frame(frame.features, map.features, map_normals, T_sensor_map,
            map_samples, theta_dl, phi_dl, map_dl, movable_probs, movable_counts);
    }

    for (uint i = 0; i < movable_probs.size(); ++i) {
        movable_probs[i] = movable_probs[i] / (movable_counts[i] + 1e-6);
        if (movable_counts[i] < 1e-6)
            movable_probs[i] = -1;
    }
    map.allocateDescriptor("movable", 1);
    uint movable_row = map.getDescriptorStartingRow("movable");
    for (uint i = 0; i < movable_probs.size(); ++i) {
        map.descriptors(movable_row, i) = movable_probs[i];
    }
    std::cout << "Adding movable descriptor to the existing map..." << std::endl;
    map.removeDescriptor("normals");  // Get rid of normals
    map.save(root + "map/map.ply");

    std::cout << "Removing movable points from the map..." << std::endl;
    uint feat_dim = map.features.rows();
    map.removeDescriptor("movable");
    uint desc_dim = map.descriptors.rows();
    uint j = 0;
    for (uint i = 0; i < movable_probs.size(); ++i) {
        if (movable_probs[i] < 0.8) {
            map.features.block(0, j, feat_dim, 1) = map.features.block(0, i, feat_dim, 1);
            map.descriptors.block(0, j, desc_dim, 1) = map.descriptors.block(0, i, desc_dim, 1);
            j++;
        }
    }
    map.conservativeResize(j);

    std::cout << "Saving the filtered map..." << std::endl;
    map.save(root + "map/" + new_map_name);
}


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

    @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_TWO_ADD_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_TWO_ADD_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

  /*!

    @ingroup group-arithmetic
    This function object computes two reals of the type of the inputs
    (in an std::pair) @c r0 and @c r1 such that:

    @code
    r0 = x + y
    r1 = r0 -(x + y)
    @endcode

    using perfect arithmetic.

    Its main usage is to be able to compute
    sum of reals and the residual error using IEEE 754 arithmetic.

    @par Header <boost/simd/function/two_add.hpp>

    @par Example:

      @snippet two_add.cpp two_add

    @par Possible output:

      @snippet two_add.txt two_add

  **/
  std::pair<Value, Value> two_add(Value const& x, Value const& y);
} }
#endif

#include <boost/simd/function/scalar/two_add.hpp>
#include <boost/simd/function/simd/two_add.hpp>

#endif


#pragma once
#include <Eigen/Dense>
#include <fstream>
#include <iostream>
#include <vector>

using namespace Eigen;

template <class Type>
Matrix<Type, Dynamic, Dynamic, RowMajor>* load_img(const std::string& path, int num_channels, int img_length) {
    using Mat = Matrix<Type, Dynamic, Dynamic, RowMajor>;

    Mat* img = new Mat[num_channels];

    int cellno = 0, c = 0;
    int img_size = img_length * img_length;

    std::ifstream in(path);
    std::string line;
    std::vector<Type> values(img_size);
    while (std::getline(in, line)) {
        std::stringstream ss(line);
        std::string cell;
        while (std::getline(ss, cell, ',')) {
            Type val = Type(std::stod(cell));
            values[cellno++] = val;

            if (cellno == img_size) {
                cellno = 0;
                img[c] = Map<Mat>(values.data(), img_length, img_length);
                values.clear();
                ++c;
            }
        }
    }

    return img;
}

template <class Type>
Matrix<Type, Dynamic, Dynamic, RowMajor>** load_conv_weights(const std::string& path, int num_channels, int num_kernels, int kernel_length) {
    using Mat = Matrix<Type, Dynamic, Dynamic, RowMajor>;

    Mat** kernels = new Mat*[num_kernels];
    for (int k = 0; k < num_kernels; ++k) {
        kernels[k] = new Mat[num_channels];
    }

    int cellno = 0, c = 0, k = 0;
    int kernel_size = kernel_length * kernel_length;

    std::ifstream in(path);
    std::string line;
    std::vector<Type> values;
    while (std::getline(in, line)) {
        std::stringstream ss(line);
        std::string cell;
        while (std::getline(ss, cell, ',')) {
            Type val = Type(std::stod(cell));
            values.push_back(val);
            ++cellno;

            if (cellno == kernel_size) {
                cellno = 0;
                kernels[k][c] = Map<Mat>(values.data(), kernel_length, kernel_length);
                values.clear();
                ++c;
                if (c == num_channels) {
                    c = 0;
                    ++k;
                }
            }
        }
    }
    return kernels;
}

template <class Type>
Matrix<Type, 1, Dynamic, RowMajor> load_conv_biases(const std::string& path, int num_kernels) {
    using RowVec = Matrix<Type, 1, Dynamic, RowMajor>;

    std::ifstream in(path);
    std::string cell;
    std::vector<Type> values;
    while (std::getline(in, cell, ',')) {
        Type val = Type(std::stod(cell));
        values.push_back(val);
    }
    return Map<RowVec>(values.data(), 1, num_kernels);
}

template <class Type>
Matrix<Type, Dynamic, Dynamic, RowMajor> load_fc(const std::string& path, int M, int N) {
    std::ifstream indata(path);
    std::string line;
    std::vector<Type> values;
    while (std::getline(indata, line)) {
        std::stringstream lineStream(line);
        std::string cell;
        while (std::getline(lineStream, cell, ',')) {
            Type val = Type(std::stod(cell));
            values.push_back(val);
        }
    }
    return Map<Matrix<Type, Dynamic, Dynamic, RowMajor>>(values.data(), M, N);
}


/*
 * DIPlib 3.0
 * This file contains definitions for functions that do matrix computations using Eigen
 *
 * (c)2016-2017, Cris Luengo.
 * Based on original DIPlib code: (c)1995-2014, Delft University of Technology.
 *
 * 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 <array>
#include "diplib/library/numeric.h"

#if defined(__GNUG__) || defined(__clang__)
// For this file, turn off -Wsign-conversion, Eigen is really bad at this!
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wsign-conversion"
#if __GNUC__ == 11
#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
#endif
#endif

#include <Eigen/Eigenvalues>
#include <Eigen/SVD>

namespace dip {

namespace {

template< class T >
struct GreaterMagnitude {
   bool operator()( T const& a, T const& b ) const {
      return std::abs( a ) > std::abs( b );
   }
};

} // namespace

void SymmetricEigenDecomposition(
      dip::uint n,
      ConstSampleIterator< dfloat > input,
      SampleIterator< dfloat > lambdas,
      SampleIterator< dfloat > vectors
) {
   DIP_ASSERT( input.Stride() >= 0 ); // TODO: (here and other asserts in this file) Eigen doesn't support negative strides, but there's a ticket for that: https://gitlab.com/libeigen/eigen/-/issues/747
   Eigen::Map< Eigen::MatrixXd const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), n, n, Eigen::InnerStride<>( input.Stride() ));
   if( vectors ) {
      Eigen::SelfAdjointEigenSolver< Eigen::MatrixXd > eigensolver( matrix );
      //if( eigensolver.info() != Eigen::Success ) { abort(); }
      Eigen::VectorXd const& eigenvalues = eigensolver.eigenvalues();
      Eigen::MatrixXd const& eigenvectors = eigensolver.eigenvectors();
      std::vector< dip::uint > indices( n );
      std::iota( indices.begin(), indices.end(), 0 );
      std::sort( indices.begin(), indices.end(), [ & ]( dip::uint a, dip::uint b ) { return std::abs( eigenvalues[ b ] ) < std::abs( eigenvalues[ a ] ); } );
      for( dip::uint ii = 0; ii < n; ++ii ) {
         dip::uint kk = indices[ ii ];
         lambdas[ ii ] = eigenvalues[ kk ];
         dip::uint offset = ii * n;
         for( dip::uint jj = 0; jj < n; ++jj ) {
            vectors[ jj + offset ] = eigenvectors( jj, kk );
         }
      }
   } else {
      DIP_ASSERT( lambdas.Stride() >= 0 );
      Eigen::Map< Eigen::VectorXd, 0, Eigen::InnerStride<> > eigenvalues( lambdas.Pointer(), n, Eigen::InnerStride<>( lambdas.Stride() ));
      eigenvalues = matrix.selfadjointView< Eigen::Lower >().eigenvalues();
      std::sort( lambdas, lambdas + n, GreaterMagnitude< dfloat >() );
   }
}

void SymmetricEigenDecomposition2(
      ConstSampleIterator< dfloat > input,
      SampleIterator< dfloat > lambdas,
      SampleIterator< dfloat > vectors
) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::Matrix2d const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), Eigen::InnerStride<>( input.Stride() ));
   if( vectors ) {
      Eigen::SelfAdjointEigenSolver< Eigen::Matrix2d > eigensolver( matrix );
      //if( eigensolver.info() != Eigen::Success ) { abort(); }
      Eigen::Vector2d const& eigenvalues = eigensolver.eigenvalues();
      Eigen::Matrix2d const& eigenvectors = eigensolver.eigenvectors();
      dip::uint indices0 = 0;
      dip::uint indices1 = 1;
      if( std::abs( eigenvalues[ 0 ] ) < std::abs( eigenvalues[ 1 ] )) {
         indices0 = 1;
         indices1 = 0;
      }
      lambdas[ 0 ] = eigenvalues[ indices0 ];
      lambdas[ 1 ] = eigenvalues[ indices1 ];
      vectors[ 0 ] = eigenvectors( 0, indices0 );
      vectors[ 1 ] = eigenvectors( 1, indices0 );
      vectors[ 2 ] = eigenvectors( 0, indices1 );
      vectors[ 3 ] = eigenvectors( 1, indices1 );
   } else {
      DIP_ASSERT( lambdas.Stride() >= 0 );
      Eigen::Map< Eigen::Vector2d, 0, Eigen::InnerStride<> > eigenvalues( lambdas.Pointer(), Eigen::InnerStride<>( lambdas.Stride() ));
      eigenvalues = matrix.selfadjointView< Eigen::Lower >().eigenvalues();
      if( std::abs( lambdas[ 0 ] ) < std::abs( lambdas[ 1 ] )) {
         std::swap( lambdas[ 0 ], lambdas[ 1 ] );
      }
   }
}

void SymmetricEigenDecomposition3(
      ConstSampleIterator< dfloat > input,
      SampleIterator< dfloat > lambdas,
      SampleIterator< dfloat > vectors
) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::Matrix3d const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), Eigen::InnerStride<>( input.Stride() ));
   if( vectors ) {
      Eigen::SelfAdjointEigenSolver< Eigen::Matrix3d > eigensolver( matrix );
      //if( eigensolver.info() != Eigen::Success ) { abort(); }
      Eigen::Vector3d const& eigenvalues = eigensolver.eigenvalues();
      Eigen::Matrix3d const& eigenvectors = eigensolver.eigenvectors();
      std::array< dip::uint, 3 > indices{{ 0, 1, 2 }};
      std::sort( indices.begin(), indices.end(), [ & ]( dip::uint a, dip::uint b ) { return std::abs( eigenvalues[ b ] ) < std::abs( eigenvalues[ a ] ); } );
      for( dip::uint ii = 0; ii < 3; ++ii ) {
         dip::uint kk = indices[ ii ];
         lambdas[ ii ] = eigenvalues[ kk ];
         dip::uint offset = ii * 3;
         for( dip::uint jj = 0; jj < 3; ++jj ) {
            vectors[ jj + offset ] = eigenvectors( jj, kk );
         }
      }
   } else {
      DIP_ASSERT( lambdas.Stride() >= 0 );
      Eigen::Map< Eigen::Vector3d, 0, Eigen::InnerStride<> > eigenvalues( lambdas.Pointer(), Eigen::InnerStride<>( lambdas.Stride() ));
      eigenvalues = matrix.selfadjointView< Eigen::Lower >().eigenvalues();
      std::sort( lambdas, lambdas + 3, GreaterMagnitude< dfloat >() );
   }
}

void LargestEigenvector(
      dip::uint n,
      ConstSampleIterator< dfloat > input,
      SampleIterator< dfloat > vector
) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXd const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), n, n, Eigen::InnerStride<>( input.Stride() ));
   Eigen::SelfAdjointEigenSolver< Eigen::MatrixXd > eigensolver( matrix );
   //if( eigensolver.info() != Eigen::Success ) { abort(); }
   Eigen::VectorXd const& eigenvalues = eigensolver.eigenvalues();
   Eigen::MatrixXd const& eigenvectors = eigensolver.eigenvectors();
   std::vector< dip::uint > indices( n );
   std::iota( indices.begin(), indices.end(), 0 );
   std::sort( indices.begin(), indices.end(), [ & ]( dip::uint a, dip::uint b ) { return std::abs( eigenvalues[ b ] ) < std::abs( eigenvalues[ a ] ); } );
   dip::uint kk = indices[ 0 ];
   for( dip::uint jj = 0; jj < n; ++jj ) {
      vector[ jj ] = eigenvectors( jj, kk );
   }
}

void SmallestEigenvector(
      dip::uint n,
      ConstSampleIterator< dfloat > input,
      SampleIterator< dfloat > vector
) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXd const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), n, n, Eigen::InnerStride<>( input.Stride() ));
   Eigen::SelfAdjointEigenSolver< Eigen::MatrixXd > eigensolver( matrix );
   //if( eigensolver.info() != Eigen::Success ) { abort(); }
   Eigen::VectorXd const& eigenvalues = eigensolver.eigenvalues();
   Eigen::MatrixXd const& eigenvectors = eigensolver.eigenvectors();
   std::vector< dip::uint > indices( n );
   std::iota( indices.begin(), indices.end(), 0 );
   std::sort( indices.begin(), indices.end(), [ & ]( dip::uint a, dip::uint b ) { return std::abs( eigenvalues[ b ] ) < std::abs( eigenvalues[ a ] ); } );
   dip::uint kk = indices.back();
   for( dip::uint jj = 0; jj < n; ++jj ) {
      vector[ jj ] = eigenvectors( jj, kk );
   }
}

void EigenDecomposition(
      dip::uint n,
      ConstSampleIterator< dfloat > input,
      SampleIterator< dcomplex > lambdas,
      SampleIterator< dcomplex > vectors
) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXd const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), n, n, Eigen::InnerStride<>( input.Stride() ));
   if( vectors ) {
      Eigen::EigenSolver< Eigen::MatrixXd > eigensolver( matrix );
      Eigen::VectorXcd const& eigenvalues = eigensolver.eigenvalues();
      Eigen::MatrixXcd const& eigenvectors = eigensolver.eigenvectors();
      std::vector< dip::uint > indices( n );
      std::iota( indices.begin(), indices.end(), 0 );
      std::sort( indices.begin(), indices.end(), [ & ]( dip::uint a, dip::uint b ) { return std::abs( eigenvalues[ b ] ) < std::abs( eigenvalues[ a ] ); } );
      for( dip::uint ii = 0; ii < n; ++ii ) {
         dip::uint kk = indices[ ii ];
         lambdas[ ii ] = eigenvalues[ kk ];
         dip::uint offset = ii * n;
         for( dip::uint jj = 0; jj < n; ++jj ) {
            vectors[ jj + offset ] = eigenvectors( jj, kk );
         }
      }
   } else {
      DIP_ASSERT( lambdas.Stride() >= 0 );
      Eigen::Map< Eigen::VectorXcd, 0, Eigen::InnerStride<> > eigenvalues( lambdas.Pointer(), n, Eigen::InnerStride<>( lambdas.Stride() ));
      eigenvalues = matrix.eigenvalues();
      std::sort( lambdas, lambdas + n, GreaterMagnitude< dcomplex >() );
   }
}

void EigenDecomposition(
      dip::uint n,
      ConstSampleIterator< dcomplex > input,
      SampleIterator< dcomplex > lambdas,
      SampleIterator< dcomplex > vectors
) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXcd const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), n, n, Eigen::InnerStride<>( input.Stride() ));
   if( vectors ) {
      Eigen::ComplexEigenSolver< Eigen::MatrixXcd > eigensolver( matrix );
      Eigen::VectorXcd const& eigenvalues = eigensolver.eigenvalues();
      Eigen::MatrixXcd const& eigenvectors = eigensolver.eigenvectors();
      std::vector< dip::uint > indices( n );
      std::iota( indices.begin(), indices.end(), 0 );
      std::sort( indices.begin(), indices.end(), [ & ]( dip::uint a, dip::uint b ) { return std::abs( eigenvalues[ b ] ) < std::abs( eigenvalues[ a ] ); } );
      for( dip::uint ii = 0; ii < n; ++ii ) {
         dip::uint kk = indices[ ii ];
         lambdas[ ii ] = eigenvalues[ kk ];
         dip::uint offset = ii * n;
         for( dip::uint jj = 0; jj < n; ++jj ) {
            vectors[ jj + offset ] = eigenvectors( jj, kk );
         }
      }
   } else {
      DIP_ASSERT( lambdas.Stride() >= 0 );
      Eigen::Map< Eigen::VectorXcd, 0, Eigen::InnerStride<> > eigenvalues( lambdas.Pointer(), n, Eigen::InnerStride<>( lambdas.Stride() ));
      eigenvalues = matrix.eigenvalues();
      std::sort( lambdas, lambdas + n, GreaterMagnitude< dcomplex >() );
   }
}

dfloat Determinant( dip::uint n, ConstSampleIterator< dfloat > input ) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXd const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), n, n, Eigen::InnerStride<>( input.Stride() ));
   return matrix.determinant();
}

dcomplex Determinant( dip::uint n, ConstSampleIterator< dcomplex > input ) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXcd const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), n, n, Eigen::InnerStride<>( input.Stride() ));
   return matrix.determinant();
}

void SingularValueDecomposition(
      dip::uint m,
      dip::uint n,
      ConstSampleIterator< dfloat > input,
      SampleIterator< dfloat > Sout,
      SampleIterator< dfloat > Uout,
      SampleIterator< dfloat > Vout
) {
   dip::uint p = std::min( m, n );
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXd const, 0, Eigen::InnerStride<> > M( input.Pointer(), m, n, Eigen::InnerStride<>( input.Stride() ));
   Eigen::JacobiSVD< Eigen::MatrixXd > svd( M, Eigen::ComputeThinU | Eigen::ComputeThinV );
   DIP_ASSERT( Sout.Stride() >= 0 );
   Eigen::Map< Eigen::VectorXd, 0, Eigen::InnerStride<> > S( Sout.Pointer(), p, Eigen::InnerStride<>( Sout.Stride() ));
   S = svd.singularValues();
   if( Uout && Vout ) {
      DIP_ASSERT( Uout.Stride() >= 0 );
      Eigen::Map< Eigen::MatrixXd, 0, Eigen::InnerStride<> > U( Uout.Pointer(), m, p, Eigen::InnerStride<>( Uout.Stride() ));
      DIP_ASSERT( Vout.Stride() >= 0 );
      Eigen::Map< Eigen::MatrixXd, 0, Eigen::InnerStride<> > V( Vout.Pointer(), n, p, Eigen::InnerStride<>( Vout.Stride() ));
      U = svd.matrixU();
      V = svd.matrixV();
   }
}

void SingularValueDecomposition(
      dip::uint m,
      dip::uint n,
      ConstSampleIterator< dcomplex > input,
      SampleIterator< dcomplex > Sout,
      SampleIterator< dcomplex > Uout,
      SampleIterator< dcomplex > Vout
) {
   dip::uint p = std::min( m, n );
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXcd const, 0, Eigen::InnerStride<> > M( input.Pointer(), m, n, Eigen::InnerStride<>( input.Stride() ));
   Eigen::JacobiSVD< Eigen::MatrixXcd > svd( M, Eigen::ComputeThinU | Eigen::ComputeThinV );
   DIP_ASSERT( Sout.Stride() >= 0 );
   Eigen::Map< Eigen::VectorXcd, 0, Eigen::InnerStride<> > S( Sout.Pointer(), p, Eigen::InnerStride<>( Sout.Stride() ));
   S = svd.singularValues();
   if( Uout && Vout ) {
      DIP_ASSERT( Uout.Stride() >= 0 );
      Eigen::Map< Eigen::MatrixXcd, 0, Eigen::InnerStride<> > U( Uout.Pointer(), m, p, Eigen::InnerStride<>( Uout.Stride() ));
      DIP_ASSERT( Vout.Stride() >= 0 );
      Eigen::Map< Eigen::MatrixXcd, 0, Eigen::InnerStride<> > V( Vout.Pointer(), n, p, Eigen::InnerStride<>( Vout.Stride() ));
      U = svd.matrixU();
      V = svd.matrixV();
   }
}

void Inverse( dip::uint n, ConstSampleIterator< dfloat > input, SampleIterator< dfloat > output ) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXd const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), n, n, Eigen::InnerStride<>( input.Stride() ));
   DIP_ASSERT( output.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXd, 0, Eigen::InnerStride<> > result( output.Pointer(), n, n, Eigen::InnerStride<>( output.Stride() ));
   result = matrix.inverse();
}

void Inverse( dip::uint n, ConstSampleIterator< dcomplex > input, SampleIterator< dcomplex > output ) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXcd const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), n, n, Eigen::InnerStride<>( input.Stride() ));
   DIP_ASSERT( output.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXcd, 0, Eigen::InnerStride<> > result( output.Pointer(), n, n, Eigen::InnerStride<>( output.Stride() ));
   result = matrix.inverse();
}

void PseudoInverse(
      dip::uint m,
      dip::uint n,
      ConstSampleIterator< dfloat > input,
      SampleIterator< dfloat > output,
      dfloat tolerance
) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXd const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), m, n, Eigen::InnerStride<>( input.Stride() ));
   Eigen::JacobiSVD< Eigen::MatrixXd > svd( matrix, Eigen::ComputeThinU | Eigen::ComputeThinV );
   tolerance = tolerance * static_cast< dfloat >( std::max( m, n )) * svd.singularValues().array().abs()( 0 );
   DIP_ASSERT( output.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXd, 0, Eigen::InnerStride<> > result( output.Pointer(), n, m, Eigen::InnerStride<>( output.Stride() ));
   result = svd.matrixV() *
            ( svd.singularValues().array().abs() > tolerance ).select( svd.singularValues().array().inverse(), 0 )
                                                              .matrix().asDiagonal()
            * svd.matrixU().adjoint();
}

void PseudoInverse(
      dip::uint m,
      dip::uint n,
      ConstSampleIterator< dcomplex > input,
      SampleIterator< dcomplex > output,
      dfloat tolerance
) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXcd const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), m, n, Eigen::InnerStride<>( input.Stride() ));
   Eigen::JacobiSVD< Eigen::MatrixXcd > svd( matrix, Eigen::ComputeThinU | Eigen::ComputeThinV );
   tolerance = tolerance * static_cast< dfloat >( std::max( m, n )) * svd.singularValues().array().abs()( 0 );
   DIP_ASSERT( output.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXcd, 0, Eigen::InnerStride<> > result( output.Pointer(), n, m, Eigen::InnerStride<>( output.Stride() ));
   result = svd.matrixV() *
            ( svd.singularValues().array().abs() > tolerance ).select( svd.singularValues().array().inverse(), 0 )
                                                              .matrix().asDiagonal()
            * svd.matrixU().adjoint();
}

dip::uint Rank( dip::uint m, dip::uint n, ConstSampleIterator< dfloat > input ) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXd const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), m, n, Eigen::InnerStride<>( input.Stride() ));
   Eigen::CompleteOrthogonalDecomposition< Eigen::MatrixXd > decomposition( matrix );
   return static_cast< dip::uint >( decomposition.rank() );
}

dip::uint Rank( dip::uint m, dip::uint n, ConstSampleIterator< dcomplex > input ) {
   DIP_ASSERT( input.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXcd const, 0, Eigen::InnerStride<> > matrix( input.Pointer(), m, n, Eigen::InnerStride<>( input.Stride() ));
   Eigen::CompleteOrthogonalDecomposition< Eigen::MatrixXcd > decomposition( matrix );
   return static_cast< dip::uint >( decomposition.rank() );
}

void Solve(
      dip::uint m,
      dip::uint n,
      ConstSampleIterator< dfloat > A,
      ConstSampleIterator< dfloat > b,
      SampleIterator< dfloat > output
) {
   DIP_ASSERT( A.Stride() >= 0 );
   DIP_ASSERT( b.Stride() >= 0 );
   Eigen::Map< Eigen::MatrixXd const, 0, Eigen::InnerStride<> > matrix( A.Pointer(), m, n, Eigen::InnerStride<>( A.Stride() ));
   Eigen::JacobiSVD< Eigen::MatrixXd > svd( matrix, Eigen::ComputeThinU | Eigen::ComputeThinV );
   Eigen::Map< Eigen::VectorXd const, 0, Eigen::InnerStride<> > vector( b.Pointer(), m, Eigen::InnerStride<>( b.Stride() ));
   Eigen::Map< Eigen::VectorXd, 0, Eigen::InnerStride<> > result( output.Pointer(), n, Eigen::InnerStride<>( output.Stride() ));
   result = svd.solve( vector );
}

} // namespace dip

#if defined(__GNUG__) || defined(__clang__)
#pragma GCC diagnostic pop
#endif

#ifdef DIP_CONFIG_ENABLE_DOCTEST
#include "doctest.h"

DOCTEST_TEST_CASE("[DIPlib] testing the EigenDecomposition functions") {
   // Test generic symmetric code with 2x2 matrix
   dip::dfloat matrix2[] = { 4, 8, 0 };
   dip::dfloat lambdas[ 3 ];
   dip::dfloat vectors[ 9 ];
   dip::SymmetricEigenDecompositionPacked( 2, matrix2, lambdas );
   DOCTEST_CHECK( lambdas[ 0 ] == 8 );
   DOCTEST_CHECK( lambdas[ 1 ] == 4 );
   dip::SymmetricEigenDecompositionPacked( 2, matrix2, lambdas, vectors );
   DOCTEST_CHECK( lambdas[ 0 ] == 8 );
   DOCTEST_CHECK( lambdas[ 1 ] == 4 );
   DOCTEST_CHECK( vectors[ 0 ] == 0 );
   DOCTEST_CHECK( vectors[ 1 ] == 1 );
   DOCTEST_CHECK( vectors[ 2 ] == 1 );
   DOCTEST_CHECK( vectors[ 3 ] == 0 );
   matrix2[ 0 ] = 8;
   matrix2[ 1 ] = 4;
   dip::SymmetricEigenDecompositionPacked( 2, matrix2, lambdas );
   DOCTEST_CHECK( lambdas[ 0 ] == 8 );
   DOCTEST_CHECK( lambdas[ 1 ] == 4 );
   dip::SymmetricEigenDecompositionPacked( 2, matrix2, lambdas, vectors );
   DOCTEST_CHECK( lambdas[ 0 ] == 8 );
   DOCTEST_CHECK( lambdas[ 1 ] == 4 );
   DOCTEST_CHECK( vectors[ 0 ] == 1 );
   DOCTEST_CHECK( vectors[ 1 ] == 0 );
   DOCTEST_CHECK( vectors[ 2 ] == 0 );
   DOCTEST_CHECK( vectors[ 3 ] == 1 );
   matrix2[ 0 ] = 3;
   matrix2[ 1 ] = 3;
   matrix2[ 2 ] = -1;
   dip::SymmetricEigenDecompositionPacked( 2, matrix2, lambdas );
   DOCTEST_CHECK( lambdas[ 0 ] == 4 );
   DOCTEST_CHECK( lambdas[ 1 ] == 2 );
   dip::SymmetricEigenDecompositionPacked( 2, matrix2, lambdas, vectors );
   DOCTEST_CHECK( lambdas[ 0 ] == 4 );
   DOCTEST_CHECK( lambdas[ 1 ] == 2 );
   DOCTEST_CHECK( vectors[ 0 ] == doctest::Approx(  std::cos( dip::pi/4 ))); // signs might be different here...
   DOCTEST_CHECK( vectors[ 1 ] == doctest::Approx( -std::sin( dip::pi/4 )));
   DOCTEST_CHECK( vectors[ 2 ] == doctest::Approx(  std::sin( dip::pi/4 )));
   DOCTEST_CHECK( vectors[ 3 ] == doctest::Approx(  std::cos( dip::pi/4 )));

   // Test 2x2-specific symmetric code
   dip::dfloat matrix2f[] = { 4, 0, 0, 8 };
   dip::SymmetricEigenDecomposition2( matrix2f, lambdas );
   DOCTEST_CHECK( lambdas[ 0 ] == 8 );
   DOCTEST_CHECK( lambdas[ 1 ] == 4 );
   dip::SymmetricEigenDecomposition2( matrix2f, lambdas, vectors );
   DOCTEST_CHECK( lambdas[ 0 ] == 8 );
   DOCTEST_CHECK( lambdas[ 1 ] == 4 );
   DOCTEST_CHECK( vectors[ 0 ] == 0 );
   DOCTEST_CHECK( vectors[ 1 ] == 1 );
   DOCTEST_CHECK( vectors[ 2 ] == 1 );
   DOCTEST_CHECK( vectors[ 3 ] == 0 );
   matrix2f[ 0 ] = 8;
   matrix2f[ 3 ] = 4;
   dip::SymmetricEigenDecomposition2( matrix2f, lambdas );
   DOCTEST_CHECK( lambdas[ 0 ] == 8 );
   DOCTEST_CHECK( lambdas[ 1 ] == 4 );
   dip::SymmetricEigenDecomposition2( matrix2f, lambdas, vectors );
   DOCTEST_CHECK( lambdas[ 0 ] == 8 );
   DOCTEST_CHECK( lambdas[ 1 ] == 4 );
   DOCTEST_CHECK( vectors[ 0 ] == 1 );
   DOCTEST_CHECK( vectors[ 1 ] == 0 );
   DOCTEST_CHECK( vectors[ 2 ] == 0 );
   DOCTEST_CHECK( vectors[ 3 ] == 1 );
   matrix2f[ 0 ] = 3;
   matrix2f[ 1 ] = -1;
   matrix2f[ 2 ] = -1;
   matrix2f[ 3 ] = 3;
   dip::SymmetricEigenDecomposition2( matrix2f, lambdas );
   DOCTEST_CHECK( lambdas[ 0 ] == 4 );
   DOCTEST_CHECK( lambdas[ 1 ] == 2 );
   dip::SymmetricEigenDecomposition2( matrix2f, lambdas, vectors );
   DOCTEST_CHECK( lambdas[ 0 ] == 4 );
   DOCTEST_CHECK( lambdas[ 1 ] == 2 );
   DOCTEST_CHECK( vectors[ 0 ] == doctest::Approx(  std::cos( dip::pi/4 ))); // signs might be different here...
   DOCTEST_CHECK( vectors[ 1 ] == doctest::Approx( -std::sin( dip::pi/4 )));
   DOCTEST_CHECK( vectors[ 2 ] == doctest::Approx(  std::sin( dip::pi/4 )));
   DOCTEST_CHECK( vectors[ 3 ] == doctest::Approx(  std::cos( dip::pi/4 )));

   // Test generic symmetric code with 3x3 matrix
   dip::dfloat matrix3[] = { 3, 1.5, 1.5, 0.0, 0.0, -0.5 };
   dip::SymmetricEigenDecompositionPacked( 3, matrix3, lambdas );
   DOCTEST_CHECK( lambdas[ 0 ] == 3 );
   DOCTEST_CHECK( lambdas[ 1 ] == 2 );
   DOCTEST_CHECK( lambdas[ 2 ] == 1 );
   dip::SymmetricEigenDecompositionPacked( 3, matrix3, lambdas, vectors );
   DOCTEST_CHECK( lambdas[ 0 ] == 3 );
   DOCTEST_CHECK( lambdas[ 1 ] == 2 );
   DOCTEST_CHECK( lambdas[ 2 ] == 1 );
   DOCTEST_CHECK( vectors[ 0 ] == doctest::Approx(  1.0 ));
   DOCTEST_CHECK( vectors[ 1 ] == doctest::Approx(  0.0 ));
   DOCTEST_CHECK( vectors[ 2 ] == doctest::Approx(  0.0 ));
   DOCTEST_CHECK( vectors[ 3 ] == doctest::Approx(  0.0 ));
   DOCTEST_CHECK( vectors[ 4 ] == doctest::Approx(  1.0 / std::sqrt( 2.0 )));
   DOCTEST_CHECK( vectors[ 5 ] == doctest::Approx( -1.0 / std::sqrt( 2.0 )));
   DOCTEST_CHECK( vectors[ 6 ] == doctest::Approx(  0.0 ));
   DOCTEST_CHECK( vectors[ 7 ] == doctest::Approx(  1.0 / std::sqrt( 2.0 )));
   DOCTEST_CHECK( vectors[ 8 ] == doctest::Approx(  1.0 / std::sqrt( 2.0 )));

   // Test 3x3-specific symmetric code
   dip::dfloat matrix3f[] = { 3, 0.0, 0.0, 0.0, 1.5, -0.5, 0.0, -0.5, 1.5 };
   dip::SymmetricEigenDecomposition3( matrix3f, lambdas );
   DOCTEST_CHECK( lambdas[ 0 ] == 3 );
   DOCTEST_CHECK( lambdas[ 1 ] == 2 );
   DOCTEST_CHECK( lambdas[ 2 ] == 1 );
   dip::SymmetricEigenDecomposition3( matrix3f, lambdas, vectors );
   DOCTEST_CHECK( lambdas[ 0 ] == 3 );
   DOCTEST_CHECK( lambdas[ 1 ] == 2 );
   DOCTEST_CHECK( lambdas[ 2 ] == 1 );
   DOCTEST_CHECK( vectors[ 0 ] == doctest::Approx(  1.0 ));
   DOCTEST_CHECK( vectors[ 1 ] == doctest::Approx(  0.0 ));
   DOCTEST_CHECK( vectors[ 2 ] == doctest::Approx(  0.0 ));
   DOCTEST_CHECK( vectors[ 3 ] == doctest::Approx(  0.0 ));
   DOCTEST_CHECK( vectors[ 4 ] == doctest::Approx(  1.0 / std::sqrt( 2.0 )));
   DOCTEST_CHECK( vectors[ 5 ] == doctest::Approx( -1.0 / std::sqrt( 2.0 )));
   DOCTEST_CHECK( vectors[ 6 ] == doctest::Approx(  0.0 ));
   DOCTEST_CHECK( vectors[ 7 ] == doctest::Approx(  1.0 / std::sqrt( 2.0 )));
   DOCTEST_CHECK( vectors[ 8 ] == doctest::Approx(  1.0 / std::sqrt( 2.0 )));

   // Test generic non-symmetric code with 2x2 matrix
   dip::dfloat matrix22[] = { 3, -1, -1, 3 };
   dip::dcomplex c_lambdas[ 2 ];
   dip::dcomplex c_vectors[ 4 ];
   dip::EigenDecomposition( 2, matrix22, c_lambdas, c_vectors );
   DOCTEST_CHECK( c_lambdas[ 0 ].real() == doctest::Approx( 4.0 ));
   DOCTEST_CHECK( c_lambdas[ 1 ].real() == doctest::Approx( 2.0 ));
   DOCTEST_CHECK( c_vectors[ 0 ].real() == doctest::Approx(  cos( dip::pi/4 ))); // signs might be different here...
   DOCTEST_CHECK( c_vectors[ 1 ].real() == doctest::Approx( -sin( dip::pi/4 )));
   DOCTEST_CHECK( c_vectors[ 2 ].real() == doctest::Approx(  sin( dip::pi/4 )));
   DOCTEST_CHECK( c_vectors[ 3 ].real() == doctest::Approx(  cos( dip::pi/4 )));
   DOCTEST_CHECK( c_lambdas[ 0 ].imag() == 0 );
   DOCTEST_CHECK( c_lambdas[ 1 ].imag() == 0 );
   DOCTEST_CHECK( c_vectors[ 0 ].imag() == 0 );
   DOCTEST_CHECK( c_vectors[ 1 ].imag() == 0 );
   DOCTEST_CHECK( c_vectors[ 2 ].imag() == 0 );
   DOCTEST_CHECK( c_vectors[ 3 ].imag() == 0 );

}

DOCTEST_TEST_CASE("[DIPlib] testing the SingularValueDecomposition and related functions") {
   dip::dfloat matrix22[] = { 4, 0, 0, 8 };
   dip::dfloat S[ 2 ];
   dip::dfloat U[ 4 ];
   dip::dfloat V[ 6 ];
   dip::SingularValueDecomposition( 2, 2, matrix22, S, U, V );
   DOCTEST_CHECK( S[ 0 ] == 8 );
   DOCTEST_CHECK( S[ 1 ] == 4 );
   DOCTEST_CHECK( U[ 0 ] == 0 );
   DOCTEST_CHECK( U[ 1 ] == 1 );
   DOCTEST_CHECK( U[ 2 ] == 1 );
   DOCTEST_CHECK( U[ 3 ] == 0 );
   DOCTEST_CHECK( V[ 0 ] == 0 );
   DOCTEST_CHECK( V[ 1 ] == 1 );
   DOCTEST_CHECK( V[ 2 ] == 1 );
   DOCTEST_CHECK( V[ 3 ] == 0 );

   dip::dfloat matrix23[] = { 3, 2, 2, 3, 2, -2 };
   dip::SingularValueDecomposition( 2, 3, matrix23, S, U, V );
   DOCTEST_CHECK( S[ 0 ] == doctest::Approx( 5.0 ));
   DOCTEST_CHECK( S[ 1 ] == doctest::Approx( 3.0 ));
   DOCTEST_CHECK( U[ 0 ] == doctest::Approx( -1.0 / std::sqrt( 2 ))); // signs might be different here...
   DOCTEST_CHECK( U[ 1 ] == doctest::Approx( -1.0 / std::sqrt( 2 )));
   DOCTEST_CHECK( U[ 2 ] == doctest::Approx( 1.0 / std::sqrt( 2 )));
   DOCTEST_CHECK( U[ 3 ] == doctest::Approx( -1.0 / std::sqrt( 2 )));
   DOCTEST_CHECK( V[ 0 ] == doctest::Approx( -1.0 / std::sqrt( 2 )));
   DOCTEST_CHECK( V[ 1 ] == doctest::Approx( -1.0 / std::sqrt( 2 )));
   DOCTEST_CHECK( V[ 2 ] == doctest::Approx( 0.0 ));
   DOCTEST_CHECK( V[ 3 ] == doctest::Approx( 1.0 / std::sqrt( 18 )));
   DOCTEST_CHECK( V[ 4 ] == doctest::Approx( -1.0 / std::sqrt( 18 )));
   DOCTEST_CHECK( V[ 5 ] == doctest::Approx( 4.0 / std::sqrt( 18 )));

   DOCTEST_CHECK( dip::Rank( 2, 3, matrix23 ) == 2 );

   dip::dfloat matrix32[ 6 ];
   dip::PseudoInverse( 2, 3, matrix23, matrix32 );
   DOCTEST_CHECK( matrix32[ 0 ] == doctest::Approx( 28.0 / 180.0 ));
   DOCTEST_CHECK( matrix32[ 1 ] == doctest::Approx(  8.0 / 180.0 ));
   DOCTEST_CHECK( matrix32[ 2 ] == doctest::Approx( 40.0 / 180.0 ));
   DOCTEST_CHECK( matrix32[ 3 ] == doctest::Approx(  8.0 / 180.0 ));
   DOCTEST_CHECK( matrix32[ 4 ] == doctest::Approx( 28.0 / 180.0 ));
   DOCTEST_CHECK( matrix32[ 5 ] == doctest::Approx(-40.0 / 180.0 ));

   dip::dfloat b[ 3 ] = { 44.0 / 180.0, 64.0 / 180.0, -40.0 / 180.0 };
   dip::dfloat x[ 2 ];
   dip::Solve( 3, 2, matrix32, b, x );
   DOCTEST_CHECK( x[ 0 ] == doctest::Approx( 1.0 ));
   DOCTEST_CHECK( x[ 1 ] == doctest::Approx( 2.0 ));
}

#endif // DIP_CONFIG_ENABLE_DOCTEST


/**
 * @file Sonar.cpp
 * @brief Projection functions for imaging sonar, e.g. DIDSON
 * @author: Michael Kaess
 * @date: Aug 2014
 */

#include <iostream>
#include <vector>
#include <utility>

#include <Eigen/Dense>

#include "Sonar.h"

using namespace std;
using namespace isam;
using namespace Eigen;

namespace sonar {

Sonar::Sonar(double rMin, double rMax, double bearingFov, double elevationFov,
    int numBearings, int numRanges) :
    _rMin(rMin), _rMax(rMax), _bearingFov(bearingFov), _elevationFov(
        elevationFov), _numBearings(numBearings), _numRanges(numRanges) {
  if (rMax <= rMin) {
    cout << "ERROR: rMax must be larger than rMin" << endl;
    exit(1);
  }
  if (bearingFov <= 0 || elevationFov <= 0) {
    cout << "ERROR: bearingFov and elevationFov must be positive" << endl;
    exit(1);
  }
  if (numBearings <= 0 || numRanges <= 0) {
    cout << "ERROR: numBearing and numRanges must be positive" << endl;
    exit(1);
  }
}

vector<Point3d> Sonar::getFrustum(const isam::Pose3d& pose) const {
  vector<Point3d> frustum;
  vector<double> range(2);
  range[0] = _rMin; range[1] = _rMax;
  vector<pair<double, double> > angles;
  double yaw = 0.5 * _bearingFov;
  double pitch = 0.5 * _elevationFov;
  angles.push_back(make_pair(-yaw, -pitch));
  angles.push_back(make_pair(yaw, -pitch));
  angles.push_back(make_pair(yaw, pitch));
  angles.push_back(make_pair(-yaw, pitch));

  for (int r = 0; r < 2; r++) {
    for (int a = 0; a < 4; a++) {
      yaw = angles[a].first;
      pitch = angles[a].second;
      Point3d p_local = Point3d(range[r], 0, 0);
      Point3d p_sonar = Pose3d(0, 0, 0, yaw, pitch, 0).transform_from(p_local);
      frustum.push_back(pose.transform_from(p_sonar));
    }
  }

  return frustum;
}

bool Sonar::project(const isam::Pose3d& pose, const isam::Point3d& point,
    isam::Point2d& projection) const {
  Point3d point_sonar = pose.transform_to(point);
  double range = point_sonar.vector().norm();
  double x = point_sonar.x();
  double y = point_sonar.y();
  double z = point_sonar.z();
  Vector2d direction(point_sonar.x(), point_sonar.y());
  direction.normalize();
  Vector2d mapping = direction * range;

  double u = mapping(0);
  double v = mapping(1);
  double b = (atan2(v, u) / _bearingFov + 0.5) * _numBearings;
  double r = (range - _rMin) / (_rMax - _rMin) * _numRanges;
  double elevation = atan2(z, sqrt(x * x + y * y)) / _elevationFov;

  projection = Point2d(b, r);

  // projection within field of view of sensor?
  if (b >= 0 && b < _numBearings && r >= 0 && r < _numRanges
      && elevation >= -0.5 && elevation <= 0.5) {
    return true;
  }

  return false;
}

} /* namespace sonar */


/*
 * Copyright (c) 2010-2012 Steffen Kieß
 *
 * 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 MATH_DIAGMATRIX3_HPP_INCLUDED
#define MATH_DIAGMATRIX3_HPP_INCLUDED

// Math::DiagMatrix3<T> is a 3x3 matrix where all non-diagonal entries are zero.
//
// The difference to Math::Vector3<T> is mostly the effect of operations
// (Vector3<T> * Vector3<T> is a dot product while DiagMatrix3<T> * Vector3<T>
// is a matrix-vector-multiplication / an elementwise multiplication)

#include <Core/Assert.hpp>
#include <Core/Util.hpp>

#include <Math/Forward.hpp>
#include <Math/DiagMatrix3.hpp>

#include <complex>

#include <boost/type_traits/is_convertible.hpp>
#include <boost/utility/enable_if.hpp>

namespace Math {
  template <typename T> class DiagMatrix3 {
    Vector3<T> diag_;

    class PrivateType {
      friend class DiagMatrix3;
      PrivateType () {}
    };

  public:
    DiagMatrix3 () {}
    explicit DiagMatrix3 (T v) : diag_ (v, v, v) {}
    DiagMatrix3 (T x, T y, T z) : diag_ (x, y, z) {}
    template <typename U> explicit DiagMatrix3 (const Vector3<U>& diag) : diag_ (diag) {}
#if (defined (__clang__) || GCC_VERSION_IS_ATLEAST(4, 6)) && !defined (__CUDACC__)
#pragma GCC diagnostic push
#endif
#pragma GCC diagnostic ignored "-Wconversion"
    template <typename U> DiagMatrix3 (DiagMatrix3<U> v, UNUSED typename boost::enable_if<boost::is_convertible<U, T>, PrivateType>::type dummy = PrivateType ()) : diag_ (v.diag ()) {}
    template <typename U> explicit DiagMatrix3 (DiagMatrix3<U> v, UNUSED typename boost::disable_if<boost::is_convertible<U, T>, PrivateType>::type dummy = PrivateType ()) : diag_ (v.diag ()) {}
#if (defined (__clang__) || GCC_VERSION_IS_ATLEAST(4, 6)) && !defined (__CUDACC__)
#pragma GCC diagnostic pop
#endif

    const Vector3<T>& diag () const {
      return diag_;
    }
    Vector3<T>& diag () {
      return diag_;
    }

    T m11 () const {
      return diag ().x ();
    }
    T m22 () const {
      return diag ().y ();
    }
    T m33 () const {
      return diag ().z ();
    }

    T& m11 () {
      return diag ().x ();
    }
    T& m22 () {
      return diag ().y ();
    }
    T& m33 () {
      return diag ().z ();
    }

    // off-diagonal entries (read-only, = 0)
    T m12 () const { return T (); }
    T m13 () const { return T (); }
    T m21 () const { return T (); }
    T m23 () const { return T (); }
    T m31 () const { return T (); }
    T m32 () const { return T (); }

    // Return the inverse matrix
    DiagMatrix3<T> inverse () const {
      return DiagMatrix3<T> (T (1) / m11 (), T (1) / m22 (), T (1) / m33 ());
    }

    // Access values on the diagonal
    const T& operator[] (size_t i) const {
      return diag ()[i];
    }
    T& operator[] (size_t i) {
      return diag ()[i];
    }
  };

  // Operations on DiagMatrix3

#define RTS(op) DECLTYPE ((*(T*)NULL) op (*(U*)NULL))
#define RT(op) DiagMatrix3<RTS(op)>
  template <typename T, typename U> inline RT(+) operator+ (DiagMatrix3<T> v1, DiagMatrix3<U> v2) {
    return RT(+) (v1.m11 () + v2.m11 (), v1.m22 () + v2.m22 (), v1.m33 () + v2.m33 ());
  }

  template <typename T, typename U> inline DiagMatrix3<T>& operator+= (DiagMatrix3<T>& v1, DiagMatrix3<U> v2) {
    v1.m11 () += v2.m11 (); v1.m22 () += v2.m22 (); v1.m33 () += v2.m33 ();
    return v1;
  }

  template <typename T, typename U> inline RT(-) operator- (DiagMatrix3<T> v1, DiagMatrix3<U> v2) {
    return RT(-) (v1.m11 () - v2.m11 (), v1.m22 () - v2.m22 (), v1.m33 () - v2.m33 ());
  }

  template <typename T, typename U> inline DiagMatrix3<T>& operator-= (DiagMatrix3<T>& v1, DiagMatrix3<U> v2) {
    v1.m11 () -= v2.m11 (); v1.m22 () -= v2.m22 (); v1.m33 () -= v2.m33 ();
    return v1;
  }

  template <typename T, typename U> inline RT(*) operator* (DiagMatrix3<T> v, U scalar) {
    return RT(*) (v.m11 () * scalar, v.m22 () * scalar, v.m33 () * scalar);
  }

  template <typename T, typename U> inline RT(*) operator* (T scalar, DiagMatrix3<U> v) {
    return RT(*) (scalar * v.m11 (), scalar * v.m22 (), scalar * v.m33 ());
  }

  template <typename T, typename U> inline DiagMatrix3<T>& operator*= (DiagMatrix3<T>& v, U scalar) {
    v.m11 () *= scalar; v.m22 () *= scalar; v.m33 () *= scalar;
    return v;
  }

  template <typename T, typename U> inline RT(/) operator/ (DiagMatrix3<T> v, U scalar) {
    return RT(/) (v.m11 () / scalar, v.m22 () / scalar, v.m33 () / scalar);
  }

  template <typename T, typename U> inline DiagMatrix3<T>& operator/= (DiagMatrix3<T>& v, U scalar) {
    v.m11 () /= scalar; v.m22 () /= scalar; v.m33 () /= scalar;
    return v;
  }

  template <typename T> inline bool operator== (DiagMatrix3<T> v1, DiagMatrix3<T> v2) {
    return v1.m11 () == v2.m11 () && v1.m22 () == v2.m22 () && v1.m33 () == v2.m33 ();
  }
  template <typename T> inline bool operator!= (DiagMatrix3<T> v1, DiagMatrix3<T> v2) {
    return !(v1 == v2);
  }

  /*
  template <typename T> struct Abs2Impl<DiagMatrix3<T> > {
    static DECLTYPE(Math::abs2 (*(T*)0)) apply (DiagMatrix3<T> v) {
      return Math::abs2 (v.m11 ()) + Math::abs2 (v.m22 ()) + Math::abs2 (v.m33 ());
    }
  };
  */

  // Matrix-Matrix multiplication
  template <typename T, typename U> inline RT(*) operator* (DiagMatrix3<T> v1, DiagMatrix3<U> v2) {
    return RT(*) (v1.m11 () * v2.m11 (), v1.m22 () * v2.m22 (), v1.m33 () * v2.m33 ());
  }

  // Matrix-Vector multiplication
  template <typename T, typename U> inline Vector3<RTS(*)> operator* (DiagMatrix3<T> v1, Vector3<U> v2) {
    return Vector3<RTS(*)> (v1.m11 () * v2.x (), v1.m22 () * v2.y (), v1.m33 () * v2.z ());
  }

#undef RTS
#undef RT

  // Unary +/-
  template <typename T> inline DiagMatrix3<T> operator+ (DiagMatrix3<T> v) {
    return DiagMatrix3<T> (+v.m11 (), +v.m22 (), +v.m33 ());
  }
  template <typename T> inline DiagMatrix3<T> operator- (DiagMatrix3<T> v) {
    return DiagMatrix3<T> (-v.m11 (), -v.m22 (), -v.m33 ());
  }

  template <typename F> DiagMatrix3<F> real (DiagMatrix3<std::complex<F> > v) {
    return DiagMatrix3<F> (real (v.m11 ()), real (v.m22 ()), real (v.m33 ()));
  }

  template <typename F> DiagMatrix3<F> imag (DiagMatrix3<std::complex<F> > v) {
    return DiagMatrix3<F> (imag (v.m11 ()), imag (v.m22 ()), imag (v.m33 ()));
  }
}

#endif // !MATH_DIAGMATRIX3_HPP_INCLUDED


#include <unistd.h>	// sleep
#include <armadillo>
#include <vector>
#include <deque>
#include <algorithm>    // std::min

#include "Functions.h"

using namespace arma;
using namespace std;

vec eta(const vec &v, const double tau){
	vec out(size(v),fill::zeros);

	uvec ind = find(abs(v) > tau);
	out(ind) = sign(v(ind))%(abs(v(ind)) - tau );
	return out;
}

vec eta_deriv(const vec &v, const double tau){
	vec out(size(v),fill::zeros);

	uvec ind = find(abs(v) > tau);
	out(ind).ones();
	return out;
}

vec eta(const vec &v, const vec tau){
	vec out(size(v),fill::zeros);

	uvec ind = find(abs(v) > tau);
	out(ind) = sign(v(ind))%(abs(v(ind)) - tau(ind) );
	return out;
}

vec eta_deriv(const vec &v, const vec tau){
	vec out(size(v),fill::zeros);

	uvec ind = find(abs(v) > tau);
	out(ind).ones();
	return out;
}

vec AMP(const mat &A, const vec &y, const int sparsity, const unsigned int max_iter,
	const double tol, unsigned int &num_iters, const simulation_parameters simulation_params){

	const unsigned int N = A.n_cols;   	// signal dimension
	const unsigned int M = y.n_elem;	// number of measurements
	//const double delta = double(num_measurements)/sig_dim;
	unsigned int i = 0;
	vec x_t(N,fill::zeros);
	vec z_t = y;
	double tau = .1;
	bool done = false;
	vec pseudo_data (N,fill::zeros);

	while(!done){
		i++;
		z_t = y - A*x_t + z_t * sum (eta_deriv( pseudo_data,tau) ) / M;
		pseudo_data = A.t() * z_t + x_t;
		tau = tau * sum(eta_deriv(pseudo_data,tau)) / M;
		x_t = eta(pseudo_data,tau) ;

		if (norm (y - A*x_t) < tol || i >= max_iter){
			done = true;
		}
	}

	num_iters = i;
	return x_t;
}

// J. Zhu, R. Pilgrim and D. Baron, "An overview of multi-processor approximate message passing,"
// http://ieeexplore.ieee.org/document/7926166/
vec R_MP_AMP(const mat &A, const vec &y, const int sparsity, const unsigned int max_iter,
	const double tol, unsigned int &num_iters, const simulation_parameters simulation_params){
	uvec slow_cores;
	set_slow_cores(slow_cores, simulation_params);

	const unsigned int N = A.n_cols;
	const unsigned int M = y.n_elem;
	const unsigned int P = simulation_params.num_cores;
	unsigned int i = 0;
	bool done = false;
	vec x_t(N,fill::zeros);
	double g_t = M;
	double tau = .1;

	vector <vec> pseudo_data (P,vec(N,fill::zeros));
	// parallel section of the code starts here
	#pragma omp parallel num_threads(simulation_params.num_cores)
	{
	// initializing variables in local memory
	const int p = omp_get_thread_num();
	mat A_p = A.rows(M*p/P , M*(p+1)/P -1 );
	vec y_p = y.subvec( M*p/P  , M*(p+1)/P -1 );

	vec z_t_p = y_p;
	// R_MP_AMP itearations
	while(!done){
		//AT processor p:
		z_t_p = y_p - A_p*x_t + z_t_p * g_t / M;
		pseudo_data[p] = A_p.t() * z_t_p + x_t/P;

		//slow cores sleep for  simulation_params.sleep_slow_cores microseconds
		if (any( slow_cores == omp_get_thread_num()) ){
			usleep(simulation_params.sleep_slow_cores);
		}

		#pragma omp barrier
		//AT fusion center:
		#pragma omp single
		{
			i++;
			vec pseudo_data_total(N,fill::zeros);
			for (unsigned int j = 0; j < P; j++){
				pseudo_data_total = pseudo_data_total + pseudo_data[j];
			}

			tau = tau * sum(eta_deriv(pseudo_data_total,tau)) / M;
			g_t = sum(eta_deriv(pseudo_data_total,tau));
			x_t = eta(pseudo_data_total,tau);
			if (norm (y - A*x_t) < tol || i >= max_iter){
				done = true;
			}
		}
		#pragma omp barrier
	}
	// parallel section of the code ends here

	}

	num_iters = i;
	return x_t;
}


// Copyright Matt Overby 2021.
// Distributed under the MIT License.

#ifndef MCL_LBFGS_HPP
#define MCL_LBFGS_HPP 1

#include <Eigen/Dense>
#include <vector>
#include <functional>

namespace mcl
{

// L-BFGS implementation based on Nocedal & Wright Numerical Optimization book (Section 7.2)
//
// Function pointers are meant to be used with lambdas, e.g.
//   LBFGS<MatrixXd> lbfgs;
//   lbfgs.gradient = [&](const MatrixXd &x, MatrixXd &g)->Scalar { return ... };
//   double obj = lbfgs.minimize(x);
//
// If the g arg is not sized, don't compute gradient
//
template <typename MatrixType>
class LBFGS
{
public:
    typedef typename MatrixType::Scalar Scalar;

	struct Options
	{
		int min_iters;
		int max_iters;
		Scalar abs_tol; // absolute tol if converged(...) not set
		Scalar rel_tol; // relative tol if converged(...) not set
		int M; // history window size
		Scalar gamma; // init Hessian = gamma * I
		Options() :
			min_iters(0),
			max_iters(100),
			abs_tol(1e-5),
			rel_tol(0),
			M(6),
			gamma(1)
			{}
	} options;

	LBFGS();

	// Output from the last call to minimize(...)
	int iters() const { return num_iters; }
	Scalar gamma() const { return gamma_k; }

	// Resizes buffers and sets to zero.
	// Called during minimize(...) ONLY if there is a change in dof or M.
	// Otherwise it's assumed you're picking up where you left off
	// on the previous call to minimize(...)
	void reset(int rows, int cols);

	// Required:
	// computes objective value and gradient
	//   obj = gradient(x, g)
	// If the g arg is not sized, don't compute gradient
	std::function<Scalar(const MatrixType&, MatrixType&)> gradient;

	// Optional:
	// Returns true if the solver should exit, default uses ||g||<abs_tol or ||g||<||x||rel_tol
	//   is_converged = converged(obj, x_prev, x, grad)
	std::function<bool(Scalar obj, const MatrixType&, const MatrixType&, const MatrixType&)> converged;

	// Optional:
	// Linesearch function, default uses bracketing weak wolfe (slow!)
	//   obj_k1 = (x, grad, descent, alpha)
	// Returns new objective value and updates both x AND gradient
	std::function<Scalar(MatrixType&, MatrixType&, const MatrixType&, Scalar&)> linesearch;

	// Optional:
	// Filter descent direction, p = B(p)
	// Otherwise p = gamma_k * p is used.
	std::function<void(MatrixType&)> filter;

	// Calls initialize(x) once and iterate(x) until converged
	Scalar minimize(MatrixType& x);

	// Initialize the solver
	void initialize(MatrixType& x);

	// Take an iteration
	// Returns objective
	Scalar iterate(MatrixType& x);

	// i.e. bisection with weak Wolfe conditions
	Scalar bracketing_weakwolfe(
		MatrixType& x,
		MatrixType& grad,
		const MatrixType& p,
		Scalar &alpha) const;

	// Used if converged not set
	// Returns true if:
	// grad.norm <= abs_tol
	// or
	// grad.norm() <= rel_tol * x.norm()
	bool default_converged(
		Scalar curr_obj,
		const MatrixType& xprev,
		const MatrixType& x,
		const MatrixType& grad) const;

	Scalar inner(
		const MatrixType &a,
		const MatrixType &b) const;

protected:
	bool initialized;
	int num_iters, max_iters, k;
	Scalar gamma_k, obj_0, obj_k;
	std::vector<MatrixType> s;
	std::vector<MatrixType> y;
	Eigen::Matrix<Scalar,Eigen::Dynamic,1> alpha;
	Eigen::Matrix<Scalar,Eigen::Dynamic,1> rho;
	MatrixType grad;
	MatrixType q;
	MatrixType descent;
	MatrixType grad_old;
	MatrixType x_old;
	MatrixType x_last;
	MatrixType s_temp;
	MatrixType y_temp;

}; // end class LBFGS

//
// Implementation
//

template <typename MatrixType>
LBFGS<MatrixType>::LBFGS() :
	initialized(false),
	num_iters(0),
	max_iters(0),
	k(0),
	gamma_k(1),
	obj_0(0),
	obj_k(0)
	{}

template <typename MatrixType>
void LBFGS<MatrixType>::reset(int rows, int cols)
{
    using namespace Eigen;
	num_iters = 0;
	max_iters = 0;
	k = 0;
	gamma_k = 1;
	obj_0 = std::numeric_limits<Scalar>::max();
	obj_k = std::numeric_limits<Scalar>::max();
	int M = options.M;
	s = std::vector<MatrixType>(M, MatrixType::Zero(rows,cols));
	y = std::vector<MatrixType>(M, MatrixType::Zero(rows,cols));
	alpha = VectorXd::Zero(M);
	rho = VectorXd::Zero(M);
	grad = MatrixType::Zero(rows, cols);
	q = MatrixType::Zero(rows, cols); // inv descent
	descent = q;
	grad_old = MatrixType::Zero(rows, cols);
	x_old = MatrixType::Zero(rows, cols);
	x_last = MatrixType::Zero(rows, cols);
	s_temp = MatrixType::Zero(rows, cols);
	y_temp = MatrixType::Zero(rows, cols);
}

// Returns number of iterations used
template <typename MatrixType>
typename LBFGS<MatrixType>::Scalar
LBFGS<MatrixType>::minimize(MatrixType &x)
{
	initialize(x);

	// Did we start at the initializer?
	if (num_iters >= options.min_iters &&
		converged(obj_k,x_last,x,grad))
	{
		num_iters = 1;
		return obj_k;
	}

	for (; k<max_iters; ++k)
	{
		iterate(x);
	} // end loop lbfgs iters

	return obj_k;

} // end minimize

template <typename MatrixType>
void LBFGS<MatrixType>::initialize(MatrixType& x)
{
	initialized = true;

	if (gradient == nullptr) {
	    throw std::runtime_error("no gradient function");
	}

	if (converged == nullptr)
	{
		using namespace std::placeholders;
		converged = std::bind(&LBFGS::default_converged, this, _1, _2, _3, _4);
	}

	if (linesearch == nullptr)
	{
		using namespace std::placeholders;
		linesearch = std::bind(&LBFGS::bracketing_weakwolfe, this, _1, _2, _3, _4);
	}

	// Resize/resize variables?
	if (alpha.rows() != options.M ||
		grad.rows() != x.rows() ||
		grad.cols() != x.cols()) {
		reset(x.rows(), x.cols());
	}

	num_iters = 0;
	max_iters = std::max(options.min_iters, options.max_iters);
	gamma_k = options.gamma;
	obj_0 = gradient(x, grad);
	obj_k = obj_0;
	k = 0;
}

template <typename MatrixType>
typename LBFGS<MatrixType>::Scalar
LBFGS<MatrixType>::iterate(MatrixType& x)
{
	if (!initialized) {
		throw std::runtime_error("not initialized");
	}

	x_old = x;
	grad_old = grad;
	q = grad;
	num_iters++;

	//
	// Two-loop recursion
	//
	{
		// L-BFGS first - loop recursion
		int iter = std::min(options.M, k);
		for(int i = iter - 1; i >= 0; --i)
		{
			Scalar denom = inner(s[i], y[i]);
			if (std::abs(denom) <= 0.0)
			{
				rho(i) = 0;
				alpha(i) = 0;
				continue;
			}
			rho(i) = 1.0 / denom;
			alpha(i) = rho(i)*inner(s[i], q);
			q -= alpha(i) * y[i];
		}

		if (filter != nullptr) { filter(q); }
		else { q = gamma_k*q; }

		// L-BFGS second - loop recursion
		for(int i = 0; i < iter; ++i)
		{
			Scalar beta = rho(i)*inner(q, y[i]);
			q += (alpha(i) - beta)*s[i];
		}
	}

	//
	// Perform step
	//
	{
		// If our hess approx is bad and we start going in
		// the wrong direction, restart memory
		Scalar step_size = 1.0;
		Scalar dir = inner(q, grad);
		if (dir <= 0)
		{
			q = grad;
			max_iters -= k; // Restart memory
			k = 0;
			step_size = std::min(1.0, 1.0 / grad.template lpNorm<Eigen::Infinity>() );
		}

		// We've hit local minima, we have to exit
		if (q.squaredNorm() <= 0.0) {
			return obj_k;
		}

		descent = -q;
		x_last = x;
		obj_k = linesearch(x, grad, descent, step_size);
		if (num_iters >= options.min_iters && converged(obj_k,x_last,x,grad)) {
			return obj_k;
		}
	}

	//
	// Correction term
	//
	{
		s_temp = x - x_old;
		y_temp = grad - grad_old;

		// update the history
		if (k < options.M)
		{
			s[k] = s_temp;
			y[k] = y_temp;
		}
		else
		{
			for (int i=0; i<options.M - 1; ++i)
			{
				s[i] = s[i+1];
				y[i] = y[i+1];
			}
			s.back() = s_temp;
			y.back() = y_temp;
		}

		Scalar denom = inner(y_temp, y_temp);
		if (std::abs(denom) > 0.0)
		{
			gamma_k = inner(s_temp, y_temp) / denom;
		}
	}

	return obj_k;
}

template <typename MatrixType>
typename LBFGS<MatrixType>::Scalar
LBFGS<MatrixType>::bracketing_weakwolfe(
		MatrixType& x,
		MatrixType& g,
		const MatrixType &drt,
		Scalar &step) const
{
    using namespace Eigen;

	Scalar c1 = 10e-4;
	Scalar c2 = 0.9;
	if (step <= 0.0) { step = 1.0; }
    int x_cols = x.cols();

	g = MatrixType::Zero(x.rows(), x_cols);
	MatrixType xp = x;
	const Scalar fx_init = gradient(x,g);
	Scalar fx = fx_init;

	Scalar dg_init = inner(g, drt);
    if (dg_init > 0) {
        throw std::runtime_error("direction increases objective");
    }

	const Scalar test_decr = c1 * dg_init;
	Scalar lower = 0;
	Scalar upper = std::numeric_limits<Scalar>::infinity();
	int maxiter = 200;
	int iter = 0;
	for (iter = 0; iter < maxiter; iter++)
	{
		x = xp + step * drt;
		fx = gradient(x, g);
		if (fx > fx_init + step * test_decr) { // Armijo rule
			upper = step;
		}
		else
		{
			Scalar dg = inner(g, drt);
			if(dg < c2 * dg_init){ // Weak wolfe
				lower = step;
			} else {
				break; // both met
			}
		}
		step = std::isinf(upper) ? 2*step : lower/2 + upper/2;
	}
	return fx;

} // end linesearch

template <typename MatrixType>
bool LBFGS<MatrixType>::default_converged(
		Scalar curr_obj,
		const MatrixType& xprev,
		const MatrixType& x,
		const MatrixType& g) const
{
	(void)(curr_obj);
	(void)(xprev);
	Scalar gnorm = g.norm();
	if (gnorm <= options.abs_tol) {
	    return true;
	}
	Scalar xnorm = x.norm();
	if (gnorm < options.rel_tol * xnorm) {
	    return true;
	}
	return false;
}

template <typename MatrixType>
typename LBFGS<MatrixType>::Scalar
LBFGS<MatrixType>::inner(
    const MatrixType &a,
    const MatrixType &b) const
{
    int cols = std::min(a.cols(), b.cols());
    Scalar dot = 0;
    for (int i=0; i<cols; ++i) {
        dot += a.col(i).dot(b.col(i));
    }
    return dot;
} // end inner product

} // end ns mcl

#endif