#include "stereo_rectification.hpp"

void StereoRectification::initConfigOptions(dv::RuntimeConfig &config) {
	config.add(
		"fitMorePixels", dv::ConfigOption::floatOption("How many input pixels to fit (alpha value in cv::stereoRectify "
													   "function), 0: maximize the image, 1: keep all source pixels.",
							 0, 0, 1));
	config.add("calibrationFile",
		dv::ConfigOption::fileOpenOption(
			"The name of the file from which to load the calibration settings for stereo rectification.",
			dv::portable_get_user_home_directory() + "/calibration_camera.xml", "xml"));

	config.setPriorityOptions({"calibrationFile"});
}

StereoRectification::StereoRectification() {
	for (size_t i = 0; i < NUM_CAMERAS; i++) {
		const auto inputName = "camera" + std::to_string(i);

		// Generate camera ID from origin.
		id[i] = getCameraID(inputs.infoNode(inputName).getString("source"));

		// Get input dimensions.
		size[i] = cv::Size{inputs.infoNode(inputName).getInt("sizeX"), inputs.infoNode(inputName).getInt("sizeY")};
	}

	// The two inputs must be from two different origins/cameras.
	if (id[0] == id[1]) {
		throw std::runtime_error("Both inputs originate from the same camera.");
	}

	// The two cameras must have the same resolution for stereo rectification.
	if (size[0] != size[1]) {
		throw std::runtime_error("Inputs do not have the same size.");
	}
}

void StereoRectification::configUpdate() {
	// Any changes to configuration mean the calibration has to be
	// reloaded and reinitialized. On file-related failure, we stop.
	if (!loadStereoRectificationMatrices(config.getString("calibrationFile"))) {
		throw std::runtime_error("Could not load calibration data.");
	}
}

bool StereoRectification::cvExists(const cv::FileNode &fn) {
	if (fn.type() == cv::FileNode::NONE) {
		return (false);
	}

	return (true);
}

std::string StereoRectification::getCameraID(const std::string &originDescription) {
	// Camera origin descriptions are fine, never start with a digit or have spaces
	// or other special characters in them that fail with OpenCV's FileStorage.
	// So we just clean/escape the string for possible other sources.
	auto str = std::regex_replace(originDescription, FILENAME_CLEANUP_REGEX, "_");

	// FileStorage node names can't start with - or 0-9.
	// Prefix with an underscore in that case.
	if ((str[0] == '-') || (std::isdigit(str[0]) != 0)) {
		str = "_" + str;
	}

	return str;
}

bool StereoRectification::loadStereoRectificationMatrices(const std::string &filename) {
	// Load stereo camera calibration file.
	if (filename.empty()) {
		log.error << "No stereo calibration file specified." << std::endl;
		return (false);
	}

	cv::FileStorage fs(filename, cv::FileStorage::READ);

	if (!fs.isOpened()) {
		log.error << "Impossible to load the stereo calibration file: " << filename << std::endl;
		return (false);
	}

	// Load matrices from new calibration file format only, as only
	// the new one supports stereo calibration.
	auto typeNode = fs["type"];

	if (!cvExists(typeNode) || !typeNode.isString() || (typeNode.string() != "stereo")) {
		log.error << "Invalid stereo calibration file: " << filename << std::endl;
		return (false);
	}

	log.info << "New-style camera calibration file found." << std::endl;

	CameraInfo loaded[NUM_CAMERAS];

	for (size_t i = 0; i < NUM_CAMERAS; i++) {
		auto cameraNode = fs[id[i]];

		if (!cvExists(cameraNode) || !cameraNode.isMap() || !cvExists(cameraNode["camera_matrix"])
			|| !cvExists(cameraNode["distortion_coefficients"]) || !cvExists(cameraNode["image_width"])
			|| !cvExists(cameraNode["image_height"])) {
			log.error.format("Calibration data for camera {:s} not present in file: {:s}", id[i], filename);
			return (false);
		}

		cameraNode["camera_matrix"] >> loaded[i].cameraMatrix;
		cameraNode["distortion_coefficients"] >> loaded[i].distCoeffs;

		cameraNode["image_width"] >> loaded[i].width;
		cameraNode["image_height"] >> loaded[i].height;

		if ((loaded[i].width != size[i].width) || (loaded[i].height != size[i].height)) {
			log.error.format(
				"Size of camera {:s} does not match image size in calibration file: {:s}", id[i], filename);
			return (false);
		}
	}

	// Load stereo matrices.
	if (!cvExists(fs["R"]) || !cvExists(fs["T"])) {
		log.error.format("Stereo calibration data not present in file: {:s}", filename);
		return (false);
	}

	cv::Mat R, T;

	fs["R"] >> R;
	fs["T"] >> T;

	log.info.format("Loaded required stereo calibration matrices from file: {:s}", filename);

	const double alpha = static_cast<double>(config.getFloat("fitMorePixels"));

	initStereoRectification(loaded, R, T, alpha);

	return (true);
}
