#include "cpu/image_proc.h"

#include <Eigen/Dense>
#include <map>
#include <string>
#include <cmath>
#include <stdexcept>

namespace helper {

bool in_bounds(Eigen::Vector2f p, int h, int w) {
	return p.x() >= 0.0 && p.x() < static_cast<float>(w) && p.y() >= 0.0 && p.y() < static_cast<float>(h);
}

bool valid_flow_at(const Eigen::Vector2f& p, int h, int w, py::array_t<float>& flow_image, Eigen::Vector2f& flow) {
	if (!in_bounds(p, h, w)) {
		return false;
	}

	flow.x() = *flow_image.data(p.y(), p.x(), 0);
	flow.y() = *flow_image.data(p.y(), p.x(), 1);

	if (!flow.allFinite()) {
		return false;
	}

	return true;
}
} // namespace helper

namespace image_proc {

using Vec2f = Eigen::Vector2f;

template<class V, class L = std::less<std::string>,
		class A = Eigen::aligned_allocator<std::pair<const std::string, V>>>
using aligned_dict = std::map<std::string, V, L, A>;

py::array_t<float> compute_augmented_flow_from_rotation(py::array_t<float>& flow_image_rot_sa2so,
                                                        py::array_t<float>& flow_image_so2to,
                                                        py::array_t<float>& flow_image_rot_to2ta,
                                                        const int height, const int width) {
	// TODO: change to runtime asserts
	// assert(flow_image_rot_sa2so.ndim() == 3);
	// assert(flow_image_rot_sa2so.shape(0) == 2);
	// assert(flow_image_rot_sa2so.shape(1) == height);
	// assert(flow_image_rot_sa2so.shape(2) == width);

	// assert(flow_image_so2to.ndim() == 3);
	// assert(flow_image_so2to.shape(0) == 2);
	// assert(flow_image_so2to.shape(1) == height);
	// assert(flow_image_so2to.shape(2) == width);

	// assert(flow_image_rot_to2ta.ndim() == 3);
	// assert(flow_image_rot_to2ta.shape(0) == 2);
	// assert(flow_image_rot_to2ta.shape(1) == height);
	// assert(flow_image_rot_to2ta.shape(2) == width);

	// allocate memory for output array
	py::array_t<float> flow_image_rot_sa2ta = py::array_t<float>(flow_image_rot_sa2so.request().size);

	// reshape array to match input shape
	flow_image_rot_sa2ta.resize({height, width, 2});

	for (int y = 0; y < height; y++) {
		for (int x = 0; x < width; x++) {

			// update output flow image
			*flow_image_rot_sa2ta.mutable_data(y, x, 0) = -std::numeric_limits<float>::infinity();
			*flow_image_rot_sa2ta.mutable_data(y, x, 1) = -std::numeric_limits<float>::infinity();

			Vec2f p_sa(x, y);

			/////////////////////////////////////////////////////////////////////////////////
			// 1. SOURCE AUGMENTED TO SOURCE ORIGINAL
			/////////////////////////////////////////////////////////////////////////////////

			// flow from source augmented to source original
			Vec2f flow_sa2so(*flow_image_rot_sa2so.data(y, x, 0), *flow_image_rot_sa2so.data(y, x, 1));

			// flow_sa2so should be dense and w/o any invalid value
			if (!flow_sa2so.allFinite()) {
				throw std::runtime_error("flow_sa2so should be dense and w/o any invalid residuals!");
			}

			// compute warped location on source original (so we're going from source augmented to source original)
			Vec2f p_so = p_sa + flow_sa2so;

			// init flow_sa2ta with the first contribution, i.e, flow_sa2so
			Vec2f flow_sa2ta = flow_sa2so;

			/////////////////////////////////////////////////////////////////////////////////
			// 2. SOURCE ORIGINAL TO TARGET ORIGINAL
			/////////////////////////////////////////////////////////////////////////////////
			int u0 = std::floor(p_so.x());
			int u1 = u0 + 1;
			int v0 = std::floor(p_so.y());
			int v1 = v0 + 1;

			Vec2f p00(u0, v0);
			Vec2f p01(u0, v1);
			Vec2f p10(u1, v0);
			Vec2f p11(u1, v1);

			aligned_dict<Vec2f> valid_coords;
			aligned_dict<Vec2f> valid_flows;

			Vec2f flow_00_so2to;
			if (helper::valid_flow_at(p00, height, width, flow_image_so2to, flow_00_so2to)) {
				valid_coords["p00"] = p00;
				valid_flows["p00"] = flow_00_so2to;
			}

			Vec2f flow_01_so2to;
			if (helper::valid_flow_at(p01, height, width, flow_image_so2to, flow_01_so2to)) {
				valid_coords["p01"] = p01;
				valid_flows["p01"] = flow_01_so2to;
			}

			Vec2f flow_10_so2to;
			if (helper::valid_flow_at(p10, height, width, flow_image_so2to, flow_10_so2to)) {
				valid_coords["p10"] = p10;
				valid_flows["p10"] = flow_10_so2to;
			}

			Vec2f flow_11_so2to;
			if (helper::valid_flow_at(p11, height, width, flow_image_so2to, flow_11_so2to)) {
				valid_coords["p11"] = p11;
				valid_flows["p11"] = flow_11_so2to;
			}

			// Depending on how many valid flows we have, do bilinear interpolation or nearest neighbor:
			Vec2f flow_so2to;

			if (valid_coords.empty()) {
				continue;
			} else if (valid_coords.size() == 4) {
				// Bilinear interpolation
				float du = p_so.x() - static_cast<float>(u0);
				float dv = p_so.y() - static_cast<float>(v0);

				float w00 = (1 - du) * (1 - dv);
				float w01 = (1 - du) * dv;
				float w10 = du * (1 - dv);
				float w11 = du * dv;

				flow_so2to = w00 * valid_flows["p00"] +
				             w01 * valid_flows["p01"] +
				             w10 * valid_flows["p10"] +
				             w11 * valid_flows["p11"];
			} else {
				// Nearest Neighbor
				std::string nn = "None";
				float min_dist = std::numeric_limits<float>::max();

				for (const auto& valid_coord : valid_coords) {
					const std::string k = valid_coord.first;
					const Vec2f& p = valid_coord.second;

					float dist = (p_so - p).norm();
					if (dist < min_dist) {
						min_dist = dist;
						nn = k;
					}
				}

				if (nn == "None") {
					throw std::runtime_error("Neighrest Neighbor 'nn' was not assigned...");
				}

				flow_so2to = valid_flows[nn];
			}

			// compute warped location on target original (so we're going from source original to target original)
			Vec2f p_to = p_so + flow_so2to;

			// add flow_so2to to flow_sa2ta
			flow_sa2ta += flow_so2to;

			/////////////////////////////////////////////////////////////////////////////////
			// 3. TARGET ORIGINAL TO TARGET AUGMENTED
			/////////////////////////////////////////////////////////////////////////////////
			u0 = std::floor(p_to.x());
			u1 = u0 + 1;
			v0 = std::floor(p_to.y());
			v1 = v0 + 1;

			p00 = Vec2f(u0, v0);
			p01 = Vec2f(u0, v1);
			p10 = Vec2f(u1, v0);
			p11 = Vec2f(u1, v1);

			valid_coords.clear();
			valid_flows.clear();

			Vec2f flow_00_to2ta;
			if (helper::valid_flow_at(p00, height, width, flow_image_rot_to2ta, flow_00_to2ta)) {
				valid_coords["p00"] = p00;
				valid_flows["p00"] = flow_00_to2ta;
			}

			Vec2f flow_01_to2ta;
			if (helper::valid_flow_at(p01, height, width, flow_image_rot_to2ta, flow_01_to2ta)) {
				valid_coords["p01"] = p01;
				valid_flows["p01"] = flow_01_to2ta;
			}

			Vec2f flow_10_to2ta;
			if (helper::valid_flow_at(p10, height, width, flow_image_rot_to2ta, flow_10_to2ta)) {
				valid_coords["p10"] = p10;
				valid_flows["p10"] = flow_10_to2ta;
			}

			Vec2f flow_11_to2ta;
			if (helper::valid_flow_at(p11, height, width, flow_image_rot_to2ta, flow_11_to2ta)) {
				valid_coords["p11"] = p11;
				valid_flows["p11"] = flow_11_to2ta;
			}

			// Depending on how many valid flows we have, do bilinear interpolation or nearest neighbor:
			Vec2f flow_to2ta;

			if (valid_coords.empty()) {
				continue;
			} else if (valid_coords.size() == 4) {
				// Bilinear interpolation
				float du = p_to.x() - static_cast<float>(u0);
				float dv = p_to.y() - static_cast<float>(v0);

				float w00 = (1 - du) * (1 - dv);
				float w01 = (1 - du) * dv;
				float w10 = du * (1 - dv);
				float w11 = du * dv;

				flow_to2ta = w00 * valid_flows["p00"] +
				             w01 * valid_flows["p01"] +
				             w10 * valid_flows["p10"] +
				             w11 * valid_flows["p11"];
			} else {
				// Nearest Neighbor
				std::string nn = "None";
				float min_dist = std::numeric_limits<float>::max();

				for (const auto& valid_coord : valid_coords) {
					const std::string k = valid_coord.first;
					const Vec2f& p = valid_coord.second;

					float dist = (p_to - p).norm();
					if (dist < min_dist) {
						min_dist = dist;
						nn = k;
					}
				}

				if (nn == "None") {
					throw std::runtime_error("Neighrest Neighbor 'nn' was not assigned...");
				}

				flow_to2ta = valid_flows[nn];
			}

			// add flow_to2ta to flow_sa2ta
			flow_sa2ta += flow_to2ta;

			// update output flow image
			*flow_image_rot_sa2ta.mutable_data(y, x, 0) = flow_sa2ta.x();
			*flow_image_rot_sa2ta.mutable_data(y, x, 1) = flow_sa2ta.y();
		}
	}

	return flow_image_rot_sa2ta;
}

static inline bool has_immediate_true_z_neighbors(const int& index, const int& z, const int& z_bound, const py::array_t<bool>& array) {
	auto z_in_range = [&z_bound](int neighbor_z) { return neighbor_z > 0 && neighbor_z < z_bound; };
	auto value_if_in_range_or_zero = [&z, &index, &z_in_range, &array](int delta_z) {
		return z_in_range(z + delta_z) ? *array.data(index, z + delta_z) : 0;
	};
	return value_if_in_range_or_zero(-1) || value_if_in_range_or_zero(0) || value_if_in_range_or_zero(1);
}

int count_tp1(py::array_t<bool>& p, py::array_t<bool>& gt) {
	assert(p.ndim() == 2);
	assert(gt.ndim() == 2);
	const int n_batch = p.shape(0);
	const int dimz = p.shape(1);

	auto& ptr = p;
	int counter = 0;
	for (int i = 0; i < n_batch; i++) {
		for (int z = 0; z < dimz; z++) {
			if (*gt.data(i, z)) {
				counter += has_immediate_true_z_neighbors(i, z, dimz, p);
			}
		}
	}

	return counter;
}

int count_tp2(py::array_t<bool>& p, py::array_t<bool>& gt) {
	assert(p.ndim() == 4);
	assert(gt.ndim() == 4);
	const int n_batch = p.shape(0);
	assert(p.shape(1) == 1);
	const int height = p.shape(2);
	const int width = p.shape(3);

	auto& ptr = p;
	int counter = 0;
	for (int i = 0; i < n_batch; i++)
		for (int y = 0; y < height; y++)
			for (int x = 0; x < width; x++) {
				if (*gt.data(i, 0, y, x)) {
					counter += CHECK2();
				}
			}
	return counter;
}

int count_tp3(py::array_t<bool>& p, py::array_t<bool>& gt) {
	assert(p.ndim() == 5);
	assert(gt.ndim() == 5);
	const int n_batch = p.shape(0);
	assert(p.shape(1) == 1);
	const int dimz = p.shape(2);
	const int dimy = p.shape(3);
	const int dimx = p.shape(4);

	auto& ptr = p;
	int counter = 0;
	for (int i = 0; i < n_batch; i++)
		for (int z = 0; z < dimz; z++)
			for (int y = 0; y < dimy; y++)
				for (int x = 0; x < dimx; x++) {
					if (*gt.data(i, 0, z, y, x)) {
						counter += CHECK3();
						//printf("i %d x %d y %d z %d in %d\n", i, x, y, z, res);
					}
				}
	return counter;
}

void extend3(py::array_t<bool>& in, py::array_t<bool>& out) {
	assert(in.ndim() == 5);
	assert(out.ndim() == 5);
	int n_batch = in.shape(0);
	assert(in.shape(1) == 1);
	int dimz = in.shape(2);
	int dimy = in.shape(3);
	int dimx = in.shape(4);

	auto& ptr = in;
	for (int i = 0; i < n_batch; i++) {
		for (int z = 1; z < dimz - 1; z++) {
			for (int y = 1; y < dimy - 1; y++) {
				for (int x = 1; x < dimx - 1; x++) {
					*out.mutable_data(i, 0, z, y, x) = CHECK3();
				}
			}
		}
	}
}

void backproject_depth_ushort(py::array_t<unsigned short>& image_in, py::array_t<float>& point_image_out,
                              const float fx, const float fy, const float cx, const float cy, const float normalizer) {
	assert(image_in.ndim() == 2);
	assert(point_image_out.ndim() == 3);

	const int width = image_in.shape(1);
	const int height = image_in.shape(0);

	assert(point_image_out.shape(0) == height);
	assert(point_image_out.shape(1) == width);
	assert(point_image_out.shape(2) == 3);

#pragma omp parallel for default(none) shared(image_in, point_image_out) firstprivate(height, width, fx, fy, cx, cy, normalizer)
	for (int y = 0; y < height; y++) {
		for (int x = 0; x < width; x++) {
			float depth = float(*image_in.data(y, x)) / normalizer;

			if (depth > 0) {
				float pos_x = depth * (static_cast<float>(x) - cx) / fx;
				float pos_y = depth * (static_cast<float>(y) - cy) / fy;
				float pos_z = depth;

				*point_image_out.mutable_data(y, x, 0) = pos_x;
				*point_image_out.mutable_data(y, x, 1) = pos_y;
				*point_image_out.mutable_data(y, x, 2) = pos_z;
			}
		}
	}
}

py::array_t<float> backproject_depth_ushort(py::array_t<unsigned short>& image_in, float fx, float fy, float cx, float cy, float normalizer) {
	py::array_t<float> point_image_out({image_in.shape(0), image_in.shape(1), static_cast<ssize_t>(3)});
	memset(point_image_out.mutable_data(0, 0, 0), 0, point_image_out.size() * sizeof(float));
	backproject_depth_ushort(image_in, point_image_out, fx, fy, cx, cy, normalizer);
	return point_image_out;
}

void backproject_depth_float(py::array_t<float>& image_in, py::array_t<float>& point_image_out,
                             float fx, float fy, float cx, float cy) {
	assert(image_in.ndim() == 2);
	assert(point_image_out.ndim() == 3);

	int width = image_in.shape(1);
	int height = image_in.shape(0);
	assert(point_image_out.shape(0) == 3);
	assert(point_image_out.shape(1) == height);
	assert(point_image_out.shape(2) == width);

#pragma omp parallel for default(none) shared(image_in, point_image_out) firstprivate(height, width, fx, fy, cx, cy)
	for (int y = 0; y < height; y++) {
		for (int x = 0; x < width; x++) {
			float depth = *image_in.data(y, x);

			if (depth > 0) {
				float pos_x = depth * (static_cast<float>(x) - cx) / fx;
				float pos_y = depth * (static_cast<float>(y) - cy) / fy;
				float pos_z = depth;

				*point_image_out.mutable_data(y, x, 0) = pos_x;
				*point_image_out.mutable_data(y, x, 1) = pos_y;
				*point_image_out.mutable_data(y, x, 2) = pos_z;
			}
		}
	}
}

void compute_mesh_from_depth(const py::array_t<float>& point_image_in, float max_triangle_edge_distance,
                             py::array_t<float>& vertex_positions_out, py::array_t<int>& vertex_pixels_out,
                             py::array_t<int>& face_indices_out) {
	assert(point_image_in.ndim() == 3);
	assert(point_image_in.shape(2) == 3);

	int width = static_cast<int>(point_image_in.shape(1));
	int height = static_cast<int>(point_image_in.shape(0));

	// Compute valid pixel vertices and faces.
	// We also need to compute the pixel -> vertex index mapping for
	// computation of faces.
	// We connect neighboring pixels on the square into two triangles.
	// We only select valid triangles, i.e. with all valid vertices and
	// not too far apart.
	// Important: The triangle orientation is set such that the normals
	// point towards the camera.
	std::vector<Eigen::Vector3f> vertices;
	std::vector<Eigen::Vector3i> faces;
	std::vector<Eigen::Vector2i> pixels;

	int vertexIdx = 0;
	std::vector<int> mapPixelToVertexIdx(width * height, -1);

	for (int y = 0; y < height - 1; y++) {
		for (int x = 0; x < width - 1; x++) {
			Eigen::Vector3f obs00(*point_image_in.data(y + 0, x + 0, 0), *point_image_in.data(y + 0, x + 0, 1), *point_image_in.data(y + 0, x + 0, 2));
			Eigen::Vector3f obs01(*point_image_in.data(y + 1, x + 0, 0), *point_image_in.data(y + 1, x + 0, 1), *point_image_in.data(y + 1, x + 0, 2));
			Eigen::Vector3f obs10(*point_image_in.data(y + 0, x + 1, 0), *point_image_in.data(y + 0, x + 1, 1), *point_image_in.data(y + 0, x + 1, 2));
			Eigen::Vector3f obs11(*point_image_in.data(y + 1, x + 1, 0), *point_image_in.data(y + 1, x + 1, 1), *point_image_in.data(y + 1, x + 1, 2));

			int idx00 = y * width + x;
			int idx01 = (y + 1) * width + x;
			int idx10 = y * width + (x + 1);
			int idx11 = (y + 1) * width + (x + 1);

			bool valid00 = obs00.z() > 0;
			bool valid01 = obs01.z() > 0;
			bool valid10 = obs10.z() > 0;
			bool valid11 = obs11.z() > 0;

			if (valid00 && valid01 && valid10) {
				float d0 = (obs00 - obs01).norm();
				float d1 = (obs00 - obs10).norm();
				float d2 = (obs01 - obs10).norm();

				if (d0 <= max_triangle_edge_distance && d1 <= max_triangle_edge_distance && d2 <= max_triangle_edge_distance) {
					int vIdx0 = mapPixelToVertexIdx[idx00];
					int vIdx1 = mapPixelToVertexIdx[idx01];
					int vIdx2 = mapPixelToVertexIdx[idx10];

					if (vIdx0 == -1) {
						vIdx0 = vertexIdx;
						mapPixelToVertexIdx[idx00] = vertexIdx;
						vertices.push_back(obs00);
						vertexIdx++;
						pixels.emplace_back(x, y);
					}
					if (vIdx1 == -1) {
						vIdx1 = vertexIdx;
						mapPixelToVertexIdx[idx01] = vertexIdx;
						vertices.push_back(obs01);
						vertexIdx++;
						pixels.emplace_back(x, y + 1);
					}
					if (vIdx2 == -1) {
						vIdx2 = vertexIdx;
						mapPixelToVertexIdx[idx10] = vertexIdx;
						vertices.push_back(obs10);
						vertexIdx++;
						pixels.emplace_back(x + 1, y);
					}

					faces.emplace_back(vIdx0, vIdx1, vIdx2);
				}
			}

			if (valid01 && valid10 && valid11) {
				float d0 = (obs10 - obs01).norm();
				float d1 = (obs10 - obs11).norm();
				float d2 = (obs01 - obs11).norm();

				if (d0 <= max_triangle_edge_distance && d1 <= max_triangle_edge_distance && d2 <= max_triangle_edge_distance) {
					int vIdx0 = mapPixelToVertexIdx[idx11];
					int vIdx1 = mapPixelToVertexIdx[idx10];
					int vIdx2 = mapPixelToVertexIdx[idx01];

					if (vIdx0 == -1) {
						vIdx0 = vertexIdx;
						mapPixelToVertexIdx[idx11] = vertexIdx;
						vertices.push_back(obs11);
						vertexIdx++;
						pixels.emplace_back(x + 1, y + 1);
					}
					if (vIdx1 == -1) {
						vIdx1 = vertexIdx;
						mapPixelToVertexIdx[idx10] = vertexIdx;
						vertices.push_back(obs10);
						vertexIdx++;
						pixels.emplace_back(x + 1, y);
					}
					if (vIdx2 == -1) {
						vIdx2 = vertexIdx;
						mapPixelToVertexIdx[idx01] = vertexIdx;
						vertices.push_back(obs01);
						vertexIdx++;
						pixels.emplace_back(x, y + 1);
					}

					faces.emplace_back(vIdx0, vIdx1, vIdx2);
				}
			}
		}
	}

	// Convert to numpy array.
	int vertex_count = vertices.size();
	int face_count = faces.size();

	if (vertex_count > 0 && face_count > 0) {
		// Reference check should be set to false otherwise there is a runtime
		// error. Check why that is the case.
		vertex_positions_out.resize({vertex_count, 3}, false);
		face_indices_out.resize({face_count, 3}, false);
		vertex_pixels_out.resize({vertex_count, 2}, false);

		for (int i = 0; i < vertex_count; i++) {
			*vertex_positions_out.mutable_data(i, 0) = vertices[i].x();
			*vertex_positions_out.mutable_data(i, 1) = vertices[i].y();
			*vertex_positions_out.mutable_data(i, 2) = vertices[i].z();

			*vertex_pixels_out.mutable_data(i, 0) = pixels[i].x();
			*vertex_pixels_out.mutable_data(i, 1) = pixels[i].y();
		}

		for (int i = 0; i < face_count; i++) {
			*face_indices_out.mutable_data(i, 0) = faces[i].x();
			*face_indices_out.mutable_data(i, 1) = faces[i].y();
			*face_indices_out.mutable_data(i, 2) = faces[i].z();
		}
	}
}

void compute_mesh_from_depth(
		const py::array_t<float>& point_image_in, float max_triangle_edge_distance,
		py::array_t<float>& vertex_positions_out, py::array_t<int>& face_indices_out
) {
	py::array_t<int> vertex_pixels;
	compute_mesh_from_depth(point_image_in, max_triangle_edge_distance, vertex_positions_out, vertex_pixels, face_indices_out);
}

py::tuple compute_mesh_from_depth(const py::array_t<float>& point_image_in, float max_triangle_edge_distance) {
	py::array_t<float> vertex_positions_out;
	py::array_t<int> face_indices_out;
	py::array_t<int> vertex_pixels_out;
	compute_mesh_from_depth(point_image_in, max_triangle_edge_distance, vertex_positions_out, vertex_pixels_out, face_indices_out);
	return py::make_tuple(vertex_positions_out, vertex_pixels_out, face_indices_out);
}

void compute_mesh_from_depth_and_color(
		const py::array_t<float>& point_image, const py::array_t<int>& color_image, float max_triangle_edge_distance,
		py::array_t<float>& vertex_positions, py::array_t<int>& vertex_colors, py::array_t<int>& face_indices
) {
	int width = static_cast<int>(point_image.shape(1));
	int height = static_cast<int>(point_image.shape(0));

	// Compute valid pixel vertices and faces.
	// We also need to compute the pixel -> vertex index mapping for
	// computation of faces.
	// We connect neighboring pixels on the square into two triangles.
	// We only select valid triangles, i.e. with all valid vertices and
	// not too far apart.
	// Important: The triangle orientation is set such that the normals
	// point towards the camera.
	std::vector<Eigen::Vector3f> vertices;
	std::vector<Eigen::Vector3i> colors;
	std::vector<Eigen::Vector3i> faces;

	int vertexIdx = 0;
	std::vector<int> mapPixelToVertexIdx(width * height, -1);

	for (int y = 0; y < height - 1; y++) {
		for (int x = 0; x < width - 1; x++) {
			Eigen::Vector3f obs00(*point_image.data(y + 0, x + 0, 0), *point_image.data(y + 0, x + 0, 1), *point_image.data(y + 0, x + 0, 2));
			Eigen::Vector3f obs01(*point_image.data(y + 1, x + 0, 0), *point_image.data(y + 1, x + 0, 1), *point_image.data(y + 1, x + 0, 2));
			Eigen::Vector3f obs10(*point_image.data(y + 0, x + 1, 0), *point_image.data(y + 0, x + 1, 1), *point_image.data(y + 0, x + 1, 2));
			Eigen::Vector3f obs11(*point_image.data(y + 1, x + 1, 0), *point_image.data(y + 1, x + 1, 1), *point_image.data(y + 1, x + 1, 2));

			Eigen::Vector3i color00(*color_image.data(y + 0, x + 0, 0), *color_image.data(y + 0, x + 0, 1), *color_image.data(y + 0, x + 0, 2));
			Eigen::Vector3i color01(*color_image.data(y + 1, x + 0, 0), *color_image.data(y + 1, x + 0, 1), *color_image.data(y + 1, x + 0, 2));
			Eigen::Vector3i color10(*color_image.data(y + 0, x + 1, 0), *color_image.data(y + 0, x + 1, 1), *color_image.data(y + 0, x + 1, 2));
			Eigen::Vector3i color11(*color_image.data(y + 1, x + 1, 0), *color_image.data(y + 1, x + 1, 1), *color_image.data(y + 1, x + 1, 2));

			// find linear indices
			int idx00 = y * width + x;
			int idx01 = (y + 1) * width + x;
			int idx10 = y * width + (x + 1);
			int idx11 = (y + 1) * width + (x + 1);

			bool valid00 = obs00.z() > 0;
			bool valid01 = obs01.z() > 0;
			bool valid10 = obs10.z() > 0;
			bool valid11 = obs11.z() > 0;

			// region ======= LOWER LEFT TRIANGLE =======
			if (valid00 && valid01 && valid10) {
				float d0 = (obs00 - obs01).norm();
				float d1 = (obs00 - obs10).norm();
				float d2 = (obs01 - obs10).norm();

				if (d0 <= max_triangle_edge_distance && d1 <= max_triangle_edge_distance && d2 <= max_triangle_edge_distance) {
					int vIdx0 = mapPixelToVertexIdx[idx00];
					int vIdx1 = mapPixelToVertexIdx[idx01];
					int vIdx2 = mapPixelToVertexIdx[idx10];

					if (vIdx0 == -1) {
						vIdx0 = vertexIdx;
						mapPixelToVertexIdx[idx00] = vertexIdx;
						vertices.push_back(obs00);
						colors.push_back(color00);
						vertexIdx++;
					}
					if (vIdx1 == -1) {
						vIdx1 = vertexIdx;
						mapPixelToVertexIdx[idx01] = vertexIdx;
						vertices.push_back(obs01);
						colors.push_back(color01);
						vertexIdx++;
					}
					if (vIdx2 == -1) {
						vIdx2 = vertexIdx;
						mapPixelToVertexIdx[idx10] = vertexIdx;
						vertices.push_back(obs10);
						colors.push_back(color10);
						vertexIdx++;
					}

					faces.emplace_back(vIdx0, vIdx1, vIdx2);
				}
			}
			// endregion
			// region ======= UPPER RIGHT TRIANGLE =======
			if (valid01 && valid10 && valid11) {
				float d0 = (obs10 - obs01).norm();
				float d1 = (obs10 - obs11).norm();
				float d2 = (obs01 - obs11).norm();

				if (d0 <= max_triangle_edge_distance && d1 <= max_triangle_edge_distance && d2 <= max_triangle_edge_distance) {
					int vIdx0 = mapPixelToVertexIdx[idx11];
					int vIdx1 = mapPixelToVertexIdx[idx10];
					int vIdx2 = mapPixelToVertexIdx[idx01];

					if (vIdx0 == -1) {
						vIdx0 = vertexIdx;
						mapPixelToVertexIdx[idx11] = vertexIdx;
						vertices.push_back(obs11);
						colors.push_back(color11);
						vertexIdx++;
					}
					if (vIdx1 == -1) {
						vIdx1 = vertexIdx;
						mapPixelToVertexIdx[idx10] = vertexIdx;
						vertices.push_back(obs10);
						colors.push_back(color10);
						vertexIdx++;
					}
					if (vIdx2 == -1) {
						vIdx2 = vertexIdx;
						mapPixelToVertexIdx[idx01] = vertexIdx;
						vertices.push_back(obs01);
						colors.push_back(color01);
						vertexIdx++;
					}

					faces.emplace_back(vIdx0, vIdx1, vIdx2);
				}
			}
			// endregion
		}
	}

	// Convert to numpy array.
	int vertex_count = vertices.size();
	int face_count = faces.size();

	if (vertex_count > 0 && face_count > 0) {
		// Reference check should be set to false otherwise there is a runtime
		// error. Check why that is the case.
		vertex_positions.resize({vertex_count, 3}, false);
		vertex_colors.resize({vertex_count, 3}, false);
		face_indices.resize({face_count, 3}, false);

		for (int i = 0; i < vertex_count; i++) {
			*vertex_positions.mutable_data(i, 0) = vertices[i].x();
			*vertex_positions.mutable_data(i, 1) = vertices[i].y();
			*vertex_positions.mutable_data(i, 2) = vertices[i].z();

			*vertex_colors.mutable_data(i, 0) = colors[i].x();
			*vertex_colors.mutable_data(i, 1) = colors[i].y();
			*vertex_colors.mutable_data(i, 2) = colors[i].z();
		}

		for (int i = 0; i < face_count; i++) {
			*face_indices.mutable_data(i, 0) = faces[i].x();
			*face_indices.mutable_data(i, 1) = faces[i].y();
			*face_indices.mutable_data(i, 2) = faces[i].z();
		}
	}
}

void compute_mesh_from_depth_and_flow(
		const py::array_t<float>& point_image_in, const py::array_t<float>& flow_image_in,
		float max_triangle_edge_distance, py::array_t<float>& vertex_positions_out,
		py::array_t<float>& vertex_flows_out, py::array_t<int>& vertex_pixels_out, py::array_t<int>& face_indices_out
) {
	int width = point_image_in.shape(1);
	int height = point_image_in.shape(0);

	// Compute valid pixel vertices and faces
	// a pixel is considered valid if the corresponding point's z coordinate is greater than zero and x, y, and z of the flow vector are all finite

	// We also need to compute the pixel -> vertex index mapping for computation of faces.
	// TODO: why are there no vertex indices > 400, if the total image has 640x480=307200 pixels?

	// For every 2x2 pixel area, we connect neighboring pixels into two triangles.
	// [ ]-------[ ]
	//  | ╲       |
	//  |    ╲    |
	//  |       ╲ |
	// [ ]-------[ ]

	// We only select valid triangles, i.e. with all valid vertices, and vertices not too far apart.

	// Important: The triangle orientation is set such that the normals point towards the camera.

	std::vector<Eigen::Vector3f> vertices;
	std::vector<Eigen::Vector3f> flows;
	std::vector<Eigen::Vector2i> pixels;
	std::vector<Eigen::Vector3i> faces;

	int vertex_index = 0;
	std::vector<int> pixel_to_vertex_index_map(width * height, -1);

	for (int y = 0; y < height - 1; y++) {
		for (int x = 0; x < width - 1; x++) {
			Eigen::Vector3f obs00(*point_image_in.data(y, x, 0), *point_image_in.data(y, x, 1), *point_image_in.data(y, x, 2));
			Eigen::Vector3f obs01(*point_image_in.data(y + 1, x, 0), *point_image_in.data(y + 1, x, 1), *point_image_in.data(y + 1, x, 2));
			Eigen::Vector3f obs10(*point_image_in.data(y, x + 1, 0), *point_image_in.data(y, x + 1, 1), *point_image_in.data(y, x + 1, 2));
			Eigen::Vector3f obs11(*point_image_in.data(y + 1, x + 1, 0), *point_image_in.data(y + 1, x + 1, 1), *point_image_in.data(y + 1, x + 1, 2));

			Eigen::Vector3f flow00(*flow_image_in.data(y, x, 0), *flow_image_in.data(y, x, 1), *flow_image_in.data(y, x, 2));
			Eigen::Vector3f flow01(*flow_image_in.data(y + 1, x, 0), *flow_image_in.data(y + 1, x, 1), *flow_image_in.data(y + 1, x, 2));
			Eigen::Vector3f flow10(*flow_image_in.data(y, x + 1, 0), *flow_image_in.data(y, x + 1, 1), *flow_image_in.data(y, x + 1, 2));
			Eigen::Vector3f flow11(*flow_image_in.data(y + 1, x + 1, 0), *flow_image_in.data(y + 1, x + 1, 1), *flow_image_in.data(y + 1, x + 1, 2));

			// linear indices of the four pixels
			int idx00 = y * width + x;
			int idx01 = (y + 1) * width + x;
			int idx10 = y * width + (x + 1);
			int idx11 = (y + 1) * width + (x + 1);

			// determine pixel validity
			bool valid00 = obs00.z() > 0 && std::isfinite(flow00.x()) && std::isfinite(flow00.y()) && std::isfinite(flow00.z());
			bool valid01 = obs01.z() > 0 && std::isfinite(flow01.x()) && std::isfinite(flow01.y()) && std::isfinite(flow01.z());
			bool valid10 = obs10.z() > 0 && std::isfinite(flow10.x()) && std::isfinite(flow10.y()) && std::isfinite(flow10.z());
			bool valid11 = obs11.z() > 0 && std::isfinite(flow11.x()) && std::isfinite(flow11.y()) && std::isfinite(flow11.z());

			// region ======= LOWER LEFT TRIANGLE =============
			if (valid00 && valid01 && valid10) {
				float d0 = (obs00 - obs01).norm();
				float d1 = (obs00 - obs10).norm(); // hypotenuse when triangle projected to image plane
				float d2 = (obs01 - obs10).norm();

				if (d0 <= max_triangle_edge_distance && d1 <= max_triangle_edge_distance && d2 <= max_triangle_edge_distance) {
					int vertex_index_0 = pixel_to_vertex_index_map[idx00];
					int vertex_index_1 = pixel_to_vertex_index_map[idx01];
					int vertex_index_2 = pixel_to_vertex_index_map[idx10];

					if (vertex_index_0 == -1) {
						vertex_index_0 = vertex_index;
						pixel_to_vertex_index_map[idx00] = vertex_index;
						vertices.push_back(obs00);
						flows.push_back(flow00);
						pixels.emplace_back(x, y);
						vertex_index++;
					}
					if (vertex_index_1 == -1) {
						vertex_index_1 = vertex_index;
						pixel_to_vertex_index_map[idx01] = vertex_index;
						vertices.push_back(obs01);
						flows.push_back(flow01);
						pixels.emplace_back(x, y + 1);
						vertex_index++;
					}
					if (vertex_index_2 == -1) {
						vertex_index_2 = vertex_index;
						pixel_to_vertex_index_map[idx10] = vertex_index;
						vertices.push_back(obs10);
						flows.push_back(flow10);
						pixels.emplace_back(x + 1, y);
						vertex_index++;
					}

					faces.emplace_back(vertex_index_0, vertex_index_1, vertex_index_2);
				}
			}
			// endregion
			// region ======= UPPER RIGHT TRIANGLE ==================
			if (valid01 && valid10 && valid11) {
				float d0 = (obs10 - obs01).norm();
				float d1 = (obs10 - obs11).norm();
				float d2 = (obs01 - obs11).norm();

				if (d0 <= max_triangle_edge_distance && d1 <= max_triangle_edge_distance && d2 <= max_triangle_edge_distance) {
					int vertex_index_0 = pixel_to_vertex_index_map[idx11];
					int vertex_index_1 = pixel_to_vertex_index_map[idx10];
					int vertex_index_2 = pixel_to_vertex_index_map[idx01];

					if (vertex_index_0 == -1) {
						vertex_index_0 = vertex_index;
						pixel_to_vertex_index_map[idx11] = vertex_index;
						vertices.push_back(obs11);
						flows.push_back(flow11);
						pixels.emplace_back(x + 1, y + 1);
						vertex_index++;
					}
					if (vertex_index_1 == -1) {
						vertex_index_1 = vertex_index;
						pixel_to_vertex_index_map[idx10] = vertex_index;
						vertices.push_back(obs10);
						flows.push_back(flow10);
						pixels.emplace_back(x + 1, y);
						vertex_index++;
					}
					if (vertex_index_2 == -1) {
						vertex_index_2 = vertex_index;
						pixel_to_vertex_index_map[idx01] = vertex_index;
						vertices.push_back(obs01);
						flows.push_back(flow01);
						pixels.emplace_back(x, y + 1);
						vertex_index++;
					}

					faces.emplace_back(vertex_index_0, vertex_index_1, vertex_index_2);
				}
			}
			// endregion
		}
	}

	// Convert to numpy array.
	int vertex_count = vertices.size();
	int face_count = faces.size();

	if (vertex_count > 0 && face_count > 0) {
		// Reference check should be set to false otherwise there is a runtime
		// error. Check why that is the case.
		vertex_positions_out.resize({vertex_count, 3}, false);
		vertex_flows_out.resize({vertex_count, 3}, false);
		vertex_pixels_out.resize({vertex_count, 2}, false);
		face_indices_out.resize({face_count, 3}, false);

		for (int i = 0; i < vertex_count; i++) {
			*vertex_positions_out.mutable_data(i, 0) = vertices[i].x();
			*vertex_positions_out.mutable_data(i, 1) = vertices[i].y();
			*vertex_positions_out.mutable_data(i, 2) = vertices[i].z();

			*vertex_flows_out.mutable_data(i, 0) = flows[i].x();
			*vertex_flows_out.mutable_data(i, 1) = flows[i].y();
			*vertex_flows_out.mutable_data(i, 2) = flows[i].z();

			*vertex_pixels_out.mutable_data(i, 0) = pixels[i].x();
			*vertex_pixels_out.mutable_data(i, 1) = pixels[i].y();
		}

		for (int i = 0; i < face_count; i++) {
			*face_indices_out.mutable_data(i, 0) = faces[i].x();
			*face_indices_out.mutable_data(i, 1) = faces[i].y();
			*face_indices_out.mutable_data(i, 2) = faces[i].z();
		}
	}
}

py::tuple compute_mesh_from_depth_and_flow(const py::array_t<float>& point_image_in, const py::array_t<float>& flow_image_in,
                                           float max_triangle_edge_distance) {
	py::array_t<float> vertex_positions_out;
	py::array_t<float> vertex_flows_out;
	py::array_t<int> vertex_pixels_out;
	py::array_t<int> face_indices_out;
	compute_mesh_from_depth_and_flow(point_image_in, flow_image_in, max_triangle_edge_distance, vertex_positions_out, vertex_flows_out, vertex_pixels_out, face_indices_out);
	return py::make_tuple(vertex_positions_out, vertex_flows_out, vertex_pixels_out, face_indices_out);
}

void filter_depth(py::array_t<unsigned short>& depth_image_in, py::array_t<unsigned short>& depth_image_out, int radius) {
	assert(depth_image_in.ndim() == 2);
	assert(depth_image_out.ndim() == 2);
	unsigned kernel_size = 2 * radius + 1;
	unsigned window_size = kernel_size * kernel_size;

	int width = depth_image_in.shape(1);
	int height = depth_image_in.shape(0);
	assert(depth_image_out.shape(0) == height);
	assert(depth_image_out.shape(1) == width);

	// #pragma omp parallel for
	for (int y = 0; y < height; y++) {
		for (int x = 0; x < width; x++) {
			// Get all residuals in the median window.
			int x_min = std::max(x - radius, 0);
			int x_max = std::min(x + radius, int(width) - 1);
			int y_min = std::max(y - radius, 0);
			int y_max = std::min(y + radius, int(height) - 1);

			std::vector<unsigned short> window_values;
			window_values.reserve(window_size);

			for (int y_near = y_min; y_near <= y_max; y_near++) {
				for (int x_near = x_min; x_near <= x_max; x_near++) {
					unsigned short depth = *depth_image_in.data(y_near, x_near);
					if (depth > 0) {
						window_values.push_back(depth);
					}
				}
			}

			// Sort the residuals and pick the median as the middle element.
			unsigned element_count = window_values.size();
			std::sort(window_values.begin(), window_values.end());

			unsigned middle_index = std::floor(element_count / 2);
			unsigned short median = window_values[middle_index];

			// Write out the median value.
			*depth_image_out.mutable_data(y, x) = median;
		}
	}
}

py::array_t<unsigned short> filter_depth(py::array_t<unsigned short>& depth_image_in, int radius) {
	py::array_t<unsigned short> depth_image_out({depth_image_in.shape(0), depth_image_in.shape(1)});
	memset(depth_image_out.mutable_data(0, 0), 0, depth_image_out.size() * sizeof(unsigned short));
	filter_depth(depth_image_in, depth_image_out, radius);
	return depth_image_out;
}

py::array_t<float> warp_flow(const py::array_t<float>& image, const py::array_t<float>& flow, const py::array_t<float>& mask) {
	// We assume:
	//      image shape (3, h, w)
	//      flow shape  (2, h, w)
	//      mask shape  (2, h, w)

	int width = image.shape(2);
	int height = image.shape(1);

	py::array_t<float> imageWarped = py::array_t<float>({3, height, width});
	py::array_t<float> weightsWarped = py::array_t<float>({1, height, width});

	// Initialize to zero.
	for (int v = 0; v < height; v++) {
		for (int u = 0; u < width; u++) {
			*imageWarped.mutable_data(0, v, u) = 0.0;
			*imageWarped.mutable_data(1, v, u) = 0.0;
			*imageWarped.mutable_data(2, v, u) = 0.0;
			*weightsWarped.mutable_data(0, v, u) = 0.0;
		}
	}

	// Compute image residuals and interpolation weights.
	for (int v = 0; v < height; v++) {
		for (int u = 0; u < width; u++) {
			// Check if pixel is inside the mask.
			if (*mask.data(0, v, u) <= 0 || *mask.data(1, v, u) <= 0) continue;

			// Compute the warped pixel.
			float u_warped = static_cast<float>(u) + *flow.data(0, v, u);
			float v_warped = static_cast<float>(v) + *flow.data(1, v, u);

			int u0 = std::floor(u_warped);
			int u1 = u0 + 1;
			int v0 = std::floor(v_warped);
			int v1 = v0 + 1;

			if (u0 < 0 || u1 >= width || v0 < 0 || v1 >= height) continue;

			// Interpolate the color contributions.
			float du = u_warped - u0;
			float dv = v_warped - v0;

			float w00 = (1 - du) * (1 - dv);
			float w01 = (1 - du) * dv;
			float w10 = du * (1 - dv);
			float w11 = du * dv;

			float c0 = *image.data(0, v, u);
			float c1 = *image.data(1, v, u);
			float c2 = *image.data(2, v, u);

			*imageWarped.mutable_data(0, v0, u0) += w00 * c0;
			*imageWarped.mutable_data(1, v0, u0) += w00 * c1;
			*imageWarped.mutable_data(2, v0, u0) += w00 * c2;
			*imageWarped.mutable_data(0, v1, u0) += w01 * c0;
			*imageWarped.mutable_data(1, v1, u0) += w01 * c1;
			*imageWarped.mutable_data(2, v1, u0) += w01 * c2;
			*imageWarped.mutable_data(0, v0, u1) += w10 * c0;
			*imageWarped.mutable_data(1, v0, u1) += w10 * c1;
			*imageWarped.mutable_data(2, v0, u1) += w10 * c2;
			*imageWarped.mutable_data(0, v1, u1) += w11 * c0;
			*imageWarped.mutable_data(1, v1, u1) += w11 * c1;
			*imageWarped.mutable_data(2, v1, u1) += w11 * c2;

			*weightsWarped.mutable_data(0, v0, u0) += w00;
			*weightsWarped.mutable_data(0, v1, u0) += w01;
			*weightsWarped.mutable_data(0, v0, u1) += w10;
			*weightsWarped.mutable_data(0, v1, u1) += w11;
		}
	}

	// Normalize image.
	for (int v = 0; v < height; v++) {
		for (int u = 0; u < width; u++) {
			float w = *weightsWarped.data(0, v, u);
			if (w > 0) {
				*imageWarped.mutable_data(0, v, u) /= w;
				*imageWarped.mutable_data(1, v, u) /= w;
				*imageWarped.mutable_data(2, v, u) /= w;
			} else {
				*imageWarped.mutable_data(0, v, u) = 1.0;
				*imageWarped.mutable_data(1, v, u) = 1.0;
				*imageWarped.mutable_data(2, v, u) = 1.0;
			}
		}
	}

	return imageWarped;
}

py::array_t<float> warp_rigid(
		const py::array_t<float>& rgbxyz_image,
		const py::array_t<float>& rotation,
		const py::array_t<float>& translation,
		float fx, float fy, float cx, float cy
) {
	// We assume:
	//      rgbd shape (6, h, w)
	//      rotation shape  (9)
	//      translation shape  (2)

	int width = rgbxyz_image.shape(2);
	int height = rgbxyz_image.shape(1);

	float r00 = *rotation.data(0);
	float r01 = *rotation.data(1);
	float r02 = *rotation.data(2);
	float r10 = *rotation.data(3);
	float r11 = *rotation.data(4);
	float r12 = *rotation.data(5);
	float r20 = *rotation.data(6);
	float r21 = *rotation.data(7);
	float r22 = *rotation.data(8);
	float t0 = *translation.data(0);
	float t1 = *translation.data(1);
	float t2 = *translation.data(2);

	py::array_t<float> image_warped = py::array_t<float>({3, height, width});
	py::array_t<float> weights_warped = py::array_t<float>({1, height, width});

	// Initialize to zero.
	for (int v = 0; v < height; v++) {
		for (int u = 0; u < width; u++) {
			*image_warped.mutable_data(0, v, u) = 0.0;
			*image_warped.mutable_data(1, v, u) = 0.0;
			*image_warped.mutable_data(2, v, u) = 0.0;
			*weights_warped.mutable_data(0, v, u) = 0.0;
		}
	}

	// Compute image residuals and interpolation weights.
	for (int v = 0; v < height; v++) {
		for (int u = 0; u < width; u++) {
			// Compute the warped pixel.
			float x = *rgbxyz_image.data(3, v, u);
			float y = *rgbxyz_image.data(4, v, u);
			float z = *rgbxyz_image.data(5, v, u);
			if (z <= 0) continue;

			float x_def = r00 * x + r01 * y + r02 * z + t0;
			float y_def = r10 * x + r11 * y + r12 * z + t1;
			float z_def = r20 * x + r21 * y + r22 * z + t2;
			if (z_def <= 0) continue;

			float u_warped = fx * x_def / z_def + cx;
			float v_warped = fy * y_def / z_def + cy;

			int u0 = std::floor(u_warped);
			int u1 = u0 + 1;
			int v0 = std::floor(v_warped);
			int v1 = v0 + 1;

			if (u0 < 0 || u1 >= width || v0 < 0 || v1 >= height) continue;

			// Interpolate the color contributions.
			float du = u_warped - u0;
			float dv = v_warped - v0;

			float w00 = (1 - du) * (1 - dv);
			float w01 = (1 - du) * dv;
			float w10 = du * (1 - dv);
			float w11 = du * dv;

			float c0 = *rgbxyz_image.data(0, v, u);
			float c1 = *rgbxyz_image.data(1, v, u);
			float c2 = *rgbxyz_image.data(2, v, u);

			*image_warped.mutable_data(0, v0, u0) += w00 * c0;
			*image_warped.mutable_data(1, v0, u0) += w00 * c1;
			*image_warped.mutable_data(2, v0, u0) += w00 * c2;
			*image_warped.mutable_data(0, v1, u0) += w01 * c0;
			*image_warped.mutable_data(1, v1, u0) += w01 * c1;
			*image_warped.mutable_data(2, v1, u0) += w01 * c2;
			*image_warped.mutable_data(0, v0, u1) += w10 * c0;
			*image_warped.mutable_data(1, v0, u1) += w10 * c1;
			*image_warped.mutable_data(2, v0, u1) += w10 * c2;
			*image_warped.mutable_data(0, v1, u1) += w11 * c0;
			*image_warped.mutable_data(1, v1, u1) += w11 * c1;
			*image_warped.mutable_data(2, v1, u1) += w11 * c2;

			*weights_warped.mutable_data(0, v0, u0) += w00;
			*weights_warped.mutable_data(0, v1, u0) += w01;
			*weights_warped.mutable_data(0, v0, u1) += w10;
			*weights_warped.mutable_data(0, v1, u1) += w11;
		}
	}

	// Normalize image.
	for (int v = 0; v < height; v++) {
		for (int u = 0; u < width; u++) {
			float w = *weights_warped.data(0, v, u);
			if (w > 0) {
				*image_warped.mutable_data(0, v, u) /= w;
				*image_warped.mutable_data(1, v, u) /= w;
				*image_warped.mutable_data(2, v, u) /= w;
			} else {
				*image_warped.mutable_data(0, v, u) = 1.0;
				*image_warped.mutable_data(1, v, u) = 1.0;
				*image_warped.mutable_data(2, v, u) = 1.0;
			}
		}
	}

	return image_warped;
}

py::array_t<float>
warp_3d(const py::array_t<float>& rgbxyz_image, const py::array_t<float>& points,
        const py::array_t<int>& mask, float fx, float fy, float cx, float cy) {

	// We assume:
	//      image shape             (6, h, w)
	//      points shape            (3, h, w)
	//      mask shape    (h, w)

	int width = rgbxyz_image.shape(2);
	int height = rgbxyz_image.shape(1);

	py::array_t<float> image_warped = py::array_t<float>({3, height, width});
	py::array_t<float> weights_warped = py::array_t<float>({1, height, width});

	// Initialize to zero.
	for (int v = 0; v < height; v++) {
		for (int u = 0; u < width; u++) {
			*image_warped.mutable_data(0, v, u) = 0.0;
			*image_warped.mutable_data(1, v, u) = 0.0;
			*image_warped.mutable_data(2, v, u) = 0.0;
			*weights_warped.mutable_data(0, v, u) = 0.0;
		}
	}

	// Compute image residuals and interpolation weights.
	for (int v = 0; v < height; v++) {
		for (int u = 0; u < width; u++) {
			// Compute the warped pixel.
			if (*mask.data(v, u) <= 0) continue;

			float z = *rgbxyz_image.data(5, v, u);
			if (z <= 0) continue;

			float x_def = *points.data(0, v, u);
			float y_def = *points.data(1, v, u);
			float z_def = *points.data(2, v, u);
			if (z_def <= 0) continue;

			float u_warped = fx * x_def / z_def + cx;
			float v_warped = fy * y_def / z_def + cy;

			int u0 = std::floor(u_warped);
			int u1 = u0 + 1;
			int v0 = std::floor(v_warped);
			int v1 = v0 + 1;

			if (u0 < 0 || u1 >= width || v0 < 0 || v1 >= height) continue;

			// Interpolate the color contributions.
			float du = u_warped - u0;
			float dv = v_warped - v0;

			float w00 = (1 - du) * (1 - dv);
			float w01 = (1 - du) * dv;
			float w10 = du * (1 - dv);
			float w11 = du * dv;

			float c0 = *rgbxyz_image.data(0, v, u);
			float c1 = *rgbxyz_image.data(1, v, u);
			float c2 = *rgbxyz_image.data(2, v, u);

			*image_warped.mutable_data(0, v0, u0) += w00 * c0;
			*image_warped.mutable_data(1, v0, u0) += w00 * c1;
			*image_warped.mutable_data(2, v0, u0) += w00 * c2;
			*image_warped.mutable_data(0, v1, u0) += w01 * c0;
			*image_warped.mutable_data(1, v1, u0) += w01 * c1;
			*image_warped.mutable_data(2, v1, u0) += w01 * c2;
			*image_warped.mutable_data(0, v0, u1) += w10 * c0;
			*image_warped.mutable_data(1, v0, u1) += w10 * c1;
			*image_warped.mutable_data(2, v0, u1) += w10 * c2;
			*image_warped.mutable_data(0, v1, u1) += w11 * c0;
			*image_warped.mutable_data(1, v1, u1) += w11 * c1;
			*image_warped.mutable_data(2, v1, u1) += w11 * c2;

			*weights_warped.mutable_data(0, v0, u0) += w00;
			*weights_warped.mutable_data(0, v1, u0) += w01;
			*weights_warped.mutable_data(0, v0, u1) += w10;
			*weights_warped.mutable_data(0, v1, u1) += w11;
		}
	}

	// Normalize image.
	for (int v = 0; v < height; v++) {
		for (int u = 0; u < width; u++) {
			float w = *weights_warped.data(0, v, u);
			if (w > 0) {
				*image_warped.mutable_data(0, v, u) /= w;
				*image_warped.mutable_data(1, v, u) /= w;
				*image_warped.mutable_data(2, v, u) /= w;
			} else {
				*image_warped.mutable_data(0, v, u) = 1.0;
				*image_warped.mutable_data(1, v, u) = 1.0;
				*image_warped.mutable_data(2, v, u) = 1.0;
			}
		}
	}

	return image_warped;
}

} //namespace image_proc


