#include "../../include/dv-sdk/data/event.hpp"
#include "../../include/dv-sdk/data/frame.hpp"
#include "../../include/dv-sdk/module.hpp"

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

#include <regex>

class Undistort : public dv::ModuleBase {
private:
	static inline const std::regex filenameCleanupRegex{"[^a-zA-Z-_\\d]"};

	std::string cameraID;
	cv::Size eventSize;
	cv::Size frameSize;

	std::vector<cv::Point2i> undistortEventMap;
	cv::Mat undistortFrameRemap1;
	cv::Mat undistortFrameRemap2;

public:
	static const char *initDescription() {
		return ("Remove distortion from lens in both frames and events.");
	}

	static void initInputs(dv::InputDefinitionList &in) {
		in.addEventInput("events", true);
		in.addFrameInput("frames", true);
	}

	static void initOutputs(dv::OutputDefinitionList &out) {
		out.addEventOutput("undistortedEvents");
		out.addFrameOutput("undistortedFrames");
	}

	static void initConfigOptions(dv::RuntimeConfig &config) {
		config.add("fitMorePixels",
			dv::ConfigOption::floatOption("How many input pixels to fit (alpha value in cv::getOptimalCameraMatrix "
										  "function), 0: maximize the image, 1: keep all source pixels.",
				0, 0, 1));
		config.add(
			"calibrationFile", dv::ConfigOption::fileOpenOption(
								   "The name of the file from which to load the calibration settings for undistortion.",
								   dv::portable_get_user_home_directory() + "/calibration_camera.xml", "xml"));

		config.setPriorityOptions({"calibrationFile"});
	}

