//This is part of MSU Video Quality Measurement Tool (VQMT) back SDK
//For MSU VQMT 12 BETA Premium
//MSU G&M Lab. Video Group 2019

#include "vqmt_sdk.h"

#include <fstream>
#include <iostream>
#include <functional>
#include <vector>
#include <map>
#include <string>
#include <cstring>

void printMetricList() {
	std::cout << "METRIC" << "\t" << "[VARIATION]" << std::endl;
	for (const VQMT_MetricInfo* info = vqmt_getMetricList(); info->handler; info++) {
		std::cout << info->group << "\t" << info->variation << std::endl;
	}
}

//find metric by specified group and variation
const VQMT_MetricInfo* findMetric(const std::string& group, const std::string& variation) {
	for (const VQMT_MetricInfo* info = vqmt_getMetricList(); info->handler; info++) {
		if (info->group == group && info->variation == variation) return info;
	}

	return nullptr;
}

//return true iff metric supports specified colorspace
bool metricSupports(const VQMT_MetricInfo* i, VQMT_MetricCC cc) {
	for (auto c = i->supportedCCs; *c != VQMT_MetricCC::unused; ++c) {
		if (*c == cc) return true;
	}

	return false;
}

//this is universal C-style callback that will call arbitrary std::function, specified by cbparam
void callerFunction(int frame, int valueIdx, float value, void* cbparam, void* payload) {
	auto function = (std::function<void(int, int, float)>*)(cbparam);
	(*function)(frame, valueIdx, value);
}

void readY(VQMT_RawBuffer* dst, std::istream& src) {
	if (!dst) return;
	VQMT_RawImage image = vqmt_getBufferData(dst);
	for (int y = 0; y < image.height[0]; ++y) {
		char* row = (char*)((intptr_t)image.planes[0] + image.strides[0] * y);
		src.read(row, image.width[0]);
	}
}

