#include "sample_face_search.h"
#include <cv_face.h>
#include <MaterialBuffer.h>
#include "STFaceCache.h"
#include <logger.h>

#include <vector>
#include <stdio.h>

#include <opencv2/opencv.hpp>
#include <libyuv/convert.h>
#include <libyuv/convert_from_argb.h>
#define SUBSAMPLE(v, a) ((((v) + (a) - 1)) / (a))

#define MAX_FACE_IMAGE_WIDTH 640
#define MAX_FACE_IMAGE_HEIGHT 480

using namespace std;
using namespace cv;

//static cv_handle_t handle_verify = nullptr;
//static cv_handle_t handle_db = nullptr;
//static cv_handle_t handle_detect = nullptr;

#define _MAX_PATH 260

cv_feature_t *stface_extract_feature(stface_handles& handles, const char *image_path) {
	Mat bgr_image = imread(image_path);  // CV_PIX_FMT_BGR888
	if (!bgr_image.data) {
		return nullptr;
	}
	
	cv_feature_t *p_feature = nullptr;
	cv_face_t *p_face = nullptr;
	int face_count = 0;
	cv_result_t result = CV_OK;
	result = cv_face_detect(handles.handle_detect, bgr_image.data, CV_PIX_FMT_BGR888,
				bgr_image.cols, bgr_image.rows, bgr_image.step,
				CV_FACE_UP, &p_face, &face_count);
	if (face_count >= 1) {
		result = cv_verify_get_feature(handles.handle_verify,
						(unsigned char *)bgr_image.data, CV_PIX_FMT_BGR888,
						bgr_image.cols, bgr_image.rows, bgr_image.step,
						p_face, &p_feature, nullptr);
		if (result != CV_OK) {
			LOGP(DEBUG, "cv_verify_get_feature failed, error code %d", result);
		}
	} else {
		LOGP(DEBUG, "can't find face in %s", image_path);
	}
	// release the memory of face
	cv_face_release_detector_result(p_face, face_count);
	return p_feature;
}

