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

#include "color_calibration.hpp"

using namespace dv::color_calibration;

class FrameColorCalibrationModule : public dv::ModuleBase {
public:
	static void initInputs(dv::InputDefinitionList &in) {
		in.addFrameInput("frames");
	}

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

	static const char *initDescription() {
		return ("Color calibration module.");
	}

	static void initConfigOptions(dv::RuntimeConfig &config) {
		config.add("fileIn", dv::ConfigOption::fileOpenOption("Path to previous calibration.",
								 dv::portable_get_user_home_directory() + "/path_to_color_calibration.yml", "yml"));
		config.add("fileOut", dv::ConfigOption::fileSaveOption("Output path for saving calibration.",
								  dv::portable_get_user_home_directory() + "/path_to_color_calibration.yml", "yml"));
		config.add("runColorCorrection", dv::ConfigOption::boolOption("Run color correction.", false));
		config.add("save", dv::ConfigOption::buttonOption("Save color calibration.", "SAVE"));
		config.add("open", dv::ConfigOption::buttonOption("Open color calibration.", "OPEN"));
		config.add(
			"openDavis", dv::ConfigOption::buttonOption("Open default DAVIS color calibration.", "OPEN DEFAULT DAVIS"));
		config.add("reset", dv::ConfigOption::buttonOption("Reset color calibration matrix.", "RESET"));
		config.add("calibrate",
			dv::ConfigOption::buttonOption("Run color calibration and create color calibration", "CALIBRATE"));
		config.add("addColorChecker",
			dv::ConfigOption::buttonOption("Search for color calibration matrix.", "ADD COLOR CHECKER"));
		config.add("colorChecker", dv::ConfigOption::listOption("Color checker to be used for calibration.",
									   "SPYDER_CHECKER", {"SPYDER_CHECKER", "MACBETH_CHECKER"}));
		config.setPriorityOptions({"addColorChecker", "calibrate", "runColorCorrection", "fileIn", "fileOut", "save",
			"open", "openDavis", "reset"});
	}

	FrameColorCalibrationModule() {
		outputs.getFrameOutput("frames").setup(inputs.getFrameInput("frames"));
		initModel();
	}

	void configUpdate() override {
		if (config.getBool("reset")) {
			config.setBool("reset", false);
			config.setBool("runColorCorrection", false);
			initModel();
			log.info("Calibration reset.");
		}
		if (config.getBool("save")) {
			config.setBool("save", false);
			try {
				const std::filesystem::path path{config.getString("fileOut")};
				mCalibrationModel.saveCalibration(path);
				log.info("Calibration saved.");
			}
			catch (dv::exceptions::Exception &e) {
				log.warning(e.what());
				log.error("Could not save calibration.");
				return;
			}
		}
		if (config.getBool("open")) {
			config.setBool("open", false);
			try {
				const std::filesystem::path path{config.getString("fileIn")};
				mCalibrationModel.openCalibration(path);
				log.info("Previous calibration loaded.");
				config.setBool("runColorCorrection", true);
			}
			catch (dv::exceptions::Exception &e) {
				log.warning(e.what());
				log.error("Could not open calibration.");
				return;
			}
		}
		if (config.getBool("openDavis")) {
			config.setBool("openDavis", false);
			mCalibrationModel.openDavisDefaultCalibration();
			log.info("Default Davis color calibration loaded.");
			config.setBool("runColorCorrection", true);
		}
		if (config.getBool("calibrate")) {
			config.setBool("calibrate", false);
			try {
				const double loss = mCalibrationModel.prepareColorCorrection();
				log.info("Calibration success! Loss: " + std::to_string(loss));
				config.setBool("runColorCorrection", true);
			}
			catch (dv::exceptions::Exception &e) {
				log.warning(e.what());
				log.error("Calibration failed!");
				return;
			}
		}
		if (config.getBool("addColorChecker")) {
			config.setBool("addColorChecker", false);
			const auto input = inputs.getFrameInput("frames");
			if (input.data()) {
				const bool success = mCalibrationModel.findColorChecker(input.frame().getMatCopy());
				if (success) {
					log.info("Color checker found.");
				}
				else {
					log.warning("No color checker found in the frame.");
				}
			}
		}
	}

	void run() override {
		const auto input = inputs.getFrameInput("frames");
		auto output      = outputs.getFrameOutput("frames");

		if (input.data()) {
			const cv::Mat uncalibratedInput = input.frame().getMatCopy();
			if (config.getBool("runColorCorrection")) {
				try {
					const cv::Mat correctedInput = mCalibrationModel.correctColor(uncalibratedInput);
					output << correctedInput << dv::commit;
				}
				catch (dv::exceptions::Exception &e) {
					log.warning(e.what());
					log.error("Model not ready for inference. Run calibration first.");
					config.setBool("runColorCorrection", false);
				}
				return;
			}
			output << uncalibratedInput << dv::commit;
		}
	}

private:
	void initModel() {
		ColorCheckerName checker;
		if (config.getString("colorChecker") == "SPYDER_CHECKER") {
			checker = ColorCheckerName::SPYDER;
		}
		else if (config.getString("colorChecker") == "MACBETH_CHECKER") {
			checker = ColorCheckerName::MACBETH;
		}
		else {
			log.error("Unknown color checker supplied. Calibration not initialised.");
			return;
		}
		mCalibrationModel = ColorCalibration(checker);
	}

	ColorCalibration mCalibrationModel;
};

registerModuleClass(FrameColorCalibrationModule)
