////
//// Created by giovanni on 25/11/21.
////

#include <dv-processing/camera/calibration_set.hpp>
#include <dv-processing/io/mono_camera_recording.hpp>
#include <dv-processing/kinematics/motion_compensator.hpp>
#include <dv-processing/optimization/contrast_maximization_translation_and_depth.hpp>
#include <dv-processing/optimization/contrast_maximization_wrapper.hpp>

#include <CLI/CLI.hpp>
#include <fmt/core.h>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>

#include <unsupported/Eigen/NonLinearOptimization>

void visualize(dv::Frame &image1, dv::Frame &image2, dv::Frame &image3);
void feedDataMotionCompensator(std::unique_ptr<dv::kinematics::MotionCompensator<>> &mc, int64_t timestampInit,
	int64_t timestampEnd, const Eigen::VectorXf &translation, float depth);

int main(int ac, char **av) {
	using namespace std::chrono_literals;

	std::string calibrationFilepath;
	std::string aedat4Path;
	float framerate;
	// The contribution value influence the result of the optimization since the contrast of the event image is directly
	// related to the contribution. This value depends on the texture in the scene and the length of the time window of
	// events used for the optimization. A proper tuning of this parameter is needed to get good optimization results.
	float contribution;
	int inputDim = 4;
	// numMeasurement represents the number of function evaluations performed to compute gradient. Since gradient is
	// computed indipendently for each variable to be optimized, the number of measurements should be >= inputDim
	int numMeasurement = inputDim;
	// most important variable defining where will the optimization converge to. Since the optimization is a non linear
	// one, the initial guess has a big impact on the final results. Please read our document about contrast
	// maximization (docs/source/motion_compensation.md) to know more about importance of initial guess.
	std::vector<float> initTranslationDepth = {0.f, 0.f, 0.f, 3.f};
	// this is a very important variable that has a big impact on the final optimization result: it defines the
	// searching space for the variable to be optimized. If learning rate is too small, the optimization is stuck around
	// initial guess. On the other hand, if learning rate is too big, optimization will jump between values having a
	// hard time converging to a final result. Unfortunately, its optimal value depends on the optimization itself
	// (number of events/time interval and variable to be optimized), so it is strongly encouraged to fine tune this
	// value to obtain good results.
	float learningRate = float(1e-1);

	CLI::App app{"Command-line for imu compensation script."};
	app.add_option("-c,--calibration", calibrationFilepath, "Path to xml calibration file.")
		->required()
		->check(CLI::ExistingFile);
	app.add_option("-i,--input", aedat4Path, "Path to an input aedat4 file.")->required()->check(CLI::ExistingFile);
	app.add_option("-t,--init-translation-depth", initTranslationDepth, "Gyroscope angle offset in radians.");
	app.add_option(
		   "-v,--contribution", contribution, "Contribution of each event to the intensity of reconstructed frame.")
		->default_val(0.02f);
	app.add_option("-f,--framerate", framerate, "Image generation framerate.")->default_val(20.f);

	try {
		app.parse(ac, av);
	}
	catch (const CLI::ParseError &e) {
		return app.exit(e);
	}

	// camera info
	auto camera = std::make_shared<dv::camera::CameraGeometry>(
		dv::camera::CalibrationSet::LoadFromFile(calibrationFilepath).getCameraCalibration("C0")->getCameraGeometry());

	// Initial guess of translation and depth. This contrast maximization sample optimize the camera translation (x, y,
	// z) and scene depth to maximize the contrast of the event image generated by motion compensating events with
	// camera translation and scene depth information. Note that camera position and scene depth are tightly couples
	// inforamtion. Since contrast maximization is based on a non-linear optimization, the initial guess for each
	// variable to optimize has a strong impact on the final result.
	Eigen::VectorXf initTranslationDepthEig
		= Eigen::Map<Eigen::VectorXf, Eigen::Unaligned>(initTranslationDepth.data(), initTranslationDepth.size());

	// accumulator is used to get event image without performing any compensation
	dv::EdgeMapAccumulator accumulator(camera->getResolution(), contribution);
	dv::EventStore store;
	std::vector<dv::EventStore> eventSlices;
	dv::Duration interval(static_cast<int64_t>(1e+6f / framerate));
	dv::Duration imuTimeOverhead = 5ms;
	int64_t nextTimestamp        = -1;

	dv::io::DataReadHandler handler;
	handler.mEventHandler = [&store](const dv::EventStore &events) {
		store.add(events);
	};

	// although we do not directly optimize imuSamples here, it is convenient to slice event chunks based on imu data.
	// Similarly, a slicer based on events could have also been used.
	std::vector<dv::IMU> imuSamples;
	handler.mImuHandler = [&](const std::vector<dv::IMU> &imuData) {
		// Integrate the IMU measured rotation to slice event chunks by time
		for (const auto &imu : imuData) {
			int64_t generationTime = imu.timestamp - imuTimeOverhead.count();
			if (nextTimestamp < 0) {
				nextTimestamp = generationTime + interval.count();
			}
			else if (generationTime >= nextTimestamp) {
				auto events = store.sliceTime(generationTime - interval.count(), generationTime);
				eventSlices.push_back(events);
				nextTimestamp = generationTime + interval.count();
			}
		}
	};

	// Construct the camera reader
	dv::io::MonoCameraRecording reader(aedat4Path);
	std::unique_ptr<dv::kinematics::MotionCompensator<>> mcCompensate;

	cv::namedWindow("Preview", cv::WINDOW_NORMAL);
	cv::setWindowTitle("Preview", "compensated (L) vs uncompensated (R) images");

	while (reader.handleNext(handler)) {
		for (const auto &events : eventSlices) {
			// Optimize average translation and depth of the scene. Here we construct a ContrastMaximizationWrapper
			// based on TranslationLossFunctor. This functor creates a motion compensator object at each optimization
			// iteration. The motion compenator is constructed in such a way that the camera translational motion is
			// assumed to be constant in the given time window. This can be achieved by setting as initial
			// transformation at the lowest event time an identiy and the actual translation value at the highest event
			// timestamp. Motion compensator will internally interpolate assuming linear translation between initial and
			// end time.
			// Contrast maximization wrapper is initialized specifying the functor type to be used, in this case
			// "dv::optimization::TranslationLossFunctor", which implements the cost evaluation in the "int operator()"
			// method. Any new functor created to optimize a motion model, should also overload this method and inherit
			// basic "OptimizationFunctor" class such that non-linear optimization can properly be handled by
			// dv::optimization::ContrastMaximizationWrapper.
			// For more information about the contrast maximization optimization assuming translational camera motion
			// please check "contrast_maximization_translation_and_depth.hpp".
			dv::optimization::ContrastMaximizationWrapper<dv::optimization::TranslationLossFunctor>
				contrastMaximization(std::make_unique<dv::optimization::TranslationLossFunctor>(
										 camera, events, contribution, inputDim, numMeasurement),
					learningRate);
			// Contrast maximization is optimized using initial translation and depth guess. The output returns a struct
			// with information about successfull optimization, number of iterations till convergence, and value of
			// optimized variables.
			auto optimizationOutput           = contrastMaximization.optimize(initTranslationDepthEig);
			Eigen::VectorXf optimizedVariable = optimizationOutput.optimizedVariable;
			// motion compensate with optimized translation and depth
			Eigen::Vector3f translationOpt = optimizedVariable.block<3, 1>(0, 0);
			float depthOpt                 = optimizedVariable.w();
			std::unique_ptr<dv::kinematics::MotionCompensator<>> mcCompensateContrastMax
				= std::make_unique<dv::kinematics::MotionCompensator<>>(
					camera, std::make_unique<dv::EdgeMapAccumulator>(camera->getResolution(), contribution));
			feedDataMotionCompensator(
				mcCompensateContrastMax, events.getLowestTime(), events.getHighestTime(), translationOpt, depthOpt);

			// motion compensate with initial guess of translation and depth
			mcCompensate = std::make_unique<dv::kinematics::MotionCompensator<>>(
				camera, std::make_unique<dv::EdgeMapAccumulator>(camera->getResolution(), contribution));
			Eigen::Vector3f translationInitGuess = initTranslationDepthEig.block<3, 1>(0, 0);
			feedDataMotionCompensator(mcCompensate, events.getLowestTime(), events.getHighestTime(),
				translationInitGuess, initTranslationDepthEig.w());

			// add events
			accumulator.accept(events);
			mcCompensate->accept(events);
			mcCompensateContrastMax->accept(events);
			// get frames
			dv::Frame uncompensatedImage          = accumulator.generateFrame();
			dv::Frame compensatedImageInitGuess   = mcCompensate->generateFrame();
			dv::Frame compensatedImageContrastMax = mcCompensateContrastMax->generateFrame();

			int64_t generationTime = events.getHighestTime();
			nextTimestamp          = generationTime + interval.count();

			// compare uncompensated event image vs compensated image with initial guess of parameters vs compensated
			// image with optimized variables
			visualize(uncompensatedImage, compensatedImageInitGuess, compensatedImageContrastMax);
			cv::waitKey(0);
		}
		eventSlices.clear();
	}

	reader.run(handler);
	return EXIT_SUCCESS;
}