cv_feature_t *stface_extract_feature(stface_handles& handles, const STFaceImage& image)
{
	cv_pixel_format stimgfmt = CV_PIX_FMT_GRAY8;
	
	if (image.width > MAX_FACE_IMAGE_WIDTH || image.height > MAX_FACE_IMAGE_HEIGHT)
	{
		LOG_WARN << "image too big" << LOG_ENDL;
		return nullptr;
	}
	
	uint8_t imgbuf[MAX_FACE_IMAGE_WIDTH * MAX_FACE_IMAGE_HEIGHT * 4];
	size_t imgbufSize = 0;
	
	if (image.mb_type == MB_Frame::MBFT_Y8)
	{
		memcpy(imgbuf, image.buff, image.size);//#todo avoid mem cpy
		
		imgbufSize = image.height * image.width;
		stimgfmt = CV_PIX_FMT_GRAY8;
	}
	else if (image.mb_type == MB_Frame::MBFT_NV12)
	{
		memcpy(imgbuf, image.buff, image.size);//#todo avoid mem cpy
		
		imgbufSize = image.height * image.width * 1.5;
		stimgfmt = CV_PIX_FMT_NV12;
	}
	else if (image.mb_type == MB_Frame::MBFT_RGB565)
	{
		//#todo libyuv can not deal no-normal-size image
		uint8* dst_y = (uint8*)(imgbuf);
        uint8* dst_u = (uint8*)(dst_y + (image.height * image.width));
        uint8* dst_v = (uint8*)(dst_u + (image.height * image.width / 4));
		
		int ret = libyuv::RGB565ToI420(
			image.buff, image.width * 2,//#todo test 
			dst_y, image.width, 
			dst_u, SUBSAMPLE(image.width, 2), 
			dst_v, SUBSAMPLE(image.width, 2), 
			image.width, image.height
			);
		
		imgbufSize = image.height * image.width * 1.5;
		stimgfmt = CV_PIX_FMT_YUV420P;
	}
	else if (image.mb_type == MB_Frame::MBFT_RGB888)
	{
		//int ret = libyuv::ARGBToBGRA(
		//	image.buff, image.width * 4,
		//	imgbuf, image.width * 4,
		//	image.width, image.height
		//	);
		//
		////memcpy(imgbuf, image.buff, image.size);
		//
		//imgbufSize = image.height * image.width * 4;
		//stimgfmt = CV_PIX_FMT_BGRA8888;

		int ret = libyuv::ABGRToI400(
				image.buff, image.width * 4,
				imgbuf, image.width,
				image.width, image.height
			);
		
		//memcpy(imgbuf, image.buff, image.size);
		
		imgbufSize = image.height * image.width;
		stimgfmt = CV_PIX_FMT_GRAY8;
		
		//{
		//	static int f = 0;
		//	++f;
		//	
		//	char fname[50];
		//	sprintf(fname, "st-%d-w%d-h%d.y8", f, image.width, image.height);
		//	FILE *pFile = fopen(fname, "wb");
		//	fwrite(imgbuf, 1, imgbufSize, pFile);
		//	fclose(pFile);
		//}
	}
	else if (image.mb_type == MB_Frame::MBFT_JPEG)
	{
		Mat matTmp = imdecode(_InputArray(image.buff, image.size), CV_LOAD_IMAGE_COLOR);
		
		//imwrite("aaa.jpg", matTmp);

		imgbufSize = matTmp.total() * matTmp.elemSize();
		memcpy(imgbuf, matTmp.ptr(), imgbufSize);

		stimgfmt = CV_PIX_FMT_BGR888;
	}
	else
	{
		LOG_WARN << "mb frame type not support" << LOG_ENDL;
		return nullptr;
	}

	//{
    //    static int f = 0;
	//	++f;
	//	
    //    char fname[50];
    //    sprintf(fname, "st-%d.img", f);
    //    FILE *pFile = fopen(fname, "wb");
    //    fwrite(image.buff, 1, image.size, pFile);
    //    fclose(pFile);
	//	
	//	sprintf(fname, "st-%d.img", f);
    //    pFile = fopen(fname, "wb");
    //    fwrite(imgbuf, 1, imgbufSize, pFile);
    //    fclose(pFile);
    //}

	cv_feature_t *p_feature = nullptr;
	cv_face_t *p_face = nullptr;
	int face_count = 0;
	cv_result_t result = CV_OK;
	result = cv_face_detect(handles.handle_detect, (unsigned char *)imgbuf, stimgfmt, image.width, image.height, image.width, CV_FACE_UP, &p_face, &face_count);
	if (face_count >= 1)
	{
		result = cv_verify_get_feature(handles.handle_verify, (unsigned char *)imgbuf, stimgfmt, image.width, image.height, image.width, p_face, &p_feature, nullptr);
		if (result != CV_OK)
		{
			LOGP(DEBUG, "cv_verify_get_feature failed, error code %d", result);
		}
	}
	else
	{
		LOGP(DEBUG, "can't find face in STFaceImage");
	}

	// release the memory of face
	cv_face_release_detector_result(p_face, face_count);
	return p_feature;
}

float stface_compare(stface_handles& handles, const STFaceImage& image1, const STFaceImage& image2)
{
	cv_feature_t* f1 = stface_extract_feature(handles, image1);
	if (f1 == nullptr)
	{
		LOGP(INFO, "can't find face in image1");
		return -1.0;
	}
	
	cv_feature_t* f2 = stface_extract_feature(handles, image2);
	if (f2 == nullptr)
	{
		cv_verify_release_feature(f1);
		LOGP(INFO, "can't find face in image2");
		return -1.0;
	}

	float score = -1.0;
	cv_result_t cv_result = cv_verify_compare_feature(handles.handle_verify, f1, f2, &score);
	if (cv_result != CV_OK) {
		LOGP(DEBUG, "cv_verify_compare_feature failed, error code %d", cv_result);
		score = -1.0;
	}
	
	cv_verify_release_feature(f1);
	cv_verify_release_feature(f2);
	return score;
}

