#pragma once
#include "color_util.hpp"
#include "dv-processing/exception/exceptions/file_exceptions.hpp"
#include "dv-processing/exception/exceptions/generic_exceptions.hpp"

#include <opencv2/core.hpp>
#include <opencv2/mcc.hpp>

#include <iostream>

namespace dv::color_calibration {
/**
 * Class for color calibration of event cameras using a color checker.
 * This class uses OpenCV's color correction model to calibrate the colors of the event camera. It uses a color checker
 * to find a mapping between the raw colors and the true colors of the color checker. This mapping is then used to
 * calibrate the colors of the event camera.
 */
class ColorCalibration {
public:
	ColorCalibration();

	/**
	 * Constructor for color calibration
	 * @param checker color checker to use
	 */
	explicit ColorCalibration(const ColorCheckerName checker);

	/**
	 * Constructor for color calibration from a saved calibration
	 * @param calibrationPath path to calibration
	 * @param checker color checker to use
	 */
	ColorCalibration(const std::filesystem::path &calibrationPath, const ColorCheckerName checker);

	/**
	 * Find color checker in a frame and adds it to detector if present
	 * @param image frame input
	 * @return false if no color chart was found - ture if success
	 */
	[[nodiscard]] bool findColorChecker(const cv::Mat &image);

	/**
	 * Calculate color calibration from given color chart
	 * @returns loss of the calibration
	 */
	[[nodiscard]] double prepareColorCorrection();

	/**
	 * Apply the color correction to an image
	 * @param uncalibratedImage image to be corrected
	 * @returns corrected image
	 */
	[[nodiscard]] cv::Mat correctColor(const cv::Mat &uncalibratedImage);

	/**
	 * Reset the detector and color correction
	 */
	void resetDetector();

	/**
	 * Save the calibration to a file
	 * @param path path to save the calibration
	 */
	void saveCalibration(const std::filesystem::path &path) const;

	/**
	 * Open a calibration from a file
	 * @param path path to the calibration
	 */
	void openCalibration(const std::filesystem::path &path);

	/**
	 * Open a default calibration for DAVIS
	 */
	void openDavisDefaultCalibration();

	/**
	 * Check if the color correction matrix is ready
	 * @returns true if ready
	 */
	[[nodiscard]] bool inferenceReady() const;

	/**
	 * Get the color correction matrix
	 * @returns color correction matrix
	 */
	[[nodiscard]] cv::Mat getColorCorrectionMatrix() const;

private:
	/**
	 * Run the color calibration
	 * @returns loss of the calibration
	 */
	double runCalibration();

	/**
	 * Set the default model parameters
	 */
	void setModelParameters();

	// Detection and Calibration objects
	cv::Ptr<cv::mcc::CCheckerDetector> mDetector = cv::mcc::CCheckerDetector::create();
	ColorCheckerName mUsedChecker                = ColorCheckerName::SPYDER;
	cv::ccm::ColorCorrectionModel mModel;

	// Calibration data
	cv::Mat mChartRGB;
	cv::Mat mTrueRgb;

