#include "stereocalibration.hpp"

#include <chrono>
#include <thread>

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

	setCameraID(input1.getOriginDescription(), 0);
	setCameraID(input2.getOriginDescription(), 1);

	camera[0].imageSize = input1.size();
	camera[1].imageSize = input2.size();

	// Fisheye lenses not supported for stereo. Ensure it's disabled.
	if (config->getBool("useFisheyeModel")) {
		config->setBool("useFisheyeModel", false);
		log.warning("Fish-eye lenses are not supported for stereo calibration, disabling option.");
	}

	// load single camera calibration from xml
	loadCalibrationStereo(config->getString("inputStereoCalibrationFile"));

	// if all three file are present the cameraMatrix1, distCoeffs1, cameraMatrix2, distCoeffs2 are updated
	loadCalibrationCamera(config->getString("input1CalibrationFile"), 0);
	loadCalibrationCamera(config->getString("input2CalibrationFile"), 1);
}

void StereoCalibration::loadCalibrationStereo(const std::string &filename) {
	// load stereo calibration file
	if (filename.empty()) {
		return;
	}

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

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

	auto typeNode = fs["type"];

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

	for (size_t i = 0; i < camera.size(); i++) {
		auto cameraNode = fs[camera[i].cameraID];

		if (!cvExists(cameraNode) || !cameraNode.isMap()) {
			log.warning.format(
				"Calibration data for camera {:s} not present in file: {:s}", camera[i].cameraID, filename);
			continue;
		}

		if (!cvExists(cameraNode["camera_matrix"]) || !cvExists(cameraNode["distortion_coefficients"])) {
			log.warning.format(
				"Calibration data for camera {:s} not present in file: {:s}", camera[i].cameraID, filename);
			continue;
		}

		cameraNode["camera_matrix"] >> camera[i].loadedCameraMatrix;
		cameraNode["distortion_coefficients"] >> camera[i].loadedDistCoeffs;
		log.info.format("Loaded camera matrix and distortion coefficients for camera {:s} from file: {:s}",
			camera[i].cameraID, filename);

		camera[i].cameraCalibrated = true;
		camera[i].cameraMatrix     = camera[i].loadedCameraMatrix.clone();
		camera[i].distCoeffs       = camera[i].loadedDistCoeffs.clone();
	}

	// Stereo parameters.
	if (!cvExists(fs["R"]) || !cvExists(fs["T"]) || !cvExists(fs["E"]) || !cvExists(fs["F"])) {
		log.warning.format("Stereo parameters not present in file: {:s}", filename);
		return;
	}

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

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

	stereoCalibrated = true;
}

bool StereoCalibration::isCalibrated() {
	return stereoCalibrated;
}

bool StereoCalibration::getInput() {
	auto it1 = camera[0].input.begin();
	auto it2 = camera[1].input.begin();

	// for intermediate saving of results
	dv::InputDataWrapper<dv::Frame> c1{nullptr}, c2{nullptr};
	bool waiting = false;

	while (!waiting && !camera[0].input.empty() && !camera[1].input.empty()) {
		if (it1->timestamp() >= it2->timestamp()) {
			// iterate until it2.timestamp > it1.timestamp
			while ((it2 != camera[1].input.end()) && (it1->timestamp() >= it2->timestamp())) {
				it2++;
			}

			if (it2 == camera[1].input.end()) {
				waiting = true;
			}
			else {
				// align with frame in input2 where difference is smallest (0 if
				// perfectly aligned)
				if ((it1->timestamp() - std::prev(it2)->timestamp()) < (it2->timestamp() - it1->timestamp())) {
					it2--;
				}
			}
		}
		else {
			// iterate until it1.timestamp > it2.timestamp
			while ((it1 != camera[0].input.end()) && (it2->timestamp() > it1->timestamp())) {
				it1++;
			}

			if (it1 == camera[0].input.end()) {
				waiting = true;
			}
			else {
				// align with frame in input1 where difference is smallest (0 if
				// perfectly aligned)
				if ((it2->timestamp() - std::prev(it1)->timestamp()) < (it1->timestamp() - it2->timestamp())) {
					it1--;
				}
			}
		}

		if (!waiting) {
			c1 = *it1;
			c2 = *it2;

			// then erase everything including current positions
			if (it1 != camera[0].input.end()) {
				it1++;
			}
			if (it2 != camera[1].input.end()) {
				it2++;
			}

			it1 = camera[0].input.erase(camera[0].input.begin(), it1);
			it2 = camera[1].input.erase(camera[1].input.begin(), it2);
		}
	}

	if ((c1.getBasePointer() != nullptr) && (c2.getBasePointer() != nullptr)) {
		// save both iterator positions as currentImage1 and 2, since they are
		// deleted from vector
		convertInputToGray(c1, camera[0].currentImage);
		convertInputToGray(c2, camera[1].currentImage);

		currentTimestamp = std::max(c1.timestamp(), c2.timestamp());

		return true;
	}

	return false;
}

