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

#include "calibration.hpp"
#include "cameracalibration.hpp"
#include "stereocalibration.hpp"

class CalibrationModule : public dv::ModuleBase {
private:
	std::unique_ptr<Calibration> calibration;

public:
	static void initInputs(dv::InputDefinitionList &in) {
		in.addFrameInput("input1");
		// second input is optional
		in.addFrameInput("input2", true);
	}

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

	static const char *initDescription() {
		return ("Perform lens calibration for one camera or stereo calibration for two cameras.");
	}

	static void initConfigOptions(dv::RuntimeConfig &config) {
		config.add("useDefaultFilename",
			dv::ConfigOption::boolOption("Use the default file names, calibration_camera.xml or "
										 "calibration_stereo.xml, without the camera ID and date/time information in "
										 "the file name, which may override existing files with the same name.",
				false, false));
		config.add("consecDetects",
			dv::ConfigOption::intOption("Number of consecutive detections for valid image.", 10, 1, 40));
		config.add("minDetections",
			dv::ConfigOption::intOption("Minimum number of valid detections to use for calibration.", 30, 20, 100));

		config.add(
			"calibrationPattern", dv::ConfigOption::listOption("Pattern to run calibration with.", 0,
									  std::vector<std::string>{"chessboard", "circlesGrid", "asymmetricCirclesGrid"}));
		config.add("useDefaultPattern",
			dv::ConfigOption::buttonOption("Set board sizes to the proper default ones for the selected pattern as "
										   "provided in iniVation's documentation.",
				"Use Default Pattern"));
		config.add("boardWidth", dv::ConfigOption::intOption("Number of pattern points in the x-axis.", 9, 1, 64));
		config.add("boardHeight", dv::ConfigOption::intOption("Number of pattern points in the y-axis.", 6, 1, 64));
		config.add("boardSquareSize",
			dv::ConfigOption::floatOption(
				"Size of one square in your defined unit (point, millimeter, etc.).", 30.0F, 0.0F, 1000.0F));
		config.add("useFisheyeModel",
			dv::ConfigOption::boolOption("Use fisheye lens model for single camera calibration.", false));
		config.add("maxReprojectionError",
			dv::ConfigOption::floatOption("Maximum total reprojection error allowed (in pixels).", 1.0F, 0.0F, 10.0F));
		config.add("maxEpipolarLineError",
			dv::ConfigOption::floatOption(
				"Maximum mean epipolar line error allowed, epipolar line error is the sum of the distance between the "
				"point and the epipolar line (in pixels) in the two images.",
				5.0F, 0.0F, 10.0F));

		config.add("highlightArea",
			dv::ConfigOption::boolOption("Set to true if already covered area should be highlighted in green.", true));
		config.add(
			"drawEpipolarLines", dv::ConfigOption::boolOption("Draw epipolar lines after stereo calibration.", false));
		config.add("undistortOutput", dv::ConfigOption::boolOption("Undistort output images after calibration.", true));

		config.add("keep", dv::ConfigOption::buttonOption("Press to keep image for calibration.", "Keep"));
		config.add("discard", dv::ConfigOption::buttonOption("Press to discard image for calibration.", "Discard"));

		config.add(
			"checkImages", dv::ConfigOption::buttonOption("Check all images collected until now.", "Check Images"));
		config.add(
			"clearImages", dv::ConfigOption::buttonOption("Remove all images collected until now.", "Clear Images"));
		config.add("calibrateNow", dv::ConfigOption::buttonOption("Force calibration execution now.", "Calibrate Now"));
		config.add("saveAnyway",
			dv::ConfigOption::buttonOption("Save current calibration data regardless of errors.", "Save Calibration"));

		config.add("outputCalibrationDirectory",
			dv::ConfigOption::directoryOption(
				"Specify directory to save the calibration settings in.", dv::portable_get_user_home_directory()));
		config.add("saveImages", dv::ConfigOption::boolOption("Save the images used for calibration.", false));

		config.add("input1CalibrationFile",
			dv::ConfigOption::fileOpenOption("Camera 1 calibration.", Calibration::FILE_FORMAT));
		config.add("input2CalibrationFile",
			dv::ConfigOption::fileOpenOption("Camera 2 calibration", Calibration::FILE_FORMAT));
		config.add("inputStereoCalibrationFile",
			dv::ConfigOption::fileOpenOption("Stereo calibration for both cameras.", Calibration::FILE_FORMAT));

		config.add("info/foundPoints", dv::ConfigOption::statisticOption("Number of point sets found so far."));
		config.add("info/calibrated", dv::ConfigOption::boolOption("Calibration completed successfuly?", false, true));

		config.setPriorityOptions(std::vector<std::string>{
			"info/foundPoints", "info/calibrated", "checkImages", "keep", "discard", "outputCalibrationDirectory"});
	}