#include "poisson_fem.hpp"
#include <boost/python.hpp>
//#include "vectorize.hpp"

namespace gd {

using namespace boost::python;
using namespace std;
using namespace dealii;

void py_export_poisson_fem() {
	class_<PoissonFEM<2>, boost::noncopyable >("PoissonFEM2d", init<int>())
		.def("make_grid", &PoissonFEM<2>::make_grid)
		.def("setup_system", &PoissonFEM<2>::setup_system)
		.def("refine_grid", &PoissonFEM<2>::refine_grid)
		.def("assemble_system", &PoissonFEM<2>::assemble_system)
		.def("solve", &PoissonFEM<2>::solve)
		.def("output_results", &PoissonFEM<2>::output_results)
		.def("output_grid", &PoissonFEM<2>::output_grid)
		.def("eval", (&PoissonFEM<2>::eval))
		.def("n_active_cells", (&PoissonFEM<2>::n_active_cells))
		.def("global_refine", (&PoissonFEM<2>::global_refine))
	;
	class_<PoissonFEM<1>, boost::noncopyable >("PoissonFEM1d", init<int>())
		.def("make_grid", &PoissonFEM<1>::make_grid)
		.def("setup_system", &PoissonFEM<1>::setup_system)
		.def("refine_grid", &PoissonFEM<1>::refine_grid)
		.def("assemble_system", &PoissonFEM<1>::assemble_system)
		.def("solve", &PoissonFEM<1>::solve)
		.def("output_results", &PoissonFEM<1>::output_results)
		.def("output_grid", &PoissonFEM<1>::output_grid)
		.def("eval", (&PoissonFEM<1>::eval))
		.def("getpoints", (&PoissonFEM<1>::getpoints))
		.def("n_active_cells", (&PoissonFEM<1>::n_active_cells))
		.def("grad", (&PoissonFEM<1>::grad))
		.def("global_refine", (&PoissonFEM<1>::global_refine))

	;
}

template <int dim>
class RightHandSide : public Function<dim>
{
public:
	RightHandSide () : Function<dim>() {}

