#include <dv-sdk/module.hpp>

class EventVisualizer : public dv::ModuleBase {
private:
	static int getColorHex(const std::string &hexColor) {
		auto c = std::stoi(hexColor, nullptr, 16);
		if ((c < 0) || (c > 255)) {
			c = 0;
		}
		return c;
	}

	dv::EventStreamSlicer slicer;
	cv::Mat frame;
	cv::Vec3b backgroundColor;
	cv::Vec3b positiveColor;
	cv::Vec3b negativeColor;
	int64_t lastTimestamp{0};

	void renderFrame(const dv::EventStore &store) {
		frame = backgroundColor;

		for (const auto &event : store) {
			frame.at<cv::Vec3b>(event.y(), event.x()) = event.polarity() ? positiveColor : negativeColor;
		}

		lastTimestamp = (store.isEmpty()) ? (lastTimestamp) : (store.getLowestTime());

		auto frameOutput = outputs.getFrameOutput("frames").frame();
		frameOutput.setTimestamp(lastTimestamp);
		frameOutput.setExposure(std::chrono::milliseconds{static_cast<int64_t>(1000.0f / 30.0f)});
		frameOutput.setPosition(0, 0);
		frameOutput.setSource(dv::FrameSource::VISUALIZATION);
		frameOutput.setMat(frame);
		frameOutput.commit();
	}

public:
	static void initInputs(dv::InputDefinitionList &in) {
		in.addEventInput("events");
	}

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

	static const char *initDescription() {
		return ("Simple event visualizer.");
	}

	static void initConfigOptions(dv::RuntimeConfig &config) {
		config.add("backgroundColor",
			dv::ConfigOption::stringOption("Background color in hex format #RRGGBB.", "FFFFFF", 6, 6));
		config.add("positiveColor",
			dv::ConfigOption::stringOption("Positive (ON) event color in hex format #RRGGBB.", "005DB7", 6, 6));
		config.add("negativeColor",
			dv::ConfigOption::stringOption("Negative (OFF) event color in hex format #RRGGBB.", "2B2B2B", 6, 6));
	}

	EventVisualizer() {
		slicer.doEveryTimeInterval(std::chrono::milliseconds{static_cast<int64_t>(1000.0f / 30.0f)},
			std::function<void(const dv::EventStore &)>(
				std::bind(&EventVisualizer::renderFrame, this, std::placeholders::_1)));

		frame = cv::Mat{inputs.getEventInput("events").size(), CV_8UC3};

		outputs.getFrameOutput("frames").setup(inputs.getEventInput("events"));
	}

	void run() override {
		slicer.accept(inputs.getEventInput("events").events());
	}

	void configUpdate() override {
		backgroundColor(0) = getColorHex(config.getString("backgroundColor").substr(4, 2));
		backgroundColor(1) = getColorHex(config.getString("backgroundColor").substr(2, 2));
		backgroundColor(2) = getColorHex(config.getString("backgroundColor").substr(0, 2));

		positiveColor(0) = getColorHex(config.getString("positiveColor").substr(4, 2));
		positiveColor(1) = getColorHex(config.getString("positiveColor").substr(2, 2));
		positiveColor(2) = getColorHex(config.getString("positiveColor").substr(0, 2));

		negativeColor(0) = getColorHex(config.getString("negativeColor").substr(4, 2));
		negativeColor(1) = getColorHex(config.getString("negativeColor").substr(2, 2));
		negativeColor(2) = getColorHex(config.getString("negativeColor").substr(0, 2));
	}
};

registerModuleClass(EventVisualizer)