	static void advancedStaticInit(dvModuleData mData) {
		auto mNode = dv::Cfg::Node(mData->moduleNode);

		mNode.addAttributeListener(nullptr, &inputFileChangeListener);
	}

	static void inputFileChangeListener(dvConfigNode node, void *userData, enum dvConfigAttributeEvents event,
		const char *changeKey, enum dvConfigAttributeType changeType, union dvConfigAttributeValue changeValue) {
		UNUSED_ARGUMENT(userData);

		const auto mNode         = dv::Cfg::Node(node);
		const auto moduleRunning = (mNode.getBool("isRunning") || mNode.getBool("running"));
		const auto changeStr     = std::string(changeKey);

		if (moduleRunning && (event == DVCFG_ATTRIBUTE_MODIFIED) && (changeType == DVCFG_TYPE_STRING)
			&& ((changeStr == "input1CalibrationFile") || (changeStr == "input2CalibrationFile")
				|| (changeStr == "inputStereoCalibrationFile"))) {
			dv::Log(dv::logLevel::WARNING,
				"Calibration file '{:s}' is only loaded at module startup, please make sure to "
				"restart the calibration module to load it.",
				changeValue.string);
		}
	}

	CalibrationModule() {
		// Setup the outputs and initialize calibration according to number of inputs.
		auto input1 = inputs.getFrameInput("input1");
		outputs.getFrameOutput("calibrated1").setup(input1);

		if (inputs.isConnected("input2")) {
			auto input2 = inputs.getFrameInput("input2");
			outputs.getFrameOutput("calibrated2").setup(input2);

			// Stereo calibration mode.
			calibration = std::make_unique<StereoCalibration>(&config, &outputs, input1, input2);
		}
		else {
			// Camera calibration mode.
			calibration = std::make_unique<CameraCalibration>(&config, &outputs, input1);

			// Remove unused output.
			dvModuleRegisterOutput(moduleData, "calibrated2", "REMOVE");
		}

		// Reset statistic at startup.
		config.set<dv::CfgType::LONG>("info/foundPoints", 0, true);
		config.set<dv::CfgType::BOOL>("info/calibrated", false, true);

		// Reset buttons to default state.
		config.setBool("checkImages", false);
		config.setBool("clearImages", false);
		config.setBool("calibrateNow", false);
		config.setBool("saveAnyway", false);

		// Disable unless checking images.
		config.setBool("keep", true);
		config.setBool("discard", true);
	}

	~CalibrationModule() override {
		// Restore unused outputs.
		if (!inputs.isConnected("input2")) {
			dvModuleRegisterOutput(moduleData, "calibrated2", dv::Frame::TableType::identifier);
		}
	}