	virtual double value (const Point<dim>   &p,
						const unsigned int  component = 0) const;
};
template <int dim>
double RightHandSide<dim>::value (const Point<dim> &p,
								const unsigned int /* component */) const
{
	double rsq = 0;
	for (unsigned int i=0; i<dim; ++i)
		rsq += std::pow(p(i), 2);
	double r = sqrt(rsq);
	double x = p(0);
	double y = p(1);
	double phi = fmod(atan2(y, x)+2*M_PI, 2*M_PI);
	//return -exp(-pow(r/4-phi, 2));
	//return ((r > 10) && (r < 15.0)) * 1.; //3/(4*M_PI) * pow(1+pow(r,2), -5./2);
	//return -1;
	//return -1/(1+x/20.);
	return -3/(4*M_PI) * pow(1+pow(r/20,2), -5./2);
}

template <int dim>
PoissonFEM<dim>::PoissonFEM (int order) :
                 fe (order),
                dof_handler (triangulation)
 {}

template <int dim>
void PoissonFEM<dim>::make_grid()
{
	const Point<dim> center;

	if(dim == 1) {
		GridGenerator::hyper_cube(triangulation, 0, 40) ;//, 10);
	} else {
		GridGenerator::hyper_ball(triangulation, center, 40.) ;//, 10);
		//const HyperShellBoundary<dim> boundary_description(center);
		//triangulation.set_boundary (0, boundary_description);
	}
	//const HyperShellBoundary<dim> boundary_description(center);
	//const HyperShellBoundary<2> boundary_description2(center);
	//triangulation.set_boundary (0, boundary_description);
	//triangulation.set_boundary (1, boundary_description2);

	triangulation.refine_global (2);
	//triangulation.set_boundary (0);

	std::cout << "Number of active cells: "
				<< triangulation.n_active_cells()
				<< std::endl;
	std::cout << "Total number of cells: "
				<< triangulation.n_cells()
				<< std::endl;

}

template <int dim>
void  PoissonFEM<dim>::setup_system() {
	dof_handler.distribute_dofs (fe);
	std::cout << "Number of degrees of freedom: "
			<< dof_handler.n_dofs()
			<< std::endl;

	sparsity_pattern.reinit (dof_handler.n_dofs(),
							dof_handler.n_dofs(),
							dof_handler.max_couplings_between_dofs());
	DoFTools::make_sparsity_pattern (dof_handler, sparsity_pattern);
	solution.reinit (dof_handler.n_dofs());
	system_rhs.reinit (dof_handler.n_dofs());

	hanging_node_constraints.clear ();
	DoFTools::make_hanging_node_constraints (dof_handler,
                                            hanging_node_constraints);
	hanging_node_constraints.condense (sparsity_pattern);

	hanging_node_constraints.close ();

   sparsity_pattern.compress();

   system_matrix.reinit (sparsity_pattern);

}

template <int dim>
void PoissonFEM<dim>::assemble_system ()
{
	QGauss<dim>  quadrature_formula(2);
	FEValues<dim> fe_values (fe, quadrature_formula,
							update_values | update_gradients | update_JxW_values | update_quadrature_points);

	const RightHandSide<dim> right_hand_side;

	const unsigned int   dofs_per_cell = fe.dofs_per_cell;
	const unsigned int   n_q_points    = quadrature_formula.size();

	FullMatrix<double>   cell_matrix (dofs_per_cell, dofs_per_cell);
	Vector<double>       cell_rhs (dofs_per_cell);

	std::vector<unsigned int> local_dof_indices (dofs_per_cell);

	typename DoFHandler<dim>::active_cell_iterator
		cell = dof_handler.begin_active(),
		endc = dof_handler.end();
	for (; cell!=endc; ++cell)  {
		fe_values.reinit (cell);

		cell_matrix = 0;
		cell_rhs = 0;

	   for (unsigned int i=0; i<dofs_per_cell; ++i) {
			for (unsigned int j=0; j<dofs_per_cell; ++j) {
				for (unsigned int q_point=0; q_point<n_q_points; ++q_point) {
					//cout << "q: " << fe_values.quadrature_point (q_point) << endl;
					double r = fe_values.quadrature_point (q_point)(0);
					cell_matrix(i,j) += ( fe_values.shape_grad (i, q_point) *
									fe_values.shape_grad (j, q_point) *
					                      fe_values.JxW (q_point)) ;/* r * r; //*/
					cout << "-> " <<fe_values.JxW (q_point) << " " << r << " " << fe_values.shape_grad (i, q_point) << endl;
				}
			}
		}

		for (unsigned int i=0; i<dofs_per_cell; ++i) {
			for (unsigned int q_point=0; q_point<n_q_points; ++q_point) {
				double r = fe_values.quadrature_point (q_point)(0);
				cell_rhs(i) += (fe_values.shape_value (i, q_point) *
							right_hand_side.value (fe_values.quadrature_point (q_point)) *
				                fe_values.JxW (q_point)) ;/* r * r; //*/

				cell->get_dof_indices (local_dof_indices);

			}
		}

		for (unsigned int i=0; i<dofs_per_cell; ++i)
		for (unsigned int j=0; j<dofs_per_cell; ++j)
			system_matrix.add (local_dof_indices[i],
								local_dof_indices[j],
								cell_matrix(i,j));

		for (unsigned int i=0; i<dofs_per_cell; ++i)
		system_rhs(local_dof_indices[i]) += cell_rhs(i);
		}
	//cout << system_matrix << endl;
	system_matrix.print(cout);
	cout << "next" << endl;
	system_rhs.print(cout);
	hanging_node_constraints.condense (system_matrix);
	hanging_node_constraints.condense (system_rhs);

	std::map<unsigned int,double> boundary_values;
	//boundary_values[0] = 0;
	if(dim == 1) {
		VectorTools::interpolate_boundary_values (dof_handler,
												1,
												ConstantFunction<dim>(0),
												boundary_values);
	} else {
		VectorTools::interpolate_boundary_values (dof_handler,
				0,
				ConstantFunction<dim>(0),
				boundary_values);
	}
	MatrixTools::apply_boundary_values (boundary_values,
										system_matrix,
										solution,
										system_rhs);
	cout << "boundary:" << endl;
	system_matrix.print(cout);
	cout << "next" << endl;
	system_rhs.print(cout);
}

template <int dim>
Point<dim> pointr(double r);

template<>
Point<1> pointr<1>(double r) {
	return Point<1>(r);
}

template<>
Point<2> pointr<2>(double r) {
	return Point<2>(r, 0);
}

template<>
Point<3> pointr(double r) {
	return Point<3>(r, 0, 0);
}

template <int dim>
 void PoissonFEM<dim>::solve ()
 {
   SolverControl           solver_control (5000, 1e-12);
   SolverCG<>              cg (solver_control);

   cg.solve (system_matrix, solution, system_rhs,
            PreconditionIdentity());
hanging_node_constraints.distribute (solution);

	std::cout << VectorTools::point_value (dof_handler, solution, pointr<dim>(0)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, pointr<dim>(5)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, pointr<dim>(10)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, pointr<dim>(15)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, pointr<dim>(20)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, pointr<dim>(40)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, pointr<dim>(0.1)) << std::endl;
	/*std::cout << VectorTools::point_value (dof_handler, solution, Point<2>(40, 0)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, Point<2>(0, 40)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, Point<2>(0, 20)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, Point<2>(1,0)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, Point<2>(0,1)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, Point<2>(2,0)) << std::endl;
	double a = VectorTools::point_value (dof_handler, solution, Point<2>(1,0))-VectorTools::point_value (dof_handler, solution, Point<2>(0,0));
	double b = VectorTools::point_value (dof_handler, solution, Point<2>(2,0))-VectorTools::point_value (dof_handler, solution, Point<2>(1,0));
	std::cout << "a = " << a << std::endl;
	std::cout << "b = " << b << std::endl;
	std::cout << "a/b = " << (a/b) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, Point<2>(0, 0)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, Point<2>(0, 5)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, Point<2>(0, 10)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, Point<2>(0, 15)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, Point<2>(0, 20)) << std::endl;
	std::cout << VectorTools::point_value (dof_handler, solution, Point<2>(0, 25)) << std::endl;*/

}

template <int dim>
void PoissonFEM<dim>::getpoints(double_vector p) // for 1d only
{
	double* pp = p.data().begin();
	/*Triangulation<2>::active_line_iterator
		cell = triangulation.begin_active_line(),
		endc = triangulation.end_line();

	double nextx;
	const Point<dim> center;
	cout << "starting.. " << (cell==endc) << endl;
	for (; cell!=endc; ++cell) {
		cout << "-> " << cell->vertex(0)(0) << endl;
		*pp++ = (cell->vertex(0))(0);
		nextx =  (cell->vertex(1))(0);
	}
	*pp++ = nextx; // only include last point of last cell
	*/
	auto vertices = triangulation.get_vertices();
	auto used = triangulation.get_used_vertices();
	for(int i = 0; i < vertices.size(); i++) {
		if(used[i]) {
			*pp++ = vertices[i](0);
		}
	}
}

template <int dim>
double PoissonFEM<dim>::eval(double r) {
	Functions::FEFieldFunction<dim> ff(dof_handler, solution);
	return ff.value(pointr<dim>(r));
	//return VectorTools::point_value (dof_handler, solution, pointr<dim>(r));
}

template <int dim>
double PoissonFEM<dim>::grad(double r) {
	Functions::FEFieldFunction<dim> ff(dof_handler, solution);
	auto tensor = ff.gradient(pointr<dim>(r));
	return tensor[0];
}

template <int dim>
void PoissonFEM<dim>::refine_grid ()
{
   Vector<float> estimated_error_per_cell (triangulation.n_active_cells());

   KellyErrorEstimator<dim>::estimate (dof_handler,
                                       QGauss<dim-1>(3),
                                       typename FunctionMap<dim>::type(),
                                       solution,
                                       estimated_error_per_cell);
   GridRefinement::refine_and_coarsen_fixed_number (triangulation,
                                                    estimated_error_per_cell,
                                                    0.3, 0.03);
triangulation.execute_coarsening_and_refinement ();

}

template <int dim>
void PoissonFEM<dim>::output_results () const
 {
   DataOut<dim> data_out;
   data_out.attach_dof_handler (dof_handler);
   data_out.add_data_vector (solution, "solution");
   data_out.build_patches ();

   std::ofstream output (dim == 2 ? "solution2d.gpl" : "solution3d.gpl" );
   data_out.write_gnuplot (output);

	std::ofstream outgrid (dim == 2 ? "grid2d.eps" : "grid3d.eps");
	GridOut grid_out;
	grid_out.write_eps (triangulation, outgrid);

   std::ofstream output2(dim == 2 ? "solution2d.vtk" : "solution3d.vtk" );
   data_out.write_vtk (output2);

 }
template <int dim>
void PoissonFEM<dim>::output_grid(char* filename) const
{
	//DataOut<dim> data_out;
	//data_out.attach_dof_handler (dof_handler);
	//data_out.add_data_vector (solution, "solution");
	//data_out.build_patches ();

	std::ofstream outgrid(filename);
	GridOut grid_out;
	grid_out.write_eps(triangulation, outgrid);

}

template <int dim>
int PoissonFEM<dim>::n_active_cells() const
{
	return triangulation.n_active_cells();
}
int refines = 1;

template <int dim>
void PoissonFEM<dim>::global_refine(int refine)
{
	return triangulation.refine_global(refine);
}

/*
template <int dim>
void PoissonFEM<dim>::run ()
{
	make_grid();
	setup_system();
	assemble_system ();
	solve ();

	for(int i = 0; i < refines; i++) {

		refine_grid();
		setup_system();
		assemble_system ();
		solve ();
	}
	output_results ();
 }
*/
 /*int main (int argc, char** argv)
 {
	if(argc < 2) {
		printf("usage: %s <refinements>\n", argv[0]);
		exit(-1);
	}
	refines = atoi(argv[1]);
	LaplaceProblem<2> laplace_problem;
	laplace_problem.run ();

	LaplaceProblem<3> laplace_problem3d;
	laplace_problem3d.run ();
   return 0;
 }*/

}


