#include "../../include/dv-processing/visualization/pose_visualizer.hpp"

#include "CLI/CLI.hpp"
#include "boost/ut.hpp"

#include <opencv2/highgui.hpp>

namespace dvv = dv::visualization;
namespace dvk = dv::kinematics;

std::vector<dvk::Transformationf> generateTrajectory(int64_t nPositions, int64_t maxTime);

int main(int ac, char **av) {
	using namespace boost::ut;

	bool showPreview = false;
	bool interactive = false;

	CLI::App app{"Image feature tracker test"};

	app.add_flag("-p,--preview", showPreview, "Display preview image");
	app.add_flag("-i,--interactive", interactive, "Interactive mode to test pose visualizer features.");
	try {
		app.parse(ac, av);
	}
	catch (const CLI::ParseError &e) {
		return app.exit(e);
	}

	if (showPreview) {
		cv::namedWindow("Preview", cv::WINDOW_NORMAL);
	}

	"simple_view"_test = [showPreview] {
		const cv::Size dims(640, 480);
		auto poses = generateTrajectory(50, 20000000);
		dvv::PoseVisualizer visualizer(dims);

		visualizer.setViewMode(dvv::PoseVisualizer::ViewMode::VIEW_XY);

		if (showPreview) {
			std::cout << "Testing visualization of different camera poses" << std::endl;
		}

		for (const auto &pose : poses) {
			visualizer.accept(pose);
			auto img = visualizer.generateFrame();

			expect(eq(img.image.size(), dims));

			if (showPreview) {
				cv::imshow("Preview", img.image);
				cv::waitKey(33);
			}
		}

		visualizer.reset();

		auto timestampOffset = dv::now();
		for (auto &pose : poses) {
			visualizer.accept(dvk::Transformationf(timestampOffset + pose.getTimestamp(), pose.getTransform()));
			auto img = visualizer.generateFrame();

			expect(eq(img.image.size(), dims));

			if (showPreview) {
				cv::imshow("Preview", img.image);
				cv::waitKey(33);
			}
		}

		for (int thickness = 1; thickness < 10; thickness++) {
			visualizer.setLineThickness(thickness);
			auto img = visualizer.generateFrame();
			expect(eq(img.image.size(), dims));
			if (showPreview) {
				cv::imshow("Preview", img.image);
				cv::waitKey(33);
			}
		}

		if (showPreview) {
			cv::imshow("Preview", visualizer.generateFrame().image);
			cv::waitKey(0);
		}
	};

	"markers"_test = [&showPreview] {
		const cv::Size dims(640, 480);
		dvv::PoseVisualizer visualizer(dims);

		auto poses = generateTrajectory(50, 20000000);

		visualizer.setViewMode(dvv::PoseVisualizer::ViewMode::VIEW_XY);

		dv::LandmarksPacket landmarks;

		std::vector<dv::Observation> obs;
		std::vector<int8_t> dsc;
		std::vector<float> cov;
		int64_t timestamp = poses.front().getTimestamp();

		landmarks.elements.emplace_back(dv::Point3f(.1f, .1f, 1.f), 1, timestamp, dsc, "", cov, obs);
		landmarks.elements.emplace_back(dv::Point3f(-.1f, .1f, 1.f), 2, timestamp, dsc, "", cov, obs);
		landmarks.elements.emplace_back(dv::Point3f(-.1f, -.1f, 1.f), 3, timestamp, dsc, "", cov, obs);
		landmarks.elements.emplace_back(dv::Point3f(.1f, -.1f, 1.f), 4, timestamp, dsc, "", cov, obs);
		landmarks.elements.emplace_back(dv::Point3f(0.f, 0.f, 1.f), 5, timestamp, dsc, "", cov, obs);

		if (showPreview) {
			std::cout << "Testing visualization of different landmarks" << std::endl;
		}

		// Just updating the landmarks
		for (const auto &pose : poses) {
			expect(nothrow([&pose, &landmarks, &visualizer] {
				visualizer.accept(pose);
				visualizer.accept(landmarks);
			}));
			auto img = visualizer.generateFrame();
			expect(eq(img.image.size(), dims));
			expect(eq(visualizer.getNumberOfLandmarks(), landmarks.elements.size()));

			if (showPreview) {
				cv::imshow("Preview", img.image);
				cv::waitKey(33);
			}
		}

		visualizer.setLandmarkLimit(10);
		timestamp += poses.back().getTimestamp() + 1000;
		landmarks.elements.clear();
		landmarks.elements.emplace_back(dv::Point3f(.2f, .2f, 1.f), 6, timestamp, dsc, "", cov, obs);
		landmarks.elements.emplace_back(dv::Point3f(-.2f, .2f, 1.f), 7, timestamp, dsc, "", cov, obs);
		landmarks.elements.emplace_back(dv::Point3f(-.2f, -.2f, 1.f), 8, timestamp, dsc, "", cov, obs);
		landmarks.elements.emplace_back(dv::Point3f(.2f, -.2f, 1.f), 9, timestamp, dsc, "", cov, obs);
		landmarks.elements.emplace_back(dv::Point3f(0.f, 0.f, 1.f), 10, timestamp, dsc, "", cov, obs);

		visualizer.accept(dv::kinematics::Transformationf(timestamp, poses.back().getTransform()));
		visualizer.accept(landmarks);
		expect(eq(visualizer.getNumberOfLandmarks(), 10));
		auto img = visualizer.generateFrame();
		if (showPreview) {
			cv::imshow("Preview", img.image);
			cv::waitKey(0);
		}
		landmarks.elements.clear();
		landmarks.elements.emplace_back(dv::Point3f(.3f, .3f, 1.f), 11, timestamp, dsc, "", cov, obs);
		visualizer.accept(landmarks);
		expect(le(visualizer.getNumberOfLandmarks(), 10));
		img = visualizer.generateFrame();
		if (showPreview) {
			cv::imshow("Preview", img.image);
			cv::waitKey(0);
		}
	};

	"resize_image"_test = [showPreview] {
		const cv::Size dims(640, 480);
		dvk::Transformationf pose;
		dvv::PoseVisualizer visualizer(dims);
		visualizer.accept(pose);

		if (showPreview) {
			std::cout << "Testing image resize during visualization" << std::endl;
		}

		dv::Frame output = visualizer.generateFrame();
		expect(eq(output.image.size(), dims));
		if (showPreview) {
			cv::imshow("Preview", output.image);
			cv::waitKey(33);
		}

		const cv::Size smaller(320, 240);
		visualizer.setResolution(smaller);
		output = visualizer.generateFrame();
		expect(eq(output.image.size(), smaller));
		if (showPreview) {
			cv::imshow("Preview", output.image);
			cv::waitKey(33);
		}

		const cv::Size larger(1280, 960);
		visualizer.setResolution(larger);
		output = visualizer.generateFrame();
		expect(eq(output.image.size(), larger));
		if (showPreview) {
			cv::imshow("Preview", output.image);
			cv::waitKey(33);
		}

		if (showPreview) {
			cv::imshow("Preview", visualizer.generateFrame().image);
			cv::waitKey(0);
		}
	};

	"changing_colors"_test = [showPreview] {
		const cv::Size dims(640, 480);
		dvk::Transformationf pose;
		dvv::PoseVisualizer visualizer(dims);
		visualizer.accept(pose);

		if (showPreview) {
			std::cout << "Testing changing of colors during visualization" << std::endl;
		}

		dv::Frame output = visualizer.generateFrame();
		expect(eq(output.image.size(), dims));
		if (showPreview) {
			cv::imshow("Preview", output.image);
			cv::waitKey(33);
		}

		cv::Vec<uint8_t, 3> newColor(255, 0, 0);
		visualizer.setBackgroundColor(newColor);
		output = visualizer.generateFrame();
		expect(eq(output.image.at<cv::Vec<uint8_t, 3>>(5, 5), newColor));
		if (showPreview) {
			cv::imshow("Preview", output.image);
			cv::waitKey(33);
		}

		newColor = cv::Vec<uint8_t, 3>(0, 255, 0);
		visualizer.setBackgroundColor(newColor);
		output = visualizer.generateFrame();
		expect(eq(output.image.at<cv::Vec<uint8_t, 3>>(5, 5), newColor));
		if (showPreview) {
			cv::imshow("Preview", output.image);
			cv::waitKey(33);
		}

		newColor = cv::Vec<uint8_t, 3>(0, 0, 255);
		visualizer.setBackgroundColor(newColor);
		output = visualizer.generateFrame();
		expect(eq(output.image.at<cv::Vec<uint8_t, 3>>(5, 5), newColor));
		if (showPreview) {
			cv::imshow("Preview", output.image);
			cv::waitKey(33);
		}

		// I'm not sure how to properly sample grid color. Maybe setting grid width
		// to some absurd size that the grid would consistently overwrite the background and then
		// it would be possible to easily sample the color?
		visualizer.setGridColor(newColor);
		output = visualizer.generateFrame();
		if (showPreview) {
			cv::imshow("Preview", output.image);
			cv::waitKey(0);
		}
	};

	if (showPreview) {
		cv::destroyWindow("Preview");
	}

	if (!interactive) {
		return EXIT_SUCCESS;
	}

	const cv::Size dims(640, 480);
	dvv::PoseVisualizer visualizer(dims);

	auto poses = generateTrajectory(50, 20000000);
	for (const auto &pose : poses) {
		visualizer.accept(pose);
	}

	cv::namedWindow("Trajectory", cv::WINDOW_NORMAL);

	// Define trackbars for camera pose
	auto cameraPosition                = visualizer.getCameraPosition<cv::Point3i>();
	auto cameraOrientation             = visualizer.getCameraOrientation<cv::Point3i>();
	cv::TrackbarCallback valueCallback = [](int newVal, void *oldVal) {
		auto *value = reinterpret_cast<int *>(oldVal);
		*value      = newVal;
	};
	cv::createTrackbar("x", "Trajectory", nullptr, 500, valueCallback, reinterpret_cast<void *>(&cameraPosition.x));
	cv::setTrackbarMin("x", "Trajectory", -500);
	cv::setTrackbarPos("x", "Trajectory", cameraPosition.x);

	cv::createTrackbar("y", "Trajectory", nullptr, 500, valueCallback, reinterpret_cast<void *>(&cameraPosition.y));
	cv::setTrackbarMin("y", "Trajectory", -500);
	cv::setTrackbarPos("y", "Trajectory", cameraPosition.y);

	cv::createTrackbar("z", "Trajectory", nullptr, 500, valueCallback, reinterpret_cast<void *>(&cameraPosition.z));
	cv::setTrackbarMin("z", "Trajectory", -500);
	cv::setTrackbarPos("z", "Trajectory", cameraPosition.z);

	cv::createTrackbar(
		"yaw [deg]", "Trajectory", nullptr, 180, valueCallback, reinterpret_cast<void *>(&cameraOrientation.x));
	cv::setTrackbarMin("yaw [deg]", "Trajectory", -180);
	cv::setTrackbarPos("yaw [deg]", "Trajectory", cameraOrientation.x);

	cv::createTrackbar(
		"pitch [deg]", "Trajectory", nullptr, 180, valueCallback, reinterpret_cast<void *>(&cameraOrientation.y));
	cv::setTrackbarMin("pitch [deg]", "Trajectory", -180);
	cv::setTrackbarPos("pitch [deg]", "Trajectory", cameraOrientation.y);

	cv::createTrackbar(
		"roll [deg]", "Trajectory", nullptr, 180, valueCallback, reinterpret_cast<void *>(&cameraOrientation.z));
	cv::setTrackbarMin("roll [deg]", "Trajectory", -180);
	cv::setTrackbarPos("roll [deg]", "Trajectory", cameraOrientation.z);

	// Define trackbars for enums (view mode and grid plane)
	dvv::PoseVisualizer::ViewMode cameraViewMode = visualizer.getViewMode();
	cv::TrackbarCallback viewModeCallback        = [](int newVal, void *oldVal) {
		auto *value = reinterpret_cast<dvv::PoseVisualizer::ViewMode *>(oldVal);
		*value      = dvv::PoseVisualizer::ViewMode(newVal);
	};
	cv::createTrackbar(
		"ViewMode [0-6]", "Trajectory", nullptr, 6, viewModeCallback, reinterpret_cast<void *>(&cameraViewMode));
	cv::setTrackbarPos("ViewMode [0-6]", "Trajectory", static_cast<int>(cameraViewMode));

	dvv::PoseVisualizer::GridPlane gridPlane = visualizer.getGridPlane();
	cv::TrackbarCallback gridPlaneCallback   = [](int newVal, void *oldVal) {
		auto *value = reinterpret_cast<dvv::PoseVisualizer::GridPlane *>(oldVal);
		*value      = dvv::PoseVisualizer::GridPlane(newVal);
	};
	cv::createTrackbar(
		"GridPlane [0-3]", "Trajectory", nullptr, 3, gridPlaneCallback, reinterpret_cast<void *>(&gridPlane));
	cv::setTrackbarPos("GridPlane [0-3]", "Trajectory", static_cast<int>(gridPlane));

	// Define trackbar for frame size
	int frameSize = visualizer.getObjectCoordinateSize();
	cv::createTrackbar("FrameSize", "Trajectory", nullptr, 10, valueCallback, reinterpret_cast<void *>(&frameSize));
	cv::setTrackbarPos("FrameSize", "Trajectory", frameSize);

	// Define trackbars for boolean options
	bool enableAutoScaling = visualizer.isAutoScalingEnabled();
	bool legendEnabled     = visualizer.isLegendVisualizationEnabled();

	cv::TrackbarCallback boolCallback = [](int newVal, void *oldVal) {
		auto *value = reinterpret_cast<bool *>(oldVal);
		*value      = newVal;
	};
	cv::createTrackbar(
		"Autoscale", "Trajectory", nullptr, 1, boolCallback, reinterpret_cast<void *>(&enableAutoScaling));
	cv::setTrackbarPos("Autoscale", "Trajectory", static_cast<int>(enableAutoScaling));

	cv::createTrackbar(
		"Enable legend", "Trajectory", nullptr, 1, boolCallback, reinterpret_cast<void *>(&legendEnabled));
	cv::setTrackbarPos("Enable legend", "Trajectory", static_cast<int>(legendEnabled));

	while (true) {
		visualizer.setCameraPosition(cameraPosition);
		visualizer.setCameraOrientation(cameraOrientation);
		visualizer.setViewMode(cameraViewMode);
		visualizer.setGridPlane(gridPlane);
		visualizer.setObjectCoordinateSize(frameSize);
		visualizer.setAutoScaling(enableAutoScaling);
		visualizer.setLegendVisualization(legendEnabled);

		const dv::Frame frame = visualizer.generateFrame();
		cv::imshow("Trajectory", frame.image);
		const int key = cv::waitKey(1);
		if (key > 0) {
			break;
		}
	}

	return EXIT_SUCCESS;
}

std::vector<dvk::Transformationf> generateTrajectory(const int64_t nPositions, const int64_t maxTime) {
	std::vector<dvk::Transformationf> poses;
	for (int64_t i = 0; i < nPositions; ++i) {
		int64_t ts = maxTime / nPositions * i;
		float tss  = static_cast<float>(ts) / 1e6f;
		// Feed some data to the plot
		auto r = Eigen::Vector3f(sin(tss) * (tss / 5 + 1), cos(tss) * (tss / 5 + 1), tss / 8);
		auto q = Eigen::Quaternionf(1, 0, 0, 0);
		poses.emplace_back(ts, r, q);
	}
	return poses;
}