	void run() override {
		// Get candidate inputs, store them.
		auto in1 = inputs.getFrameInput("input1").data();
		auto in2 = inputs.getFrameInput("input2").data();

		while (in1) {
			calibration->addInput(in1, 0);
			dvModuleInputAdvance(moduleData, "input1");
			in1 = inputs.getFrameInput("input1").data();
		}
		while (in2) {
			calibration->addInput(in2, 1);
			dvModuleInputAdvance(moduleData, "input2");
			in2 = inputs.getFrameInput("input2").data();
		}

		// Is there anything new to work on?
		bool newData      = calibration->getInput();
		bool patternAdded = false;

		if (newData) {
			// Search for pattern, output current results.
			patternAdded = calibration->searchPattern();
		}

		// If calibrated, the output is already part of the previous call above.
		if (calibration->isCalibrated()) {
			config.set<dv::CfgType::BOOL>("info/calibrated", true, true);
			return;
		}

		// If not, then we either follow the instructions of the user (buttons)
		// or attempt auto-calibration when all given conditions are met.
		if (config.getBool("clearImages")) {
			config.setBool("clearImages", false); // Reset button.

			calibration->clearImages();
		}

		if ((patternAdded && calibration->enoughImagesCollected()) || config.getBool("checkImages")) {
			config.setBool("checkImages", false); // Reset button.

			// Issue #360: CPU usage remains high due to previous activity while
			// we're off here checking images. This confuses users who may think
			// the module is off doing something complex, but it's actually waiting
			// on their input. Momentarily reset usage to 1% as we're doing nothing.
			moduleNode.updateReadOnly<dv::CfgType::FLOAT>("cpuUsage", 1.0f);

			// Check images first, so only valid ones are left after.
			calibration->checkImages();
		}

		bool calibrationSuccessful = false;

		if ((patternAdded && calibration->enoughImagesCollected()) || config.getBool("calibrateNow")) {
			config.setBool("calibrateNow", false); // Reset button.

			// Issue #360: CPU usage remains high due to previous activity while
			// we're off here checking images. This confuses users who may think
			// the module is off doing something complex, but it's actually waiting
			// on their input. Momentarily reset usage to 1% as we're doing nothing.
			moduleNode.updateReadOnly<dv::CfgType::FLOAT>("cpuUsage", 1.0f);

			// Ensure images have been checked. If already done above, will do nothing.
			calibration->checkImages();

			// Calibration is always heavy.
			moduleNode.updateReadOnly<dv::CfgType::FLOAT>("cpuUsage", 99.0f);

			// If enough images survived checking, or the user wants it,
			// execute a calibration pass.
			calibrationSuccessful = calibration->calibrate();

			if (!calibrationSuccessful) {
				// Failed calibration, advise user on how to continue.
				log.error("Calibration attempt failed; try again after having collected more patterns, clear the "
						  "accumulated images to get new ones, or restart the module.");
			}
		}

		if (calibrationSuccessful || config.getBool("saveAnyway")) {
			config.setBool("saveAnyway", false); // Reset button.

			// Save calibration data to XML.
			calibration->saveCalibrationData();
		}
	}

	void configUpdate() override {
		// If config changed (pattern, width/height), regenerate the points to find.
		calibration->regeneratePatternPoints();

		// Check that the selected pattern size is correct and follows
		// our recommendations (see documentation).
		const auto boardWidth  = config.getInt("boardWidth");
		const auto boardHeight = config.getInt("boardHeight");

		if (boardWidth == boardHeight) {
			log.warning << "A square pattern size is not recommended, please check our documentation for recommended "
						   "patterns and sizes."
						<< std::endl;
		}

		if ((boardWidth < 5) && (boardHeight < 4)) {
			log.warning << "A pattern size smaller than 4x5 is not recommended, please check our documentation for "
						   "recommended patterns and sizes."
						<< std::endl;
		}

		if ((((boardWidth % 2) == 0) && ((boardHeight % 2) == 0))
			|| (((boardWidth % 2) == 1) && ((boardHeight % 2) == 1))) {
			log.warning << "A symmetric pattern (dimensions both even or both odd) is not recommended, please check "
						   "our documentation for recommended patterns and sizes."
						<< std::endl;
		}
	}
};

registerModuleClass(CalibrationModule)