#include <iostream>
#include <fstream>
#include <xtensor/xmath.hpp>
#include "../src/distribution/distribution.cpp"
#define BOOST_TEST_MODULE "Distribution test"
#include <boost/test/unit_test.hpp>

using namespace std;
namespace utf = boost::unit_test;

BOOST_AUTO_TEST_CASE(TestLoadDistributions, * utf::tolerance(0.00001))
{
	ifstream f("../fixtures/distribution.csv");

	Distributions distributions(f);

	BOOST_TEST(distributions.size() == 10);

	auto keys = distributions.Keys();

	BOOST_TEST(keys.size() == 10);
	BOOST_TEST(keys[4] == "GCA_001735525.1");

	auto dist = distributions["GCA_001735525.1"];

	BOOST_TEST(dist.size() == 64);
	BOOST_TEST(dist[1] == 0.018183);
	BOOST_TEST(xt::sum(dist)[0] == 1.0);

	auto values = distributions.Values();
	BOOST_TEST(values(4, 1) == 0.018183);
	BOOST_TEST(xt::sum(xt::row(values, 4))[0] == 1.0);
}


///////////////////////////////////////////////////////////////////////////////
// count_less_equal_than.hpp                                                 //
//                                                                           //
//  Copyright 2010 Erwann Rogard. Distributed under the Boost                //
//  Software License, Version 1.0. (See accompanying file                    //
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)         //
///////////////////////////////////////////////////////////////////////////////
#ifndef BOOST_STATISTICS_DETAIL_NON_PARAMETRIC_EMPIRICAL_DISTRIBUTION_CUMULATIVE_COUNT_HPP_ER_2010
#define BOOST_STATISTICS_DETAIL_NON_PARAMETRIC_EMPIRICAL_DISTRIBUTION_CUMULATIVE_COUNT_HPP_ER_2010