bool StereoCalibration::searchPattern() {
	std::vector<std::vector<cv::Point2f>> points(camera.size());
	std::vector<bool> found(camera.size(), false);

	// try to find calibration pattern in both images
	for (size_t i = 0; i < camera.size(); i++) {
		found[i] = findPattern(camera[i].currentImage, points[i]);

		if ((i == 1) && (static_cast<size_t>(std::count(found.begin(), found.end(), true)) == camera.size())) {
			checkPointsOrder(points);
		}

		// Simple output drawing.
		Calibration::updateCurrentOutput(camera[i].currentImage, points[i], found[i], i, true);
	}

	for (size_t i = 0; i < camera.size(); i++) {
		extraOutputVisualization(points[i], found[i], i);
	}

	for (size_t i = 0; i < camera.size(); i++) {
		sendCurrentOutput(i);
	}

	if (static_cast<size_t>(std::count(found.begin(), found.end(), true)) == camera.size()) {
		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;

		for (size_t i = 0; i < camera.size(); i++) {
			camera[i].images.push_back(camera[i].currentImage.clone());
			camera[i].imagePoints.push_back(points[i]);
		}

		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 StereoCalibration::calibrate() {
	// Exit condition: module stopped.
	if (!config->getBool("running")) {
		return false;
	}

	log.info("Calibrating ...");
	totalReprojectionError = calibrateStereo();

	totalEpipolarError = epipolarLinesError();

	if ((totalEpipolarError < static_cast<double>(config->getFloat("maxEpipolarLineError")))
		&& (totalReprojectionError < static_cast<double>(config->getFloat("maxReprojectionError")))) {
		stereoCalibrated = true;

		// Single camera calibration is also calculated as part of stereo calibration.
		for (auto &cam : camera) {
			cam.cameraCalibrated = true;
		}

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

		return true;
	}
	else {
		log.info.format("Calibration unsuccessful with reprojection error = {:.4f} and epipolar line error = {:.4f}",
			totalReprojectionError, totalEpipolarError);

		return false;
	}
}

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

	saveCalibration(totalReprojectionError);

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

void StereoCalibration::extraOutputVisualization(
	const std::vector<cv::Point2f> &points, const bool found, const size_t source) {
	// Draw on the other camera output.
	const size_t otherCamera = otherCameraID(source);

	// Produce output to verify calibration quality.
	if (config->getBool("drawEpipolarLines") && isCalibrated() && found) {
		drawEpipolarLines(camera[otherCamera].currentOutput, points, source);
	}
}

void StereoCalibration::checkPointsOrder(std::vector<std::vector<cv::Point2f>> &points) {
	// check the order of the found points for the second camera
	bool reverse = false;

	if (points[0].front().x < points[0].back().x) {
		if (points[1].front().x > points[1].back().x) {
			reverse = true;
		}
	}
	else {
		if (points[1].front().x < points[1].back().x) {
			reverse = true;
		}
	}

	if (points[0].front().y < points[0].back().y) {
		if (points[1].front().y > points[1].back().y) {
			reverse = true;
		}
	}
	else {
		if (points[1].front().y < points[1].back().y) {
			reverse = true;
		}
	}

	if (reverse) {
		std::reverse(points[1].begin(), points[1].end());
	}
}

void StereoCalibration::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
	std::vector<std::vector<cv::Mat>::iterator> iIts;
	std::vector<std::vector<std::vector<cv::Point2f>>::iterator> pIts;

	for (size_t i = 0; i < camera.size(); i++) {
		iIts.push_back(camera[i].images.begin());
		pIts.push_back(camera[i].imagePoints.begin());
	}

	size_t imageCtr = 0;

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

		const auto checkStr
			= fmt::format("Checking image #{:03d}/{:03d}", imageCtr + 1, config->getLong("info/foundPoints"));

		// Show current image for camera 1 to user.
		Calibration::updateCurrentOutput(*iIts[0], *pIts[0], true, 0, false, checkStr);
		sendCurrentOutput(0);

		// Show current image for camera 2 to user.
		Calibration::updateCurrentOutput(*iIts[1], *pIts[1], true, 1, false, checkStr);
		sendCurrentOutput(1);

		// 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 images ...");
			imageCtr++;
			for (size_t i = 0; i < iIts.size(); i++) {
				iIts[i]++;
				pIts[i]++;
			}
			config->setBool("keep", false);
		}
		// Discard button selected.
		else if (config->getBool("discard")) {
			log.info("Discarding images ...");
			for (size_t i = 0; i < iIts.size(); i++) {
				iIts[i] = camera[i].images.erase(iIts[i]);
				pIts[i] = camera[i].imagePoints.erase(pIts[i]);
			}
			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);
}

