/***************************************************************************
 *   Copyright (C) 2007 by Reed A. Cartwright                              *
 *   reed@scit.us                                                          *
 *                                                                         *
 *   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 BE LIABLE FOR ANY CLAIM, DAMAGES OR     *
 *   OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, *
 *   ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR *
 *   OTHER DEALINGS IN THE SOFTWARE.                                       *
 ***************************************************************************/

#include <boost/math/special_functions/zeta.hpp>

#include "models.h"

inline double zeta(double z) { return boost::math::zeta<double>(z); }

extern const int g_pupy[] = {0,1,0,1,-1};
extern const int g_atgc[] = {0,1,1,0,-1};

const model_k2p_zeta::freq_type model_k2p_zeta::k2p_freqs = {0.25,0.25,0.25,0.25};
const model_k2p_geo::freq_type model_k2p_geo::k2p_freqs = {0.25,0.25,0.25,0.25};

/***************************************************************************
 * class model_k2p_zeta                                                    *
 ***************************************************************************/

void model_k2p_zeta::preallocate(size_t maxa, size_t maxd)
{
	size_t sz_max = std::max(maxa, maxd)+1;
	p_indel_size.resize(sz_max, 0.0);
}

void model_k2p_zeta::expectation_setup(const params_type &params, const freq_type &freq)
{
	p_ts = 0.25-0.5*exp(-params[pT]*(2.0*params[pK]+1.0)/(params[pK]+1.0))
			+ 0.25*exp(-2.0*params[pT]/(params[pK]+1.0));
	p_tv = 0.5-0.5*exp(-2.0*params[pT]/(params[pK]+1.0));
	p_match = 1.0-p_ts-p_tv;

	p_end = 1.0/(params[pA]+1.0);
	p_2h =  (1.0-p_end)*exp(-2.0*params[pR]*params[pT]);
	p_2g =  (1.0-p_end)*0.5*(1.0-exp(-2.0*params[pR]*params[pT]));

	double nuc_scale2 = 0.25*p_2h
		*pow(p_match, p_match)
		*pow(p_ts, p_ts)
		*pow(0.5*p_tv, p_tv);
	nuc_scale = sqrt(nuc_scale2);
	amb_scale = 4.0;

	for(size_t i=0;i<nN;++i)
	{
		p_substitution[i][i] = 0.25*p_2h*p_match/nuc_scale2;

		for(size_t j=i+1;j<nN;++j)
		{
			p_substitution[j][i] = p_substitution[i][j] = p_2h*
				((g_pupy[i] == g_pupy[j]) ? 0.25*p_ts : 0.125*p_tv)
				/ nuc_scale2;
		}
		p_substitution[nN][i] = p_substitution[i][nN] = 0.0625*p_2h / nuc_scale2;
	}
	p_substitution[nN][nN] = 0.0625*p_2h / nuc_scale2;

	double z = params[pZ];
	double Z = p_2g/zeta(z);
	double s = 0.25/nuc_scale;
	double ss = s;
	for(size_t u = 1; u < p_indel_size.size(); ++u)
	{
		p_indel_size[u] = pow((double)u, -z)*Z*ss;
		ss *= s;
	}
}

/***************************************************************************
 * class model_k2p_geo                                                     *
 ***************************************************************************/

void model_k2p_geo::preallocate(size_t maxa, size_t maxd)
{
	size_t sz_max = std::max(maxa, maxd)+1;
	p_indel_size.resize(sz_max, 0.0);
}

void model_k2p_geo::expectation_setup(const params_type &params, const freq_type &freq)
{
	p_ts = 0.25-0.5*exp(-params[pT]*(2.0*params[pK]+1.0)/(params[pK]+1.0))
			+ 0.25*exp(-2.0*params[pT]/(params[pK]+1.0));
	p_tv = 0.5-0.5*exp(-2.0*params[pT]/(params[pK]+1.0));
	p_match = 1.0-p_ts-p_tv;

	p_end = 1.0/(params[pA]+1.0);
	p_2h =  (1.0-p_end)*exp(-2.0*params[pR]*params[pT]);
	p_2g =  (1.0-p_end)*0.5*(1.0-exp(-2.0*params[pR]*params[pT]));

	double nuc_scale2 = 0.25*p_2h
		*pow(p_match, p_match)
		*pow(p_ts, p_ts)
		*pow(0.5*p_tv, p_tv);
	nuc_scale = sqrt(nuc_scale2);
	amb_scale = 4.0;

	for(size_t i=0;i<nN;++i)
	{
		p_substitution[i][i] = 0.25*p_2h*p_match/nuc_scale2;

		for(size_t j=i+1;j<nN;++j)
		{
			p_substitution[j][i] = p_substitution[i][j] = p_2h*
				((g_pupy[i] == g_pupy[j]) ? 0.25*p_ts : 0.125*p_tv)
				/ nuc_scale2;
		}
		p_substitution[nN][i] = p_substitution[i][nN] = 0.0625*p_2h / nuc_scale2;
	}
	p_substitution[nN][nN] = 0.0625*p_2h / nuc_scale2;

	double s = 0.25/nuc_scale;
	double qq = 1.0/params[pQ];
	double dp = qq*s;
	for(size_t u = 1; u < p_indel_size.size(); ++u)
	{
		p_indel_size[u] = p_2g*dp;
		dp *= (1.0-qq)*s;
	}
	p_open = qq*p_2g*s;
	p_extend = (1.0-qq)*s;
}


/*
 * Superpixel.cpp
 *
 *  Created on: Jan 26, 2012
 *      Author: david
 */

#include "Superpixels.hpp"
#include "impl/RepairDepth.hpp"
#include "impl/Sampling.hpp"
#include "impl/Clustering.hpp"
#include <density/Smooth.hpp>
#include <pds/PDS.hpp>
#define DANVIL_ENABLE_BENCHMARK
#include <Danvil/Tools/Benchmark.h>
#include <Danvil/Tools/MoreMath.h>
#include <Danvil/Color.h>
#include <Danvil/Color/LAB.h>
#include <Danvil/Color/HSV.h>
#include <eigen3/Eigen/Eigenvalues>
#include <boost/math/constants/constants.hpp>
#include <queue>
#include <set>
#include <fstream>

//#define CLUSTER_UPDATE_MULTITHREADING
#define CREATE_DEBUG_IMAGES