#include <boost/range.hpp>
#include <boost/numeric/conversion/converter.hpp>

#include <boost/accumulators/framework/extractor.hpp>
#include <boost/accumulators/framework/accumulator_base.hpp>
#include <boost/accumulators/framework/parameters/sample.hpp>
#include <boost/accumulators/framework/parameters/accumulator.hpp>
#include <boost/accumulators/framework/depends_on.hpp>
#include <boost/accumulators/statistics/count.hpp>
#include <boost/statistics/detail/non_parametric/empirical_distribution/ordered_sample.hpp>

namespace boost {
namespace statistics{
namespace detail{
namespace empirical_distribution{
namespace impl{

    // T can be an integer or a float
    template<typename T,typename Comp = std::less<T> >
	class cumulative_count : public boost::accumulators::accumulator_base
    {
        typedef Comp comp_;
        typedef std::size_t size_;
        typedef boost::accumulators::dont_care dont_care_;

        public:

        typedef size_ result_type;
        typedef T sample_type;

        cumulative_count(dont_care_){}

        void operator()(dont_care_){}

        template<typename Args>
        result_type result(const Args& args)const{
            namespace ns = statistics::detail::empirical_distribution;
            typedef ns::tag::ordered_sample tag_;
            return this->result_impl(
                boost::accumulators::extract_result<tag_>(
                    args[ boost::accumulators::accumulator ]
                ),
                args[ boost::accumulators::sample ]
            );
        }

