#include "stereo_rectification_events.hpp"

#include <opencv2/calib3d.hpp>
#include <opencv2/imgproc.hpp>

const char *StereoRectificationEvents::initDescription() {
	return ("Perform stereo rectification on events.");
}

void StereoRectificationEvents::initInputs(dv::InputDefinitionList &in) {
	in.addEventInput("camera0");
	in.addEventInput("camera1");
}

void StereoRectificationEvents::initOutputs(dv::OutputDefinitionList &out) {
	out.addEventOutput("rectifiedCamera0");
	out.addEventOutput("rectifiedCamera1");
}

StereoRectificationEvents::StereoRectificationEvents() {
	for (size_t i = 0; i < NUM_CAMERAS; i++) {
		// Populate event output info node, keep same as input info node.
		outputs.getEventOutput("rectifiedCamera" + std::to_string(i))
			.setup(inputs.getEventInput("camera" + std::to_string(i)));
	}

	// Load calibration from file.
	if (!loadStereoRectificationMatrices(config.getString("calibrationFile"))) {
		throw std::runtime_error("Could not load calibration data.");
	}
}

void StereoRectificationEvents::run() {
	for (size_t i = 0; i < NUM_CAMERAS; i++) {
		const auto inputEvents = inputs.getEventInput("camera" + std::to_string(i)).events();

		if (inputEvents) {
			auto outputEvents = outputs.getEventOutput("rectifiedCamera" + std::to_string(i)).events();
			stereoRectifyEvents(inputEvents, outputEvents, i);
		}
	}
}

void StereoRectificationEvents::initStereoRectification(
	const CameraInfo loaded[NUM_CAMERAS], const cv::Mat R, const cv::Mat T, const double alpha) {
	cv::Mat RN[NUM_CAMERAS], CN[NUM_CAMERAS], Q;

	cv::stereoRectify(loaded[0].cameraMatrix, loaded[0].distCoeffs, loaded[1].cameraMatrix, loaded[1].distCoeffs,
		size[0], R, T, RN[0], RN[1], CN[0], CN[1], Q, cv::CALIB_ZERO_DISPARITY, alpha, size[0], nullptr, nullptr);

	for (size_t i = 0; i < NUM_CAMERAS; i++) {
		// Allocate undistort events maps.
		std::vector<cv::Point2f> undistortEventInputMap;
		std::vector<cv::Point2f> undistortEventOutputMap;

		undistortEventInputMap.reserve(static_cast<size_t>(size[i].area()));
		undistortEventOutputMap.reserve(static_cast<size_t>(size[i].area()));

		// Populate undistort events input map with all possible (x, y) address combinations.
		for (int y = 0; y < size[i].height; y++) {
			for (int x = 0; x < size[i].width; x++) {
				// Use center of pixel to get better approximation, since we're using floats anyway.
				undistortEventInputMap.push_back(
					cv::Point2f(static_cast<float>(x) + 0.5f, static_cast<float>(y) + 0.5f));
			}
		}

		cv::undistortPoints(undistortEventInputMap, undistortEventOutputMap, loaded[i].cameraMatrix,
			loaded[i].distCoeffs, RN[i], CN[i]);

		// Convert undistortEventOutputMap to integer from float for faster calculations later on.
		remapEvents[i].clear();
		remapEvents[i].reserve(undistortEventOutputMap.size());

		for (const auto &val : undistortEventOutputMap) {
			remapEvents[i].push_back(val);
		}

		remapEvents[i].shrink_to_fit();
	}
}

void StereoRectificationEvents::stereoRectifyEvents(const dv::InputVectorDataWrapper<dv::EventPacket, dv::Event> &in,
	dv::OutputVectorDataWrapper<dv::EventPacket, dv::Event> &out, size_t cam) {
	for (const auto &evt : in) {
		// Get new coordinates at which event shall be remapped.
		size_t mapIdx                        = static_cast<size_t>((evt.y() * size[cam].width) + evt.x());
		cv::Point2i eventStereoRectification = remapEvents[cam].at(mapIdx);

		// Check that new coordinates are still within view boundary. If yes, use new remapped coordinates.
		if (eventStereoRectification.x >= 0 && eventStereoRectification.x < size[cam].width
			&& eventStereoRectification.y >= 0 && eventStereoRectification.y < size[cam].height) {
			out.emplace_back(evt.timestamp(), static_cast<int16_t>(eventStereoRectification.x),
				static_cast<int16_t>(eventStereoRectification.y), evt.polarity());
		}
	}

	out.commit();
}

registerModuleClass(StereoRectificationEvents)