//------------------------------------------------------------------------------
namespace dasp {
//------------------------------------------------------------------------------

std::map<std::string,slimage::ImagePtr> sDebugImages;

Parameters::Parameters()
{
	camera.cx = 318.39f;
	camera.cy = 271.99f;
	camera.focal = 528.01f;
	camera.z_slope = 0.001f;
	color_space = ColorSpaces::RGB;
	density_mode = DensityModes::DASP;
	weight_color = 2.0f;
	weight_spatial = 1.0f;
	weight_normal = 3.0f;
	iterations = 5;
	coverage = 1.7f;
	base_radius = 0.024f;
	count = 0;
	seed_mode = SeedModes::SimplifiedPDS;
	gradient_adaptive_density = true;
	use_density_depth = true;
	ignore_pixels_with_bad_visibility = false;
	is_conquer_enclaves = true;
	segment_threshold = 1.0f;
	is_repair_depth = false;
	is_smooth_depth = false;
	is_improve_seeds = false;
	is_smooth_density = false;
	enable_clipping = false;
	clip_x_min = -3.0f; clip_x_max = +3.0f;
	clip_y_min = -3.0f; clip_y_max = -3.0f;
	clip_z_min = 0.0f; clip_z_max = 3.0f;
	enable_roi_2d = false;
	roi_2d_x_min = 0; roi_2d_x_max = 640;
	roi_2d_y_min = 0; roi_2d_y_max = 480;
}

void Cluster::UpdateCenter(const ImagePoints& points, const Parameters& opt)
{
	assert(isValid());

	if(!is_fixed) {
		// update cluster means (position, color) and update screen position and depth
		Eigen::Vector3f mean_color = Eigen::Vector3f::Zero();
		Eigen::Vector3f mean_world = Eigen::Vector3f::Zero();
		for(unsigned int i : pixel_ids) {
			const Point& p = points[i];
			assert(p.is_valid);
			mean_color += p.color;
			mean_world += p.position;
		}
		center.color = mean_color / float(pixel_ids.size());
		center.position = mean_world / float(pixel_ids.size());
		// compute screen position
		if(opt.density_mode == DensityModes::ASP_RGB) {
			// compute by averaging
			Eigen::Vector2f mean_p = Eigen::Vector2f::Zero();
			for(unsigned int i : pixel_ids) {
				const Point& p = points[i];
				mean_p += Eigen::Vector2f{ p.px, p.py };
			}
			mean_p /= static_cast<float>(pixel_ids.size());
			center.px = mean_p[0];
			center.py = mean_p[1];
		}
		else {
			// compute by projection
			Eigen::Vector2f pixel = opt.camera.project(center.position);
			center.px = static_cast<float>(pixel.x() + 0.5f);
			center.py = static_cast<float>(pixel.y() + 0.5f);
		}
		// // compute color covariance
		// color_covariance = Eigen::Matrix3f::Zero();
		// for(unsigned int i : pixel_ids) {
		// 	const Point& p = points[i];
		// 	auto d = p.color - mean_color;
		// 	color_covariance += d * d.transpose();
		// }
		// color_covariance *= 1.0f / static_cast<float>(pixel_ids.size() - 1);
	}

//	// FIXME change or not change? (SLIC mode does not allow change!)
//	//center.cluster_radius_px = opt.base_radius * opt.camera.scala(center.depth_i16);
//	if(opt.use_density_depth) {
//		center.spatial_normalizer = 1.0f;
//	}
//	else {
//		float rpx = std::sqrt(static_cast<float>(points.width()*points.height())/static_cast<float>(4*opt.count));
//		center.spatial_normalizer = rpx / (opt.base_radius * opt.camera.scala(center.depth_i16)) / center.circularity;
//	}

	if(opt.density_mode == DensityModes::ASP_RGB) {
		return;
	}

	if(pixel_ids.size() >= 3) {
		cov = PointCovariance(pixel_ids,
			[this,&points](unsigned int i) {
				return points[i].position - center.position;
			});
	}
	else {
//		std::cerr << "Not enough points!" << std::endl;
		cov = opt.base_radius * opt.base_radius * Eigen::Matrix3f::Identity();
	}

	Eigen::SelfAdjointEigenSolver<Eigen::Matrix3f> solver;
	solver.compute(cov);
	ew = solver.eigenvalues();
	ev = solver.eigenvectors();
	// std::cout << std::sqrt(ew[0]) << " " << std::sqrt(ew[1]) << " " << std::sqrt(ew[2]) << std::endl;

	if(!is_fixed) {
		center.setNormal(ev.col(0));
	}

//	center.normal = points(center.pos).normal;

	// FIXME all of the following code does not use the fixed normal!

	// eigensystem defines an ellipsoid
	// eigenvalues are the square of the standard deviation!
	// eigenvalues have to be multiplied by a factor to get the probable size
	// eigenvalues are sorted smallest to largest (d, b, a)

	// thickness: smallest diameter of ellipsoid
	thickness = cSigmaScale * std::sqrt(std::abs(ew(0))) * 2.0f;

	// eccentricity: \sqrt{1 - \frac{b^2}{a^2}}
	eccentricity = std::sqrt(1.0f - std::abs(ew(1) / ew(2)));

	// flatness: \sqrt{1 - \frac{d^2}{a^2}}
	flatness = std::sqrt(1.0f - std::abs(ew(0) / ew(2)));

	float area_base = cSigmaScale * cSigmaScale * std::sqrt(std::abs(ew(1) * ew(2)));
	area = area_base * boost::math::constants::pi<float>();

	// area_actual / area_expected = (a*b)/(R*R)
	area_quotient = area_base / (opt.base_radius * opt.base_radius);

	// shape

	if(pixel_ids.size() >= 6) {
		Eigen::Matrix<float,6,1> s = Shape(pixel_ids,
			[this,&points](unsigned int i) {
				return points[i].position - center.position;
			});
		shape_0 = s[0];
		shape_x = s[1];
		shape_y = s[2];
		shape_xy = s[3];
		shape_xx = s[4];
		shape_yy = s[5];
	}

}

void Cluster::ComputeExt(const ImagePoints& points, const Parameters& opt)
{
	std::vector<float> dist;
	dist.reserve(pixel_ids.size());
	for(unsigned int id : pixel_ids) {
		Eigen::Vector3f p = points[id].position - center.position;
		float d = p.dot(center.normal);
		dist.push_back(d);
	}
	unsigned int nth = static_cast<unsigned int>(cPercentage * static_cast<float>(dist.size()));
	std::nth_element(dist.begin(), dist.begin() + nth, dist.end());
	thickness_plane = (*(dist.begin() + nth));

	{
		area_actual = 0.0f;
		area_expected = 0.0f;
		float error_area = 0;
		int cx = center.px;
		int cy = center.py;
		int R = int(center.cluster_radius_px * opt.coverage * 1.5f);
		unsigned int xmin = std::max(0, cx - R);
		unsigned int xmax = std::min(int(points.width()-1), cx + R);
		unsigned int ymin = std::max(0, cy - R);
		unsigned int ymax = std::min(int(points.height()-1), cy + R);
		for(unsigned int y=ymin; y<=ymax; y++) {
			for(unsigned int x=xmin; x<=xmax; x++) {
				unsigned int i = points.index(x,y);
				const Point& p = points[i];
				Eigen::Vector3f t = p.position - center.position;
				float lam = t.dot(center.normal);
				float d2 = t.squaredNorm() - lam*lam;
				bool expected = lam < thickness_plane && d2 < opt.base_radius*opt.base_radius;
				bool actual = std::find(pixel_ids.begin(), pixel_ids.end(), i) != pixel_ids.end();
				float size_of_a_px = p.depth() / opt.camera.focal;
				float area_of_a_px = size_of_a_px*size_of_a_px / p.computeCircularity();
				if(actual) {
					area_actual += area_of_a_px;
				}
				if(expected) {
					area_expected += area_of_a_px;
				}
				if(actual != expected) {
					error_area += area_of_a_px;
				}
			}
		}
//		std::cout << 10000.0f*error_area << " - " << 10000.0f*expected_area << " - " << 10000.0f*real_expected_area << std::endl;
		coverage_error = (area_expected == 0) ? 0 : error_area / area_expected;
		area_expected_global = M_PI * opt.base_radius * opt.base_radius;
	}

	// we delay the multiplication with two, so that we do not need to divide
	// by two in the computation of the area error
	thickness_plane *= 2.0f;

}

//------------------------------------------------------------------------------

Superpixels::Superpixels()
{

}

std::vector<Seed> Superpixels::getClusterCentersAsSeeds() const
{
	std::vector<Seed> seeds(cluster.size());
	for(std::size_t i=0; i<cluster.size(); i++) {
		const Cluster& c = cluster[i];
		seeds[i] = Seed{c.center.px, c.center.py, c.center.cluster_radius_px, c.seed_id, c.is_fixed};
	}
	return seeds;
}

std::vector<Eigen::Vector2f> Superpixels::getClusterCentersAsPoints() const
{
	std::vector<Eigen::Vector2f> pnts(cluster.size());
	for(unsigned int i=0; i<pnts.size(); i++) {
		pnts[i] = Eigen::Vector2f(cluster[i].center.px, cluster[i].center.py);
	}
	return pnts;
}

void Superpixels::CreatePoints(const slimage::Image3ub& image, const slimage::Image1ui16& depth, const slimage::Image3f& normals)
{
	color_raw = image.clone();

	assert(normals.isNull());

	const float cTempBaseRadius = 0.025f;

	if(opt.camera.focal == 0.0f) {
		std::cerr << "Invalid focal length!" << std::endl;
	}

	// compute base radius from desired super pixel count
	if(opt.count > 0) {
		opt.base_radius = cTempBaseRadius;
	}

	unsigned int width = image.width();
	unsigned int height = image.height();
	assert(width == depth.width() && height == depth.height());
	if(!normals.isNull()) {
		assert(width == normals.width() && height == normals.height());
		// FIXME this case is disabled!
	}

	points = ImagePoints(width, height);

	const bool is_clipping_2d = opt.enable_roi_2d
		&& (opt.roi_2d_x_min < opt.roi_2d_x_max)
		&& (opt.roi_2d_y_min < opt.roi_2d_y_max);

	const bool is_clipping_3d = opt.enable_clipping
		&& (opt.clip_x_min < opt.clip_x_max)
		&& (opt.clip_y_min < opt.clip_y_max)
		&& (opt.clip_z_min < opt.clip_z_max);

	unsigned int i=0;
	for(unsigned int y=0; y<height; y++) {
		for(unsigned int x=0; x<width; x++, i++) {
			Point& p = points[i];
			// write point pixel coordinate
			p.px = x;
			p.py = y;
			// convert color
			{
				const auto& cub = image[i];
				// p.color[0] = float(cub[0]) / 255.0f;
				// p.color[1] = float(cub[1]) / 255.0f;
				// p.color[2] = float(cub[2]) / 255.0f;
				p.color = Eigen::Vector3f(
					float(cub[0]),
					float(cub[1]),
					float(cub[2])) / 255.0f;
				if(opt.color_space != ColorSpaces::RGB)
					p.color = ColorFromRGB(p.color);
			}
			// clip points which are outside of 2D ROI
			if(is_clipping_2d) {
				if(    x < opt.roi_2d_x_min || opt.roi_2d_x_max < x
					|| y < opt.roi_2d_y_min || opt.roi_2d_y_max < y
				) {
					// point is clipped
					p.is_valid = false;
					goto label_pixel_invalid_omg;
				}
			}
			if(opt.density_mode == DensityModes::ASP_RGB) {
				p.is_valid = true;
				p.position = Eigen::Vector3f::Zero();
				p.cluster_radius_px = 0.0f;
			}
			else {
				// p.pixel[0] = static_cast<float>(x);
				// p.pixel[1] = py;
				p.is_valid = false;
				// get depth and compute z/f
				const uint16_t depth_i16 = depth[i];
				const float z_over_f = opt.camera.convertKinectToMeter(depth_i16) / opt.camera.focal;
				// if depth is 0 the point is invalid
				p.is_valid = (depth_i16 != 0);
				if(!p.is_valid) goto label_pixel_invalid_omg;
				// compute position
				p.position = opt.camera.unprojectImpl(
					static_cast<float>(p.px), static_cast<float>(p.py),
					z_over_f); // will give 0 if depth==0
				// clip points which are outside of 3D ROI
				if(is_clipping_3d) {
					if(    p.position.x() < opt.clip_x_min || opt.clip_x_max < p.position.x()
						|| p.position.y() < opt.clip_y_min || opt.clip_y_max < p.position.y()
						|| p.position.z() < opt.clip_z_min || opt.clip_z_max < p.position.z()
					) {
						// point is clipped
						p.is_valid = false;
						goto label_pixel_invalid_omg;
					}
				}
				// compute cluster radius [px]
				p.cluster_radius_px = opt.base_radius / z_over_f;
				// compute normal
				{
					// FIXME in count mode the gradient is computed using a default radius of 0.02
					// FIXME regardless of the radius chosen later
					Eigen::Vector2f gradient = LocalDepthGradient(depth, x, y, z_over_f, 0.5f*p.cluster_radius_px, opt.camera);
					p.setNormalFromGradient(gradient);
					// limit minimal circularity such that the maximum angle is 80 deg
					// FIXME limit normal angle
				}
				continue;
			}
label_pixel_invalid_omg:
			p.cluster_radius_px = 0.0f; // FIXME why do we have to set this?
			p.normal = Eigen::Vector3f(0,0,-1);
		}
	}

	DANVIL_BENCHMARK_START(density)
	if(opt.density_mode == DensityModes::ASP_RGB || opt.density_mode == DensityModes::ASP_RGBD) {
		// constant density and constant cluster radius
		float rho;
		float cluster_radius_px;
		if(opt.count == 0) {
			// estimate cluster radius [px]
			constexpr float M_TO_PX = 400.0f; // random constant to convert m to px
			cluster_radius_px = opt.base_radius * M_TO_PX;
			// compute density
			rho = 1.0f / (cluster_radius_px*cluster_radius_px*boost::math::constants::pi<float>());
		}
		else {
                        if (!opt.width || !opt.height || (opt.width <= 0 && opt.height <= 0)) {
                                opt.width = 640.f;
                                opt.height = 480.f;
                        }

			rho = static_cast<float>(opt.count) / (opt.width * opt.height);
			cluster_radius_px = 1.0f / std::sqrt(boost::math::constants::pi<float>() * rho);
		}
		// constant density
		density = Eigen::MatrixXf::Constant(points.width(), points.height(), rho);
		// set cluster_radius_px
		const float* p_density = density.data();
		for(unsigned int i=0; i<points.size(); i++) {
			points[i].cluster_radius_px = cluster_radius_px;
		}
	}
	else if(opt.density_mode == DensityModes::DASP) {
		// compute depth-adaptive density
		density = ComputeDepthDensity(points, opt);
		// depth-adaptive smooth
		if(opt.is_smooth_density) {
			density = density::DensityAdaptiveSmooth(density);
		}
		// adapt density to fit number of superpixels
		if(opt.count > 0) {
			// sum density image
			const float p = static_cast<float>(opt.count) / density.sum();
			density *= p;
			// correct base radius
			float scl = 1.0f / std::sqrt(p);
			opt.base_radius *= scl;
			// adapt cluster radius to density
			for(unsigned int i=0; i<points.size(); i++) {
				points[i].cluster_radius_px *= scl;
			}
		}
	}
	else {
		std::cerr << "Invalid density mode type" << std::endl;
		throw 0;
	}
	DANVIL_BENCHMARK_STOP(density)

	if(opt.ignore_pixels_with_bad_visibility) {
		// for ignore_pixels_with_bad_visibility
		constexpr unsigned int cNmin = 5;
		// following two lines are for the case where the normal is not used
	//	const float cAlphaMax = 60.0f / 180.0f * boost::math::constants::pi<float>();
	//	float threshold = std::sqrt(1 + 2.0f * std::pow(std::tan(cAlphaMax), 2)) * std::sqrt(static_cast<float>(cNmin) / boost::math::constants::pi<float>()) / opt.camera.focal;
		float threshold = std::sqrt(static_cast<float>(cNmin) / boost::math::constants::pi<float>()) / opt.camera.focal;
		for(unsigned int i=0; i<points.size(); i++) {
			Point& p = points[i];
			if(p.is_valid) {
				if(p.position[2] * threshold > opt.base_radius * p.computeCircularity()) {
					p.is_valid = false;
				}
			}
		}
	}
}

void Superpixels::ComputeSuperpixels(const std::vector<Seed>& seeds)
{

	SetRandomNumberSeed(opt.random_seed);

//	std::cout << "DASP START" << std::endl;
//	std::cout << "Density total = " << density.sum() << ", seed points = " << seeds.size() << std::endl;
	CreateClusters(seeds);

//	std::cout << "0: n=" << cluster.size() << std::endl;
//	for(unsigned int i=0; i<cluster.size(); i++) {
//		std::cout << cluster[i].pixel_ids.size() << " ";
//	}
//	std::cout << std::endl;
	for(unsigned int i=0; i<opt.iterations; i++) {
		MoveClusters();
//		std::cout << i+1 << ": n=" << cluster.size() << std::endl;
//		for(unsigned int j=0; j<cluster.size(); j++) {
//			std::cout << cluster[j].pixel_ids.size() << " ";
//		}
	}
//	std::cout << std::endl;
	if(opt.is_conquer_enclaves) {
		//ConquerEnclaves();
		ConquerMiniEnclaves();
	}
	// delete empty superpixels
	PurgeInvalidClusters();
//	std::cout << "final: n=" << cluster.size() << std::endl;

	opt.count_actual = cluster.size();
}

namespace SegmentExtraction
{
	struct Point {
		int x, y;
	};

	struct Segment {
		int label;
		std::vector<unsigned int> points_indices;
		std::set<int> border_labels;
		bool has_nan_neighbor;
	};

	Segment ExtractSegment(const slimage::Image1i& labels, const slimage::Image1i& visited, const Point& start) {
		// we need to temporarily store if border pixels have been visited
		slimage::Image1i border_visited(visited.width(), visited.height(), slimage::Pixel1i{0});
		// init segment
		Segment result;
		int base_label = labels(start.x, start.y);
		result.label = base_label;
		result.has_nan_neighbor = false;
		// do a flood fill by managing a list of "open" pixels which need investigation
		std::queue<Point> open;
		open.push(start);
		while(!open.empty()) {
			// get the next pixel which needs investigation
			Point p = open.front();
			open.pop();
			// skip pixels which are out of bounds
			if(!labels.isValidIndex(p.x, p.y)) {
				continue;
			}
			// do not visit pixels twice in this run
			if(visited(p.x, p.y) == base_label || border_visited(p.x, p.y) == 1) {
				continue;
			}
			// check label of current point
			int current_label = labels(p.x, p.y);
			if(current_label == base_label) {
				// point belongs to label -> investigate neighbors
				unsigned int index = static_cast<unsigned int>(p.x) + static_cast<unsigned int>(p.y) * labels.width();
				result.points_indices.push_back(index);
				// add neighbors
				open.push(Point{p.x-1, p.y});
				open.push(Point{p.x+1, p.y});
				open.push(Point{p.x, p.y-1});
				open.push(Point{p.x, p.y+1});
				// mark pixel as visited in this run (visited is re-used in later runs!)
				visited(p.x, p.y) = base_label;
			}
			else {
				// point does not belong to label -> add to neighbors
				if(current_label == -1) {
					result.has_nan_neighbor = true;
				}
				else {
					result.border_labels.insert(current_label);
				}
				// mark pixel temporarily as border-visited
				border_visited(p.x, p.y) = 1;
			}
		}
		return result;
	}

