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

#include <opencv2/imgproc.hpp>

#include <algorithm>
#include <vector>

class FrameHistogram : public dv::ModuleBase {
private:
	/**
	 * Returns the color for the histogram curve with index i, white if channels is 1
	 */
	inline cv::Scalar getColor(int i, int channels) {
		if (channels <= 1) {
			return cv::Scalar(255, 255, 255);
		}
		switch (i) {
			case 0:
				return cv::Scalar(255, 0, 0);
			case 1:
				return cv::Scalar(0, 255, 0);
			case 2:
				return cv::Scalar(0, 0, 255);
		}
		return cv::Scalar(255, 255, 255);
	}

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

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

	static const char *initDescription() {
		return ("Display a histogram of the incoming frame");
	}

	static void initConfigOptions(dv::RuntimeConfig &config) {
		config.add("numBins", dv::ConfigOption::intOption("Number of bins in which to divide values up.", 256, 4, 256));
	}

	FrameHistogram() {
		// Populate frame output info node. Must have generated statistics histogram frame
		// Populate frame output info node. Must have generated statistics histogram frame
		// maximum size. Max size is 256 x 128 due to max number of bins being 256.
		outputs.getFrameOutput("histograms").setup(256, 128, inputs.getFrameInput("frames").getOriginDescription());
	}

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

		hist_out.setTimestamp(frame_in.timestamp()); // Only set main timestamp.
		hist_out.setSource(dv::FrameSource::VISUALIZATION);

		auto numBins = config.get<dv::CfgType::INT>("numBins");

		// Calculate histogram, full uint8 range.
		const float range[]    = {0, 256};
		const float *histRange = {range};

		auto inMat = frame_in.getMatPointer();
		std::vector<cv::Mat> hists;
		switch (inMat->channels()) {
			case 1: {
				hists.resize(1);
				cv::calcHist(inMat.get(), 1, nullptr, cv::Mat(), hists[0], 1, &numBins, &histRange, true, false);
				break;
			}
			case 3: {
				std::vector<cv::Mat> bgr_planes;
				cv::split(*inMat, bgr_planes);
				hists.resize(3);
				cv::calcHist(&bgr_planes[0], 1, nullptr, cv::Mat(), hists[0], 1, &numBins, &histRange, true, false);
				cv::calcHist(&bgr_planes[1], 1, nullptr, cv::Mat(), hists[1], 1, &numBins, &histRange, true, false);
				cv::calcHist(&bgr_planes[2], 1, nullptr, cv::Mat(), hists[2], 1, &numBins, &histRange, true, false);
				break;
			}
		}

		// Generate histogram image, with N x N/2 pixels.
		cv::Mat outMat(numBins / 2, numBins, CV_8UC3, cv::Scalar(0, 0, 0));

		// Normalize the result to [0, histImage.rows].
		for (auto &hist : hists) {
			cv::normalize(hist, hist, 0, outMat.rows, cv::NORM_MINMAX, -1, cv::Mat());
		}

		// Draw the histogram.
		for (int i = 1; i < outMat.cols; i++) {
			for (int h = 0; h < static_cast<int>(hists.size()); h++) {
				const auto &hist = hists[h];
				cv::line(outMat, cv::Point(i - 1, outMat.rows - cvRound(hist.at<float>(i - 1))),
					cv::Point(i, outMat.rows - cvRound(hist.at<float>(i))), getColor(h, static_cast<int>(hists.size())),
					1, cv::LINE_8, 0);
			}
		}

		// Resize histogram to fit 256x128 output
		cv::resize(outMat, outMat, cv::Size(256, 128), cv::INTER_NEAREST);

		// Send histogram out.
		hist_out.setMat(std::move(outMat));
		hist_out.commit();
	}
};

registerModuleClass(FrameHistogram)