	// Misc variables
	bool mInferenceReady = false;
};

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
/// Implementation
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
ColorCalibration::ColorCalibration() :
	mUsedChecker(ColorCheckerName::SPYDER),
	mModel(cv::Mat(), cv::Mat(), cv::ccm::COLOR_SPACE_sRGB) {
}

ColorCalibration::ColorCalibration(const ColorCheckerName checker) :
	mUsedChecker(checker),
	mModel(cv::Mat(), cv::Mat(), cv::ccm::COLOR_SPACE_sRGB) {
}

ColorCalibration::ColorCalibration(const std::filesystem::path &calibrationPath, const ColorCheckerName checker) :
	mUsedChecker(checker),
	mModel(cv::Mat(), cv::Mat(), cv::ccm::COLOR_SPACE_sRGB) {
	openCalibration(calibrationPath);
}

bool ColorCalibration::findColorChecker(const cv::Mat &image) {
	if (image.empty()) {
		return false;
	}
	return mDetector->process(image, cv::mcc::TYPECHART::MCC24, 1, true);
}

double ColorCalibration::prepareColorCorrection() {
	// Iterate detected checkers and average the color chart
	const auto checkers = mDetector->getListColorChecker();
	if (checkers.empty()) {
		throw dv::exceptions::LengthError("Preparation failed: Couldn't find any added color checkers.");
	}
	// Average the color charts
	cv::Mat sumChartsRGB(24, 1, CV_64FC3, cv::Scalar(0, 0, 0));
	for (const auto &cc : checkers) {
		cv::Mat chartRGB = cc->getChartsRGB().col(1).clone().reshape(3, 24);
		chartRGB.convertTo(chartRGB, CV_64FC3, 1. / std::numeric_limits<uint8_t>::max());
		sumChartsRGB += chartRGB;
	}
	mChartRGB = sumChartsRGB / static_cast<double>(checkers.size());

	// Prepare the true colors
	cv::Mat trueColors;
	switch (mUsedChecker) {
		case ColorCheckerName::SPYDER:
			mTrueRgb = ColorCheckerTrueColors::SPYDER;
			break;
		case ColorCheckerName::MACBETH:
			mTrueRgb = ColorCheckerTrueColors::MACBETH;
			break;
	}
	return runCalibration();
}

cv::Mat ColorCalibration::correctColor(const cv::Mat &uncalibratedImage) {
	if (!mInferenceReady) {
		throw dv::exceptions::RuntimeError("Correction failed: No color calibration matrix to apply. "
										   "Run prepareColorCorrection() first.");
	}
	else if (uncalibratedImage.empty()) {
		throw dv::exceptions::RuntimeError("Correction failed: Empty input image.");
	}
	cv::Mat calibratedImage = uncalibratedImage.clone();

	// Perform color correction
	cv::cvtColor(calibratedImage, calibratedImage, cv::COLOR_BGR2RGB);
	calibratedImage.convertTo(calibratedImage, CV_64FC3, 1. / std::numeric_limits<uint8_t>::max());
	calibratedImage = mModel.infer(calibratedImage);
	calibratedImage.convertTo(calibratedImage, CV_8UC3, std::numeric_limits<uint8_t>::max());
	cv::cvtColor(calibratedImage, calibratedImage, cv::COLOR_RGB2BGR);
	return calibratedImage;
}

void ColorCalibration::resetDetector() {
	mDetector.reset();
}

void ColorCalibration::saveCalibration(const std::filesystem::path &path) const {
	if (!(path.extension() == ".yml" || path.extension() == ".yaml")) {
		throw dv::exceptions::FileWriteError("Wrong file extension. Must be yml/yaml file format.");
	}
	else if (!mInferenceReady) {
		throw dv::exceptions::RuntimeError("No color calibration matrix to save. First make sure to run "
										   "prepareColorCorrection() and then save the calibration.");
	}

	cv::FileStorage fs(path.string(), cv::FileStorage::WRITE);
	if (!fs.isOpened()) {
		throw dv::exceptions::FileError("Could not open file for writing.");
	}

	if (fs.isOpened()) {
		fs << "normChartRgb" << mChartRGB;
		fs << "trueColors" << mTrueRgb;
		fs.release();
	}
}

void ColorCalibration::openCalibration(const std::filesystem::path &path) {
	if (path.empty()) {
		throw dv::exceptions::FileNotFound("Empty Path.");
	}
	else if (!(path.extension() == ".yml" || path.extension() == ".yaml")) {
		throw dv::exceptions::FileReadError("Wrong file extension. Must be yml/yaml file format.");
	}

	cv::FileStorage file(path.string(), cv::FileStorage::READ);
	if (!file.isOpened()) {
		throw dv::exceptions::FileNotFound(path);
	}
	file["normChartRgb"] >> mChartRGB;
	file["trueColors"] >> mTrueRgb;
	runCalibration();
}

void ColorCalibration::openDavisDefaultCalibration() {
	mChartRGB = DAVIS_CHART_RGB;
	mTrueRgb  = ColorCheckerTrueColors::SPYDER;
	runCalibration();
}

bool ColorCalibration::inferenceReady() const {
	return mInferenceReady;
}

cv::Mat ColorCalibration::getColorCorrectionMatrix() const {
	return mModel.getCCM();
}

double ColorCalibration::runCalibration() {
	mModel = cv::ccm::ColorCorrectionModel(mChartRGB, mTrueRgb, cv::ccm::COLOR_SPACE_sRGB);
	setModelParameters();
	mModel.run();
	mInferenceReady = true;
	return mModel.getLoss();
}

void ColorCalibration::setModelParameters() {
	// The parameters below are values from experiments to get a good calibration for the DAVIS COLOR 346 sensor
	static constexpr double MODEL_GAMMA          = 2.6;
	static constexpr double MODEL_THRESHOLD_LOW  = 0.25;
	static constexpr double MODEL_THRESHOLD_HIGH = 1;

	mModel.setCCM_TYPE(cv::ccm::CCM_3x3);
	mModel.setDistance(cv::ccm::DISTANCE_RGB);
	mModel.setLinear(cv::ccm::LINEARIZATION_GAMMA);
	mModel.setLinearGamma(MODEL_GAMMA);
	mModel.setSaturatedThreshold(MODEL_THRESHOLD_LOW, MODEL_THRESHOLD_HIGH);
}

} // namespace dv::color_calibration