        private:

        template<typename Map>
        result_type result_impl(
            Map& map,
            const sample_type& x
        )const{
           return std::for_each(
                boost::const_begin(map),
                this->bound(map,x),
                accumulator()
           ).value;
        }

        template<typename Map>
        typename boost::range_iterator<const Map>::type
        bound(
            const Map& map,
            const sample_type& x
        )const{
            return map.upper_bound(x);
        }

        struct accumulator{
            mutable size_ value;

            accumulator():value(0){}
            accumulator(const accumulator& that)
            	:value(that.value){}

            template<typename Data>
            void operator()(const Data& data)const{
            	value += data.second;
            }

        };
    };

}
namespace tag
{
    struct cumulative_count: boost::accumulators::depends_on<
        empirical_distribution::tag::ordered_sample
    >
    {
        struct impl{
            template<typename T,typename W>
            struct apply{
                typedef empirical_distribution::impl::cumulative_count<T> type;
            };
        };
    };
}
namespace result_of{

    template<typename AccSet>
    struct cumulative_count{
    	typedef empirical_distribution::tag::cumulative_count tag_;
        typedef typename
            boost::accumulators::detail::template
            	extractor_result<AccSet,tag_>::type type;
    };

}
namespace extract
{

    template<typename AccSet,typename T>
    typename detail::empirical_distribution::result_of::template
        cumulative_count<AccSet>::type
  	cumulative_count(AccSet const& acc,const T& x)
    {
        namespace ns = detail::empirical_distribution;
    	typedef ns::tag::cumulative_count tag_;
        return boost::accumulators::extract_result<tag_>(
            acc,
            (boost::accumulators::sample = x)
        );
  	}

}// extract
}// empirical_distribution
}// detail
}// statistics
}// boost

#endif


