#include "../../include/dv-sdk/module.hpp"

#include <opencv2/imgproc.hpp>

#include <numbers>

class FlipRotate : public dv::ModuleBase {
private:
	bool eventsFlipVertically   = false;
	bool eventsFlipHorizontally = false;
	bool framesFlipVertically   = false;
	bool framesFlipHorizontally = false;

	float degreeOfRotation = 0.0;
	float cosOfDegree      = 0.0;
	float sinOfDegree      = 0.0;

	cv::Size eventSize;
	cv::Point2f eventCenter;

	cv::Size frameSize;
	cv::Point2f frameCenter;

public:
	static const char *initDescription() {
		return "Flip and rotate events and frames horizontally, vertically or by setting a degree of rotation.";
	}

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

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

	static void initConfigOptions(dv::RuntimeConfig &config) {
		config.add("eventFlipVertically", dv::ConfigOption::boolOption("Flip events vertically"));
		config.add("eventFlipHorizontally", dv::ConfigOption::boolOption("Flip events horizontally"));
		config.add("frameFlipVertically", dv::ConfigOption::boolOption("Flip frames vertically"));
		config.add("frameFlipHorizontally", dv::ConfigOption::boolOption("Flip frames horizontally"));
		config.add("degreeOfRotation",
			dv::ConfigOption::floatOption("Degree of rotation from -180° to 180°", 0.0f, -180.0f, 180.0f));

		config.setPriorityOptions({"degreeOfRotation", "eventFlipVertically", "eventFlipHorizontally",
			"frameFlipVertically", "frameFlipHorizontally"});
	}

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

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

			eventSize   = inputs.getEventInput("events").size();
			eventCenter = cv::Point2f{eventSize} / 2;
		}
		else {
			// Remove unused output.
			dvModuleRegisterOutput(moduleData, "events", "REMOVE");
		}

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

			frameSize   = inputs.getFrameInput("frames").size();
			frameCenter = cv::Point2f{frameSize} / 2;
		}
		else {
			// Remove unused output.
			dvModuleRegisterOutput(moduleData, "frames", "REMOVE");
		}
	}

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

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

	void run() override {
		//------------------------------------------------------------------------
		// Events management
		const auto events = inputs.getEventInput("events").events();

		if (events) {
			auto outEvent = outputs.getEventOutput("events").events();

			for (const auto &event : events) {
				int16_t x;
				int16_t y;

				// flip of events
				if (eventsFlipHorizontally) {
					x = static_cast<int16_t>(eventSize.width - 1 - event.x());
				}
				else {
					x = event.x();
				}

				if (eventsFlipVertically) {
					y = static_cast<int16_t>(eventSize.height - 1 - event.y());
				}
				else {
					y = event.y();
				}

				// rotation of events
				if (double(degreeOfRotation) != 0.0) {
					float dx = static_cast<float>(x - eventCenter.x);
					float dy = static_cast<float>(y - eventCenter.y);

					x = static_cast<int16_t>((dx * cosOfDegree) - (dy * sinOfDegree) + eventCenter.x);
					y = static_cast<int16_t>((dx * sinOfDegree) + (dy * cosOfDegree) + eventCenter.y);

					if ((x < 0) || (x >= eventSize.width)) {
						continue;
					}
					if ((y < 0) || (y >= eventSize.height)) {
						continue;
					}
				}

				outEvent.emplace_back(event.timestamp(), x, y, event.polarity());
			}

			outEvent.commit();
		}

		//------------------------------------------------------------------------
		// Frames management
		const auto frame = inputs.getFrameInput("frames").frame();

		if (frame) {
			auto outFrame = outputs.getFrameOutput("frames").frame();

			// Setup output frame.
			outFrame.setPosition(frame.position());
			outFrame.setTimestamp(frame.timestamp());
			outFrame.setExposure(frame.exposure());
			outFrame.setSource(frame.source());

			// flip of frame
			cv::Mat frameFlip;

			if (framesFlipVertically && framesFlipHorizontally) {
				cv::flip(*frame.getMatPointer(), frameFlip, -1);
			}
			else if (framesFlipVertically && !framesFlipHorizontally) {
				cv::flip(*frame.getMatPointer(), frameFlip, 0);
			}
			else if (!framesFlipVertically && framesFlipHorizontally) {
				cv::flip(*frame.getMatPointer(), frameFlip, 1);
			}
			else {
				frameFlip = *frame.getMatPointer();
			}

			// rotation of frame
			cv::Mat frameRotate;

			if (degreeOfRotation != 0) {
				cv::Mat frameRotateInput;

				if (frameFlip.size() != frameSize) {
					// ROI frame. Only the actual pixels are in data(), but cropping
					// is based on the full original size frame as given during setup.
					// So if ROI is active, we make an original sized black frame and
					// copy the ROI region to it at the right position, consistent
					// with the event output.
					frameRotateInput = cv::Mat(frameSize, frameFlip.type(), cv::Scalar{0});
					frameFlip.copyTo(frameRotateInput(frame.roi()));
				}
				else {
					// Use input frame directly if full size.
					frameRotateInput = frameFlip;
				}

				const auto M = cv::getRotationMatrix2D(frameCenter, static_cast<double>(-degreeOfRotation), 1);
				cv::warpAffine(frameRotateInput, frameRotate, M, frameSize);

				// Frame is always expanded for rotation to max input size,
				// so we reset the ROI position to the origin 0,0.
				outFrame.setPosition(0, 0);
			}
			else {
				frameRotate = frameFlip;
			}

			outFrame.setMat(frameRotate);
			outFrame.commit();
		}
	}

	void configUpdate() override {
		eventsFlipVertically   = config.getBool("eventFlipVertically");
		eventsFlipHorizontally = config.getBool("eventFlipHorizontally");
		framesFlipVertically   = config.getBool("frameFlipVertically");
		framesFlipHorizontally = config.getBool("frameFlipHorizontally");
		degreeOfRotation       = config.getFloat("degreeOfRotation");

		cosOfDegree = std::cos(degreeOfRotation * std::numbers::pi_v<float> / 180.0f);
		sinOfDegree = std::sin(degreeOfRotation * std::numbers::pi_v<float> / 180.0f);
	}
};

registerModuleClass(FlipRotate)