double StereoCalibration::calibrateStereo() {
	// Reset matrices when recalibrating.
	for (auto &cam : camera) {
		cam.cameraMatrix = cam.loadedCameraMatrix.clone();
		cam.distCoeffs   = cam.loadedDistCoeffs.clone();
	}

	// If cameraSizes are different, calibrate each camera individually and then stereoCalibrate them
	// otherwise just call the cv function. You can pre-calibrate each camera for better results.
	if (camera[0].imageSize != camera[1].imageSize) {
		for (size_t i = 0; i < camera.size(); i++) {
			if (!camera[i].cameraCalibrated) {
				const double totalReprojError = calibrateCamera(i);

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

					log.info.format("Independent camera {:d} calibration successful with reprojection error = {:.4f}",
						i, totalReprojError);
				}
				else {
					log.warning.format(
						"Independent camera {:d} calibration unsuccessful with reprojection error = {:.4f}", i,
						totalReprojError);
				}
			}
		}
	}

	uint32_t flags = cv::CALIB_FIX_ASPECT_RATIO | cv::CALIB_ZERO_TANGENT_DIST | cv::CALIB_RATIONAL_MODEL;

	// If both cameras are already individually calibrated, use their parameters,
	// instead of trying to guess them as part of stereo calibration.
	if (camera[0].cameraCalibrated && camera[1].cameraCalibrated) {
		flags |= cv::CALIB_FIX_INTRINSIC;
	}
	else {
		flags |= cv::CALIB_FIX_K3 | cv::CALIB_FIX_K4 | cv::CALIB_FIX_K5 | cv::CALIB_USE_INTRINSIC_GUESS;
	}

	std::vector<std::vector<cv::Point3f>> objectPoints{camera[0].imagePoints.size(), patternPoints};

	const double error
		= cv::stereoCalibrate(objectPoints, camera[0].imagePoints, camera[1].imagePoints, camera[0].cameraMatrix,
			camera[0].distCoeffs, camera[1].cameraMatrix, camera[1].distCoeffs, camera[0].imageSize, R, T, E, F,
			static_cast<int>(flags), cv::TermCriteria(cv::TermCriteria::COUNT | cv::TermCriteria::EPS, 100, 1e-5));

	log.info.format("Stereo calibration completed with error: {:.4f}", error);

	return error;
}