int stface_db_add(stface_handles& handles, const char *image_path) {
	cv_feature_t *p_feature = stface_extract_feature(handles, image_path);
	if (!p_feature) {
		return -1;
	}
	int idx;
	cv_result_t cv_result = cv_verify_add_face(handles.handle_db, p_feature, &idx);
	if (cv_result != CV_OK) {
		LOGP(DEBUG, "cv_verify_add_face failed, error code %d", cv_result);
	}
	cv_verify_release_feature(p_feature);
	return idx;
}

int stface_db_add(stface_handles& handles, const STFaceImage& image)
{
	cv_feature_t *p_feature = stface_extract_feature(handles, image);
	if (!p_feature)
	{
		return -1;
	}
	int idx;
	cv_result_t cv_result = cv_verify_add_face(handles.handle_db, p_feature, &idx);
	if (cv_result != CV_OK)
	{
		LOGP(DEBUG, "cv_verify_add_face failed, error code %d", cv_result);
	}
	cv_verify_release_feature(p_feature);
	return idx;
}

bool stface_db_del(stface_handles& handles, int idx) {
	if (idx < 0) {
		LOGP(DEBUG, "invalid idx!");
		return false;
	}
	cv_result_t cv_result = CV_OK;
	cv_result = cv_verify_delete_face(handles.handle_db, idx);
	if (cv_result != CV_OK) {
		LOGP(DEBUG, "cv_verify_delete_face failed, error code %d", cv_result);
	}
	else {
		LOGP(DEBUG, "delete succeed");
	}
}

bool stface_db_save(stface_handles& handles, const char *db_path) {
	cv_result_t cv_result = CV_OK;
	cv_result = cv_verify_save_db(handles.handle_db, db_path);
	if (cv_result != CV_OK) {
		LOGP(DEBUG, "cv_verify_save_db failed, error code %d", cv_result);
		return false;
	}
	else {
		LOGP(DEBUG, "save done!");
	}

	return true;
}

bool stface_db_load(stface_handles& handles, const char *db_path) {
	cv_result_t cv_result = CV_OK;
	
	if (handles.handle_db == nullptr)
	{
		cv_result = cv_verify_create_db(&handles.handle_db);
		if (cv_result != CV_OK)
		{
			LOGP(DEBUG, "cv_verify_create_db failed, error code %d", cv_result);
			return false;
		}
	}
	
	cv_result = cv_verify_load_db(handles.handle_db, db_path);
	if (cv_result != CV_OK)
	{
		LOGP(DEBUG, "cv_verify_load_db failed, error code %d", cv_result);
		return false;
	}
	else
	{
		LOGP(DEBUG, "load done!");
	}

	return true;
}

bool stface_db_create(stface_handles& handles, const char *db_path)
{
	cv_result_t cv_result = CV_OK;

	if (handles.handle_db == nullptr)
	{
		cv_result = cv_verify_create_db(&handles.handle_db);
		if (cv_result != CV_OK)
		{
			LOGP(DEBUG, "cv_verify_create_db failed, error code %d", cv_result);
			return false;
		}
	}
	
	// add a bias image so that stfacesdk can save db
	int idx = stface_db_add(handles, "./bias.jpg");
	cout << "idx = " << idx << endl;
	
	cv_result = cv_verify_save_db(handles.handle_db, db_path);
	if (cv_result != CV_OK)
	{
		LOG_WARN << "cv_verify_save_db return false" << LOG_ENDL;
		return false;
	}
	
	return true;
}