	std::vector<Segment> FindAllSegments(const slimage::Image1i& labels) {
		std::vector<Segment> segments;
		slimage::Image1i visited(labels.width(), labels.height(), slimage::Pixel1i{-1});
		for(unsigned int y=0; y<labels.height(); y++) {
			for(unsigned int x=0; x<labels.width(); x++) {
				// skip invalid pixels
				if(labels(x,y) == -1) {
					continue;
				}
				// skip pixels which are already visited in their run
				if(visited(x,y) != -1) {
					continue;
				}
				// find segment
				Segment result = ExtractSegment(labels, visited, Point{x, y});
				segments.push_back(result);
			}
		}
		return segments;
	}
}

void Superpixels::ConquerEnclaves()
{
	// compute labels for every pixel
	std::vector<int> labels_v = ComputePixelLabels();
	slimage::Image1i labels(width(), height(), slimage::Buffer<int>(width()*height(), labels_v.begin().base()));
	// find all segments (i.e. connected regions with the same label)
	std::vector<SegmentExtraction::Segment> segments = SegmentExtraction::FindAllSegments(labels);
	// count number of segments per label
	std::vector<unsigned int> label_segment_count(cluster.size(), 0);
	for(const SegmentExtraction::Segment& x : segments) {
		assert(x.label < cluster.size());
		assert(x.label != -1);
		label_segment_count[x.label] ++;
	}
	// sort segments by size
	std::sort(segments.begin(), segments.end(),
			[](const SegmentExtraction::Segment& a, const SegmentExtraction::Segment& b) {
					return a.points_indices.size() < b.points_indices.size();
			});
	// remove enclaves and exclaves
	for(SegmentExtraction::Segment& x : segments) {
		// do not delete last segment of a label
		if(label_segment_count[x.label] == 1) {
			continue;
		}
		label_segment_count[x.label] --;
		// can not delete if no neighbor (i.e. only nan neighbors)
		if(x.border_labels.size() == 0) {
			continue;
		}
		// find neighbor with largest segment
		// FIXME possibly use a big enough neighbor which matches color and normal best
		unsigned int neighbor_best_size = 0;
		unsigned int neighbor_best_label = 0;
		for(int nlab : x.border_labels) {
			unsigned int nsize = cluster[nlab].pixel_ids.size();
			if(nsize > neighbor_best_size) {
				neighbor_best_size = nsize;
				neighbor_best_label = nlab;
			}
		}
		// assign segment to one of the neighbors
		std::cout << "Conquering enclave with " << x.points_indices.size() << " points" << std::endl;
		// remove pixels from old cluster
		std::vector<unsigned int>& original_pixel_ids = cluster[x.label].pixel_ids;
		std::sort(original_pixel_ids.begin(), original_pixel_ids.end());
		std::sort(x.points_indices.begin(), x.points_indices.end());
		std::vector<unsigned int> new_pixel_ids(original_pixel_ids.size());
		auto it = std::set_difference(original_pixel_ids.begin(), original_pixel_ids.end(),
				x.points_indices.begin(), x.points_indices.end(),
				new_pixel_ids.begin());
		original_pixel_ids = std::vector<unsigned int>(new_pixel_ids.begin(), it);
		// add pixels to new cluster
		cluster[neighbor_best_label].pixel_ids.insert(cluster[neighbor_best_label].pixel_ids.begin(), x.points_indices.begin(), x.points_indices.end());
	}
}

void Superpixels::ConquerMiniEnclaves()
{
	// edge neighbors are checked later
	int offset[4] = {
			-1, +1, -static_cast<int>(width()), +static_cast<int>(width())
	};
	// compute pixel labels
	std::vector<int> labels_v = ComputePixelLabels();
	slimage::Image1i labels(width(), height(), slimage::Buffer<int>(width()*height(), labels_v.begin().base()));
	// iterate over all pixels
	for(unsigned int y=1; y+1<labels.height(); y++) {
		for(unsigned int x=1; x+1<labels.width(); x++) {
			// compute pixel index and get label
			unsigned int index = x + y * labels.width();
			int lab = labels(x,y);
			// skip invalid pixels
			if(lab == -1) {
				continue;
			}
			// check neighbors
			int neighbors[4]; // will contain the labels of neighbor pixels
			bool are_all_different = true; // will be true if all neighbors have a different label
			bool are_all_nan = true; // will be true if all neighbors are invalid
			for(unsigned int i=0; i<4; i++) {
				int x = labels[index + offset[i]];
				neighbors[i] = x;
				are_all_different = are_all_different && (x != lab);
				are_all_nan = are_all_nan && (x == -1);
			};
			// relabel pixel if isolated and at least one neighbors is valid
			if(are_all_different && !are_all_nan) {
				// remove from old cluster
				cluster[lab].removePixel(index);
				// find best new cluster (valid and nearest in world coordinates)
				int best_lab = 0;
				float best_dist = 1e9;
				for(unsigned int i=0; i<4; i++) {
					if(neighbors[i] == -1) {
						continue;
					}
					float d2 = (points[index].position - points[index + offset[i]].position).squaredNorm();
					if(d2 < best_dist) {
						best_lab = neighbors[i];
						best_dist = d2;
					}
				}
				// add to new cluster
				assert(best_lab != -1);
				cluster[best_lab].addPixel(index);
				labels(x,y) = best_lab;
			}
		}
	}
}

std::vector<Seed> CreateSeedPoints(
	const ImagePoints& img,
	const std::vector<Eigen::Vector2f>& pnts,
	std::vector<int>* origin=0)
{
	std::vector<Seed> pnts_final;
	pnts_final.reserve(pnts.size());
	std::vector<int> origin_final;
	if(origin) {
		origin_final.reserve(pnts.size());
		assert(pnts.size() == origin->size());
	}
	for(size_t i=0; i<pnts.size(); i++) {
		const Eigen::Vector2f& p = pnts[i];
		int x = static_cast<int>(std::round(p[0]));
		int y = static_cast<int>(std::round(p[1]));
		if(0 <= x && x < static_cast<int>(img.width())
			&& 0 <= y && y < static_cast<int>(img.height())
			&& img(x,y).is_valid
		) {
			pnts_final.push_back(
				Seed::Dynamic(x, y, img(x,y).cluster_radius_px));
			if(origin) {
				origin_final.push_back((*origin)[i]);
			}
		}
	}
	if(origin) {
		*origin = origin_final;
	}
	return pnts_final;
}

std::vector<Seed> Superpixels::FindSeeds()
{
	switch(opt.seed_mode) {
	case SeedModes::Random:
		return CreateSeedPoints(points,
			pds::Random(density));
	case SeedModes::Grid:
		return CreateSeedPoints(points,
			pds::RectGrid(density));
	case SeedModes::SimplifiedPDS_Old:
		return CreateSeedPoints(points,
			pds::SimplifiedPDSOld(density));
	case SeedModes::SimplifiedPDS:
		return CreateSeedPoints(points,
			pds::SimplifiedPDS(density));
	case SeedModes::FloydSteinberg:
		return CreateSeedPoints(points,
			pds::FloydSteinberg(density));
	case SeedModes::FloydSteinbergExpo:
		return CreateSeedPoints(points,
			pds::FloydSteinbergExpo(density));
	case SeedModes::FloydSteinbergMultiLayer:
		return CreateSeedPoints(points,
			pds::FloydSteinbergMultiLayer(density));
	case SeedModes::Fattal:
		return CreateSeedPoints(points,
			pds::Fattal(density));
	case SeedModes::Delta: {
		std::vector<Eigen::Vector2f> pnts_prev(cluster.size());
		for(std::size_t i=0; i<cluster.size(); i++) {
			const Cluster& c = cluster[i];
			pnts_prev[i] = Eigen::Vector2f(c.center.px, c.center.py);
		}
		std::vector<int> seed_origin;
		std::vector<Eigen::Vector2f> pnts = pds::DeltaDensitySampling(density, pnts_prev, &seed_origin);
		std::vector<Seed> seeds = CreateSeedPoints(points, pnts, &seed_origin);
		assert(seeds.size() != seed_origin.size());
		if(seeds.size() != seed_origin.size()) {
			std::cerr << "ERROR with DDS: invalid point!" << std::endl;
		}
		for(size_t i=0; i<seeds.size(); ++i) {
			seeds[i].label = seed_origin[i];
		}
		return seeds;
	}
	default:
		assert(false && "FindSeeds: Unkown mode!");
	};
}

std::vector<int> Superpixels::ComputePixelLabels() const
{
	std::vector<int> labels(points.size(), -1);
	for(unsigned int j=0; j<cluster.size(); j++) {
		for(unsigned int i : cluster[j].pixel_ids) {
			labels[i] = static_cast<int>(j);
		}
	}
	return labels;
}

slimage::Image1i Superpixels::ComputeLabels() const
{
	slimage::Image1i img(width(), height(), slimage::Pixel1i{-1});
	for(unsigned int j=0; j<cluster.size(); j++) {
		for(unsigned int i : cluster[j].pixel_ids) {
			img[i] = static_cast<int>(j);
		}
	}
	return img;
}

Partition Superpixels::ComputePartition() const
{
	Partition p;
	p.segments.resize(cluster.size());
	for(unsigned int i=0; i<cluster.size(); i++) {
		p.segments[i] = cluster[i].pixel_ids;
	}
	return p;
}

void Superpixels::CreateClusters(const std::vector<Seed>& seeds)
{
	// create clusters
	cluster.clear();
	cluster.reserve(seeds.size());
	for(unsigned int k=0; k<seeds.size(); k++) {
		const Seed& p = seeds[k];
		Cluster c;
		c.seed_id = p.label; // k
		c.is_fixed = p.is_fixed;
		c.center.px = p.x;
		c.center.py = p.y;
		c.center.is_valid = true;
		c.center.cluster_radius_px = p.scala;
		if(c.is_fixed) {
			// use fixed world position (and compute depth from it)
			c.center.position = p.fixed_world;
			c.center.color = p.fixed_color;
			c.center.normal = p.fixed_normal;
		}
		// assign points (use only half the radius)
		int R = static_cast<int>(std::ceil(c.center.cluster_radius_px * 0.35f));
		R = std::min(2, R);
		assert(R >= 0 && "CreateClusters: Invalid radius!");
		c.pixel_ids.reserve((2*R + 1)*(2*R + 1));
		unsigned int xmin = std::max<int>(p.x - R, 0);
		unsigned int xmax = std::min<int>(p.x + R, int(points.width()) - 1);
		unsigned int ymin = std::max<int>(p.y - R, 0);
		unsigned int ymax = std::min<int>(p.y + R, int(points.height()) - 1);
		for(unsigned int yi=ymin; yi<=ymax; yi++) {
			for(unsigned int xi=xmin; xi<=xmax; xi++) {
				unsigned int index = points.index(xi, yi);
				if(!points(xi, yi).is_valid) {
					// omit invalid points
					continue;
				}
				c.pixel_ids.push_back(index);
			}
		}
		// update center
		if(c.isValid()) {
			c.UpdateCenter(points, opt);
			cluster.push_back(c);
		}
	}
}

slimage::Image1f Superpixels::ComputeEdges()
{
	// FIXME metric needs central place!
	if(opt.density_mode == DensityModes::ASP_RGB) {
		DensityAdaptiveMetric_UxRGB metric(opt.weight_spatial, opt.weight_color);
		return dasp::ComputeEdges(points, metric);
	}
	else if(opt.density_mode == DensityModes::ASP_RGBD) {
		DensityAdaptiveMetric_UxRGBxD metric(opt.weight_spatial, opt.weight_color, opt.weight_normal);
		return dasp::ComputeEdges(points, metric);
	}
	else if(opt.density_mode == DensityModes::DASP) {
		// dasp
		DepthAdaptiveMetric metric(opt.weight_spatial, opt.weight_color, opt.weight_normal, opt.base_radius);
		return dasp::ComputeEdges(points, metric);
	}
	else {
		// error!
		std::cerr << "Invalid depth mode" << std::endl;
	}
}

void Superpixels::ImproveSeeds(std::vector<Seed>& seeds, const slimage::Image1f& edges)
{
	const unsigned int width = points.width();
	const unsigned int height = points.height();

	const int dx8[8] = {-1, -1,  0,  1, 1, 1, 0, -1};
	const int dy8[8] = { 0, -1, -1, -1, 0, 1, 1,  1};

	for(Seed& seed : seeds) {
		if(seed.is_fixed) {
			continue;
		}
		int sx = seed.x;
		int sy = seed.y;
		int bestid = sy*width + sx;
		for(unsigned int i=0; i<8; i++) {
			int nx = sx + dx8[i];
			int ny = sy + dy8[i];
			if(nx >= 0 && nx < int(width) && ny >= 0 && ny < int(height)) {
				int nind = ny*width + nx;
				if(edges[nind] < edges[bestid]) {
					bestid = nind;
				}
			}
		}
		seed.x = bestid % width;
		seed.y = bestid / width;
	}
}

void Superpixels::PurgeInvalidClusters()
{
	auto it = std::remove_if(cluster.begin(), cluster.end(),
		[](const Cluster& c) {
			return !c.isValid();
		});
	cluster.resize(it - cluster.begin());
}

void Superpixels::MoveClusters()
{
	// compute next iteration of cluster labeling
	slimage::Image1i labels;
	// FIXME metric needs central place!
	if(opt.density_mode == DensityModes::ASP_RGB) {
		DensityAdaptiveMetric_UxRGB metric(opt.weight_spatial, opt.weight_color);
		labels = dasp::IterateClusters(cluster, points, opt, metric);
	}
	else if(opt.density_mode == DensityModes::ASP_RGBD) {
		DensityAdaptiveMetric_UxRGBxD metric(opt.weight_spatial, opt.weight_color, opt.weight_normal);
		labels = dasp::IterateClusters(cluster, points, opt, metric);
	}
	else if(opt.density_mode == DensityModes::DASP) {
		DepthAdaptiveMetric metric(opt.weight_spatial, opt.weight_color, opt.weight_normal, opt.base_radius);
		labels = dasp::IterateClusters(cluster, points, opt, metric);
	}
	else {
		// error!
		std::cerr << "Invalid depth mode" << std::endl;
	}
	// delete old clusters assignments
	for(Cluster& c : cluster) {
		unsigned int n = c.pixel_ids.size();
		c.pixel_ids.clear();
		c.pixel_ids.reserve(n);
	}
	// assign points to clusters
	for(unsigned int i=0; i<points.size(); i++) {
		int label = labels[i];
		if(label >= 0) {
			cluster[label].pixel_ids.push_back(i);
		}
	}
	// remove invalid clusters
	PurgeInvalidClusters();
	// update remaining (valid) clusters
	for(Cluster& c : cluster) {
		c.UpdateCenter(points, opt);
	}
}

ClusterGroupInfo Superpixels::ComputeClusterGroupInfo(unsigned int n, float max_thick)
{
	ClusterGroupInfo cgi;
	cgi.hist_thickness = Histogram<float>(n, 0, max_thick);
	cgi.hist_area_quotient = Histogram<float>(n, 0.5f, 2.0f);
	cgi.hist_coverage_error = Histogram<float>(n, 0, 0.10f);
	for(const Cluster& c : cluster) {
		cgi.hist_thickness.add(c.thickness);
		cgi.hist_area_quotient.add(c.area_quotient);
		cgi.hist_coverage_error.add(c.coverage_error);
//		std::cout << ci.t << " " << ci.b << " " << ci.a << std::endl;
	}
	return cgi;
}

Eigen::Vector3f Superpixels::ColorToRGB(const Eigen::Vector3f& source) const
{
	Eigen::Vector3f target;
	switch(opt.color_space) {
	case ColorSpaces::HSV: {
		Danvil::convert_hsv_2_rgb(source[0], source[1], source[2], target[0], target[1], target[2]);
	} break;
	case ColorSpaces::LAB: {
		Danvil::color_lab_to_rgb(100.0f*source[0], 100.0f*source[1], 100.0f*source[2], target[0], target[1], target[2]);
	} break;
	case ColorSpaces::HN: {
		float ra = source[0];
		float ga = source[1];
		float a = source[2];
		if(a > 0) {
			target[0] = ra * a;
			target[1] = ga * a;
			target[2] = a - (target[0] + target[1]);
		}
		else {
			// fixme
			target[0] = 0;
			target[1] = 0;
			target[2] = 0;
		}
	} break;
	default: case ColorSpaces::RGB: {
		target[0] = source[0];
		target[1] = source[1];
		target[2] = source[2];
	} break;
	}
	return target;
}

Eigen::Vector3f Superpixels::ColorFromRGB(const Eigen::Vector3f& source) const
{
	switch(opt.color_space) {
	case ColorSpaces::HSV: {
		Eigen::Vector3f target;
		Danvil::convert_rgb_2_hsv(source[0], source[1], source[2], target[0], target[1], target[2]);
		return target;
	}
	case ColorSpaces::LAB: {
		Eigen::Vector3f target;
		Danvil::color_rgb_to_lab(source[0], source[1], source[2], target[0], target[1], target[2]);
		return 0.01f * target;
	}
	case ColorSpaces::HN: {
		float r = source[0];
		float g = source[1];
		float b = source[2];
		float a = r + g + b;
		if(a > 0.05f) {
			return Eigen::Vector3f(r / a, g / a, a * 0.1f);
		}
		else {
			// FIXME
			return Eigen::Vector3f::Zero();
		}
	}
	default: case ColorSpaces::RGB:
		return source;
	}
}

Superpixels ComputeSuperpixels(const slimage::Image3ub& color, const slimage::Image1ui16& depth, const Parameters& opt)
{
	Superpixels clustering;
	clustering.opt = opt;
	ComputeSuperpixelsIncremental(clustering, color, depth);
	return clustering;
}

void ComputeSuperpixelsIncremental(Superpixels& clustering, const slimage::Image3ub& color, const slimage::Image1ui16& depth)
{
	if(clustering.opt.is_repair_depth) {
		DANVIL_BENCHMARK_START(dasp_repair)
		RepairDepth(depth, color);
		DANVIL_BENCHMARK_STOP(dasp_repair)
	}

	if(clustering.opt.is_smooth_depth) {
		DANVIL_BENCHMARK_START(dasp_smooth)
		SmoothDepth(depth, color);
		DANVIL_BENCHMARK_STOP(dasp_smooth)
	}

	// compute normals only if necessary
	slimage::Image3f normals;
//	if(super_params_ext.weight_normal > 0.0f) {
//		DANVIL_BENCHMARK_START(normals)
//		slimage::Image3f kinect_points = dasp::PointsAndNormals::ComputePoints(kinect_depth, slimage::ThreadingOptions::UsePool(thread_pool_index_));
//		slimage::Image3f kinect_normals = dasp::PointsAndNormals::ComputeNormals(kinect_depth, kinect_points, slimage::ThreadingOptions::UsePool(thread_pool_index_));
//	//	dasp::PointsAndNormals::ComputeNormalsFast(kinect_depth, kinect_points, kinect_normals);
//		DANVIL_BENCHMARK_STOP(normals)
//	}

	// prepare super pixel points
	DANVIL_BENCHMARK_START(dasp_points)
	clustering.CreatePoints(color, depth, normals);
	DANVIL_BENCHMARK_STOP(dasp_points)

	// compute super pixel seeds
	DANVIL_BENCHMARK_START(dasp_seeds)
	clustering.seeds_previous = clustering.seeds;
	clustering.seeds = clustering.FindSeeds();
//	std::cout << "Seeds: " << seeds.size() << std::endl;
	DANVIL_BENCHMARK_STOP(dasp_seeds)

	// compute super pixel point edges and improve seeds with it
	if(clustering.opt.is_improve_seeds) {
		DANVIL_BENCHMARK_START(dasp_improve)
		slimage::Image1f edges = clustering.ComputeEdges();
		clustering.ImproveSeeds(clustering.seeds, edges);
		DANVIL_BENCHMARK_STOP(dasp_improve)
	}

	// compute clusters
	DANVIL_BENCHMARK_START(dasp_clusters)
	clustering.ComputeSuperpixels(clustering.seeds);
	DANVIL_BENCHMARK_STOP(dasp_clusters)
}

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


/**
 * @file sdirk.cc
 * @brief NPDE homework SDIRK code
 * @author Unknown, Oliver Rietmann
 * @date 31.03.2021
 * @copyright Developed at ETH Zurich
 */

#include "sdirk.h"

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

#include "../../../lecturecodes/helperfiles/polyfit.h"

namespace SDIRK {

/* SAM_LISTING_BEGIN_0 */
Eigen::Vector2d SdirkStep(const Eigen::Vector2d &z0, double h, double gamma) {
  Eigen::Vector2d res;
  // Compute one timestep of the SDIRK implicit RK-SSM for the linear ODE
  //====================
  // Your code goes here
  //====================
  return res;
}
/* SAM_LISTING_END_0 */

/* SAM_LISTING_BEGIN_1 */
std::vector<Eigen::Vector2d> SdirkSolve(const Eigen::Vector2d &z0,
                                        unsigned int M, double T,
                                        double gamma) {
  // Solution vector
  std::vector<Eigen::Vector2d> res(M + 1);
  // Solve the ODE with uniform timesteps using the SDIRK method
  //====================
  // Your code goes here
  //====================
  return res;
}
/* SAM_LISTING_END_1 */

/* SAM_LISTING_BEGIN_2 */
double CvgSDIRK() {
  double conv_rate;
  // Study the convergence rate of the method.
  //====================
  // Your code goes here
  //====================
  return conv_rate;
}
/* SAM_LISTING_END_2 */

}  // namespace SDIRK


#include <qlo/objects/obj_credit.hpp>

#include <qlo/enumerations/factories/termstructuresfactory.hpp>

#include <ql/instruments/stock.hpp>
#include <ql/quote.hpp>
#include <ql/currencies/europe.hpp>
#include <ql/time/daycounter.hpp>
#include <ql/time/daycounters/actual360.hpp>
#include <ql/termstructures/yield/piecewiseyieldcurve.hpp>
#include <ql/termstructures/credit/interpolatedhazardratecurve.hpp>
#include <ql/termstructures/credit/flathazardrate.hpp>
#include <ql/termstructures/credit/piecewisedefaultcurve.hpp>
#include <ql/math/interpolations/backwardflatinterpolation.hpp>
#include <ql/math/interpolations/loginterpolation.hpp>
#include <ql/math/interpolations/bilinearinterpolation.hpp>
#include <ql/math/interpolations/bicubicsplineinterpolation.hpp>
#include <ql/math/solver1d.hpp>
#include <ql/math/solvers1d/brent.hpp>

#include <ql/pricingengines/credit/midpointcdsengine.hpp>

#include <ql/instruments/creditdefaultswap.hpp>
#include <ql/experimental/credit/riskybond.hpp>
#include <ql/experimental/credit/syntheticcdo.hpp>
#include <ql/experimental/credit/midpointcdoengine.hpp>
#include <ql/experimental/credit/nthtodefault.hpp>
#include <ql/experimental/credit/integralntdengine.hpp>
#include <ql/experimental/credit/basecorrelationstructure.hpp>
#include <ql/experimental/credit/cdsoption.hpp>
#include <ql/experimental/credit/blackcdsoptionengine.hpp>

#include <boost/algorithm/string/case_conv.hpp>
#include <boost/make_shared.hpp>

#include <ql/settings.hpp>

QuantLibAddin::CreditDefaultSwap::CreditDefaultSwap(
    const boost::shared_ptr<reposit::ValueObject>& properties,
            // BEGIN typemap rp_tm_default
            QuantLib::Protection::Side BuyerSeller,
            QuantLib::Real Notional,
            QuantLib::Rate Upfront,
            QuantLib::Rate Spread,
            QuantLib::Schedule const &PremiumSchedule,
            QuantLib::BusinessDayConvention PaymentConvention,
            QuantLib::DayCounter const &DayCounter,
            bool SettlesAccrual,
            bool PayAtDefault,
            QuantLib::Date const &ProtectionStart,
            QuantLib::Date const &UpfrontDate,
            // END   typemap rp_tm_default
    bool permanent)
: Instrument(properties, permanent) {
		// dirty way to decide if this is constructed through a run only version
        if(UpfrontDate == QuantLib::Null<QuantLib::Date>() && Upfront == 0.) {
            libraryObject_ = boost::shared_ptr<QuantLib::CreditDefaultSwap>(
                new QuantLib::CreditDefaultSwap(
                    BuyerSeller,
                    Notional,
                    Spread,
                    PremiumSchedule,
                    PaymentConvention,
                    DayCounter,
                    SettlesAccrual,
                    PayAtDefault,
                    ProtectionStart,
                    boost::shared_ptr<QuantLib::Claim>()));
        }else{
                libraryObject_ = boost::shared_ptr<QuantLib::CreditDefaultSwap>(
                    new QuantLib::CreditDefaultSwap(
                        BuyerSeller,
                        Notional,
                        Upfront,
                        Spread,
                        PremiumSchedule,
                        PaymentConvention,
                        DayCounter,
                        SettlesAccrual,
                        PayAtDefault,
                        ProtectionStart,
                        UpfrontDate,
                        boost::shared_ptr<QuantLib::Claim>()));
        }
}

QuantLibAddin::HazardRateCurve::HazardRateCurve(
    const boost::shared_ptr<reposit::ValueObject>& properties,
            // BEGIN typemap rp_tm_default
            std::vector< QuantLib::Date > const &CurveDates,
            std::vector< QuantLib::Rate > const &CurveRates,
            QuantLib::DayCounter const &DayCounter,
            // END   typemap rp_tm_default
    bool permanent)
: DefaultProbabilityTermStructure(properties, permanent) {
        QL_REQUIRE(!CurveDates.empty(), "no input dates given");
        QL_REQUIRE(CurveDates.size() == CurveRates.size(),
                   "vector sizes differ");
        libraryObject_ = boost::shared_ptr<QuantLib::Extrapolator>(
        new QuantLib::InterpolatedHazardRateCurve<QuantLib::BackwardFlat>(
 				 CurveDates, CurveRates, DayCounter));
}

QuantLibAddin::PiecewiseHazardRateCurve::PiecewiseHazardRateCurve(
    const boost::shared_ptr<reposit::ValueObject>& properties,
            // BEGIN typemap rp_tm_default
            std::vector< boost::shared_ptr< QuantLib::DefaultProbabilityHelper > > const &Helpers,
            QuantLib::DayCounter const &DayCounter,
            QuantLib::Calendar const &Calendar,
            std::string const &Interpolation,
            QuantLib::Real Accuracy,
            // END   typemap rp_tm_default
    bool permanent)
: DefaultProbabilityTermStructure(properties, permanent) {
        if(Interpolation == std::string("LINEAR")){
            libraryObject_ = boost::shared_ptr<QuantLib::Extrapolator>(new
                   QuantLib::PiecewiseDefaultCurve<QuantLib::HazardRate,
                        QuantLib::Linear>(
                            0,
                            Calendar,
                            Helpers,
                            DayCounter));
        }else if(Interpolation == std::string("BACKWARDFLAT")) {
            libraryObject_ = boost::shared_ptr<QuantLib::Extrapolator>(new
                   QuantLib::PiecewiseDefaultCurve<QuantLib::HazardRate,
                        QuantLib::BackwardFlat>(
                            0,
                            Calendar,
                            Helpers,
                            DayCounter));
        }else{
            QL_FAIL("Unrecognised interpolator");
        }

        libraryObject_->enableExtrapolation();
}

const std::vector<QuantLib::Date>& QuantLibAddin::PiecewiseHazardRateCurve::dates() const {
    typedef QuantLib::PiecewiseDefaultCurve<QuantLib::HazardRate, QuantLib::BackwardFlat> flat_curve;
    typedef QuantLib::PiecewiseDefaultCurve<QuantLib::HazardRate, QuantLib::Linear> lin_curve;
    boost::shared_ptr<flat_curve> ptrBF =
        boost::dynamic_pointer_cast<flat_curve>(libraryObject_);
    if(ptrBF) return ptrBF->dates();
    boost::shared_ptr<lin_curve> ptrLIN =
        boost::dynamic_pointer_cast<lin_curve>(libraryObject_);
    if(ptrLIN) return ptrLIN->dates();
    QL_FAIL("Unable to cast default probability term structure.");
}

const std::vector<QuantLib::Real>& QuantLibAddin::PiecewiseHazardRateCurve::data() const {
    typedef QuantLib::PiecewiseDefaultCurve<QuantLib::HazardRate, QuantLib::BackwardFlat> flat_curve;
    typedef QuantLib::PiecewiseDefaultCurve<QuantLib::HazardRate, QuantLib::Linear> lin_curve;
    boost::shared_ptr<flat_curve> ptrBF =
        boost::dynamic_pointer_cast<flat_curve>(libraryObject_);
    if(ptrBF) return ptrBF->data();
    boost::shared_ptr<lin_curve> ptrLIN =
        boost::dynamic_pointer_cast<lin_curve>(libraryObject_);
    if(ptrLIN) return ptrLIN->data();
    QL_FAIL("Unable to cast default probability term structure.");
}

QuantLibAddin::PiecewiseFlatForwardCurve::PiecewiseFlatForwardCurve(
    const boost::shared_ptr<reposit::ValueObject>& properties,
            // BEGIN typemap rp_tm_default
            QuantLib::Date const &ReferenceDate,
            std::vector< boost::shared_ptr< QuantLib::RateHelper > > const &RateHelpers,
            QuantLib::DayCounter const &DayCounter,
            QuantLib::Real Accuracy,
            // END   typemap rp_tm_default
    bool permanent)
: YieldTermStructure(properties, permanent) {
        libraryObject_ = boost::shared_ptr<QuantLib::Extrapolator>(new
               QuantLib::PiecewiseYieldCurve<QuantLib::Discount,QuantLib::LogLinear>(ReferenceDate, RateHelpers, DayCounter));
}

QuantLibAddin::RiskyFixedBond::RiskyFixedBond(
    const boost::shared_ptr<reposit::ValueObject>& properties,
            // BEGIN typemap rp_tm_default
            std::string Bondname,
            QuantLib::Currency Currency,
            QuantLib::Real Recovery,
            QuantLib::Handle< QuantLib::DefaultProbabilityTermStructure > const &DefaultCurve,
            QuantLib::Schedule const &Schedule,
            QuantLib::Real Rate,
            QuantLib::DayCounter DayCounter,
            QuantLib::BusinessDayConvention PaymentConvention,
            QuantLib::Real Notional,
            QuantLib::Handle< QuantLib::YieldTermStructure > const &DiscountingCurve,
            QuantLib::Date PricingDate,
            // END   typemap rp_tm_default
    bool permanent)
: Instrument(properties, permanent) {
        std::vector<QuantLib::Real> notionals(1,Notional);

        libraryObject_ = boost::shared_ptr<QuantLib::RiskyFixedBond>(
            new QuantLib::RiskyFixedBond(
                    Bondname,Currency,Recovery,DefaultCurve,Schedule,Rate,DayCounter,
                    PaymentConvention,notionals,DiscountingCurve///, PricingDate
                                       ));
}

QuantLibAddin::Issuer::Issuer(
    const boost::shared_ptr<reposit::ValueObject>& properties,
            // BEGIN typemap rp_tm_default
            boost::shared_ptr< QuantLib::DefaultProbabilityTermStructure > const &DefaultCurves,
            boost::shared_ptr< QuantLib::DefaultEventSet > const &DefaultEvents,
            // END   typemap rp_tm_default
    bool permanent)
: reposit::LibraryObject<QuantLib::Issuer>(properties, permanent) {
        std::vector<QuantLib::Issuer::key_curve_pair> curves(1, std::make_pair(
            QuantLib::NorthAmericaCorpDefaultKey(QuantLib::EURCurrency(),
                                                     QuantLib::SeniorSec,
                                                     QuantLib::Period(),
                                                     1. // amount threshold
                                                     ),
            QuantLib::Handle<QuantLib::DefaultProbabilityTermStructure>(DefaultCurves)
        ));
        libraryObject_ = boost::shared_ptr<QuantLib::Issuer>(new QuantLib::Issuer(curves, *DefaultEvents));
}

QuantLibAddin::DefaultEventSet::DefaultEventSet(
    const boost::shared_ptr<reposit::ValueObject>& properties,
            // BEGIN typemap rp_tm_default
            std::string const &EventType,
            QuantLib::Date const &EventDate,
            QuantLib::Currency const &Currency,
            QuantLib::Seniority Seniority,
            QuantLib::Date const &SettlementDate,
            QuantLib::Real SettledRecovery,
            // END   typemap rp_tm_default
    bool permanent)
: reposit::LibraryObject<QuantLib::DefaultEventSet>(properties, permanent) {
        // if no match return empty set
        libraryObject_ = boost::shared_ptr<QuantLib::DefaultEventSet>(new QuantLib::DefaultEventSet());

        // only one recovery parsed by now; bankruptcy events need the whole
        //  set or they fail to construct.
        std::map<QuantLib::Seniority, QuantLib::Real> rrs;
        rrs.insert(std::pair<QuantLib::Seniority, QuantLib::Real>(Seniority, SettledRecovery));
        if(EventType==std::string("FailureToPayEvent")) {
            libraryObject_->insert(boost::shared_ptr<QuantLib::FailureToPayEvent> (
                new QuantLib::FailureToPayEvent(EventDate, Currency, Seniority, 1.e7,
                //implSettlemt,
                SettlementDate,
                rrs)));
        }else if(EventType==std::string("BankruptcyEvent")){
            libraryObject_->insert(boost::shared_ptr<QuantLib::BankruptcyEvent> (
                new QuantLib::BankruptcyEvent(EventDate, Currency, Seniority,
                SettlementDate,
                rrs)));
        }
}

QuantLibAddin::BaseCorrelationTermStructure::BaseCorrelationTermStructure(
    const boost::shared_ptr<reposit::ValueObject>& properties,
            // BEGIN typemap rp_tm_default
            std::string const &InterpolatorType,
            QuantLib::Natural SettlementDays,
            QuantLib::Calendar const &Calendar,
            QuantLib::BusinessDayConvention Convention,
            std::vector< QuantLib::Period > const &Tenors,
            std::vector< QuantLib::Real > const &LossLevel,
            std::vector< std::vector< QuantLib::Handle< QuantLib::Quote > > > const &Correlations,
            QuantLib::DayCounter const &DayCounter,
            // END   typemap rp_tm_default
    bool permanent)
: CorrelationTermStructure(properties, permanent) {
        // select 2D interpolator:
        //  Another option is to write one class per interpolator.
        if(InterpolatorType == std::string("BILIN")) {
            libraryObject_ = boost::shared_ptr<QuantLib::BaseCorrelationTermStructure<QuantLib::BilinearInterpolation> >(new
                QuantLib::BaseCorrelationTermStructure<QuantLib::BilinearInterpolation>(SettlementDays, Calendar, Convention, Tenors, LossLevel, Correlations, DayCounter));
        //}else if(InterpolatorType == std::string("")) {
        }else if(InterpolatorType == std::string("BICUBIC")) {
            libraryObject_ = boost::shared_ptr<QuantLib::BaseCorrelationTermStructure<QuantLib::BicubicSpline> >(new
                QuantLib::BaseCorrelationTermStructure<QuantLib::BicubicSpline>(SettlementDays, Calendar, Convention, Tenors, LossLevel, Correlations, DayCounter));
        }else{
            QL_FAIL("Can't determine base correlation surface interpolator.");
        }
}

QuantLib::Real QuantLibAddin::BaseCorrelationTermStructure::correlation(const QuantLib::Date& d, QuantLib::Real lossLevel) {
    if(interpolType_ == std::string("BILIN")) {
        return boost::dynamic_pointer_cast<QuantLib::BaseCorrelationTermStructure<QuantLib::BilinearInterpolation> >
            (libraryObject_)->correlation(d, lossLevel);
    }else if(interpolType_ == std::string("BICUBIC")) {
        return boost::dynamic_pointer_cast<QuantLib::BaseCorrelationTermStructure<QuantLib::BicubicSpline> >
            (libraryObject_)->correlation(d, lossLevel);
    }else{
        QL_FAIL("unknown 2D interpolator");
    }
}


#include "run.hpp"
#include "crit_crauser.hpp"
#include "crit_dijkstra.hpp"
#include "crit_heuristic.hpp"
#include "crit_oracle.hpp"
#include "crit_paper_in.hpp"
#include "crit_paper_out.hpp"
#include "crit_traff_bridge.hpp"
#include "dijkstra.hpp"
#include "generate_edges.hpp"
#include "generate_positions.hpp"
#include "math.hpp"
#include <boost/algorithm/string.hpp>
#include <cctype>
#include <fstream>
#include <random>

#ifndef DISABLE_CAIRO
#include "draw_graph.hpp"
#include <cairomm/cairomm.h>
#endif

using namespace boost::algorithm;

const std::string sssp::dijkstra_result_csv_header("node_count,phase,relaxed,fringe_size");

void sssp::execute_run(const arguments& args, std::ostream* out, std::ostream* err) {
    size_t start_node = 0;
    node_map<vec2> positions;
    graph graph;

    if (args.graph_file.empty()) {
        std::mt19937_64 rng(args.seed);
        std::uniform_int_distribution<int> uniform_seed(INT_MIN, INT_MAX);
        int position_seed = uniform_seed(rng);
        int cost_seed = uniform_seed(rng);
        int edge_seed = uniform_seed(rng);

        if (args.edge_gen.algorithm != edge_algorithm::kronecker) {
            switch (args.position_gen.algorithm) {
                case position_algorithm::poisson:
                    positions = generate_poisson_disc_positions(
                        position_seed, args.position_gen.poisson.min_distance, args.position_gen.poisson.max_reject);
                    break;
                case position_algorithm::uniform:
                    positions = generate_uniform_positions(position_seed, args.position_gen.uniform.count);
                    break;
                default:
                    BOOST_ASSERT(false);
                    break;
            }

            for (size_t i = 0; i < positions.size(); ++i) {
                graph.add_node();
            }
        }

        edge_cost_fn edge_cost_fn;
        switch (args.cost_gen.algorithm) {
            case cost_algorithm::uniform:
                edge_cost_fn = [rng = std::mt19937(cost_seed)](const line& line) mutable {
                    return std::uniform_real_distribution<double>(0.0, 1.0)(rng);
                };
                break;
            case cost_algorithm::one:
                edge_cost_fn = [](const line& line) { return 1.0; };
                break;
            case cost_algorithm::euclidean:
                edge_cost_fn = [](const line& line) { return distance(line.start, line.end); };
                break;
            default:
                BOOST_ASSERT(false);
                break;
        }

        switch (args.edge_gen.algorithm) {
            case edge_algorithm::planar:
                generate_planar_edges(edge_seed, args.edge_gen.planar.probability, edge_cost_fn, graph, positions);
                break;
            case edge_algorithm::uniform:
                generate_uniform_edges(edge_seed, args.edge_gen.uniform.probability, edge_cost_fn, graph, positions);
                break;
            case edge_algorithm::layered:
                generate_layered_edges(edge_seed,
                                       args.edge_gen.layered.probability,
                                       args.edge_gen.layered.count,
                                       edge_cost_fn,
                                       graph,
                                       positions);
                break;
            case edge_algorithm::kronecker:
                generate_kronecker_graph(edge_seed,
                                         args.edge_gen.kronecker.initiator,
                                         args.edge_gen.kronecker.k,
                                         edge_cost_fn,
                                         graph,
                                         positions);
                break;
            default:
                BOOST_ASSERT(false);
                break;
        }
    } else {
        std::ifstream in(args.graph_file);
        std::string line;
        std::vector<std::string> columns;
        std::set<std::tuple<std::string, std::string, double>> edges;
        std::unordered_map<std::string, size_t> nodes;
        size_t node_index = 0;
        while (std::getline(in, line)) {
            trim(line);
            if (line.empty() || starts_with(line, "#") || starts_with(line, "//") || starts_with(line, "--")) {
                continue;
            }

            columns.clear();
            split(columns, line, [](char c) { return std::isspace(static_cast<int>(c)); }, token_compress_on);
            if (columns.size() < 2) {
                if (err) {
                    (*err) << "Parse error in " << args.graph_file << "!\n";
                }
                return;
            }

            double cost = 1.0;
            if (columns.size() >= 3) {
                try {
                    cost = boost::lexical_cast<double>(columns[2]);
                } catch (const boost::bad_lexical_cast& ex) {
                    if (err) {
                        (*err) << "Parse error in " << args.graph_file << ": " << ex.what() << "\n";
                    }
                    return;
                }
            }

            edges.emplace(columns[0], columns[1], cost);
            if (nodes.find(columns[0]) == nodes.end()) {
                nodes[columns[0]] = node_index;
                node_index += 1;
                graph.add_node();
            }
            if (nodes.find(columns[1]) == nodes.end()) {
                nodes[columns[1]] = node_index;
                node_index += 1;
                graph.add_node();
            }
        }
        if (!in.eof()) {
            if (err) {
                (*err) << "Read error: " << args.graph_file << "\n";
            }
            return;
        }

        if (graph.node_count() == 0) {
            return; // nothing to do ...
        }

        for (const auto& edge : edges) {
            graph.add_edge(nodes[std::get<0>(edge)], nodes[std::get<1>(edge)], std::get<double>(edge));
        }
        positions = graph.make_node_map([](size_t) { return vec2(0.0, 0.0); });
    }

    if (graph.node_count() == 0) {
        (*err) << "The generated graph is empty!\n";
        return;
    }

    boost::base_collection<criteria> criteria;
    for (sssp_algorithm crit : args.algorithms) {
        switch (crit) {
            case sssp_algorithm::crauser_in:
                criteria.insert(crauser_in(&graph, start_node, false));
                break;
            case sssp_algorithm::crauser_in_dyn:
                criteria.insert(crauser_in(&graph, start_node, true));
                break;
            case sssp_algorithm::crauser_out:
                criteria.insert(crauser_out(&graph, start_node, false));
                break;
            case sssp_algorithm::crauser_out_dyn:
                criteria.insert(crauser_out(&graph, start_node, true));
                break;
            case sssp_algorithm::dijkstra:
                criteria.insert(smallest_tentative_distance(&graph, start_node));
                break;
            case sssp_algorithm::heuristic:
                criteria.insert(heuristic(
                    &graph, start_node, [&](size_t node) { return distance(positions[start_node], positions[node]); }));
                break;
            case sssp_algorithm::oracle:
                criteria.insert(oracle(&graph, start_node));
                break;
            case sssp_algorithm::traff:
                criteria.insert(traff_bridge(&graph, start_node));
                break;
            case sssp_algorithm::paper_in:
                criteria.insert(paper_in(&graph, start_node));
                break;
            case sssp_algorithm::paper_out:
                criteria.insert(paper_out(&graph, start_node));
                break;
            default:
                BOOST_ASSERT(false);
                break;
        }
    }

    node_map<dijkstra_result> result = dijkstra(graph, start_node, criteria);

    size_t reachable = 0;
    int max_phase = 0;
    std::vector<size_t> relaxed(graph.node_count(), 0);
    std::vector<size_t> fringe_size(graph.node_count(), 0);
    for (const auto& r : result) {
        if (r.settled()) {
            reachable += 1;
            max_phase = std::max(max_phase, r.relaxation_phase);
            relaxed[r.relaxation_phase] += 1;
            fringe_size[r.relaxation_phase] = std::max(fringe_size[r.relaxation_phase], r.fringe_size);
        }
    }

    if (out) {
        for (int phase = 0; phase <= max_phase; ++phase) {
            *out << arguments_csv_values(args) << "," << graph.node_count() << "," << phase << "," << relaxed[phase]
                 << "," << fringe_size[phase] << "\n";
        }
    }

#ifndef DISABLE_CAIRO
    if (args.image.size() > 0) {
        int layers = 1;
        std::vector<size_t> count_by_layer(1, graph.node_count());
        std::vector<size_t> placed_by_layer(1, 0);
        if (args.edge_gen.algorithm == edge_algorithm::layered) {
            layers = args.edge_gen.layered.count;
            count_by_layer = std::vector<size_t>(layers, 0);
            placed_by_layer = std::vector<size_t>(layers, 0);
            for (size_t node = 0; node < graph.node_count(); ++node) {
                count_by_layer[y_bucket(layers, positions[node].y)] += 1;
            }
        }

        node_map<node_style> node_styles = graph.make_node_map([&](size_t i) {
            node_style style;

            style.position = positions[i];
            if (result[i].relaxation_phase == -1) {
                style.color = rgb(1.0, 0.5, 0.5);
            } else {
                double c = std::min(1.0, std::pow(static_cast<double>(result[i].relaxation_phase) / 15, 1.0 / 2.2));
                style.color = rgb(c, c, 1.0);
                style.text = std::to_string(result[i].relaxation_phase);
            }

            // Override values if disabled
            if (!args.image_node_labels) {
                style.text = "";
            }
            if (!args.image_node_colors) {
                style.color = rgb(1.0, 1.0, 1.0);
            }
            if (args.edge_gen.algorithm == edge_algorithm::layered) {
                int layer = y_bucket(layers, positions[i].y);
                style.position.x = (1.0 / (count_by_layer[layer] + 1)) * (placed_by_layer[layer] + 1);
                style.position.y = (1.0 / (layers + 1)) * (layer + 1);
                placed_by_layer[layer] += 1;
            }

            return style;
        });

        edge_map<edge_style> edge_styles = graph.make_edge_map([&](size_t source, size_t destination) {
            edge_style style;

            for (const edge_info& edge : graph.outgoing_edges(source)) {
                if (edge.destination == destination) {
                    style.text = std::to_string(int(std::round(edge.cost * 1000))) + ">";
                    break;
                }
            }

            if (result[destination].predecessor == source) {
                style.line_width *= 2;
                style.color = rgb(0.25, 0.25, 1.0);
                style.foreground = true;
            }

            // Override values if disabled
            if (!args.image_edge_colors) {
                style.color = rgb(0.0, 0.0, 0.0);
                style.line_width = edge_style().line_width;
                style.foreground = edge_style().foreground;
            }
            if (!args.image_edge_labels) {
                style.text = "";
            }

            return style;
        });

        try {
            int width = 800;
            int height = 800;
            auto surface = Cairo::PdfSurface::create(args.image, width, height);
            auto cr = Cairo::Context::create(surface);
            cr->translate(25, 25);
            cr->scale(width - 50, height - 50);
            cr->save();
            cr->set_source_rgb(1.0, 1.0, 1.0);
            cr->paint();
            cr->restore();
            draw_graph(cr, graph, node_styles, edge_styles);
        } catch (const std::ios_base::failure& ex) {
            if (err) {
                *err << "Could not write file " << args.image << ": " << ex.what() << "\n";
            }
        }
    }
#endif
}


#include <boost/random/non_central_chi_squared_distribution.hpp>


// Copyright (C) 2019 The Regents of the University of California (Regents).
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above
//       copyright notice, this list of conditions and the following
//       disclaimer in the documentation and/or other materials provided
//       with the distribution.
//
//     * Neither the name of The Regents or University of California nor the
//       names of its contributors may be used to endorse or promote products
//       derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
// Please contact the author of this library if you have any questions.
// Author: Chris Sweeney (cmsweeney@cs.ucsb.edu)

// This file was created by Steffen Urban (urbste@googlemail.com) or
// company address (steffen.urban@zeiss.com)
// December 2018

#include <glog/logging.h>
#include <Eigen/Cholesky>
#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/Eigenvalues>
#include <Eigen/Geometry>
#include <Eigen/LU>
#include <Eigen/QR>
#include <Eigen/SVD>
#include <random>

#include "theia/sfm/pose/four_point_focal_length_radial_distortion.h"
#include "theia/sfm/pose/four_point_focal_length_radial_distortion_helper.h"
#include "theia/util/random.h"

namespace theia {

double sgn(double val) { return (0.0 < val) - (val < 0.0); }

using Matrix34d = Eigen::Matrix<double, 3, 4>;
using Matrix24d = Eigen::Matrix<double, 2, 4>;
using Matrix42d = Eigen::Matrix<double, 4, 2>;
using Vector8d  = Eigen::Matrix<double, 8, 1>;
using Vector5d  = Eigen::Matrix<double, 5, 1>;
using Eigen::Map;
using Eigen::Vector3d;
using Eigen::Vector4d;
using Eigen::Matrix4d;
using Eigen::Matrix3d;
using Eigen::MatrixXd;
using Eigen::Matrix;

bool FourPointsPoseFocalLengthRadialDistortion(
    const std::vector<Eigen::Vector2d>& feature_vectors,
    const std::vector<Eigen::Vector3d>& world_points,
    std::vector<Eigen::Matrix3d>* rotations,
    std::vector<Eigen::Vector3d>* translations,
    std::vector<double>* radial_distortions,
    std::vector<double>* focal_lengths) {
  // check that input size of features and world points is 4
  CHECK_GE(feature_vectors.size(), 4);
  CHECK_EQ(feature_vectors.size(), world_points.size());

  Vector4d d;
  Matrix34d world_points_;
  Matrix34d u;  // image points. Will be normalized.
  for (int i = 0; i < 4; ++i) {
    d[i] = feature_vectors[i].squaredNorm();
    world_points_.col(i) = world_points[i];
    u.col(i) = feature_vectors[i].homogeneous();
  }

  const Vector3d t0 = world_points_.rowwise().mean();
  Matrix34d t0_mat;
  t0_mat << t0, t0, t0, t0;

  Matrix4d U;
  U.topRows<3>() = world_points_ - t0_mat;
  U.bottomRows<1>() << 1.0, 1.0, 1.0, 1.0;

  Eigen::JacobiSVD<Eigen::MatrixXd> svd(
      U.topRows<3>(), Eigen::ComputeFullV | Eigen::ComputeFullU);
  Matrix3d R0 = svd.matrixU();

  if (sgn(R0.determinant()) < 0.0) {
    R0.col(0) *= -1;
  }
  R0.transposeInPlace();

  U.topRows<3>() = R0 * U.topRows<3>();
  const double scale =
      U.topRows<3>().array().pow(2).colwise().sum().sqrt().mean();

  U.topRows<3>() /= scale;

  // rescale image points
  const double f0 = u.topRows<2>().array().pow(2).colwise().sum().sqrt().mean();
  u.topRows<2>() /= f0;

  const double k0 = d.array().mean();
  d /= k0;

  Matrix<double, 5, 8> M;
  M.fill(0.0);
  M.row(0).leftCols<4>() = U.col(0);
  M.row(1).rightCols<4>() = U.col(0);
  for (int k = 1; k < 4; ++k) {
    M.row(k + 1).leftCols<4>() = u.row(1).col(k) * U.col(k).transpose();
    M.row(k + 1).rightCols<4>() = -u.row(0).col(k) * U.col(k).transpose();
  }

  Vector5d b;
  b.fill(0.0);
  b.topRows<2>() = u.col(0).topRows<2>();
  Eigen::HouseholderQR<Eigen::MatrixXd> qr;
  qr.compute(M.transpose());

  Matrix<double, 8, 8> Q = qr.householderQ();
  Matrix<double, 8, 5> R = qr.matrixQR().triangularView<Eigen::Upper>();

  Matrix<double, 8, 4> N;
  N.fill(0.0);
  N.leftCols<3>() = Q.rightCols<3>();

  // this random rotation is supposed to make the solver more stable
  static RandomNumberGenerator random_number_gen(42);
  Vector3d rot_vec(random_number_gen.RandDouble(-0.5, 0.5),
                   random_number_gen.RandDouble(-0.5, 0.5),
                   random_number_gen.RandDouble(-0.5, 0.5));
  Eigen::AngleAxisd random_rot(rot_vec.norm(), rot_vec);
  N.leftCols<3>() *= random_rot.toRotationMatrix();
  Matrix<double, 8, 1> x0 =
      Q.leftCols<5>() * (R.topRows<5>().transpose().fullPivLu().solve(b));
  N.rightCols<1>() = x0;

  Matrix<double, 6, 3> C;
  C.fill(0.0);
  Matrix34d UN1 = U.rightCols<3>().transpose() * N.topRows<4>();
  Matrix34d UN2 = U.rightCols<3>().transpose() * N.bottomRows<4>();
  Matrix<double, 6, 9> B;
  B.fill(0.0);

  B.topLeftCorner<3, 3>() = UN1.leftCols<3>();
  B.bottomLeftCorner<3, 3>() = UN2.leftCols<3>();
  B.topRightCorner<3, 1>() = UN1.rightCols<1>();
  B.bottomRightCorner<3, 1>() = UN2.rightCols<1>();

  B.block<3, 4>(0, 3) =
      d.bottomRows<3>().transpose().replicate(4, 1).transpose().cwiseProduct(
          UN1);
  B.block<3, 4>(3, 3) =
      d.bottomRows<3>().transpose().replicate(4, 1).transpose().cwiseProduct(
          UN2);

  B.col(7).topRows<3>() = -u.row(0).rightCols<3>().transpose().cwiseProduct(
      U.row(2).rightCols<3>().transpose());
  B.col(7).bottomRows<3>() = -u.row(1).rightCols<3>().transpose().cwiseProduct(
      U.row(2).rightCols<3>().transpose());

  // fill these guys
  Matrix3d Utmp;
  Utmp.row(0) = U.row(0).rightCols<3>();
  Utmp.row(1) = U.row(1).rightCols<3>();
  Utmp.row(2) = U.row(3).rightCols<3>();

  Matrix3d u1temp;
  u1temp.row(0) = u.row(0).rightCols<3>();
  u1temp.row(1) = u.row(0).rightCols<3>();
  u1temp.row(2) = u.row(0).rightCols<3>();
  Matrix3d u2temp;
  u2temp.row(0) = u.row(1).rightCols<3>();
  u2temp.row(1) = u.row(1).rightCols<3>();
  u2temp.row(2) = u.row(1).rightCols<3>();

  C.block<3, 3>(0, 0) = Utmp.transpose().cwiseProduct(u1temp.transpose());
  C.block<3, 3>(3, 0) = Utmp.transpose().cwiseProduct(u2temp.transpose());

  Matrix<double, 3, 9> D = C.colPivHouseholderQr().solve(B);
  Map<Eigen::RowVectorXd> N_(N.data(), N.size());
  Map<Eigen::RowVectorXd> D_(D.data(), D.size());

  Matrix<double, 64, 1> data;
  data(0) = 0.0;  // used to keep matlab indices, just an index offset of 1
  data.block<32, 1>(1, 0) = N_;
  data.block<27, 1>(33, 0) = D_;
  data(60, 0) = d(0, 0);
  data.bottomRows<3>() = U.topRows<3>().col(0);

  std::vector<Vector5d> valid_solutions;
  FourPointsPoseFocalLengthRadialDistortionSolver(data, &valid_solutions);

  rotations->resize(valid_solutions.size());
  translations->resize(valid_solutions.size());
  radial_distortions->resize(valid_solutions.size());
  focal_lengths->resize(valid_solutions.size());

  for (int i = 0; i < valid_solutions.size(); ++i) {
    const double k = valid_solutions[i][3];
    const double P33 = valid_solutions[i][4];
    Eigen::Vector4d alpha(valid_solutions[i][0], valid_solutions[i][1],
                          valid_solutions[i][2], 1.0);

    Vector8d P12_ = N * alpha;
    Map<Matrix42d> P12(P12_.data(), 4, 2);

    Matrix<double, 9, 1> tmp;
    tmp(0, 0) = alpha[0];
    tmp(1, 0) = alpha[1];
    tmp(2, 0) = alpha[2];
    tmp(3, 0) = k * alpha(0);
    tmp(4, 0) = k * alpha(1);
    tmp(5, 0) = k * alpha(2);
    tmp(6, 0) = k;
    tmp(7, 0) = P33;
    tmp(8, 0) = 1.0;

    const Vector3d P3_124 = D * tmp;

    Vector4d P3;
    P3(0) = P3_124(0);
    P3(1) = P3_124(1);
    P3(2) = P33;
    P3(3) = P3_124(2);

    Matrix34d P;
    P.topRows<2>() = P12.transpose();
    P.bottomRows<1>() = P3;
    P /= P.bottomRows<1>().leftCols<3>().norm();
    const double f = P.topRows<1>().leftCols<3>().norm();
    Matrix3d K = Matrix3d::Identity();
    K(0, 0) = 1. / f;
    K(1, 1) = 1. / f;

    (*focal_lengths)[i] = f * f0;
    (*radial_distortions)[i] = k / k0;

    Matrix34d Rt = K * P;

    if (Rt.topLeftCorner<3, 3>().determinant() < 0.0) Rt *= -1.0;

    // scale radial distortion
    // radial_distortions[i] *= (focal_lengths[i]*focal_lengths[i]);

    Rt.col(3) = Rt.col(3) * scale - Rt.topLeftCorner<3, 3>() * R0 * t0;
    Rt.topLeftCorner<3, 3>() = Rt.topLeftCorner<3, 3>() * R0;

    (*rotations)[i] = Rt.topLeftCorner<3, 3>();
    (*translations)[i] = Rt.col(3);
  }

  return valid_solutions.size() > 0;
}
}


#include "apriltag_ros/detector_offline.h"

#include <Eigen/Eigen>
#include <opencv2/opencv.hpp>
#include <sys/time.h>

using namespace std;

namespace apriltag_ros {
/*
 * Z-Y-X Eular angles
 */
template <typename T>
Eigen::Matrix<T, 3, 1> Quaternion2Euler(const Eigen::Quaternion<T>& q) {
    const double Epsilon = 0.0009765625f;
    const double threshold = 0.5 - Epsilon;
    Eigen::Matrix<T, 3, 1> eular(0, 0, 0);
    double singular = static_cast<double>(q.w() * q.y() - q.x() * q.z());
    if (singular < -threshold || singular > threshold) {
        if (singular < -0.00001) {
            eular(2) = static_cast<T>(2 * atan2(q.x(), q.w()));
            eular(1) = static_cast<T>(-M_PI / 2.0);
        } else if (singular > 0.000001) {
            eular(2) = static_cast<T>(-2 * atan2(q.x(), q.w()));
            eular(1) = static_cast<T>(M_PI / 2.0);
        }
        eular(0) = 0.0;
    } else {
        eular(0) = static_cast<T>(
            atan2(2 * (q.y() * q.z() + q.w() * q.x()), q.w() * q.w() - q.x() * q.x() - q.y() * q.y() + q.z() * q.z()));
        eular(1) = static_cast<T>(std::asin(-2 * (q.x() * q.z() - q.w() * q.y())));
        eular(2) = static_cast<T>(
            atan2(2 * (q.x() * q.y() + q.w() * q.z()), q.w() * q.w() + q.x() * q.x() - q.y() * q.y() - q.z() * q.z()));
    }
    return eular;
}

DetectorOffline::DetectorOffline(ros::NodeHandle nh, ros::NodeHandle nh_private) : nh_(nh), nh_private_(nh_private) {
    if (nh_private_.hasParam("video_file")) {
        ROS_INFO("nh_private_ has video_file parameter");
    }
    nh_private_.param<std::string>("video_file", video_file, "");
    nh_private_.param<std::string>("bag_file", bag_file, "");
    nh_private_.param<std::string>("output_file", output_file, "/home/guoziwei/Documents/outputpose.txt");
    if (video_file.empty()) {
        ROS_WARN("Video name is not passed\n");
    } else {
        ROS_INFO("Got param: %s", video_file.c_str());
        cap = cv::VideoCapture(video_file);
    }
    if (bag_file.empty()) {
        ROS_WARN("bagfile name is not passed\n");
    } else {
        ROS_INFO("Got param: %s", bag_file.c_str());
    }
    ROS_INFO("Got param: %s", output_file.c_str());
    pub_poses = nh.advertise<visualization_msgs::Marker>("position_xy", 1000);
    tag_detector_ = std::shared_ptr<TagDetector>(new TagDetector(nh_private_));
    draw_tag_detections_image_ = getAprilTagOption<bool>(nh_private_, "publish_tag_detections_image", false);
    it_ = std::shared_ptr<image_transport::ImageTransport>(new image_transport::ImageTransport(nh_));
    tag_detections_publisher_ = nh_.advertise<AprilTagDetectionArray>("tag_detections", 1);
    if (draw_tag_detections_image_) {
        tag_detections_image_publisher_ = it_->advertise("tag_detections_image", 1);
        // tag_detections_image_publisher_.publish(cv_image_->toImageMsg());
    }
    start_time = ros::Time::now();
    header.frame_id = "/camera_link";
    header.seq = 0;
}

int DetectorOffline::readImage() {
    if (!cap.isOpened()) {
        std::cerr << "Error opening video file" << std::endl;
        return -1;
    }
    int frame_width = cap.get(CV_CAP_PROP_FRAME_WIDTH);
    int frame_height = cap.get(CV_CAP_PROP_FRAME_HEIGHT);
    int frames = cap.get(CV_CAP_PROP_FRAME_COUNT);
    ROS_INFO("frames:%d %d %d", frame_width, frame_height, frames);
    // struct timeval start, end;
    Frame frame;
    start_time = ros::Time(0.0);
    header.stamp = start_time;

    visualization_msgs::Marker poses;
    poses.header = header;
    poses.ns = "tag_position";
    poses.type = visualization_msgs::Marker::SPHERE_LIST;
    poses.action = visualization_msgs::Marker::ADD;
    poses.pose.orientation.w = 1.0;
    poses.lifetime = ros::Duration();
    poses.id = 0;
    poses.scale.x = 0.05;
    poses.scale.y = 0.05;
    poses.scale.z = 0.05;
    poses.color.r = 1.0;
    poses.color.a = 1.0;
    FILE* outfile;
    outfile = std::fopen(output_file.c_str(), "w");
    Camera_Info camera_;
    camera_.CreatefromYaml("/home/guoziwei/catkin_ws/src/apriltag_ros/apriltag_ros/config/camera_info.yaml");
    cv::Mat distortion_image;
    cv::Mat camera_matrix =
        (cv::Mat_<double>(3, 3) << camera_.fx, 0.0, camera_.cx, 0.0, camera_.fy, camera_.cy, 0.0, 0.0, 1.0);
    cv::Mat distcoeffs = (cv::Mat_<double>(5, 1) << camera_.k1, camera_.k2, 0.0, 0.0, camera_.k3);
    std::vector<cv::Point2d> imagePoints;
    while (1) {
        cap >> distortion_image;
        if (distortion_image.empty())
            break;
        cv::undistort(distortion_image, frame.frame, camera_matrix, distcoeffs);
        cv::cvtColor(frame.frame, frame.frame, CV_BGR2GRAY);
        // cv::imshow("Undistortion", frame.frame);
        // cv::waitKey(3);
        // std::cout<<"get frame"<<std::endl;
        frame.led_on = 0;
        frame.time_stamp = cap.get(CV_CAP_PROP_POS_MSEC) / 1e3;
        // gettimeofday(&start, nullptr);

        try {
            header.stamp = start_time + ros::Duration(frame.time_stamp);
            cv_image_ = boost::make_shared<cv_bridge::CvImage>(header, "bgr8", frame.frame);
        } catch (cv_bridge::Exception& e) {
            ROS_ERROR("cv_bridge exception: %s", e.what());
            return -2;
        }
        AprilTagDetectionArray tag_pose_array = tag_detector_->detectTags(cv_image_, imagePoints);
        // detect led
        // cv::Rect roi(imagePoints[2].x - 60, imagePoints[2].y - 60, 120, 120);
        // cv::Mat crroped = distortion_image(roi);
        // cv::Mat1b mask1, mask2;
        // cv::cvtColor(crroped, crroped, cv::COLOR_BGR2HSV);
        // cv::inRange(crroped, cv::Scalar(0, 70, 50), cv::Scalar(10, 255, 255), mask1);
        // cv::inRange(crroped, cv::Scalar(170, 70, 50), cv::Scalar(180, 250, 255), mask2);
        // cv::Mat1b mask = mask1 | mask2;
        // cv::Mat_<uchar>::iterator it = mask.begin();
        // cv::Mat_<uchar>::iterator it_end = mask.end();
        // for (; it != it_end; it++) {
        //     if (*it == 255)
        //         frame.led_on = 1;
        // }

        // cv::imshow("Red Detections", crroped);
        // cv::waitKey(3);
        for (unsigned int i = 0; i < tag_pose_array.detections.size(); i++) {
            geometry_msgs::PoseStamped pose;
            geometry_msgs::Point position_marker;
            pose.pose = tag_pose_array.detections[i].pose.pose.pose;
            pose.header = tag_pose_array.detections[i].pose.header;
            poses.header = tag_pose_array.detections[i].pose.header;
            poses.header.stamp = ros::Time::now();
            position_marker.x = pose.pose.position.x;
            position_marker.y = pose.pose.position.y;
            position_marker.z = pose.pose.position.z;
            poses.points.push_back(position_marker);
            std::fprintf(
                outfile, "%lf %lf %lf %lf %lf %lf %lf %lf %d\n", pose.header.stamp.toSec(), pose.pose.position.x,
                pose.pose.position.y, pose.pose.position.z, pose.pose.orientation.x, pose.pose.orientation.y,
                pose.pose.orientation.z, pose.pose.orientation.w, frame.led_on);
        }
        pub_poses.publish(poses);
        // ROS_INFO("seq:%d",header.seq);
        if (draw_tag_detections_image_) {
            tag_detector_->drawDetections(cv_image_);
            tag_detections_image_publisher_.publish(cv_image_->toImageMsg());
        }
        header.seq += 1;
        // gettimeofday(&end, nullptr);
        // double elapsed_seconds = (end.tv_sec - start.tv_sec) * 1e6;
        // elapsed_seconds += (end.tv_usec - start.tv_usec) * 1e-6;
        // std::cout<<"elapsed time: "<<std::setprecision(8)<<elapsed_seconds<<endl;
    }
    cap.release();
    std::fclose(outfile);
    return 0;
}

int DetectorOffline::readCompressedImageFromBag() {
    FILE* outfile;
    outfile = std::fopen(output_file.c_str(), "w");
    std::vector<cv::Point2d> imagePoints;
    bag.open(bag_file);
    std::vector<std::string> topics;
    topics.push_back(std::string("/ninebot/fisheye_left/compressed"));
    rosbag::View view(bag, rosbag::TopicQuery(topics));
    for (rosbag::MessageInstance const compress_image : view) {
        sensor_msgs::CompressedImageConstPtr i = compress_image.instantiate<sensor_msgs::CompressedImage>();
        header = i->header;
        cv_bridge::CvImagePtr compressed_ptr = cv_bridge::toCvCopy(i, sensor_msgs::image_encodings::RGB8);
        try {
            cv_image_ = boost::make_shared<cv_bridge::CvImage>(header, "rgb8", compressed_ptr->image);
        } catch (cv_bridge::Exception& e) {
            ROS_ERROR("cv_bridge exception: %s", e.what());
            return -2;
        }
        AprilTagDetectionArray tag_pose_array = tag_detector_->detectTags(cv_image_, imagePoints);
        for (unsigned int i = 0; i < tag_pose_array.detections.size(); i++) {
            geometry_msgs::PoseStamped pose;
            geometry_msgs::Point position_marker;
            pose.pose = tag_pose_array.detections[i].pose.pose.pose;
            pose.header = tag_pose_array.detections[i].pose.header;
            // poses.header = tag_pose_array.detections[i].pose.header;
            // poses.header.stamp = ros::Time::now();
            position_marker.x = pose.pose.position.x;
            position_marker.y = pose.pose.position.y;
            position_marker.z = pose.pose.position.z;
            // poses.points.push_back(position_marker);
            Eigen::Vector3d position(pose.pose.position.x, pose.pose.position.y, pose.pose.position.z);
            Eigen::Quaterniond attitude(
                pose.pose.orientation.w, pose.pose.orientation.x, pose.pose.orientation.y, pose.pose.orientation.z);
            attitude = attitude.conjugate();
            position = attitude * position;
            const auto eular_vec = Quaternion2Euler(attitude);
            std::fprintf(
                outfile, "%lf %lf %lf %lf %lf %lf %lf %lf %lf %lf %lf\n", pose.header.stamp.toSec(), position[0],
                position[1], position[2], attitude.w(), attitude.x(), attitude.y(), attitude.z(), eular_vec[0], eular_vec[1],
                eular_vec[2]);
        }
        if (draw_tag_detections_image_) {
            tag_detector_->drawDetections(cv_image_);
            tag_detections_image_publisher_.publish(cv_image_->toImageMsg());
        }
        // cv::imshow("image", compressed_ptr->image);
        // cv::waitKey(1);
    }
    std::fclose(outfile);
    return 0;
}
}  // namespace apriltag_ros


/* @file
 * @copyright This code is licensed under the 3-clause BSD license.
 *   Copyright ETH Zurich, Laboratory of Physical Chemistry, Reiher Group.
 *   See LICENSE.txt for details.
 */

#include <boost/test/unit_test.hpp>

#include "Molassembler/Shapes/CoordinateSystemTransformation.h"
#include "Molassembler/Shapes/InertialMoments.h"
#include "Molassembler/Shapes/ContinuousMeasures.h"
#include "Molassembler/Shapes/Shapes.h"
#include "Molassembler/Shapes/Data.h"

using namespace Scine::Molassembler;
using namespace Shapes;

// From ContinuousMeasures.cpp
extern Continuous::PositionCollection addOrigin(const Continuous::PositionCollection& vs);

void randomlyRotate(Eigen::Ref<Continuous::PositionCollection> vs) {
  vs = rotationMatrix(CoordinateSystem {}, CoordinateSystem::random()) * vs;
}

const std::vector<std::string>& topNames() {
  static const std::vector<std::string> strings {
    "Line",
    "Asymmetric",
    "Prolate",
    "Oblate",
    "Spherical"
  };

  return strings;
}

const std::string& topName(Top top) {
  return topNames().at(
    static_cast<std::underlying_type<Top>::type>(top)
  );
}

BOOST_AUTO_TEST_CASE(InertialStandardization, *boost::unit_test::label("Shapes")) {
  const std::vector<
    std::pair<Shape, Top>
  > tops {
    {Shape::Line, Top::Line},
    {Shape::Bent, Top::Asymmetric},
    {Shape::Seesaw, Top::Asymmetric},
    {Shape::TrigonalBipyramid, Top::Prolate},
    {Shape::PentagonalBipyramid, Top::Oblate},
    {Shape::Square, Top::Oblate},
    {Shape::Octahedron, Top::Spherical},
    {Shape::Tetrahedron, Top::Spherical}
  };

  for(const auto& nameTopPair : tops) {
    auto positions = addOrigin(coordinates(nameTopPair.first));

    // Apply a random coordinate transformation

    // Analyze it
    auto normalizedPositions = Continuous::normalize(positions);
    Top top = standardizeTop(normalizedPositions);
    BOOST_CHECK_MESSAGE(
      top == nameTopPair.second,
      "Top standardization failed. Expected "
      << topName(nameTopPair.second)
      << " for shape " << Shapes::name(nameTopPair.first)
      << ", got " << topName(top) << " instead."
    );
  }
}


#include <vector>
#include <algorithm>
#include <iostream>
#include <iterator>
#include <boost/property_map.hpp>
#include <boost/graph/adjacency_list.hpp>

#include "../../src/edmonds_optimum_branching.hpp"

// Define a directed graph type that associates a weight with each
// edge. We store the weights using internal properties as described
// in BGL.
typedef boost::property<boost::edge_weight_t, double>       EdgeProperty;
typedef boost::adjacency_list<boost::listS,
                              boost::vecS,
                              boost::directedS,
                              boost::no_property,
                              EdgeProperty>                 Graph;
typedef boost::graph_traits<Graph>::vertex_descriptor       Vertex;
typedef boost::graph_traits<Graph>::edge_descriptor         Edge;

int
main(int argc, char *argv[])
{
    const int N = 4;

    // Graph with N vertices
    Graph G(N);

    // Create a vector to keep track of all the vertices and enable us
    // to index them. As a side note, observe that this is not
    // necessary since Vertex is probably an integral type. However,
    // this may not be true of arbitrary graphs and I think this code
    // is a better illustration of a more general case.
    std::vector<Vertex> the_vertices;
    BOOST_FOREACH (Vertex v, vertices(G))
    {
        the_vertices.push_back(v);
    }

    // add a few edges with weights to the graph
    add_edge(the_vertices[0], the_vertices[1], 3.0, G);
    add_edge(the_vertices[0], the_vertices[2], 1.5, G);
    add_edge(the_vertices[0], the_vertices[3], 1.8, G);
    add_edge(the_vertices[1], the_vertices[2], 4.3, G);
    add_edge(the_vertices[2], the_vertices[3], 2.2, G);

    // This is how we can get a property map that gives the weights of
    // the edges.
    boost::property_map<Graph, boost::edge_weight_t>::type weights =
        get(boost::edge_weight_t(), G);

    // This is how we can get a property map mapping the vertices to
    // integer indices.
    boost::property_map<Graph, boost::vertex_index_t>::type vertex_indices =
        get(boost::vertex_index_t(), G);

    // Print the graph (or rather the edges of the graph).
    std::cout << "This is the graph:\n";
    BOOST_FOREACH (Edge e, edges(G))
    {
        std::cout << "(" << boost::source(e, G) << ", "
                  << boost::target(e, G) << ")\t"
                  << get(weights, e) << "\n";
    }

    // Find the maximum branching.
    std::vector<Edge> branching;
    edmonds_optimum_branching<true, false, false>(G,
                                                  vertex_indices,
                                                  weights,
                                                  static_cast<Vertex *>(0),
                                                  static_cast<Vertex *>(0),
                                                  std::back_inserter(branching));

    // Print the edges of the maximum branching
    std::cout << "This is the maximum branching\n";
    BOOST_FOREACH (Edge e, branching)
    {
        std::cout << "(" << boost::source(e, G) << ", "
                  << boost::target(e, G) << ")\t"
                  << get(weights, e) << "\n";
    }
    return 0;
}


///////////////////////////////////////////////////////////////////////////////
// BSD 3-Clause License
//
// Copyright (C) 2018-2020, LAAS-CNRS, New York University, Max Planck Gesellschaft,
//                          University of Edinburgh
// Copyright note valid unless otherwise stated in individual files.
// All rights reserved.
///////////////////////////////////////////////////////////////////////////////

#ifndef CROCODDYL_CORE_NUMDIFF_STATE_HPP_
#define CROCODDYL_CORE_NUMDIFF_STATE_HPP_

#include <boost/make_shared.hpp>
#include <boost/shared_ptr.hpp>
#include "crocoddyl/core/state-base.hpp"

namespace crocoddyl {

class StateNumDiff : public StateAbstract {
 public:
  explicit StateNumDiff(boost::shared_ptr<StateAbstract> state);
  ~StateNumDiff();

