#include "../../include/dv-sdk/data/frame.hpp"
#include "../../include/dv-sdk/module.hpp"

#include <opencv2/imgproc.hpp>

#include <algorithm>
#include <vector>

class FrameContrast : public dv::ModuleBase {
private:
	enum class ContrastAlgorithms {
		NORMALIZATION,
		HISTOGRAM_EQUALIZATION,
		CLAHE
	} contrastAlgo;

public:
	static void initInputs(dv::InputDefinitionList &in) {
		in.addFrameInput("frames");
	}

	static void initOutputs(dv::OutputDefinitionList &out) {
		out.addFrameOutput("frames");
	}

	static const char *initDescription() {
		return ("Enhance images by applying contrast enhancement algorithms.");
	}

	static void initConfigOptions(dv::RuntimeConfig &config) {
		config.add(
			"contrastAlgorithm", dv::ConfigOption::listOption("Contrast enhancement algorithm to apply.", 0,
									 std::vector<std::string>{"normalization", "histogram_equalization", "clahe"}));

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

	FrameContrast() {
		// setup output frame stream
		outputs.getFrameOutput("frames").setup(inputs.getFrameInput("frames"));

		// Call once to translate string into enum properly.
		configUpdate();
	}

	void run() override {
		auto frame_in  = inputs.getFrameInput("frames").frame();
		auto frame_out = outputs.getFrameOutput("frames").frame();

		// Setup output frame. Same size.
		frame_out.setPosition(frame_in.position());
		frame_out.setTimestamp(frame_in.timestamp());
		frame_out.setExposure(frame_in.exposure());
		frame_out.setSource(frame_in.source());

		// Get input OpenCV Mat. Lifetime is properly managed.
		auto input = frame_in.getMatPointer();

		CV_Assert((input->type() == CV_8UC1) || (input->type() == CV_8UC3) || (input->type() == CV_8UC4));

		// This generally only works well on grayscale intensity images.
		// So, if this is a grayscale image, good, else if its a color
		// image we convert it to YCrCb and operate on the Y channel only.
		const cv::Mat *intensity = nullptr;
		cv::Mat yCrCbPlanes[3];

		switch (input->channels()) {
			case 1: {
				intensity = input.get();
				break;
			}
			case 3: {
				cv::Mat yCrCb;
				cv::cvtColor(*input, yCrCb, cv::COLOR_BGR2YCrCb);
				cv::split(yCrCb, yCrCbPlanes);
				intensity = &yCrCbPlanes[0];
				break;
			}
			default: {
				log.error.format("Channel number of {:d} not supported", input->channels());
				throw std::runtime_error("Channel number not supported");
			}
		}

		// Apply contrast enhancement algorithm.
		cv::Mat contrastOut;

		switch (contrastAlgo) {
			case ContrastAlgorithms::NORMALIZATION:
				contrastNormalize(*intensity, contrastOut, 1.0);
				break;

			case ContrastAlgorithms::HISTOGRAM_EQUALIZATION:
				contrastEqualize(*intensity, contrastOut);
				break;

			case ContrastAlgorithms::CLAHE:
				contrastCLAHE(*intensity, contrastOut, 4.0, 8);
				break;

			default:
				// Other contrast enhancement types are not available in OpenCV.
				return;
				break;
		}

		// If original was a color frame, we have to mix the various
		// components back together into an BGR image.

		switch (input->channels()) {
			case 1: {
				frame_out.setMat(std::move(contrastOut));
				break;
			}
			case 3: {
				yCrCbPlanes[0] = contrastOut;

				cv::Mat yCrCb;
				cv::merge(yCrCbPlanes, 3, yCrCb);

				cv::Mat output;
				cv::cvtColor(yCrCb, output, cv::COLOR_YCrCb2BGR);
				frame_out.setMat(std::move(output));
				break;
			}
		}
		frame_out.commit();
	}

	void contrastNormalize(const cv::Mat &input, cv::Mat &output, float clipHistPercent) {
		CV_Assert((input.type() == CV_8UC1));
		CV_Assert((clipHistPercent >= 0) && (clipHistPercent < 100));

		// O(x, y) = alpha * I(x, y) + beta, where alpha maximizes the range
		// (contrast) and beta shifts it so lowest is zero (brightness).
		double minValue;
		double maxValue;

		if (clipHistPercent == 0) {
			// Determine minimum and maximum values.
			cv::minMaxLoc(input, &minValue, &maxValue);
		}
		else {
			// Calculate histogram.
			int histSize           = UINT8_MAX + 1;
			float hRange[]         = {0, static_cast<float>(histSize)};
			const float *histRange = {hRange};
			bool uniform           = true;
			bool accumulate        = false;

			cv::Mat hist;
			cv::calcHist(&input, 1, nullptr, cv::Mat(), hist, 1, &histSize, &histRange, uniform, accumulate);

			// Calculate cumulative distribution from the histogram.
			for (int i = 1; i < histSize; i++) {
				hist.at<float>(i) += hist.at<float>(i - 1);
			}

			// Locate points that cut at required value.
			float total      = hist.at<float>(histSize - 1);
			clipHistPercent *= (total / 100.0F); // Calculate absolute value from percent.
			clipHistPercent /= 2.0F;             // Left and right wings, so divide by two.

			// Locate left cut.
			minValue = 0;
			while (hist.at<float>(static_cast<int>(minValue)) < clipHistPercent) {
				minValue++;
			}

			// Locate right cut.
			maxValue = UINT8_MAX;
			while (hist.at<float>(static_cast<int>(maxValue)) >= (total - clipHistPercent)) {
				maxValue--;
			}
		}

		// Use min/max to calculate input range.
		double range = maxValue - minValue;

		// Calculate alpha (contrast).
		double alpha = (static_cast<double>(UINT8_MAX)) / range;

		// Calculate beta (brightness).
		double beta = -minValue * alpha;

		// Apply alpha and beta to pixels array.
		input.convertTo(output, -1, alpha, beta);
	}

	void contrastEqualize(const cv::Mat &input, cv::Mat &output) {
		CV_Assert((input.type() == CV_8UC1));
		output = cv::Mat::zeros(input.size(), CV_8UC1);

		// Calculate histogram.
		int histSize           = UINT8_MAX + 1;
		float hRange[]         = {0, static_cast<float>(histSize)};
		const float *histRange = {hRange};
		bool uniform           = true;
		bool accumulate        = false;

		cv::Mat hist;
		cv::calcHist(&input, 1, nullptr, cv::Mat(), hist, 1, &histSize, &histRange, uniform, accumulate);

		// Calculate cumulative distribution from the histogram.
		for (int i = 1; i < histSize; i++) {
			hist.at<float>(i) += hist.at<float>(i - 1);
		}

		// Total number of pixels. Must be the last value!
		float total = hist.at<float>(histSize - 1);

		// Smallest non-zero cumulative distribution value. Must be the first non-zero value!
		float min = 0;
		for (int i = 0; i < histSize; i++) {
			if (hist.at<float>(i) > 0) {
				min = hist.at<float>(i);
				break;
			}
		}

		// Calculate lookup table for histogram equalization.
		hist -= static_cast<double>(min);
		hist /= static_cast<double>(total - min);
		hist *= static_cast<double>(UINT8_MAX);

		// Apply lookup table to input image.
		int idx = 0;
		std::for_each(input.begin<uint8_t>(), input.end<uint8_t>(), [&hist, &output, &idx](const uint8_t &elem) {
			output.at<uint8_t>(idx++) = static_cast<uint8_t>(hist.at<float>(elem));
		});
	}

	void contrastCLAHE(const cv::Mat &input, cv::Mat &output, float clipLimit, int tilesGridSize) {
		CV_Assert((input.type() == CV_8UC1));
		CV_Assert((clipLimit >= 0) && (clipLimit < 100));
		CV_Assert((tilesGridSize >= 1) && (tilesGridSize <= 64));
		output = cv::Mat::zeros(input.size(), CV_8UC1);

		// Apply the CLAHE algorithm to the intensity channel (luminance).
		cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE();
		clahe->setClipLimit(static_cast<double>(clipLimit));
		clahe->setTilesGridSize(cv::Size(tilesGridSize, tilesGridSize));
		clahe->apply(input, output);
	}

	void configUpdate() override {
		// Parse available choices into enum value.
		auto selectedContrastAlgo = config.get<dv::CfgType::STRING>("contrastAlgorithm");

		if (selectedContrastAlgo == "histogram_equalization") {
			contrastAlgo = ContrastAlgorithms::HISTOGRAM_EQUALIZATION;
		}
		else if (selectedContrastAlgo == "clahe") {
			contrastAlgo = ContrastAlgorithms::CLAHE;
		}
		else {
			// Default choice.
			contrastAlgo = ContrastAlgorithms::NORMALIZATION;
		}
	}
};

registerModuleClass(FrameContrast)