	Undistort() {
		// Check if any inputs are connected.
		if (!inputs.isConnected("events") && !inputs.isConnected("frames")) {
			throw std::runtime_error("No input is connected.");
		}

		// If both are connected, they must be from the same origin.
		if (inputs.isConnected("events") && inputs.isConnected("frames")) {
			if (inputs.getEventInput("events").getOriginDescription()
				!= inputs.getFrameInput("frames").getOriginDescription()) {
				throw std::runtime_error("Both inputs are connected, but do not originate from the same camera.");
			}
		}

		if (inputs.isConnected("events")) {
			eventSize = inputs.getEventInput("events").size();

			// Populate event output info node, keep same as input info node.
			outputs.getEventOutput("undistortedEvents").setup(inputs.getEventInput("events"));

			// Generate camera ID from origin.
			setCameraID(inputs.getEventInput("events").getOriginDescription());
		}
		else {
			// Remove unused output.
			dvModuleRegisterOutput(moduleData, "undistortedEvents", "REMOVE");
		}

		if (inputs.isConnected("frames")) {
			frameSize = inputs.getFrameInput("frames").size();

			// Populate frame output info node, keep same as input info node.
			outputs.getFrameOutput("undistortedFrames").setup(inputs.getFrameInput("frames"));

			// Generate camera ID from origin.
			setCameraID(inputs.getFrameInput("frames").getOriginDescription());
		}
		else {
			// Remove unused output.
			dvModuleRegisterOutput(moduleData, "undistortedFrames", "REMOVE");
		}

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

	~Undistort() override {
		// Restore unused outputs.
		if (!inputs.isConnected("events")) {
			dvModuleRegisterOutput(moduleData, "undistortedEvents", dv::EventPacket::TableType::identifier);
		}

		if (!inputs.isConnected("frames")) {
			dvModuleRegisterOutput(moduleData, "undistortedFrames", dv::Frame::TableType::identifier);
		}
	}

	void run() override {
		// Undistortion can be applied to both frames and events.
		const auto events_in = inputs.getEventInput("events").events();
		const auto frame_in  = inputs.getFrameInput("frames").frame();

		if (events_in) {
			auto events_out = outputs.getEventOutput("undistortedEvents").events();
			undistortEvents(events_in, events_out);
		}

		if (frame_in) {
			auto frame_out = outputs.getFrameOutput("undistortedFrames").frame();
			undistortFrame(frame_in, frame_out);
		}
	}

	void configUpdate() override {
		// Any changes to configuration mean the calibration has to be
		// reloaded and reinitialized. On file-related failure, we stop.
		if (!loadUndistortMatrices(config.getString("calibrationFile"))) {
			config.setBool("running", false);
		}
	}

private:
	static bool cvExists(const cv::FileNode &fn) {
		if (fn.type() == cv::FileNode::NONE) {
			return (false);
		}

		return (true);
	}

	void setCameraID(const std::string &originDescription) {
		// Camera origin descriptions are fine, never start with a digit or have spaces
		// or other special characters in them that fail with OpenCV's FileStorage.
		// So we just clean/escape the string for possible other sources.
		auto str = std::regex_replace(originDescription, filenameCleanupRegex, "_");

		// FileStorage node names can't start with - or 0-9.
		// Prefix with an underscore in that case.
		if ((str[0] == '-') || (std::isdigit(str[0]) != 0)) {
			str = "_" + str;
		}

		cameraID = str;
	}

	bool loadUndistortMatrices(const std::string &filename) {
		// loads single camera calibration file
		if (filename.empty()) {
			log.error << "No camera calibration file specified." << std::endl;
			return (false);
		}

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

		if (!fs.isOpened()) {
			log.error << "Impossible to load the camera calibration file: " << filename << std::endl;
			return (false);
		}

		// Load matrices from old and new calibration file formats.
		cv::Mat undistortCameraMatrix;
		cv::Mat undistortDistCoeffs;

		auto typeNode = fs["type"];

		if (!cvExists(typeNode) || !typeNode.isString()) {
			log.info << "Old-style camera calibration file found." << std::endl;

			if (!cvExists(fs["camera_matrix"]) || !cvExists(fs["distortion_coefficients"])) {
				log.error.format("Calibration data not present in file: {:s}", filename);
				return (false);
			}

			fs["camera_matrix"] >> undistortCameraMatrix;
			fs["distortion_coefficients"] >> undistortDistCoeffs;

			log.info.format("Loaded camera matrix and distortion coefficients from file: {:s}", filename);
		}
		else {
			log.info << "New-style camera calibration file found." << std::endl;

			auto cameraNode = fs[cameraID];

			if (!cvExists(cameraNode) || !cameraNode.isMap()) {
				log.error.format("Calibration data for camera {:s} not present in file: {:s}", cameraID, filename);
				return (false);
			}

			if (!cvExists(cameraNode["camera_matrix"]) || !cvExists(cameraNode["distortion_coefficients"])) {
				log.error.format("Calibration data for camera {:s} not present in file: {:s}", cameraID, filename);
				return (false);
			}

			cameraNode["camera_matrix"] >> undistortCameraMatrix;
			cameraNode["distortion_coefficients"] >> undistortDistCoeffs;

			log.info.format(
				"Loaded camera matrix and distortion coefficients for camera {:s} from file: {:s}", cameraID, filename);
		}

		// Determine if fisheye lens mode should be enabled.
		bool useFisheyeModel = false;

		if (cvExists(fs["use_fisheye_model"])) {
			fs["use_fisheye_model"] >> useFisheyeModel;
		}

		const double alpha = static_cast<double>(config.getFloat("fitMorePixels"));

		if (inputs.getEventInput("events").isConnected()) {
			// Allocate undistort events maps.
			std::vector<cv::Point2f> undistortEventInputMap;
			std::vector<cv::Point2f> undistortEventOutputMap;

			undistortEventInputMap.reserve(static_cast<size_t>(eventSize.area()));
			undistortEventOutputMap.reserve(static_cast<size_t>(eventSize.area()));

			// Populate undistort events input map with all possible (x, y) address combinations.
			for (int y = 0; y < eventSize.height; y++) {
				for (int x = 0; x < eventSize.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::Mat optimalCameraMatrixEvents;

			if (useFisheyeModel) {
				cv::fisheye::estimateNewCameraMatrixForUndistortRectify(undistortCameraMatrix, undistortDistCoeffs,
					eventSize, cv::Matx33d::eye(), optimalCameraMatrixEvents, alpha);

				cv::fisheye::undistortPoints(undistortEventInputMap, undistortEventOutputMap, undistortCameraMatrix,
					undistortDistCoeffs, cv::Matx33d::eye(), optimalCameraMatrixEvents);
			}
			else {
				optimalCameraMatrixEvents
					= cv::getOptimalNewCameraMatrix(undistortCameraMatrix, undistortDistCoeffs, eventSize, alpha);
				cv::undistortPoints(undistortEventInputMap, undistortEventOutputMap, undistortCameraMatrix,
					undistortDistCoeffs, cv::Matx33d::eye(), optimalCameraMatrixEvents);
			}

			// Convert undistortEventOutputMap to integer from float for faster calculations later on.
			undistortEventMap.clear();
			undistortEventMap.reserve(static_cast<size_t>(eventSize.area()));

			for (size_t i = 0; i < undistortEventOutputMap.size(); i++) {
				const cv::Point2i eventUndistort = undistortEventOutputMap.at(i);

				// Check that new coordinates are still within view boundary. If yes, use new remapped coordinates.
				if ((eventUndistort.x >= 0) && (eventUndistort.x < eventSize.width) && (eventUndistort.y >= 0)
					&& (eventUndistort.y < eventSize.height)) {
					undistortEventMap.push_back(eventUndistort);
				}
				else {
					// Else push placeholder values.
					undistortEventMap.emplace_back(-1, -1);
				}
			}
		}

		if (inputs.getFrameInput("frames").isConnected()) {
			cv::Mat optimalCameraMatrixFrames;

			if (useFisheyeModel) {
				cv::fisheye::estimateNewCameraMatrixForUndistortRectify(undistortCameraMatrix, undistortDistCoeffs,
					frameSize, cv::Matx33d::eye(), optimalCameraMatrixFrames, alpha);

				cv::fisheye::initUndistortRectifyMap(undistortCameraMatrix, undistortDistCoeffs, cv::Matx33d::eye(),
					optimalCameraMatrixFrames, frameSize, CV_16SC2, undistortFrameRemap1, undistortFrameRemap2);
			}
			else {
				optimalCameraMatrixFrames
					= cv::getOptimalNewCameraMatrix(undistortCameraMatrix, undistortDistCoeffs, frameSize, alpha);

				cv::initUndistortRectifyMap(undistortCameraMatrix, undistortDistCoeffs, cv::Matx33d::eye(),
					optimalCameraMatrixFrames, frameSize, CV_16SC2, undistortFrameRemap1, undistortFrameRemap2);
			}
		}

		return (true);
	}

	void undistortEvents(const dv::InputVectorDataWrapper<dv::EventPacket, dv::Event> &in,
		dv::OutputVectorDataWrapper<dv::EventPacket, dv::Event> &out) {
		for (const auto &evt : in) {
			// Get new coordinates at which event shall be remapped.
			size_t mapIdx              = static_cast<size_t>((evt.y() * eventSize.width) + evt.x());
			cv::Point2i eventUndistort = undistortEventMap.at(mapIdx);

			// Jump over events that would resolve to a placeholder, signifying out-of-bounds.
			if (eventUndistort.x == -1) {
				continue;
			}

			out.emplace_back(evt.timestamp(), static_cast<int16_t>(eventUndistort.x),
				static_cast<int16_t>(eventUndistort.y), evt.polarity());
		}

		out.commit();
	}

	void undistortFrame(const dv::InputDataWrapper<dv::Frame> &in, dv::OutputDataWrapper<dv::Frame> &out) {
		// Setup output frame. Same size.
		out.setPosition(in.position());
		out.setTimestamp(in.timestamp());
		out.setExposure(in.exposure());
		out.setSource(in.source());

		// Get input OpenCV Mat. Lifetime is properly managed.
		auto input = in.getMatPointer();

		// Get output OpenCV Mat. Memory must have been allocated already.
		cv::Mat output;

		cv::remap(*input, output, undistortFrameRemap1, undistortFrameRemap2, cv::INTER_CUBIC, cv::BORDER_CONSTANT);

		out.setMat(std::move(output));
		out.commit();
	}
};

registerModuleClass(Undistort)