double StereoCalibration::epipolarLinesError() {
	// calculate the total error based on epipolar constraint violation
	double errTotal = 0;

	for (size_t i = 0; i < camera[0].imagePoints.size(); i++) {
		const double errSingleImage  = errorFunction(i);
		errTotal                    += errSingleImage;
	}

	// error normalized by the total number of projected points n_tot=(#frames)*(#points)
	return errTotal / static_cast<double>(camera[0].imagePoints.size() * camera[0].imagePoints[0].size());
}

double StereoCalibration::errorFunction(const size_t i) {
	// stereo calibration uses the epipolar geometry constraint:
	// m2^t*F*m1=0 as an error metric.
	std::vector<std::vector<cv::Vec3f>> lines;

	for (size_t j = 0; j < camera.size(); j++) {
		lines.push_back(epipolarLines(camera[j].imagePoints[i], j));
	}

	std::vector<cv::Point2f> p(camera.size());
	std::vector<double> errSinglePoint(camera.size());
	double error = 0;

	for (size_t n = 0; n < camera[0].imagePoints[i].size(); n++) {
		p[0] = camera[0].imagePoints[i][n];
		p[1] = camera[1].imagePoints[i][n];

		// Distance between a point and the relative epipolar line in the two images.
		errSinglePoint[0]
			= std::fabs(static_cast<double>((p[0].x * lines[1][n][0]) + (p[0].y * lines[1][n][1]) + lines[1][n][2]))
			/ std::sqrt(std::pow(lines[1][n][0], 2) + std::pow(lines[1][n][1], 2));
		errSinglePoint[1]
			= std::fabs(static_cast<double>((p[1].x * lines[0][n][0]) + (p[1].y * lines[0][n][1]) + lines[0][n][2]))
			/ std::sqrt(std::pow(lines[0][n][0], 2) + std::pow(lines[0][n][1], 2));

		error += errSinglePoint[0] + errSinglePoint[1];
	}

	return error;
}

std::vector<cv::Vec3f> StereoCalibration::epipolarLines(
	const std::vector<cv::Point2f> &sigleImagePoints, const size_t source) {
	// calculate the epipolare lines for one image given a set of points in the other image
	std::vector<cv::Vec3f> line;
	cv::Mat idealPoints;

	cv::undistortPoints(cv::Mat(sigleImagePoints), idealPoints, camera[source].cameraMatrix, camera[source].distCoeffs,
		cv::Mat(), camera[source].cameraMatrix);
	cv::computeCorrespondEpilines(idealPoints, static_cast<int>(source + 1), F, line);

	return line;
}

void StereoCalibration::drawEpipolarLines(
	cv::Mat &output, const std::vector<cv::Point2f> &points, const size_t source) {
	// Draw epipolar lines on the output image

	// OpenCV colormap (from function cv::drawChessboardCorners)
	static const size_t line_max                          = 7;
	static const uint8_t pattern_line_colors[line_max][4] = {
		{0,   0,   255, 0},
        {0,   128, 255, 0},
        {0,   200, 200, 0},
        {0,   255, 0,   0},
        {200, 200, 0,   0},
        {255, 0,   0,   0},
		{255, 0,   255, 0}
    };

	const auto lines = epipolarLines(points, source);

	for (size_t i = 0; i < lines.size(); i++) {
		const uint8_t *const pattern_line_color
			= &pattern_line_colors[(i / static_cast<size_t>(getBoardSize().width)) % line_max][0];
		const cv::Scalar color(
			pattern_line_color[0], pattern_line_color[1], pattern_line_color[2], pattern_line_color[3]);

		const auto p1 = cv::Point2f(0.0, -lines[i][2] / lines[i][1]);
		const auto p2 = cv::Point2f(static_cast<float>(output.cols),
			-(lines[i][2] + (lines[i][0] * static_cast<float>(output.cols))) / lines[i][1]);

		cv::line(output, p1, p2, color, 1, cv::LINE_AA);
	}
}

std::string StereoCalibration::getDefaultFileName() {
	// stereo calibration default file name
	std::string filename = "calibration_stereo";

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

	return filename;
}

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

	fs << "R" << R;
	fs << "T" << T;
	fs << "E" << E;
	fs << "F" << F;

	fs << "epipolar_error" << totalEpipolarError;

	fs << "type" << "stereo";
}
