#include "cameracalibration.hpp"

#include <chrono>
#include <thread>

CameraCalibration::CameraCalibration(
	dv::RuntimeConfig *config_, dv::RuntimeOutputs *outputs_, dv::RuntimeInput<dv::Frame> &input) {
	config  = config_;
	outputs = outputs_;

	setCameraID(input.getOriginDescription(), 0);

	camera[0].imageSize = input.size();

	if (config->getBool("useFisheyeModel")) {
		// Reset differently for fish-eye model.
		camera[0].loadedDistCoeffs = cv::Mat::zeros(4, 1, CV_64F);
	}

	loadCalibrationCamera(config->getString("input1CalibrationFile"), 0);
}

bool CameraCalibration::isCalibrated() {
	return camera[0].cameraCalibrated;
}

bool CameraCalibration::getInput() {
	if (camera[0].input.empty()) {
		return (false);
	}

	auto in = camera[0].input.back();
	camera[0].input.clear();

	convertInputToGray(in, camera[0].currentImage);
	currentTimestamp = in.timestamp();

	return (true);
}

bool CameraCalibration::searchPattern() {
	std::vector<cv::Point2f> points;
	bool found = false;

	// try to find calibration pattern in the image
	found = findPattern(camera[0].currentImage, points);

	updateCurrentOutput(camera[0].currentImage, points, found, 0, true);

	sendCurrentOutput(0);

	if (found) {
		consecFound++;
	}
	else {
		consecFound = 0;
	}

	// points are only added if calibration pattern detected for every frame for a given number of frames
	if ((!isCalibrated()) && (consecFound >= static_cast<size_t>(config->getInt("consecDetects")))) {
		consecFound = 0;

		camera[0].images.push_back(camera[0].currentImage.clone());
		camera[0].imagePoints.push_back(points);

		config->set<dv::CfgType::LONG>("info/foundPoints", static_cast<int64_t>(camera[0].images.size()), true);
		log.info << "Added point set, total number of points: " << camera[0].images.size() << std::endl;

		return true;
	}

	return false;
}

bool CameraCalibration::calibrate() {
	// Exit condition: module stopped.
	if (!config->getBool("running")) {
		return false;
	}

	log.info("Calibrating ...");
	totalReprojectionError = calibrateCamera(0);

	if (totalReprojectionError < static_cast<double>(config->getFloat("maxReprojectionError"))) {
		camera[0].cameraCalibrated = true;

		log.info.format("Calibration successful with reprojection error = {:.4f}", totalReprojectionError);

		return true;
	}
	else {
		log.warning.format("Calibration unsuccessful with reprojection error = {:.4f}", totalReprojectionError);

		return false;
	}
}

void CameraCalibration::saveCalibrationData() {
	if (totalReprojectionError < 0) {
		log.error("No calibration ever executed, cannot save it.");
		return;
	}

	saveCalibration(totalReprojectionError);

	if (config->getBool("saveImages")) {
		savePNGImages(0);
	}
}

void CameraCalibration::checkImages() {
	// Exit condition: module stopped.
	if (!config->getBool("running")) {
		return;
	}

	size_t imagesToCheck = (camera[0].images.size() - checkedImages);

	if (imagesToCheck == 0) {
		log.info("No further images to check.");
		return;
	}

	log.info.format("Checking {:d} new images to discard low quality ones ...", imagesToCheck);

	// iterate over images and decide whether image should be kept for calibration
	auto iIt = camera[0].images.begin();
	auto pIt = camera[0].imagePoints.begin();

	size_t imageCtr = 0;

	while (iIt != camera[0].images.end() && pIt != camera[0].imagePoints.end()) {
		// Skip images and points that have already been checked previously.
		if (imageCtr < checkedImages) {
			imageCtr++;
			iIt++;
			pIt++;
			continue;
		}

		// Show current image to user.
		updateCurrentOutput(*iIt, *pIt, true, 0, false,
			fmt::format("Checking image #{:03d}/{:03d}", imageCtr + 1, config->getLong("info/foundPoints")));
		sendCurrentOutput(0);

		// Update config status and set buttons to disabled.
		config->update();
		config->setBool("keep", false);
		config->setBool("discard", false);

		// Wait until one of the buttons is selected.
		while (!config->getBool("keep") && !config->getBool("discard")) {
			std::this_thread::sleep_for(std::chrono::milliseconds(100));
			config->update();

			// Exit condition: module stopped.
			if (!config->getBool("running")) {
				return;
			}
		}

		// Both buttons selected, invalid!
		if ((config->getBool("keep")) && (config->getBool("discard"))) {
			log.warning("Selecting both 'keep' and 'discard' is not allowed, doing nothing.");
			config->setBool("keep", false);
			config->setBool("discard", false);
		}
		// Keep button selected.
		else if (config->getBool("keep")) {
			log.info("Keeping image ...");
			imageCtr++;
			iIt++;
			pIt++;
			config->setBool("keep", false);
		}
		// Discard button selected.
		else if (config->getBool("discard")) {
			log.info("Discarding image ...");
			iIt = camera[0].images.erase(iIt);
			pIt = camera[0].imagePoints.erase(pIt);
			config->setBool("discard", false);
			config->set<dv::CfgType::LONG>("info/foundPoints", static_cast<int64_t>(camera[0].images.size()), true);
		}
	}

	checkedImages = camera[0].images.size();

	log.info.format("Verified {:d} new images, total good images is now {:d}.", imagesToCheck, checkedImages);

	// Disable buttons when not checking images.
	config->setBool("keep", true);
	config->setBool("discard", true);
}

std::string CameraCalibration::getDefaultFileName() {
	// camera calibration default file name
	std::string filename = "calibration_camera";

	if (!config->getBool("useDefaultFilename")) {
		filename += "_" + camera[0].cameraID;
	}

	return filename;
}

void CameraCalibration::writeToFile(cv::FileStorage &fs) {
	// write calibration result to file
	writeToFileCamera(fs, 0);

	fs << "use_fisheye_model" << config->getBool("useFisheyeModel");

	fs << "type" << "camera";
}