bool stface_db_gen(stface_handles& handles, char *image_list, char *output_db_path) {
	bool bresult = true;
	FILE *fp_path = fopen(image_list, "r");
	if(!fp_path) {
		LOGP(DEBUG, "failed to load %s", image_list);
		return false;
	}
	std::vector<cv_feature_t *> list_feature;
	list_feature.clear();
	for (;;) {
		char image_path[1024];
		int num = fscanf(fp_path, "%s", image_path);
		if (num != 1) {
			bresult = false;
			break;
		}
		LOGP(DEBUG, "extracting %s", image_path);

		// get the face feature
		cv_feature_t *p_feature = stface_extract_feature(handles, image_path);
		if (!p_feature) {
			LOGP(DEBUG, "failed to extract image: %s", image_path);
			continue;
		}
		list_feature.push_back(p_feature);
	}
	fclose(fp_path);
	cv_verify_destroy_db(handles.handle_db);
	cv_result_t cv_result = CV_OK;
	cv_verify_create_db(&handles.handle_db);
	cv_result = cv_verify_build_db(handles.handle_db, &list_feature[0], list_feature.size());
	if (cv_result != CV_OK) {
		LOGP(DEBUG, "cv_verify_build_db failed, error code %d", cv_result);
		bresult = false;
	}
	cv_verify_save_db(handles.handle_db, output_db_path);
	for (int i = 0; i < list_feature.size(); i++) {
		cv_verify_release_feature(list_feature[i]);
	}
	cout << "db gen done!" << endl;

	return bresult;
}

bool stface_search_db(stface_handles& handles, char *image_path) {
	FILE *fp_path = fopen(image_path, "r");
	if (fp_path == nullptr) {
		LOGP(DEBUG, "invalid path !");
		return false;
	}
	fclose(fp_path);

	cv_feature_t *p_feature = stface_extract_feature(handles, image_path);
	if (p_feature == nullptr) {
		LOGP(DEBUG, "extract failed !");
		return false;
	}

	int top_k = 10;
	int *top_idxs = new int[top_k];
	float *top_scores = new float[top_k];
	unsigned int result_length = 0;
	cv_result_t cv_result = cv_verify_search_face(handles.handle_verify, handles.handle_db,
		p_feature, top_k,
		top_idxs, top_scores, &result_length);
	if (cv_result == CV_OK) {
		for (unsigned int t = 0; t < result_length; t++) {
			// const cv_feature_t item = result[t].item;
			LOGP(DEBUG, "%d\t%0.2f", top_idxs[t], top_scores[t]);
		}
	}
	else {
		LOGP(DEBUG, "cv_verify_search_face failed, error code %d", cv_result);
	}
	if (top_idxs) {
		delete[]top_idxs;
	}
	if (top_scores) {
		delete[]top_scores;
	}
	cv_verify_release_feature(p_feature);

	return true;
}

bool stface_search_db(stface_handles& handles, const STFaceImage& image, top_idx_score_vect_t& result)
{
	cv_feature_t *p_feature = stface_extract_feature(handles, image);
	if (p_feature == nullptr)
	{
		LOGP(DEBUG, "extract failed !");
		return false;
	}

	int top_k = 10;
	int *top_idxs = new int[top_k];
	float *top_scores = new float[top_k];
	unsigned int result_length = 0;
	cv_result_t cv_result = cv_verify_search_face(handles.handle_verify, handles.handle_db, p_feature, top_k, top_idxs, top_scores, &result_length);
	if (cv_result == CV_OK)
	{
		for (unsigned int t = 0; t < result_length; t++)
		{
			// const cv_feature_t item = result[t].item;
			LOGP(DEBUG, "%d\t%0.2f", top_idxs[t], top_scores[t]);
			
			result.push_back(TopIdxScore(top_idxs[t],  top_scores[t]));
		}
	}
	else
	{
		LOGP(DEBUG, "cv_verify_search_face failed, error code %d", cv_result);
	}
	if (top_idxs)
	{
		delete[] top_idxs;
	}
	if (top_scores)
	{
		delete[] top_scores;
	}
	cv_verify_release_feature(p_feature);

	return true;
}

