#include "calibration.hpp"

#include "../../include/dv-sdk/cross/portable_io.h"

#include <chrono>
#include <regex>

static const std::regex filenameCleanupRegex{"[^a-zA-Z-_\\d]"};

void Calibration::setCameraID(const std::string &originDescription, const size_t source) {
	// 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, filenameCleanupRegex, "_");

	// 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;
	}

	camera[source].cameraID = str;
}

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

	return (true);
}

void Calibration::putText(const std::string &str, cv::Mat &img, const int maxWidth) {
	const auto strSize = cv::getTextSize(str, cv::FONT_HERSHEY_DUPLEX, 1.0, 2, nullptr);

	const auto strScale = static_cast<double>((img.cols >= maxWidth) ? (maxWidth) : (img.cols))
						/ static_cast<double>(strSize.width + 20);
	cv::putText(img, str, cv::Point{10, static_cast<int>(10 + (strScale * strSize.height))}, cv::FONT_HERSHEY_DUPLEX,
		strScale, CV_RGB(255, 0, 255), 2);
}

void Calibration::loadCalibrationCamera(const std::string &filename, const size_t source) {
	// loads single camera calibration file
	if (filename.empty()) {
		return;
	}

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

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

	auto typeNode = fs["type"];

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

	auto cameraNode = fs[camera[source].cameraID];

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

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

	cameraNode["camera_matrix"] >> camera[source].loadedCameraMatrix;
	cameraNode["distortion_coefficients"] >> camera[source].loadedDistCoeffs;

	log.info.format("Loaded camera matrix and distortion coefficients for camera {:s} from file: {:s}",
		camera[source].cameraID, filename);

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

void Calibration::addInput(const dv::InputDataWrapper<dv::Frame> &input, const size_t source) {
	camera[source].input.push_back(input);
}

Calibration::CalibrationPattern Calibration::getCalibrationPattern() {
	// returns the calibration pattern as set by the user
	auto pattern = config->getString("calibrationPattern");

	if (pattern == "circlesGrid") {
		return CalibrationPattern::CIRCLES_GRID;
	}
	else if (pattern == "asymmetricCirclesGrid") {
		return CalibrationPattern::ASYMMETRIC_CIRCLES_GRID;
	}
	else {
		return CalibrationPattern::CHESSBOARD;
	}
}

void Calibration::setDefaultPatternSizes() {
	if (config->getBool("useDefaultPattern")) {
		// Depending on chosen pattern, reset sizes.
		switch (getCalibrationPattern()) {
			case CalibrationPattern::CHESSBOARD:
				config->setInt("boardWidth", 9);
				config->setInt("boardHeight", 6);
				config->setFloat("boardSquareSize", 30.0f);
				break;

			case CalibrationPattern::CIRCLES_GRID:
				config->setInt("boardWidth", 9);
				config->setInt("boardHeight", 6);
				config->setFloat("boardSquareSize", 30.0f);
				break;

			case CalibrationPattern::ASYMMETRIC_CIRCLES_GRID:
				config->setInt("boardWidth", 6);
				config->setInt("boardHeight", 9);
				config->setFloat("boardSquareSize", 40.0f);
				break;
		}

		// Disarm button.
		config->setBool("useDefaultPattern", false);
	}
}

void Calibration::regeneratePatternPoints() {
	// Ensure pattern sizes are updated.
	setDefaultPatternSizes();

	auto boardSize = getBoardSize();

	// computes the same as opencv calcBoardCornerPositions and fills the object points accordingly
	patternPoints.clear();

	switch (getCalibrationPattern()) {
		case CalibrationPattern::CHESSBOARD:
		case CalibrationPattern::CIRCLES_GRID:
			for (int y = 0; y < boardSize.height; y++) {
				for (int x = 0; x < boardSize.width; x++) {
					patternPoints.push_back(cv::Point3f(static_cast<float>(x) * config->getFloat("boardSquareSize"),
						static_cast<float>(y) * config->getFloat("boardSquareSize"), 0));
				}
			}
			break;

		case CalibrationPattern::ASYMMETRIC_CIRCLES_GRID:
			for (int y = 0; y < boardSize.height; y++) {
				for (int x = 0; x < boardSize.width; x++) {
					patternPoints.push_back(cv::Point3f(
						static_cast<float>((2 * x) + (y % 2)) * (config->getFloat("boardSquareSize") / 2.0F),
						static_cast<float>(y) * (config->getFloat("boardSquareSize") / 2.0F), 0));
				}
			}
			break;
	}

	patternPoints.shrink_to_fit();
}

void Calibration::convertInputToGray(dv::InputDataWrapper<dv::Frame> &input, cv::Mat &dest) {
	// make sure we are only dealing with gray scale images
	if (input.format() == dv::FrameFormat::OPENCV_8U_C1) {
		dest = input.getMatPointer()->clone();
	}
	else if (input.format() == dv::FrameFormat::OPENCV_8U_C3) {
		cv::cvtColor(*input.getMatPointer(), dest, cv::COLOR_BGR2GRAY);
	}
	else if (input.format() == dv::FrameFormat::OPENCV_8U_C4) {
		cv::cvtColor(*input.getMatPointer(), dest, cv::COLOR_BGRA2GRAY);
	}
}

bool Calibration::enoughImagesCollected() {
	if (camera[0].images.size() < static_cast<size_t>(config->getInt("minDetections"))) {
		return false;
	}

	log.info("Enough images collected!");
	return true;
}

void Calibration::clearImages() {
	checkedImages = 0;
	config->set<dv::CfgType::LONG>("info/foundPoints", 0, true);

	for (auto &cam : camera) {
		cam.images.clear();
		cam.imagePoints.clear();
	}
}

cv::Size Calibration::getSubPixelSearchWindowHalfSize(const std::vector<cv::Point2f> &points) {
	// Find the square size in pixels for the given points, and use that to define a search window for subpixel
	// refinement
	bool initializedMinDifference = false;
	cv::Point2f minPointDifference;
	for (size_t i = 1; i < points.size(); i++) {
		const auto pointDifference = points.at(i) - points.at(i - 1);
		if (!initializedMinDifference) {
			minPointDifference       = pointDifference;
			initializedMinDifference = true;
			continue;
		}

		if (cv::norm(minPointDifference) > cv::norm(pointDifference)) {
			minPointDifference = pointDifference;
		}
	}

	const auto squareSizePixels = std::max(std::abs(minPointDifference.x), std::abs(minPointDifference.y));

	// The search window half size is set to half of the square size in pixels.
	// The minimum search window half size is set to 3x3.
	static const int minSearchWindowHalfSize = 3;
	const auto searchWindowHalfSize = std::max(static_cast<int>(squareSizePixels / 2.0f), minSearchWindowHalfSize);
	return {searchWindowHalfSize, searchWindowHalfSize};
}

bool Calibration::findPattern(const cv::Mat &image, std::vector<cv::Point2f> &points) {
	// finds the correct calibration pattern in a given image
	bool found = false;

	switch (getCalibrationPattern()) {
		case CalibrationPattern::CHESSBOARD: {
			uint32_t chessBoardFlags
				= cv::CALIB_CB_ADAPTIVE_THRESH | cv::CALIB_CB_NORMALIZE_IMAGE | cv::CALIB_CB_FILTER_QUADS;

			if (!config->getBool("useFisheyeModel")) {
				// Fast check erroneously fails with high distortions like fisheye lens.
				chessBoardFlags |= cv::CALIB_CB_FAST_CHECK;
			}

			found = cv::findChessboardCorners(image, getBoardSize(), points, static_cast<int>(chessBoardFlags));

			if (found && !points.empty()) {
				// Improve the found corners' coordinate accuracy for chessboard pattern.
				const auto searchWindow = getSubPixelSearchWindowHalfSize(points);
				cv::cornerSubPix(image, points, searchWindow, cv::Size(-1, -1),
					cv::TermCriteria(cv::TermCriteria::EPS + cv::TermCriteria::COUNT, 30, 0.0001));
			}

			break;
		}

		case CalibrationPattern::CIRCLES_GRID:
			found = cv::findCirclesGrid(image, getBoardSize(), points, cv::CALIB_CB_SYMMETRIC_GRID);
			break;

		case CalibrationPattern::ASYMMETRIC_CIRCLES_GRID:
			found = cv::findCirclesGrid(
				image, getBoardSize(), points, cv::CALIB_CB_ASYMMETRIC_GRID | cv::CALIB_CB_CLUSTERING);
			break;
	}

	return (found);
}

void Calibration::updateCurrentOutput(const cv::Mat &image, const std::vector<cv::Point2f> &points, const bool found,
	const size_t source, const bool highlight, const std::string &text) {
	// copy the most recent image to currentOutput, adding covered area and pattern points
	// convert the grayscale images back to bgr for output
	cv::cvtColor(image, camera[source].currentOutput, cv::COLOR_GRAY2BGR);

	if (highlight && config->getBool("highlightArea")) {
		highlightCoveredArea(source);
	}

	if (!text.empty()) {
		const double alpha = 0.6;
		cv::Mat oldOverlay;

		camera[source].currentOutput.copyTo(oldOverlay);

		Calibration::putText(text, oldOverlay);

		cv::addWeighted(oldOverlay, alpha, camera[source].currentOutput, 1 - alpha, 0, camera[source].currentOutput);
	}

	cv::drawChessboardCorners(camera[source].currentOutput, getBoardSize(), cv::Mat(points), found);

	// If calibration available for a specific camera, offer to undistort output.
	if (camera[source].cameraCalibrated) {
		undistortOutput(source);
	}
}

void Calibration::sendCurrentOutput(const size_t source) {
	// push currentOutput to output stream
	const auto name = "calibrated" + std::to_string(source + 1);
	outputs->getFrameOutput(name) << currentTimestamp << camera[source].currentOutput << dv::commit;
}

void Calibration::highlightCoveredArea(const size_t source) {
	// highlight the covered area by drawing a green circle around each point
	const double alpha = 0.1;
	cv::Mat oldOverlay;

	for (size_t i = 0; i < camera[source].imagePoints.size(); i++) {
		camera[source].currentOutput.copyTo(oldOverlay);

		for (size_t j = 0; j < camera[source].imagePoints[i].size(); j++) {
			cv::circle(oldOverlay, camera[source].imagePoints[i][j], camera[source].currentOutput.cols / 20,
				cv::Scalar(0, 105, 0), cv::FILLED);
		}

		cv::addWeighted(oldOverlay, alpha, camera[source].currentOutput, 1 - alpha, 0, camera[source].currentOutput);
	}
}

void Calibration::undistortOutput(const size_t source) {
	// Show reprojected points to generate an indicative output
	// that can be used to verify calibration.
	if (config->getBool("undistortOutput")) {
		cv::Mat undistortFrame;
		cv::undistort(
			camera[source].currentOutput, undistortFrame, camera[source].cameraMatrix, camera[source].distCoeffs);
		camera[source].currentOutput = undistortFrame;
	}
}

cv::Size Calibration::getBoardSize() {
	// utility function to return pattern board size
	switch (getCalibrationPattern()) {
		case CalibrationPattern::CHESSBOARD:
			// Inner corners, so -1 on each side.
			return cv::Size(config->getInt("boardWidth") - 1, config->getInt("boardHeight") - 1);

		case CalibrationPattern::CIRCLES_GRID:
		case CalibrationPattern::ASYMMETRIC_CIRCLES_GRID:
		default:
			return cv::Size(config->getInt("boardWidth"), config->getInt("boardHeight"));
	}
}

double Calibration::calibrateCamera(const size_t source) {
	// Reset matrices when recalibrating.
	camera[source].cameraMatrix = camera[source].loadedCameraMatrix.clone();
	camera[source].distCoeffs   = camera[source].loadedDistCoeffs.clone();

	double fError = 0;

	// performs calibration of a single camera
	std::vector<std::vector<cv::Point3f>> objectPoints{camera[source].imagePoints.size(), patternPoints};

	camera[source].rVecs.clear();
	camera[source].tVecs.clear();

	if (config->getBool("useFisheyeModel")) {
		cv::Mat _rvecs, _tvecs;

		fError = cv::fisheye::calibrate(objectPoints, camera[source].imagePoints, camera[source].imageSize,
			camera[source].cameraMatrix, camera[source].distCoeffs, _rvecs, _tvecs,
			cv::fisheye::CALIB_FIX_SKEW | cv::fisheye::CALIB_RECOMPUTE_EXTRINSIC | cv::fisheye::CALIB_FIX_K2
				| cv::fisheye::CALIB_FIX_K3 | cv::fisheye::CALIB_FIX_K4);

		camera[source].rVecs.reserve(static_cast<size_t>(_rvecs.rows));
		camera[source].tVecs.reserve(static_cast<size_t>(_tvecs.rows));

		for (size_t i = 0; i < objectPoints.size(); i++) {
			camera[source].rVecs.push_back(_rvecs.row(static_cast<int>(i)));
			camera[source].tVecs.push_back(_tvecs.row(static_cast<int>(i)));
		}
	}
	else {
		fError = cv::calibrateCamera(objectPoints, camera[source].imagePoints, camera[source].imageSize,
			camera[source].cameraMatrix, camera[source].distCoeffs, camera[source].rVecs, camera[source].tVecs,
			cv::CALIB_FIX_K4 | cv::CALIB_FIX_K5);
	}

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

	return fError;
}

void Calibration::saveCalibration(const double totalAvgError) {
	// checks for valid file and saves calibration
	const auto path = saveFilePath();
	cv::FileStorage fs(path, cv::FileStorage::WRITE);

	if (!fs.isOpened()) {
		log.warning << "Impossible to save the calibration file: " << path << std::endl;
		return;
	}

	writeToFile(fs);

	fs << "pattern_width" << getBoardSize().width;
	fs << "pattern_height" << getBoardSize().height;
	fs << "pattern_type" << config->getString("calibrationPattern");

	fs << "board_width" << config->getInt("boardWidth");
	fs << "board_height" << config->getInt("boardHeight");
	fs << "square_size" << config->getFloat("boardSquareSize");

	fs << "calibration_error" << totalAvgError;
	fs << "calibration_time" << saveFileTimeSuffix("{:%c}");

	log.info << "Calibration file saved in: " << path << std::endl;
}

std::string Calibration::saveFilePath() {
	// returns the filePath with correct extension
	boost::filesystem::path outputDir{config->getString("outputCalibrationDirectory")};

	if (outputDir.empty()) {
		outputDir = dv::portable_get_user_home_directory();
	}

	auto fileName = getDefaultFileName();
	if (!config->getBool("useDefaultFilename")) {
		fileName += "-" + saveFileTimeSuffix("{:%Y_%m_%d_%H_%M_%S}");
	}
	fileName += "." + FILE_FORMAT;

	const auto filePath = (outputDir / fileName).string();

	log.info.format("Saving calibration to file: {:s}", filePath);

	return (filePath);
}

std::string Calibration::saveFileTimeSuffix(const std::string &fmt) {
	// Get current time for suffix part.
	return fmt::format(fmt::runtime(fmt), dv::localTime(std::chrono::system_clock::now()));
}

void Calibration::writeToFileCamera(cv::FileStorage &fs, const size_t source) {
	// writes single camera calibration to file
	fs << camera[source].cameraID; // Node name.
	fs << "{";                     // Node start.
	fs << "camera_matrix" << camera[source].cameraMatrix;
	fs << "distortion_coefficients" << camera[source].distCoeffs;
	fs << "image_width" << camera[source].imageSize.width;
	fs << "image_height" << camera[source].imageSize.height;
	fs << "}"; // Node end.
}

void Calibration::savePNGImages(const size_t source) {
	boost::filesystem::path outputDir{config->getString("outputCalibrationDirectory")};

	if (outputDir.empty()) {
		outputDir = dv::portable_get_user_home_directory();
	}

	auto folderName  = camera[source].cameraID + "-images";
	folderName      += "-" + saveFileTimeSuffix("{:%Y_%m_%d_%H_%M_%S}");

	const auto folderPath = (outputDir / folderName);

	if (boost::filesystem::create_directory(folderPath)) {
		log.info.format("Saving calibration images to: {:s}", folderPath.string());

		size_t counter = 0;
		for (const auto &frame : camera[source].images) {
			auto fileName  = camera[source].cameraID + "-image";
			fileName      += "_" + std::to_string(counter);
			fileName      += "." + IMAGE_FORMAT;

			const auto filePath = (folderPath / fileName);

			cv::imwrite(filePath.string(), frame);
			counter++;
		}
	}
}