void feedDataMotionCompensator(std::unique_ptr<dv::kinematics::MotionCompensator<>> &mc, int64_t timestampInit,
	int64_t timestampEnd, const Eigen::VectorXf &translation, float depth) {
	dv::kinematics::Transformation<float> initTransform = {timestampInit, Eigen::Matrix4f::Identity()};

	Eigen::Matrix<float, 4, 4> transformNoOptimize               = Eigen::Matrix4f::Identity();
	transformNoOptimize.block<3, 1>(0, 3)                        = translation;
	dv::kinematics::Transformation<float> endTransformNoOptimize = {timestampEnd, transformNoOptimize};

	mc->accept(initTransform);
	mc->accept(endTransformNoOptimize);
	mc->setConstantDepth(depth);
}

void visualize(dv::Frame &image1, dv::Frame &image2, dv::Frame &image3) {
	cv::Mat preview;
	cv::hconcat(image1.image, image2.image, preview);
	cv::hconcat(preview, image3.image, preview);

	cv::Mat mean, stdUncompensated, stdCompensatedImu, stdCompensatedContrastMax;
	cv::meanStdDev(image1.image, mean, stdUncompensated);
	cv::meanStdDev(image2.image, mean, stdCompensatedImu);
	cv::meanStdDev(image3.image, mean, stdCompensatedContrastMax);

	// convert contrast to string (to show contrast results in preview title)
	cv::String title;
	std::string s1 = fmt::format("{:.2f}", stdUncompensated.at<double>(0));
	std::string s2 = fmt::format("{:.2f}", stdCompensatedImu.at<double>(0));
	std::string s3 = fmt::format("{:.2f}", stdCompensatedContrastMax.at<double>(0));

	title = "Contrast (uncompensated - imu - contrast max))" + s1 + " - " + s2 + " - " + s3;
	cv::imshow("Preview", preview);
	cv::setWindowTitle("Preview", title);
}
