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

#include <opencv2/imgproc.hpp>

class Accumulator : public dv::ModuleBase {
private:
	dv::EventStreamSlicer slicer;
	dv::Accumulator frameAccumulator;
	int sliceJob;
	int demosaic             = 0;
	int64_t accumulationTime = -1;
	int64_t currentFrameTime = -1;
	std::string sliceMethod;

public:
	static const char *initDescription() {
		return "Accumulates events into a frame. Provides various configurations to tune the integration process";
	}

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

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

	static void initConfigOptions(dv::RuntimeConfig &configuration) {
		configuration.add(
			"sliceMethod", dv::ConfigOption::listOption(
							   "Method for slicing events, elapsed time between frames or number of events per frame",
							   0, {"TIME", "NUMBER"}));
		configuration.add("rectifyPolarity", dv::ConfigOption::boolOption("All events have positive contribution"));
		configuration.add("eventContribution",
			dv::ConfigOption::floatOption("The contribution of a single event", 0.15f, 0.0f, 1.0f));
		configuration.add(
			"maxPotential", dv::ConfigOption::floatOption("Value at which to clip the integration", 1.0f, 0.0f, 1.0f));
		configuration.add("neutralPotential",
			dv::ConfigOption::floatOption("Value to which the decay tends over time", 0.0f, 0.0f, 1.0f));
		configuration.add(
			"minPotential", dv::ConfigOption::floatOption("Value at which to clip the integration", 0.0f, 0.0f, 1.0f));
		configuration.add("decayFunction", dv::ConfigOption::listOption("The decay function to be used", 2,
											   {"None", "Linear", "Exponential", "Step"}));
		configuration.add(
			"decayParam", dv::ConfigOption::doubleOption(
							  "Slope for linear decay, tau for exponential decay, time for step decay", 1e6, 0, 1e10));
		configuration.add("synchronousDecay", dv::ConfigOption::boolOption("Decay at frame generation time"));
		configuration.add(
			"accumulationTime", dv::ConfigOption::intOption("Time in ms to accumulate events over", 33, 1, 1000));
		configuration.add("accumulationNumber",
			dv::ConfigOption::intOption("Number of events to accumlate for a frame", 100000, 1, INT_MAX));
		configuration.add("colorDemosaicing", dv::ConfigOption::boolOption("Apply color demosaicing filter", false));

		configuration.setPriorityOptions(
			{"sliceMethod", "decayFunction", "decayParam", "synchronousDecay", "colorDemosaicing"});
	}

	void elaborateFrame(const dv::EventStore &events, const dv::Duration exposure) {
		frameAccumulator.accumulate(events);
		// generate frame
		auto frame = frameAccumulator.generateFrame();

		auto newFrame = outputs.getFrameOutput("frames").frame();
		newFrame.setTimestamp(currentFrameTime);
		newFrame.setExposure(exposure);
		newFrame.setSource(dv::FrameSource::ACCUMULATION);

		if (demosaic) {
			cv::Mat demosaicFrame;
			cv::cvtColor(frame.image, demosaicFrame, demosaic);
			newFrame.setMat(std::move(demosaicFrame));
			newFrame.commit();
		}
		else {
			newFrame.setMat(std::move(frame.image));
			newFrame.commit();
		}
	}

	void doPerFrameTime(const dv::EventStore &events) {
		if (currentFrameTime < 0) {
			currentFrameTime = events.getLowestTime();
		}

		elaborateFrame(events, dv::Duration{accumulationTime});
		currentFrameTime += accumulationTime;
	}

	void doPerEventNumber(const dv::EventStore &events) {
		currentFrameTime = events.getLowestTime();
		elaborateFrame(events, dv::Duration{events.getHighestTime() - events.getLowestTime()});
	}

	Accumulator() : slicer(dv::EventStreamSlicer()) {
		sliceMethod = config.getString("sliceMethod");

		outputs.getFrameOutput("frames").setup(inputs.getEventInput("events"));
		frameAccumulator = dv::Accumulator(inputs.getEventInput("events").size(), dv::Accumulator::Decay::EXPONENTIAL,
			1e6, false, 0.04f, 0.3f, 0.0f, 0.0f, false);

		if (sliceMethod == "TIME") {
			accumulationTime = config.getInt("accumulationTime") * 1000;
			sliceJob         = slicer.doEveryTimeInterval(dv::Duration{accumulationTime},
				std::function<void(const dv::EventStore &)>(
					std::bind(&Accumulator::doPerFrameTime, this, std::placeholders::_1)));
		}
		else if (sliceMethod == "NUMBER") {
			sliceJob = slicer.doEveryNumberOfElements(config.getInt("accumulationNumber"),
				std::function<void(const dv::EventStore &)>(
					std::bind(&Accumulator::doPerEventNumber, this, std::placeholders::_1)));
		}
		else {
			// This condition should be impossible.
			throw std::runtime_error("Unknown slice method selected.");
		}
	}

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

	static dv::Accumulator::Decay decayFromString(const std::string &name) {
		if (name == "Linear")
			return dv::Accumulator::Decay::LINEAR;
		if (name == "Exponential")
			return dv::Accumulator::Decay::EXPONENTIAL;
		if (name == "Step")
			return dv::Accumulator::Decay::STEP;
		return dv::Accumulator::Decay::NONE;
	}

	void configUpdate() override {
		frameAccumulator.setIgnorePolarity(config.getBool("rectifyPolarity"));
		frameAccumulator.setEventContribution(config.getFloat("eventContribution"));
		frameAccumulator.setMaxPotential(config.getFloat("maxPotential"));
		frameAccumulator.setNeutralPotential(config.getFloat("neutralPotential"));
		frameAccumulator.setMinPotential(config.getFloat("minPotential"));
		frameAccumulator.setDecayFunction(decayFromString(config.getString("decayFunction")));
		frameAccumulator.setDecayParam(config.getDouble("decayParam"));
		frameAccumulator.setSynchronousDecay(config.getBool("synchronousDecay"));

		if (sliceMethod != config.getString("sliceMethod")) {
			log.warning << "Please restart the module for 'sliceMethod' changes to take effect." << std::endl;
		}

		if (sliceMethod == "TIME") {
			accumulationTime = config.getInt("accumulationTime") * 1000;
			slicer.modifyTimeInterval(sliceJob, dv::Duration{accumulationTime});
		}
		else {
			slicer.modifyNumberInterval(sliceJob, config.getInt("accumulationNumber"));
		}

		if (config.getBool("colorDemosaicing")) {
			if (inputs.getEventInput("events").infoNode().exists<dv::CfgType::INT>("colorFilter")) {
				demosaic = cv::COLOR_BayerBG2BGR + inputs.getEventInput("events").infoNode().getInt("colorFilter");
			}
			else {
				// Turn off again if not supported.
				config.setBool("colorDemosaicing", false);
				demosaic = 0;
			}
		}
		else {
			demosaic = 0;
		}
	}
};

registerModuleClass(Accumulator)
