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

#include <opencv2/imgproc.hpp>

class CropScale : public dv::ModuleBase {
private:
	int maxWidth;
	int maxHeight;

	int cropStartX;
	int cropStartY;
	int cropEndX;
	int cropEndY;
	int cropWidth;
	int cropHeight;

	bool resize;
	int resizeWidth;
	int resizeHeight;

public:
	static const char *initDescription() {
		return "Crop Events and Frames and resize";
	}

	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("cropOffsetX",
			dv::ConfigOption::intOption("Region-of-Interst offset value for the X axis.", 0, 0, (8192 - 1)));
		config.add("cropOffsetY",
			dv::ConfigOption::intOption("Region-of-Interst offset value for the Y axis.", 0, 0, (6144 - 1)));
		config.add("cropWidth", dv::ConfigOption::intOption("Region-of-Interst width (X axis).", 8192, 1, 8192));
		config.add("cropHeight", dv::ConfigOption::intOption("Region-of-Interst height (Y axis).", 6144, 1, 6144));

		config.add("resize", dv::ConfigOption::boolOption("Enable resize. Applied after cropping."));
		config.add("resizeWidth", dv::ConfigOption::intOption("The new width. Applied after cropping.", 640, 1, 8192));
		config.add(
			"resizeHeight", dv::ConfigOption::intOption("The new height. Applied after cropping.", 480, 1, 6144));

		config.setPriorityOptions(
			{"cropOffsetX", "cropOffsetY", "cropWidth", "cropHeight", "resize", "resizeWidth", "resizeHeight"});
	}

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

		// Check if events and frames have the same size
		if (inputs.isConnected("events") && inputs.isConnected("frames")) {
			if ((inputs.getEventInput("events").sizeX() != inputs.getFrameInput("frames").sizeX()
					|| (inputs.getEventInput("events").sizeY() != inputs.getFrameInput("frames").sizeY()))) {
				throw std::runtime_error("Events and Frames don't have the same size.");
			}
		}

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

			maxWidth  = inputs.getEventInput("events").sizeX();
			maxHeight = inputs.getEventInput("events").sizeY();
		}
		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"));

			maxWidth  = inputs.getFrameInput("frames").sizeX();
			maxHeight = inputs.getFrameInput("frames").sizeY();
		}
		else {
			// Remove unused output.
			dvModuleRegisterOutput(moduleData, "frames", "REMOVE");
		}

		// Reset ranges.
		moduleNode.createAttribute<dv::CfgType::INT>("cropOffsetX", 0, {0, maxWidth - 1}, dv::CfgFlags::NORMAL,
			moduleNode.getAttributeDescription<dv::CfgType::INT>("cropOffsetX"));
		moduleNode.createAttribute<dv::CfgType::INT>("cropOffsetY", 0, {0, maxHeight - 1}, dv::CfgFlags::NORMAL,
			moduleNode.getAttributeDescription<dv::CfgType::INT>("cropOffsetY"));

		moduleNode.createAttribute<dv::CfgType::INT>("cropWidth", maxWidth, {1, maxWidth}, dv::CfgFlags::NORMAL,
			moduleNode.getAttributeDescription<dv::CfgType::INT>("cropWidth"));
		moduleNode.createAttribute<dv::CfgType::INT>("cropHeight", maxHeight, {1, maxHeight}, dv::CfgFlags::NORMAL,
			moduleNode.getAttributeDescription<dv::CfgType::INT>("cropHeight"));

		moduleNode.createAttribute<dv::CfgType::INT>("resizeWidth", maxWidth, {1, maxWidth}, dv::CfgFlags::NORMAL,
			moduleNode.getAttributeDescription<dv::CfgType::INT>("resizeWidth"));
		moduleNode.createAttribute<dv::CfgType::INT>("resizeHeight", maxHeight, {1, maxHeight}, dv::CfgFlags::NORMAL,
			moduleNode.getAttributeDescription<dv::CfgType::INT>("resizeHeight"));
	}

	void run() override {
		if (inputs.isConnected("events") && inputs.getEventInput("events").events()) {
			const auto inEvent = inputs.getEventInput("events").events();
			auto outEvent      = outputs.getEventOutput("events").events();

			for (const auto &event : inEvent) {
				// Select only events within ROI.
				if ((event.x() >= cropStartX) && (event.x() <= cropEndX) && (event.y() >= cropStartY)
					&& (event.y() <= cropEndY)) {
					// Remap start coordinates to 0,0.
					const int16_t cropX = static_cast<int16_t>(event.x() - cropStartX);
					const int16_t cropY = static_cast<int16_t>(event.y() - cropStartY);

					if (resize) {
						const int16_t resizeX
							= static_cast<int16_t>((static_cast<float>(cropX) / static_cast<float>(cropWidth))
												   * static_cast<float>(resizeWidth));
						const int16_t resizeY
							= static_cast<int16_t>((static_cast<float>(cropY) / static_cast<float>(cropHeight))
												   * static_cast<float>(resizeHeight));

						outEvent.emplace_back(event.timestamp(), resizeX, resizeY, event.polarity());
					}
					else {
						outEvent.emplace_back(event.timestamp(), cropX, cropY, event.polarity());
					}
				}
			}

			outEvent.commit();
		}

		if (inputs.isConnected("frames") && inputs.getFrameInput("frames").frame()) {
			const auto inFrame = inputs.getFrameInput("frames").frame();
			auto outFrame      = outputs.getFrameOutput("frames").frame();

			// Incoming frame could already be an ROI frame itself, so
			// smaller than the full input.
			cv::Mat img;

			if (inFrame.size() != inputs.getFrameInput("frames").size()) {
				// 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.
				img = cv::Mat(inputs.getFrameInput("frames").size(), static_cast<int>(inFrame.format()), cv::Scalar{0});
				(*inFrame.getMatPointer()).copyTo(img(inFrame.roi()));
			}
			else {
				// Use input frame directly if full size.
				img = *inFrame.getMatPointer();
			}

			cv::Mat imgROI = img(cv::Rect(cropStartX, cropStartY, cropWidth, cropHeight));

			cv::Mat imgResize;

			if (resize) {
				cv::resize(imgROI, imgResize, cv::Size(resizeWidth, resizeHeight));
			}
			else {
				imgResize = imgROI;
			}

			outFrame.setMat(imgResize); // Force copy.
			outFrame.setPosition(0, 0); // Consistent with events remapping to 0,0.

			outFrame.setTimestamp(inFrame.timestamp());
			outFrame.setExposure(inFrame.exposure());
			outFrame.setSource(inFrame.source());

			outFrame.commit();
		}
	}

	~CropScale() 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 configUpdate() override {
		cropStartX = config.getInt("cropOffsetX");
		cropStartY = config.getInt("cropOffsetY");

		cropWidth  = config.getInt("cropWidth");
		cropHeight = config.getInt("cropHeight");

		if (cropWidth > (maxWidth - cropStartX)) {
			cropWidth = (maxWidth - cropStartX);
		}
		if (cropHeight > (maxHeight - cropStartY)) {
			cropHeight = (maxHeight - cropStartY);
		}

		cropEndX = cropStartX + cropWidth - 1;
		cropEndY = cropStartY + cropHeight - 1;

		config.setInt("cropWidth", cropWidth);
		config.setInt("cropHeight", cropHeight);

		resize = config.getBool("resize");

		resizeWidth  = config.getInt("resizeWidth");
		resizeHeight = config.getInt("resizeHeight");
	}
};

registerModuleClass(CropScale)