  Eigen::VectorXd zero() const;
  Eigen::VectorXd rand() const;
  void diff(const Eigen::Ref<const Eigen::VectorXd>& x0, const Eigen::Ref<const Eigen::VectorXd>& x1,
            Eigen::Ref<Eigen::VectorXd> dxout) const;
  void integrate(const Eigen::Ref<const Eigen::VectorXd>& x, const Eigen::Ref<const Eigen::VectorXd>& dx,
                 Eigen::Ref<Eigen::VectorXd> xout) const;
  /**
   * @brief This computes the Jacobian of the diff method by finite
   * differentiation:
   * \f{equation}{
   *    Jfirst[:,k] = diff(int(x_1, dx_dist), x_2) - diff(x_1, x_2)/disturbance
   * \f}
   * and
   * \f{equation}{
   *    Jsecond[:,k] = diff(x_1, int(x_2, dx_dist)) - diff(x_1, x_2)/disturbance
   * \f}
   *
   * @param Jfirst
   * @param Jsecond
   * @param firstsecond
   */
  void Jdiff(const Eigen::Ref<const Eigen::VectorXd>& x0, const Eigen::Ref<const Eigen::VectorXd>& x1,
             Eigen::Ref<Eigen::MatrixXd> Jfirst, Eigen::Ref<Eigen::MatrixXd> Jsecond,
             Jcomponent firstsecond = both) const;
  /**
   * @brief This computes the Jacobian of the integrate method by finite
   * differentiation:
   * \f{equation}{
   *    Jfirst[:,k] = diff( int(x, d_x), int( int(x, dx_dist), dx) )/disturbance
   * \f}
   * and
   * \f{equation}{
   *    Jsecond[:,k] = diff( int(x, d_x), int( x, dx + dx_dist) )/disturbance
   * \f}
   *
   * @param Jfirst
   * @param Jsecond
   * @param firstsecond
   */
  void Jintegrate(const Eigen::Ref<const Eigen::VectorXd>& x, const Eigen::Ref<const Eigen::VectorXd>& dx,
                  Eigen::Ref<Eigen::MatrixXd> Jfirst, Eigen::Ref<Eigen::MatrixXd> Jsecond,
                  Jcomponent firstsecond = both) const;
  const double& get_disturbance() const;
  void set_disturbance(const double& disturbance);

 private:
  /**
   * @brief This is the state we need to compute the numerical differentiation
   * from.
   */
  boost::shared_ptr<StateAbstract> state_;
  /**
   * @brief This the increment used in the finite differentiation and integration.
   */
  double disturbance_;
};

}  // namespace crocoddyl

#endif  // CROCODDYL_CORE_NUMDIFF_STATE_HPP_