bool stface_search_list(stface_handles& handles, char *image_path, char *list_path) {
	cv_feature_t *p_feature = stface_extract_feature(handles, image_path);
	if (p_feature == nullptr) {
		LOGP(DEBUG, "failed to extract image: %s", image_path);
		return false;
	}

	FILE *fp_path = fopen(list_path, "r");
	if(!fp_path) {
		LOGP(DEBUG, "failed to load %s", list_path);
		return false;
	}
	std::vector<cv_feature_t *> list_feature;
	list_feature.clear();
	for (;;) {
		char image_path[_MAX_PATH];
		int num = fscanf(fp_path, "%s", image_path);
		if (num != 1) {
			break;
		}
		LOGP(DEBUG, "extracting %s", image_path);

		// get the face feature
		cv_feature_t *p_feature = stface_extract_feature(handles, image_path);
		if (!p_feature) {
			LOGP(DEBUG, "failed to extract image: %s", image_path);
			continue;
		}
		list_feature.push_back(p_feature);
	}
	fclose(fp_path);


	const int top_k = 10;
	int top_idxs[top_k];
	float top_scores[top_k];
	unsigned int result_length = 0;
	cv_result_t cv_result = cv_verify_search_face_from_list(handles.handle_verify,
				&list_feature[0], list_feature.size(),
				p_feature, top_k,
				top_idxs, top_scores, &result_length);

	if (cv_result == CV_OK) {
		for (unsigned int t = 0; t < result_length; t++) {
			LOGP(DEBUG, "%d\t%0.2f", top_idxs[t], top_scores[t]);
		}
	} else {
		LOGP(DEBUG, "search face failed");
	}
	cv_verify_release_feature(p_feature);

	for (int i = 0; i < list_feature.size(); i++) {
		cv_verify_release_feature(list_feature[i]);
	}
	LOGP(DEBUG, "list search done!");

	return true;
}

void stface_get_help() {
	LOGP(DEBUG, "Usage: help | Get cmd list");
	LOGP(DEBUG, "Usage: search p_image_colorpath | Search image in db");
	LOGP(DEBUG, "Usage: add p_image_colorpath | Add image in db, return idx in db, if idx < 0 means failed");
	LOGP(DEBUG, "Usage: del idx | Delete image in db");
	LOGP(DEBUG, "Usage: save db_file | Save current db in db_file");
	LOGP(DEBUG, "Usage: load db_file | Load db in db_file");
	LOGP(DEBUG, "Usage: gen p_image_colorlist db_file | Gen images in p_image_colorlist and save in db_file");
	LOGP(DEBUG, "Usage: exit | Exit the program");
}

