#include <opencv2/calib3d.hpp>
#include <opencv2/core.hpp>
#include <opencv2/cudaarithm.hpp>
#include <opencv2/cudafilters.hpp>
#include <opencv2/cudaoptflow.hpp>
#include <opencv2/cudawarping.hpp>
#include <opencv2/highgui.hpp>

using namespace cv;

cv::Mat makeCheckerboard() {
	int dim = 2048;
	int cell_size = dim / 8;
	constexpr int maxval = (1 << 16) - 1;
	Mat currentImage = Mat::zeros(dim, dim, CV_16UC1);
	for (int i = 0; i < dim * dim / (cell_size * cell_size); i++) {
		int y = (i * cell_size / dim);
		cv::rectangle(currentImage,
			{ ((y % 2) * cell_size + (i * cell_size * 2)) % dim,
			 y * cell_size, cell_size, cell_size },
			maxval, cv::FILLED);
	}
	return currentImage;
}

std::vector<cv::Mat> makeRotationSerie(int size, const cv::Mat& original) {
	std::vector<cv::Mat> vec;
	vec.push_back(original);

	for (int i = 1; i < size; i++) {
		auto mat = original.clone();
		auto center = original.size() / 2;
		auto rot_mat = cv::getRotationMatrix2D(
			{ (float)center.width, (float)center.height }, i, 1.0);
		cv::warpAffine(original, mat, rot_mat, original.size());
		vec.push_back(mat);
	}
	return vec;
}

int main()
{
	auto opt_flow = cv::cuda::FarnebackOpticalFlow::create();

	auto images = makeRotationSerie(10, makeCheckerboard());
	auto stream = cv::cuda::Stream();

	cuda::GpuMat init_image(images[0]);
	cuda::GpuMat init_image_resized;
	cuda::resize(init_image, init_image_resized, Size(), 0.5,
		0.5, cv::INTER_LINEAR, stream);

	Mat result_image(2048, 2048, CV_16UC1, Scalar(0));

	for (const auto& img : images) {

		cuda::GpuMat next_image(img);

		cuda::GpuMat next_image_resized;
		cuda::resize(next_image, next_image_resized, Size(), 0.5,
			0.5, cv::INTER_LINEAR, stream);

		stream.waitForCompletion(); // <------------------ Why is this line necessary ----------------------
		// Compute the optical flow
		cuda::GpuMat gpu_flow;
		opt_flow->calc(init_image_resized, next_image_resized, gpu_flow, stream);

		stream.waitForCompletion();
		Mat flow(gpu_flow);

		//estimate uniform transformation by querying at regular intervals and
		//Rescale to original image size
		std::vector<Point2f> moved_points;
		std::vector<Point2f> original_points;
		for (int x = 15; x < 1024; x += 15) {
			for (int y = 15; y < 1024; y += 15) {
				original_points.emplace_back(x / 0.5, y / 0.5);
				Point2f f = flow.at<Point2f>(y, x);
				moved_points.push_back(
					Point2f((x + f.x) / 0.5, (y + f.y) / 0.5));
			}
		}

		Mat warp_matrix = estimateAffinePartial2D(original_points, moved_points);

		cuda::GpuMat next_image_aligned;
		cuda::warpAffine(next_image, next_image_aligned, warp_matrix,
			next_image.size(), INTER_LINEAR | WARP_INVERSE_MAP, 0,
			cv::Scalar(), stream);

		stream.waitForCompletion();
		result_image += Mat(next_image_aligned) / 10;
	}

	imshow("Result", result_image);
	waitKey(0);
	imwrite("result.png", result_image);
}