int calcMetric(const VQMT_MetricInfo* metr, int files, std::istream& i1, std::istream& i2, int w, int h) {
	std::map<int, std::string> averageNames;

	//this lambda will print metric value
	std::function<void(int, int, float)> func = [&averageNames](int frame, int idx, float val) {
		if (idx != 0) return;
		if (frame >= 0)
			std::cout << "Frame " << frame << " value " << val << std::endl;
		else
			std::cout << averageNames[frame] << ": " << val << std::endl;
	};

	//create metric instance
	VQMT_MetricInst* inst = vqmt_createMetric(metr->handler, w, h, VQMT_MetricCC::Y, false, nullptr, callerFunction, &func);
	if (!inst) {
		std::cerr << "Unable to create metric" << std::endl;
		return 5;
	}

	const char* inputFormats[] = {"yuv420p", "yuv420p"};
	VQMT_MetricInst* metricList[] = { inst, nullptr };

	char errBuff[1024];
	VQMT_ImageConverter* converter = vqmt_initImageConverter(inputFormats, metricList, errBuff, sizeof(errBuff));

	if (!converter) {
		std::cerr << "Can't initialize converter: " << errBuff << std::endl;
		return 6;
	}

	VQMT_ConfiguredMetricInfo metricInfo = vqmt_getConfiguredMetricInfo(inst);
	for (int i = 0; i < metricInfo.accumulatorsCount; ++i) {
		//std::cout << metricInfo.accumulators[i].frameId
		averageNames[metricInfo.accumulators[i].frameId] = metricInfo.accumulators[i].name;
	}

	/*VQMT_ImageRequest request[2];
	std::memset(&request[0], 0, sizeof(request[0]));
	std::memset(&request[1], 0, sizeof(request[1]));

	//we have uint8_t and float data for Y-component:
	request[0].availablePlanes[(int)VQMT_ValueType::V_8u][(int)VQMT_CC::Y] = true;
	request[0].availablePlanes[(int)VQMT_ValueType::V_32f][(int)VQMT_CC::Y] = true;
	request[1].availablePlanes[(int)VQMT_ValueType::V_8u][(int)VQMT_CC::Y] = true;
	request[1].availablePlanes[(int)VQMT_ValueType::V_32f][(int)VQMT_CC::Y] = true;
	
	//ask, do we need conversion to float?
	if (vqmt_requestPlanes(inst, request) != 0) {
		std::cerr << "Unable to request planes" << std::endl;
		return 6;
	}

	bool input1Float = request[0].requestedPlanes[(int)VQMT_ValueType::V_32f][(int)VQMT_CC::Y];
	bool input2Float = request[1].requestedPlanes[(int)VQMT_ValueType::V_32f][(int)VQMT_CC::Y];*/

	std::vector<uint8_t> chBuff1(w*h/2);
	//std::vector<uint8_t> chBuff2;

	VQMT_RawBuffer* buff1;
	VQMT_RawBuffer* buff2 = nullptr; 

	buff1 = vqmt_allocateBuffer("yuv420p", w, h);
	if(!buff1) {
		std::cerr << "Can't allocate buffer for orig image" << std::endl;
		return 8;
	}

	if (files > 1) {
		buff2 = vqmt_allocateBuffer("yuv420p", w, h);
		if (!buff2) {
			std::cerr << "Can't allocate buffer for distorted image" << std::endl;
			return 8;
		}
	}
	//std::vector<float> fBuff1;
	//std::vector<float> fBuff2;

	int frame;
	for (frame = 0;; frame++) {
		//uint8_t* u8data[2] = { nullptr, nullptr };
		//float* f32data[2] = { nullptr, nullptr };

		//reading input 1
		//chBuff1.resize(w*h);
		readY(buff1, i1);
		readY(buff2, i2);

		/*for (int y = 0; y < h; ++y) {
			i1.read()
		}
		i1.read((char*)chBuff1.data(), w*h);*/
		//u8data[0] = chBuff1.data();
		//
		////optional conversion to float
		//if (input1Float) {
		//	fBuff1.resize(w*h);
		//	std::copy(chBuff1.begin(), chBuff1.end(), fBuff1.begin());
		//	f32data[0] = fBuff1.data();
		//}
		//
		//if (files > 1) {
		//	//reading input 2
		//	chBuff2.resize(w*h);
		//	i2.read((char*)chBuff2.data(), w*h);
		//	u8data[1] = chBuff2.data();
		//
		//	//optional conversion to float
		//	if (input2Float) {
		//		fBuff2.resize(w*h);
		//		std::copy(chBuff2.begin(), chBuff2.end(), fBuff2.begin());
		//		f32data[1] = fBuff2.data();
		//	}
		//}

		if (!i1 || (files > 1 && !i2)) break;

		//filling input images
		VQMT_MetricImage mi1{ w,h };
		VQMT_MetricImage mi2{ w,h };

		if (!vqmt_convertImage(buff1, converter, 0, &mi1)) {
			std::cerr << "Can't convert original image for frame " << frame << std::endl;
			return 6;
		}
		if (!vqmt_convertImage(buff2, converter, 1, &mi2)) {
			std::cerr << "Can't convert distorted image for frame " << frame << std::endl;
			return 6;
		}

		/*mi1.planes[(int)VQMT_ValueType::V_8u][(int)VQMT_CC::Y] = VQMT_PlaneDescription{ int(w * sizeof(uint8_t)), 0.f, 255.f, u8data[0] };
		mi2.planes[(int)VQMT_ValueType::V_8u][(int)VQMT_CC::Y] = VQMT_PlaneDescription{ int(w * sizeof(uint8_t)), 0.f, 255.f, u8data[1] };
		mi1.planes[(int)VQMT_ValueType::V_32f][(int)VQMT_CC::Y] = VQMT_PlaneDescription{ int(w * sizeof(float)), 0.f, 255.f, f32data[0] };
		mi2.planes[(int)VQMT_ValueType::V_32f][(int)VQMT_CC::Y] = VQMT_PlaneDescription{ int(w * sizeof(float)), 0.f, 255.f, f32data[1] };*/
		const VQMT_MetricImage* mis[] = { &mi1,&mi2 };

		//calculation of selected metric for i1 and i2 (or for i1 if it's no-reference metric)
		//vqmt_requestFrame is obligatory call before each frame calculation
		vqmt_requestCalculationMonotonic(inst, frame);
		vqmt_putFrame(inst, nullptr, frame, mis, nullptr);

		//skipping UV
		i1.read((char*)chBuff1.data(), w*h / 2);
		if (files > 1) {
			i2.read((char*)chBuff1.data(), w*h / 2);
		}
	}

	vqmt_putEof(inst, nullptr, frame);

	vqmt_freeImageConverter(converter);
	vqmt_freeBuffer(buff1);
	vqmt_freeBuffer(buff2);

	//tell metric that there are no more frames
	//and release it's memory
	vqmt_releaseMetric(inst);

	return 0;
}

int main(int argc, const char** argv) {
	if (argc == 1) {
		std::cout << "Usage: " << argv[0] << " yuv1 yuv2 format width height metric [variation]" << std::endl << "Format now can be only `yuv420p'\nSecond file can be `-' for noref metric\n";
		printMetricList();
		return 0;
	}
	if (argc != 7 && argc != 8) {
		std::cerr << "Wrong usage" << std::endl;
		return 1;
	}

	int w = std::stoi(argv[4]);
	int h = std::stoi(argv[5]);
	bool yuv420 = std::string(argv[3]) == "yuv420p";
	int files = std::string(argv[2]) == "-" ? 1 : 2;

	if (!yuv420) {
		std::cerr << "Only yuv420p is supported by this sample" << std::endl;
		return 4;
	}

	const VQMT_MetricInfo* metr = nullptr;

	//find specified metric
	metr = findMetric(argv[6], argc == 8 ? argv[7] : "");
	if (!metr) {
		std::cerr << "Unknown metric" << std::endl;
		return 1;
	}

	//check whether metric supports YYUV
	if (!metricSupports(metr, VQMT_MetricCC::Y)) {
		std::cerr << "Metric not supports Y-component" << std::endl;
		return 2;
	}

	if (metr->needReference && files < 2) {
		std::cerr << "Need at least 2 files for this metric" << std::endl;
		return 3;
	}

	std::ifstream i1(argv[1], std::ios::binary);
	if (!i1) {
		std::cerr << "Can't open file " << argv[1] << std::endl;
		return 4;
	}
	std::ifstream i2;
	if (files > 1) {
		i2.open(argv[2], std::ios::binary);
		if (!i2) {
			std::cerr << "Can't open file " << argv[2] << std::endl;
			return 4;
		}
	}

	calcMetric(metr, files, i1, i2, w, h);
}