/*    Copyright (c) 2010-2019, Delft University of Technology
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 *
 *    References
 *      Mireles James, J.D. Celestial Mechanics Notes Set 4: The Circular Restricted Three Body
 *          Problem, 2006, http://www.math.utexas.edu/users/jjames/hw4Notes.pdf,
 *          last accessed: 26 May, 2012.
 *      JPL, NASA. Astrodynamic Constants, http://ssd.jpl.nasa.gov/?constants,
 *        last updated: 13 Dec, 2012, last accessed: 19th March, 2013.
 *
 *    Notes
 *      Reference values for position Lagrange libration points are taken from (Mireles James,
 *      2006). There seems to be a bug in the computation of the L3 location! Note that this code
 *      uses the Planet class (in Bodies sub-directory), which is marked for deprecation.
 *
 */

#define BOOST_TEST_MAIN

#include <cmath>
#include <limits>

#include <boost/make_shared.hpp>
#include <memory>
#include <boost/test/floating_point_comparison.hpp>
#include <boost/test/unit_test.hpp>

#include "Tudat/Astrodynamics/BasicAstrodynamics/celestialBodyConstants.h"
#include "Tudat/Astrodynamics/Gravitation/librationPoint.h"
#include "Tudat/Mathematics/RootFinders/newtonRaphson.h"

namespace tudat
{
namespace unit_tests
{

using namespace root_finders;

BOOST_AUTO_TEST_SUITE( test_libration_points )

//! Test if computation of mass parameter is working correctly.
BOOST_AUTO_TEST_CASE( testComputationOfMassParameter )
{
    // Set expected mass parameter for Earth-Moon system.
    const double expectedMassParameter = 0.0121505811805237;

    // Set Earth gravitational parameter.
    const double earthGravitationalParameter
            = celestial_body_constants::EARTH_GRAVITATIONAL_PARAMETER;

    // Set Moon gravitational parameter (Earth/Moon mass ratio taken from (JPL, 2012).
    const double moonGravitationalParameter = earthGravitationalParameter / 81.30059;

    // Compute mass parameter.
    const double computedMassParameter = circular_restricted_three_body_problem::computeMassParameter(
                earthGravitationalParameter, moonGravitationalParameter );

    // Check if computed value corresponds to expected mass parameter.
    BOOST_CHECK_CLOSE_FRACTION( expectedMassParameter, computedMassParameter, 1.0e-14 );
}

//! Test if computation of location of L1 Lagrange libration point is working correctly.
BOOST_AUTO_TEST_CASE( testComputationOfLocationOfL1LibrationPoint )
{
    // Declare and initialize Earth-Moon mass parameter from (Mireles James, 2006).
    const double earthMoonMassParameter = 0.012277471;

    // Set expected location of L1.
    const Eigen::Vector3d expectedLocationOfL1( 0.83629259089993, 0.0, 0.0 );

    // Declare L1 libration point object with Earth-Moon mass parameter and Newton-Raphson method
    // with 1000 iterations as maximum and 1.0e-14 relative X-tolerance.
    circular_restricted_three_body_problem::LibrationPoint librationPointL1( earthMoonMassParameter,
                                            std::make_shared< NewtonRaphson >( 1.0e-14, 1000 ) );

    // Compute location of Lagrange libration point.
    librationPointL1.computeLocationOfLibrationPoint( circular_restricted_three_body_problem::LibrationPoint::l1 );

    // Determine location of libration point in Earth-Moon system.
    const Eigen::Vector3d positionOflibrationPointL1
            = librationPointL1.getLocationOfLagrangeLibrationPoint( );

    // Check if computed location of L1 matches expected location.
    BOOST_CHECK_CLOSE_FRACTION( expectedLocationOfL1.x( ),
                                positionOflibrationPointL1.x( ),
                                1.0e-14 );
    BOOST_CHECK_SMALL( positionOflibrationPointL1.y( ), std::numeric_limits< double >::min( ) );
    BOOST_CHECK_SMALL( positionOflibrationPointL1.z( ), std::numeric_limits< double >::min( ) );
}

//! Test if computation of location of L2 Lagrange libration point is working correctly.
BOOST_AUTO_TEST_CASE( testComputationOfLocationOfL2LibrationPoint )
{
    // Declare and initialize Earth-Moon mass parameter from (Mireles James, 2006).
    const double earthMoonMassParameter = 0.012277471;

    // Set expected location of L2.
    const Eigen::Vector3d expectedLocationOfL2( 1.15616816590553, 0.0, 0.0 );

    // Declare L2 libration point object with Earth-Moon mass parameter and Newton-Raphson method
    // with 1000 iterations as maximum and 1.0e-14 relative X-tolerance.
    circular_restricted_three_body_problem::LibrationPoint librationPointL2( earthMoonMassParameter,
                                            std::make_shared< NewtonRaphson >( 1.0e-14, 1000 ) );

    // Compute location of Lagrange libration point.
    librationPointL2.computeLocationOfLibrationPoint( circular_restricted_three_body_problem::LibrationPoint::l2 );

    // Determine location of libration point in Earth-Moon system.
    const Eigen::Vector3d positionOflibrationPointL2
            = librationPointL2.getLocationOfLagrangeLibrationPoint( );

    // Check if computed location of L2 matches expected location.
    BOOST_CHECK_CLOSE_FRACTION( expectedLocationOfL2.x( ),
                                positionOflibrationPointL2.x( ),
                                1.0e-14 );
    BOOST_CHECK_SMALL( positionOflibrationPointL2.y( ), std::numeric_limits< double >::min( ) );
    BOOST_CHECK_SMALL( positionOflibrationPointL2.z( ), std::numeric_limits< double >::min( ) );
}

//! Test if computation of location of L3 Lagrange libration point is working correctly.
// THERE IS A BUG IN THIS CASE!
BOOST_AUTO_TEST_CASE( testComputationOfLocationOfL3LibrationPoint )
{
    // Declare and initialize Earth-Moon mass parameter from (Mireles James, 2006).
    const double earthMoonMassParameter = 0.012277471;

    // Set expected location of L3.
    const Eigen::Vector3d expectedLocationOfL3( -1.00511551160689, 0.0, 0.0 );

    // Declare L3 libration point object with Earth-Moon mass parameter and Newton-Raphson method
    // with 1000 iterations as maximum and 1.0e-14 relative X-tolerance.
    circular_restricted_three_body_problem::LibrationPoint librationPointL3( earthMoonMassParameter,
                                            std::make_shared< NewtonRaphson >( 1.0e-14, 1000 ) );

    // Compute location of Lagrange libration point.
    librationPointL3.computeLocationOfLibrationPoint( circular_restricted_three_body_problem::LibrationPoint::l3 );

    // Determine location of libration point in Earth-Moon system.
    const Eigen::Vector3d positionOflibrationPointL3
            = librationPointL3.getLocationOfLagrangeLibrationPoint( );

    // Check if computed location of L3 matches expected location.
    BOOST_CHECK_CLOSE_FRACTION( expectedLocationOfL3.x( ),
                                positionOflibrationPointL3.x( ),
                                1.0e-2 );
    BOOST_CHECK_SMALL( positionOflibrationPointL3.y( ), std::numeric_limits< double >::min( ) );
    BOOST_CHECK_SMALL( positionOflibrationPointL3.z( ), std::numeric_limits< double >::min( ) );
}

//! Test if computation of location of L4 Lagrange libration point is working correctly.
BOOST_AUTO_TEST_CASE( testComputationOfLocationOfL4LibrationPoint )
{
    // Declare and initialize Earth-Moon mass parameter from (Mireles James, 2006).
    const double earthMoonMassParameter = 0.012277471;

    // Set expected location of L4.
    const Eigen::Vector3d expectedLocationOfL4( 0.487722529, 0.86602540378444, 0.0 );

    // Declare L4 libration point object with Earth-Moon mass parameter and Newton-Raphson method
    // with 1000 iterations as maximum and 1.0e-14 relative X-tolerance.
    circular_restricted_three_body_problem::LibrationPoint librationPointL4( earthMoonMassParameter,
                                            std::make_shared< NewtonRaphson >( 1.0e-14, 1000 ) );

    // Compute location of Lagrange libration point.
    librationPointL4.computeLocationOfLibrationPoint( circular_restricted_three_body_problem::LibrationPoint::l4 );

    // Determine location of libration point in Earth-Moon system.
    const Eigen::Vector3d positionOflibrationPointL4
            = librationPointL4.getLocationOfLagrangeLibrationPoint( );

    // Check if computed location of L4 matches expected location.
    BOOST_CHECK_CLOSE_FRACTION( expectedLocationOfL4.x( ),
                                positionOflibrationPointL4.x( ),
                                1.0e-15 );
    BOOST_CHECK_CLOSE_FRACTION( expectedLocationOfL4.y( ),
                                positionOflibrationPointL4.y( ),
                                1.0e-14 );
    BOOST_CHECK_SMALL( positionOflibrationPointL4.z( ), std::numeric_limits< double >::min( ) );
}

//! Test if computation of location of L5 Lagrange libration point is working correctly.
BOOST_AUTO_TEST_CASE( testComputationOfLocationOfL5LibrationPoint )
{
    // Declare and initialize Earth-Moon mass parameter from (Mireles James, 2006).
    const double earthMoonMassParameter = 0.012277471;

    // Set expected location of L5.
    const Eigen::Vector3d expectedLocationOfL5( 0.487722529, -0.86602540378444, 0.0 );

    // Declare L5 libration point object with Earth-Moon mass parameter and Newton-Raphson method
    // with 1000 iterations as maximum and 1.0e-14 relative X-tolerance.
    circular_restricted_three_body_problem::LibrationPoint librationPointL5( earthMoonMassParameter,
                                            std::make_shared< NewtonRaphson >( 1.0e-14, 1000 ) );

    // Compute location of Lagrange libration point.
    librationPointL5.computeLocationOfLibrationPoint( circular_restricted_three_body_problem::LibrationPoint::l5 );

    // Determine location of libration point in Earth-Moon system.
    const Eigen::Vector3d positionOflibrationPointL5
            = librationPointL5.getLocationOfLagrangeLibrationPoint( );

    // Check if computed location of L5 matches expected location.
    BOOST_CHECK_CLOSE_FRACTION( expectedLocationOfL5.x( ),
                                positionOflibrationPointL5.x( ),
                                1.0e-15 );
    BOOST_CHECK_CLOSE_FRACTION( expectedLocationOfL5.y( ),
                                positionOflibrationPointL5.y( ),
                                1.0e-14 );
    BOOST_CHECK_SMALL( positionOflibrationPointL5.z( ), std::numeric_limits< double >::min( ) );
}

BOOST_AUTO_TEST_SUITE_END( )

} // namespace unit_tests
} // namespace tudat


#ifndef YANNQ_TESTS_LAYERHELPER_HPP
#define YANNQ_TESTS_LAYERHELPER_HPP
#include <Eigen/Dense>
#include <sstream>

#include <Machines/layers/AbstractLayer.hpp>
#include <Utilities/type_traits.hpp>

template<typename T>
typename yannq::AbstractLayer<T>::Matrix
ndiff_in(yannq::AbstractLayer<T>& layer, const typename yannq::AbstractLayer<T>::Vector& input, const int outSize)
{
	using Vector = typename yannq::AbstractLayer<T>::Vector;
	using Matrix = typename yannq::AbstractLayer<T>::Matrix;
	Matrix mat(input.size(), outSize);
	Vector output1(outSize);
	Vector output2(outSize);
	const typename yannq::remove_complex<T>::type h = 1e-5;

	Vector inH = input;
	for(int i = 0; i < input.size(); i++)
	{
		inH(i) += h;
		layer.forward(inH,output1);
		inH(i) = input(i) - h;
		layer.forward(inH,output2);
		mat.row(i) = (output1 - output2)/(2*h);
		inH(i) = input(i);
	}
	return mat;
}
template<typename T>
typename yannq::AbstractLayer<T>::Matrix
ndiff_weight(yannq::AbstractLayer<T>& layer, const typename yannq::AbstractLayer<T>::Vector& input, const int outSize)
{
	using Vector = typename yannq::AbstractLayer<T>::Vector;
	using Matrix = typename yannq::AbstractLayer<T>::Matrix;
	Matrix mat(layer.paramDim(), outSize);
	Vector output1(outSize);
	Vector output2(outSize);
	const typename yannq::remove_complex<T>::type h = 1e-5;

	Vector wH = layer.getParams();
	for(int i = 0; i < wH.size(); i++)
	{
		const auto val = wH(i);

		wH(i) = val + h;
		layer.setParams(wH);
		layer.forward(input,output1);

		wH(i) = val - h;
		layer.setParams(wH);
		layer.forward(input,output2);

		mat.row(i) = (output1 - output2)/(2*h);
		wH(i) = val;
	}
	return mat;
}

