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

#include <dv-processing/camera/calibration_set.hpp>
#include <dv-processing/imu/rotation-integrator.hpp>
#include <dv-processing/io/mono_camera_recording.hpp>
#include <dv-processing/kinematics/motion_compensator.hpp>
#include <dv-processing/optimization/contrast_maximization_rotation.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);

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

	std::string calibrationFilepath;
	std::string aedat4Path;
	// Initial guess of gyroscope offset. This contrast maximization sample optimize the gyroscope offset (x, y,
	// z) which is used to update imu gyroscope data. Based on those updated data, the camera orientation is computed.
	// Then, contrast of the event image generated by motion compensating events with
	// corresponding camera orientation is calculated.
	// 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.
	std::vector<float> gyroscopeOffsetAngle = {0.f, 0.f, 0.f};
	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;
	// size of vector variables to optimize
	int inputDimension = 3;
	// 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 = inputDimension;
	// 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(
		   "-v,--contribution", contribution, "Contribution of each event to the intensity of reconstructed frame.")
		->default_val(0.02f);
	app.add_option("-g,--gyro-offset", gyroscopeOffsetAngle, "Gyroscope angle offset in radians.");
	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);
	}

	cv::FileStorage fs(calibrationFilepath, cv::FileStorage::READ);

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

	cv::Mat T_CS;
	fs["transformation_cam_imu"] >> T_CS;
	auto imuToCamTimeOffsetUs
		= static_cast<int64_t>(static_cast<float>(fs.getFirstTopLevelNode()["time_offset_cam_imu"]) * 1e+6f);
	Eigen::Matrix4f T_CS_eigen;
	cv::cv2eigen(T_CS, T_CS_eigen);
	// Compute target (camera) position wrt to sensor (imu).
	// This will be used to get integrated imu positions wrt to camera initial frame (i.e. get camera orientation in
	// time).
	const dv::kinematics::Transformationf T_S_target = dv::kinematics::Transformationf(0, T_CS_eigen).inverse();
	// move std::vector to Eigen::Vector
	Eigen::VectorXf gyroscopeOffsetAngleEig
		= Eigen::Map<Eigen::VectorXf, Eigen::Unaligned>(gyroscopeOffsetAngle.data(), gyroscopeOffsetAngle.size());
	// create rotation integrator integrating imu data using initial guess of gyroscope offset
	dv::imu::RotationIntegrator rotationIntegrator(T_S_target, imuToCamTimeOffsetUs, gyroscopeOffsetAngleEig);

	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;
	// accumulator is used to get event image without performing any compensation
	dv::EdgeMapAccumulator accumulator(camera->getResolution(), contribution);
	dv::io::DataReadHandler handler;
	std::unique_ptr<dv::kinematics::MotionCompensator<>> mcCompensateInitialGuess
		= std::make_unique<dv::kinematics::MotionCompensator<>>(
			camera, std::make_unique<dv::EdgeMapAccumulator>(camera->getResolution(), contribution));

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

	// slice events every "interval" time
	std::vector<dv::IMU> imuSamples;
	handler.mImuHandler = [&](const std::vector<dv::IMU> &imuData) {
		// Integrate the IMU measured rotation and feed it into the motion compensator
		for (const auto &imu : imuData) {
			imuSamples.push_back(imu);
			rotationIntegrator.accept(imu);
			auto transform = rotationIntegrator.getTransformation();
			mcCompensateInitialGuess->accept(transform);
			int64_t generationTime = transform.getTimestamp() - 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);

	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 gyrscope offset to maximize contrast of the event image. Here we construct a
			// ContrastMaximizationWrapper based on RotationLossFunctor. This functor creates a motion compensator
			// object at each optimization iteration. The motion compensator is constructed in such a way that the
			// camera orientation is updated using current value of gyroscope offset (x, y, z). This can be achieved by
			// integrating imu data after updating them with current gyroscope offset value.
			// E.g. new_gyroscope_value_x = initial_gyroscope_value_x + gyroscope_offset_x, where gyroscope_offset_x
			// variable is updated at each optimization iteration to maximize image contrast.
			// Contrast maximization wrapper is initialized specifying the functor type to be used, in this case
			// "dv::optimization::RotationLossFunctor", 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 rotational camera motion
			// please check "contrast_maximization_rotation.hpp".
			dv::optimization::ContrastMaximizationWrapper<dv::optimization::RotationLossFunctor> contrastMaximization(
				std::make_unique<dv::optimization::RotationLossFunctor>(camera, events, contribution, imuSamples,
					T_S_target, imuToCamTimeOffsetUs, inputDimension, numMeasurement),
				learningRate);
			auto optimizationOutput                 = contrastMaximization.optimize(gyroscopeOffsetAngleEig);
			Eigen::VectorXf gyroscopeOffsetImuFloat = optimizationOutput.optimizedVariable;
			// integrate using updated gyroscope offset
			std::unique_ptr<dv::kinematics::MotionCompensator<>> mcCompensateContrastMax
				= std::make_unique<dv::kinematics::MotionCompensator<>>(
					camera, std::make_unique<dv::EdgeMapAccumulator>(camera->getResolution(), contribution));

			// add motion to motion compensator
			dv::imu::RotationIntegrator accumulateRotationCompensator(
				T_S_target, imuToCamTimeOffsetUs, gyroscopeOffsetImuFloat);
			// integrate imu & feed transformations to motion compensator using optimized gyroscope offset
			for (const auto &imu : imuSamples) {
				accumulateRotationCompensator.accept(imu);
				auto transform = accumulateRotationCompensator.getTransformation();
				mcCompensateContrastMax->accept(transform);
			}

			// add events
			accumulator.accept(events);
			mcCompensateInitialGuess->accept(events);
			mcCompensateContrastMax->accept(events);

			// get frames
			dv::Frame uncompensatedImage          = accumulator.generateFrame();
			dv::Frame compensatedImageInitGuess   = mcCompensateInitialGuess->generateFrame();
			dv::Frame compensatedImageContrastMax = mcCompensateContrastMax->generateFrame();

			visualize(uncompensatedImage, compensatedImageInitGuess, compensatedImageContrastMax);
			cv::waitKey(0);

			int64_t generationTime = events.getHighestTime();
			nextTimestamp          = generationTime + interval.count();
		}
		eventSlices.clear();
	}
	reader.run(handler);

	return EXIT_SUCCESS;
}

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);
}