int stface_main(stface_handles& handles, int argc, char *argv[]) {
	cv_result_t cv_result = cv_face_create_detector(&handles.handle_detect, nullptr, CV_DETECT_ENABLE_ALIGN_21);
	if (cv_result != CV_OK){
		LOGP(DEBUG, "create detect handle failed, error code %d", cv_result);
	}
	cv_result = cv_verify_create_handle(&handles.handle_verify, "../../../models/verify.model");
	if (cv_result != CV_OK){
		LOGP(DEBUG, "create verify handle failed, error code %d", cv_result);
	}
	cv_result = cv_verify_create_db(&handles.handle_db);
	if (cv_result != CV_OK){
		LOGP(DEBUG, "create db handle failed, error code %d", cv_result);
	}
	
	if (handles.handle_detect != nullptr && handles.handle_verify != nullptr && handles.handle_db != nullptr) {
		// stface_db_gen("list.txt","out.db");
		LOGP(DEBUG, "Database is empty at the beginning");
		LOGP(DEBUG, "Please input 'help' to get the cmd list");
		char input_code[256];
		char image_path[_MAX_PATH];
		char db_path[_MAX_PATH];
		char command[256];
		input_code[0] = 0;
		while (1) {
			LOGP(DEBUG, ">>");
			if (!fgets(input_code, 256, stdin)) {
				LOGP(DEBUG, "read nothing");
				continue;
			}
			int input_length = strlen(input_code);
			if (input_length > 0 && input_code[input_length - 1] == '\n') {
				input_code[--input_length] = 0;
			}
			if (input_length == 0) {
				continue;
			}
			std::string str_input_code(input_code);
			if (strcmp(str_input_code.c_str(), "help") == 0) {
				stface_get_help();
			}
			else if (strcmp(str_input_code.substr(0, 3).c_str(), "add") == 0) {
				int input_number = sscanf(input_code, "%s%s", command, image_path);
				if (input_number != 2) {
					LOGP(DEBUG, "invalid! Usage: add p_image_colorpath");
					continue;
				}
				int idx = stface_db_add(handles, image_path);
				cout << "idx = " << idx << endl;
			}
			else if (strcmp(str_input_code.substr(0, 3).c_str(), "del") == 0) {
				int idx = -1;
				int input_number = sscanf(input_code, "%s%d", image_path, &idx);
				if (input_number != 2) {
					LOGP(DEBUG, "invalid! Usage: del idx(unsigned int");
					continue;
				}
				stface_db_del(handles, idx);
			}
			else if (strcmp(str_input_code.substr(0, 4).c_str(), "save") == 0) {
				int input_number = sscanf(input_code, "%s%s", command, image_path);
				if (input_number != 2) {
					LOGP(DEBUG, "invalid! Usage: save db_file");
					continue;
				}
				stface_db_save(handles, image_path);
			}
			else if (strcmp(str_input_code.substr(0, 4).c_str(), "load") == 0) {
				int input_number = sscanf(input_code, "%s%s", command, image_path);
				if (input_number != 2) {
					LOGP(DEBUG, "invalid! Usage: load db_file");
					continue;
				}
				stface_db_load(handles, image_path);
			}
			else if (strcmp(str_input_code.c_str(), "exit") == 0) {
				break;
			}
			else if (strcmp(str_input_code.substr(0, 3).c_str(), "gen") == 0) {
				int input_number = sscanf(input_code, "%s%s%s", command, image_path, db_path);
				if (input_number != 3) {
					LOGP(DEBUG, "invalid! Usage: gen p_image_colorlist_file db_file");
					continue;
				}
				stface_db_gen(handles, image_path, db_path);
			}
			else if (strcmp(str_input_code.substr(0, 6).c_str(), "search") == 0) {
				int input_number = sscanf(input_code, "%s%s", command, image_path);
				if (input_number != 2) {
					LOGP(DEBUG, "invalid! Usage: search p_image_colorpath");
					continue;
				}
				stface_search_db(handles, image_path);
			}
			else if (strcmp(str_input_code.substr(0, 10).c_str(), "listsearch") == 0) {
				char search_path[_MAX_PATH];
				int input_number = sscanf(input_code, "%s%s%s", command, search_path,
					image_path);
				if (input_number != 3) {
					LOGP(DEBUG, "invalid! Usage: listsearch p_image_colorsrcpath p_image_colorlistpath");
					continue;
				}
				stface_search_list(handles, search_path, image_path);
			}
			else {
				LOGP(DEBUG, "invalid cmd, please input 'help' to get the cmd list");
			}
		}
	}

	cv_face_destroy_detector(handles.handle_detect);
	cv_verify_destroy_db(handles.handle_db);
	cv_verify_destroy_handle(handles.handle_verify);
	return 0;
}

int stface_init_license(const char* lic_path)
{
	//加载授权证书
	cv_result_t cv_result = cv_face_public_init_license(lic_path, "license");
	if(cv_result != CV_OK)
	{
		LOGP(ERROR, "cv_face_public_init_license error %d", cv_result);
		return -1;
	}
	
	return 0;
}