#endif//YANNQ_TESTS_LAYERHELPER_HPP


#pragma once
#include <Eigen/Dense>
#include <Eigen/Geometry>
#include <utility>
#include <cmath>
#include "ray.hpp"
#include "interaction.hpp"
#include "mathext.hpp"
#include "config.hpp"

extern Config conf;

class Light
{
public:
    Light(Eigen::Vector3f pos, Eigen::Vector3f color)
        : m_Pos(std::move(pos)),
        m_Color(std::move(color))
    {
    }
    virtual ~Light() = default;

    /* Sample a point on the light's surface, if light is not delta light
     * Set PDF and the surface position
     * Return color of light
     */
    virtual Eigen::Vector3f sampleSurfacePos(float* pdf = nullptr) = 0;

    virtual float sampleSurfacePdf() = 0;

    virtual bool isHit(Ray* ray, float* hit_t = nullptr) = 0;

    virtual Eigen::Vector3f emission(Eigen::Vector3f d) = 0;

    Eigen::Vector3f m_Pos;
    Eigen::Vector3f m_Color;
};

class AreaLight:public Light
{
public:
    float radius;
    Eigen::Vector3f normal;
    Eigen::Vector3f right;

    AreaLight(Eigen::Vector3f pos, Eigen::Vector3f color,
        float r = 1.0f, Eigen::Vector3f n = Eigen::Vector3f(0.0f, -1.0f, 0.0f))
    : Light(pos, color), radius(r), normal(n.normalized())
    {
        right = mathext::rot_align_normal(normal) * Eigen::Vector3f(1.0f, 0.0f, 0.0f);
    }

    float sampleSurfacePdf()
    {
        return PI_INV / (radius * radius);
    }

    Eigen::Vector3f sampleSurfacePos(float* pdf = nullptr) override
    {
        Eigen::Vector2f sample = mathext::disk(radius);
        /* Convert polar to cartesian coordinate */
        Eigen::Vector3f sampled_light_pos = m_Pos + (Eigen::AngleAxisf(sample.y(), normal) * (sample.x() * right));

        /* Compute PDF of sampling a point in the area light */
        if (pdf) {
            *pdf = sampleSurfacePdf();
        }

        return sampled_light_pos;
    }

    Eigen::Vector3f emission(Eigen::Vector3f d) override
    {
        return conf.light_power * m_Color * abs(d.dot(normal));
    }

    /*
     * Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/minimal-ray-tracer-rendering-simple-shapes/ray-plane-and-ray-disk-intersection
     */
    bool isHit(Ray* ray, float* hit_t = nullptr) override
    {
        float d = ray->m_Dir.dot(normal);

        if (abs(d) < EPSILON) {
            return false;
        }

        float t = (m_Pos - ray->m_Ori).dot(normal) / d;
        if (t < 0) {
            return false;
        }

        Eigen::Vector3f isect = ray->getPoint(t);
        bool is_hit = (isect - m_Pos).norm() <= radius;

        if (is_hit && hit_t) {
            *hit_t = t;
        }

        return is_hit;
    }
};


//
// Created by keszocze on 15.10.18.
//

#pragma once

#include <vector>
#include <cudd/cplusplus/cuddObj.hh>
#include <boost/multiprecision/cpp_int.hpp>
#include <boost/multiprecision/cpp_dec_float.hpp>

#include "number_representation.hpp"

namespace abo::error_metrics {

    using abo::util::NumberRepresentation;

    /**
     * @brief Computes the maximal value of a function represented by a vector of BDDs
     *
     * The function is assumed to return an natural number
     *
     * @param mgr
     * @param fun The function given by a vector of BDDs
     */
    boost::multiprecision::uint256_t get_max_value(const Cudd &mgr, const std::vector<BDD> &fun);

    /**
     * @brief Computes the maximum absolute difference between the f and f_hat for any input
     * The computation is performed symbolically using BDDs
     * @param mgr The BDD object manager
     * @param f The original function
     * @param f_hat The approximated function. Must have the same number of bits as f
     * @param num_rep The number representation for f and f_hat
     * @return The maximum absolute difference
     */
    boost::multiprecision::uint256_t
        worst_case_error(const Cudd &mgr, const std::vector<BDD> &f, const std::vector<BDD> &f_hat,
                             const NumberRepresentation num_rep= NumberRepresentation::BaseTwo);

    /**
     * @brief Computes the maximum absolute difference between the f and f_hat for any input
     * divided by 2^n - 1 to normalize it to the range [0, 1] regardless of the function size (with n = f.size())
     * This is mainly intended to be a helper function to make the use of the worst case error easier
     * @param mgr The BDD object manager
     * @param f The original function
     * @param f_hat The approximated function. Must have the same number of bits as f
     * @param num_rep The number representation for f and f_hat
     * @return The maximum absolute difference normalized to [0, 1]
     */
    double worst_case_error_percent(const Cudd &mgr, const std::vector<BDD> &f, const std::vector<BDD> &f_hat,
                         const NumberRepresentation num_rep= NumberRepresentation::BaseTwo);

    /**
     * @brief Computes the maximum absolute difference between the f and f_hat for any input
     * The computation is performed using BDDs and will be typically slow compared to the BDD based variant
     * @param mgr The BDD object manager
     * @param f The original function
     * @param f_hat The approximated function. Must have the same number of bits as f
     * @param num_rep The number representation for f and f_hat
     * @return The maximum absolute difference
     */
    boost::multiprecision::uint256_t
        worst_case_error_add(const Cudd &mgr, const std::vector<BDD> &f, const std::vector<BDD> &f_hat,
                             const NumberRepresentation num_rep = NumberRepresentation::BaseTwo);

    /**
     * @brief approximate_worst_case_error
     *  Calculates the worst case error approximately, to a given relative error.
     *  The time and memory this function uses scales exponentially with the desired precision (n) (in the worst case).
     *  With m being the input size (f, f_hat), in the worst case, it lies in O(m ^ n).
     * @param mgr
     * @param f
     * @param f_hat
     * @param n The precision to calculate in number of bits, for more details see the return value. It must be greater than zero
     * @param num_rep The number representation for f and f_hat
     * @return The approximated worst case error. It is an upper bound and therefore guaranteed to be larger
     *  than the actual error. Let wc be the correct worst case error and x be the result of this function.
     *  Then it holds that wc <= x <= wc * (1 + 1 / (2 ^ (n + 1) - 1)))
     */
    boost::multiprecision::uint256_t approximate_worst_case_error(const Cudd &mgr, const std::vector<BDD> &f,
                                                                  const std::vector<BDD> &f_hat, int n,
                                                                  const NumberRepresentation num_rep = NumberRepresentation::BaseTwo);

}


//////////////////////////////////////////////////////////////////////////////
///   Copyright 2003 and onward LASMEA UMR 6602 CNRS/U.B.P Clermont-Ferrand
///   Copyright 2009 and onward LRI    UMR 8623 CNRS/Univ Paris Sud XI
///
///          Distributed under the Boost Software License, Version 1.0
///                 See accompanying file LICENSE.txt or copy at
///                     http://www.boost.org/LICENSE_1_0.txt
//////////////////////////////////////////////////////////////////////////////
#ifndef NT2_TOOLBOX_EXPONENTIAL_FUNCTION_SCALAR_IMPL_LOGS_D_LOG_HPP_INCLUDED
#define NT2_TOOLBOX_EXPONENTIAL_FUNCTION_SCALAR_IMPL_LOGS_D_LOG_HPP_INCLUDED

#include <boost/fusion/include/vector_tie.hpp>
#include <nt2/include/functions/sqr.hpp>
#include <nt2/include/functions/tofloat.hpp>
#include <nt2/include/functions/is_nan.hpp>
#include <nt2/include/functions/is_ltz.hpp>
#include <nt2/include/functions/is_eqz.hpp>
#include <nt2/include/functions/fast_frexp.hpp>
#include <nt2/include/functions/genmask.hpp>
#include <nt2/include/functions/amul.hpp>
#include <nt2/include/functions/minusone.hpp>
#include <nt2/include/functions/madd.hpp>
#include <nt2/sdk/constant/digits.hpp>
#include <nt2/sdk/constant/real.hpp>

namespace nt2
{
  namespace details
  {
    namespace internal
    {
      //////////////////////////////////////////////////////////////////////////////
      // math log functions
      //////////////////////////////////////////////////////////////////////////////

      template < class A0 >
      struct logarithm< A0, tag::not_simd_type, double>
      {
	static inline void kernel_log(const A0& a0,
				      A0& dk,
				      A0& hfsq,
				      A0& s,
				      A0& R,
				      A0& f)
	{
	  typedef typename meta::as_integer<A0, signed>::type int_type;
	  typedef typename meta::scalar_of<A0>::type               sA0;
	  A0 x;
	  int_type k;
	  boost::fusion::vector_tie(x, k) = fast_frexp(a0);
	  const int_type x_lt_sqrthf = -is_greater(Sqrt_2o_2<A0>(), x);
	  k += x_lt_sqrthf;
	  f = minusone(x+b_and(x, genmask(x_lt_sqrthf)));
	  dk = tofloat(k);
	  s = f/add(Two<A0>(),f);
	  A0 z = sqr(s);
	  A0 w = sqr(z);
	  A0 t1= w*horner<NT2_HORNER_COEFF_T(sA0, 3,
					 (0x3fc39a09d078c69fll,
					  0x3fcc71c51d8e78afll,
					  0x3fd999999997fa04ll)
					 )> (w);
	  A0 t2= z*horner<NT2_HORNER_COEFF_T(sA0, 4,
				       (0x3fc2f112df3e5244ll,
					0x3fc7466496cb03dell,
					0x3fd2492494229359ll,
					0x3fe5555555555593ll)
			)> (w);
	  R = t2+t1;
	  hfsq = mul(Half<A0>(), sqr(f));
	}

	static inline A0 log(const A0& a0)
	{
	  // ln(2)hi  =  6.93147180369123816490e-01  or  0x3fe62e42fee00000
	  // ln(2)lo  =  1.90821492927058770002e-10  or  0x3dea39ef35793c76
	  if (a0 == Inf<A0>()) return a0;
	  if (is_eqz(a0)) return Minf<A0>();
	  if (nt2::is_nan(a0)||is_ltz(a0)) return Nan<A0>();
	  A0 dk, hfsq, s, R, f;
	  kernel_log(a0, dk, hfsq, s, R, f);
	  return  mul(dk, double_constant<A0, 0x3fe62e42fee00000ll>())-
	    ((hfsq-(s*(hfsq+R)+mul(dk,double_constant<A0, 0x3dea39ef35793c76ll>())))-f);
	}
	static inline A0 log2(const A0& a0)
	{
	  if (a0 == Inf<A0>()) return a0;
	  if (is_eqz(a0)) return Minf<A0>();
	  if (nt2::is_nan(a0)||is_ltz(a0)) return Nan<A0>();
	  A0 dk, hfsq, s, R, f;
	  kernel_log(a0, dk, hfsq, s, R, f);
	  return -(hfsq-(s*(hfsq+R))-f)*Invlog_2<A0>()+dk;
	}

	static inline A0 log10(const A0& a0)
	{
	  if (a0 == Inf<A0>()) return a0;
	  if (is_eqz(a0)) return Minf<A0>();
	  if (nt2::is_nan(a0)||is_ltz(a0)) return Nan<A0>();
	  A0 dk, hfsq, s, R, f;
	  kernel_log(a0, dk, hfsq, s, R, f);
	  return -(hfsq-(s*(hfsq+R))-f)*Invlog_10<A0>()+dk*Log_2olog_10<A0>();
	}
      };
    }
  }
}

#endif

// /////////////////////////////////////////////////////////////////////////////
// End of d_log.hpp
// /////////////////////////////////////////////////////////////////////////////