#include "Inc/LivenessIRMnn.h"
#include "Inc/DeepCamDef.h"
#include <chrono>

#define M_PI float(3.1415926535898)



static unsigned char liveness_ir_model[] = {
#include "algorithm_module/models/IRLiveness_mnn.dat"
};

LivenessIRMnn* LivenessIRMnn::m_instance = nullptr;

LivenessIRMnn* LivenessIRMnn::GetInstance()
{
	if (m_instance == nullptr)
	{
		m_instance = new LivenessIRMnn();
	}
	return m_instance;
}

LivenessIRMnn::LivenessIRMnn()
{
	//const unsigned char liveness_ir_pw[] = "rushuai.liu_deepcam";
	//static bool decode = false;
	//if (!decode)
	//{
	//	for (int i = 0; i < sizeof(liveness_ir_model); i++)
	//	{
	//		liveness_ir_model[i] = liveness_ir_model[i] ^ liveness_ir_pw[i % sizeof(liveness_ir_pw)];
	//	}
	//	decode = true;
	//}
	
	m_ir_liveness = std::shared_ptr<MNN::Interpreter>(MNN::Interpreter::createFromBuffer(liveness_ir_model, sizeof(liveness_ir_model)));
	if (m_sess_liveness == nullptr) {
		MNN::ScheduleConfig config;
		config.type = MNN_FORWARD_OPENCL;
		config.numThread = 4;
		MNN::BackendConfig backendConfig;
		backendConfig.precision = MNN::BackendConfig::Precision_High;
		backendConfig.power = MNN::BackendConfig::Power_High;
		config.backendConfig = &backendConfig;

		m_sess_liveness = m_ir_liveness->createSession(config);
		m_input_tensor = m_ir_liveness->getSessionInput(m_sess_liveness, NULL);
		m_score_tensor = m_ir_liveness->getSessionOutput(m_sess_liveness, "score");
		m_ir_liveness->resizeTensor(m_input_tensor, 1, 1, 128, 128);
		m_ir_liveness->resizeSession(m_sess_liveness);

		const float mean_vals[1] = { 123.675f };
		const float norm_vals[1] = { 0.0171247538f };
		::memcpy(m_img_config.mean, mean_vals, sizeof(mean_vals));
		::memcpy(m_img_config.normal, norm_vals, sizeof(norm_vals));

		m_img_config.sourceFormat = MNN::CV::ImageFormat::GRAY;
		m_img_config.destFormat = MNN::CV::ImageFormat::GRAY;

		m_img_config.filterType = (MNN::CV::Filter)(1);
		m_img_config.wrap = (MNN::CV::Wrap)(1);
	}

}

LivenessIRMnn::~LivenessIRMnn()
{
	m_ir_liveness->releaseSession(m_sess_liveness);
}

long long LivenessIRMnn::get_cur_time()
{
	std::chrono::system_clock::duration d = std::chrono::system_clock::now().time_since_epoch();
	std::chrono::milliseconds mil = std::chrono::duration_cast<std::chrono::milliseconds>(d);
	return mil.count();

}

float LivenessIRMnn::IOU(float face0_x0, float face0_y0, float face0_x1, float face0_y1,
	float face1_x0, float face1_y0, float face1_x1, float face1_y1)
{
	float area0 = (face0_x1 - face0_x0 + 1) * (face0_y1 - face0_y0 + 1);
	float area1 = (face1_x1 - face1_x0 + 1) * (face1_y1 - face1_y0 + 1);

	float xx0 = std::max(face0_x0, face1_x0);
	float yy0 = std::max(face0_y0, face1_y0);
	float xx1 = std::min(face0_x1, face1_x1);
	float yy1 = std::min(face0_y1, face1_y1);

	float w = std::max(0.0f, xx1 - xx0 + 1);
	float h = std::max(0.0f, yy1 - yy0 + 1);

	float inter = w * h;
	return inter / (area0 + area1 - inter);
}

int LivenessIRMnn::GetLightBlurScore(const cv::Mat& face, float& light, float& blur)
{
	cv::Mat laplacian;
	cv::Laplacian(face, laplacian, CV_16S, 3);
	cv::Mat tmp_m, tmp_sd;
	cv::meanStdDev(laplacian, tmp_m, tmp_sd);
	blur = (float)tmp_sd.at<double>(0, 0);
	cv::meanStdDev(face, tmp_m, tmp_sd);
	light = (float)tmp_m.at<double>(0, 0);
	return ERR_OK;
}


int LivenessIRMnn::GetLivenessScore(const cv::Mat& face_img, float& score)
{
	std::shared_ptr<MNN::CV::ImageProcess> pretreat(MNN::CV::ImageProcess::create(m_img_config));
	pretreat->convert(face_img.data, 128, 128, face_img.step[0], m_input_tensor);

	m_ir_liveness->runSession(m_sess_liveness);

	std::shared_ptr<MNN::Tensor> tensor_score = std::make_shared<MNN::Tensor>(m_score_tensor, MNN::Tensor::CAFFE);
	m_score_tensor->copyToHostTensor(tensor_score.get());

	//std::vector<int> feature_shape = tensor_score->shape();

	float* out = tensor_score.get()->host<float>();

	float tmp_max = std::max(out[0], out[1]);
	out[0] -= tmp_max;
	out[1] -= tmp_max;
	float sum = exp(out[0]) + exp(out[1]);
	score = exp(out[1]) / sum;
	return 0;
}

int LivenessIRMnn::LivenessDetect(const cv::Mat& img, FaceInfo& face_info, float& score, bool motionFilter, int uniteFrames)
{
	static int frameCount = 0;
	static cv::Mat preGray;
	static float pre_x0 = 0.0f;
	static float pre_y0 = 0.0f;
	static float pre_x1 = 0.0f;
	static float pre_y1 = 0.0f;
	static std::vector<float> scores;
	score = 0.0f;

	static unsigned long preTime = get_cur_time();
	unsigned long curTime = get_cur_time();

	cv::Mat irFace, curGray;
	uniteFrames = std::min(5, std::max(1, uniteFrames));

	float iou = IOU(face_info.x1, face_info.y1, face_info.x2, face_info.y2, pre_x0, pre_y0, pre_x1, pre_y1);
	if (iou < 0.6 || curTime - preTime > 800000) {
		scores.clear();
	}
	preTime = curTime;
	pre_x0 = face_info.x1;
	pre_x1 = face_info.x2;
	pre_y0 = face_info.y1;
	pre_y1 = face_info.y2;

	cv::Rect max_face(face_info.x1, face_info.y1, face_info.x2 - face_info.x1, face_info.y2 - face_info.y1);
	if (max_face.x < 0) max_face.x = 0;
	if (max_face.y < 0) max_face.y = 0;
	if (max_face.x + max_face.width >= img.cols) max_face.width = img.cols - max_face.x - 1;
	if (max_face.y + max_face.height >= img.rows) max_face.height = img.rows - max_face.y - 1;

	// adjust the bounding box
	if (max_face.height > max_face.width) {
		int left = (max_face.height - max_face.width) / 2;
		max_face.y += left;
		max_face.height = max_face.width;
	}
	else {
		int left = (max_face.width - max_face.height) / 2;
		max_face.x += left;
		max_face.width = max_face.height;
	}

	cv::cvtColor(img, curGray, cv::COLOR_BGR2GRAY);

	if (preGray.empty()) {
		preGray = curGray.clone();
		return ERR_IMG;
	}

	if (motionFilter)
	{
		cv::Mat frame_binary;
		cv::absdiff(curGray(max_face), preGray(max_face), frame_binary);
		cv::threshold(frame_binary, frame_binary, 20, 255, cv::THRESH_BINARY); //change fixed value 20 to thr
		static cv::Mat kernel = getStructuringElement(cv::MORPH_RECT, cv::Size(3, 3));
		static cv::Mat kernel_ero = getStructuringElement(cv::MORPH_RECT, cv::Size(5, 5)); //the last variable can be (3,3)
		cv::morphologyEx(frame_binary, frame_binary, cv::MORPH_OPEN, kernel, cv::Point(-1, -1), 1);
		preGray = curGray.clone();
		cv::Scalar img_mean = cv::mean(frame_binary);
		if (img_mean.val[0] > 25) {
			return ERR_IMG;
		}
	}


	float pose_x = fabs(atan2((face_info.landmarks[3] - face_info.landmarks[1]), (face_info.landmarks[2] - face_info.landmarks[0])) * 180 / M_PI);
	float alpha_1 = atan2((face_info.landmarks[4] - face_info.landmarks[0]), (face_info.landmarks[5] - face_info.landmarks[1])) * 180 / M_PI;
	float alpha_2 = atan2((face_info.landmarks[2] - face_info.landmarks[4]), (face_info.landmarks[5] - face_info.landmarks[3])) * 180 / M_PI;
	// compute yaw
	float pose_y = fabs(alpha_1 - alpha_2);
	if (pose_x >= 25 || pose_y >= 25 + pose_x * 2) {
		return ERR_IMG;
	}

	irFace = FaceAlign(curGray, face_info.landmarks);
	//cv::imshow("IR", irFace);
	if (irFace.rows < 70 || irFace.cols < 70) {
		return ERR_IMG;
	}

	float light, blur;
	GetLightBlurScore(irFace, light, blur);
	if (light < 50 || light > 220 || blur < 10) {
		return ERR_IMG;
	}

	cv::resize(irFace, irFace, cv::Size(128, 128));
	GetLivenessScore(irFace, score);

	if (scores.size() >= uniteFrames) {
		scores.erase(scores.begin());
	}
	scores.push_back(score);

	for (int i = 0; i < scores.size(); i++) {
		if (scores[i] < 0.5f) {
			score = 0.0f;
		}
	}
	if (scores.size() < uniteFrames) {
		return ERR_OK;
	}
	return ERR_OK;
}

int LivenessIRMnn::LivenessDetect(const cv::Mat& img, int face_x1, int face_y1, int face_x2, int face_y2, const float* landmark, float& score, bool motionFilter, int uniteFrames)
{
	std::lock_guard<std::mutex> lock(m_mt);
	static int frameCount = 0;
	static cv::Mat preGray;
	static float pre_x0 = 0.0f;
	static float pre_y0 = 0.0f;
	static float pre_x1 = 0.0f;
	static float pre_y1 = 0.0f;
	static std::vector<float> scores;
	score = 0.0f;

	static unsigned long preTime = get_cur_time();
	unsigned long curTime = get_cur_time();

	cv::Mat irFace, curGray;
	uniteFrames = std::min(5, std::max(1, uniteFrames));

	float iou = IOU(face_x1, face_y1, face_x2, face_y2, pre_x0, pre_y0, pre_x1, pre_y1);
	if (iou < 0.6 || curTime - preTime > 800000) {
		scores.clear();
	}
	preTime = curTime;
	pre_x0 = face_x1;
	pre_x1 = face_x2;
	pre_y0 = face_y1;
	pre_y1 = face_y2;

	cv::Rect max_face(face_x1, face_y1, face_x2 - face_x1, face_y2 - face_y1);
	if (max_face.x < 0) max_face.x = 0;
	if (max_face.y < 0) max_face.y = 0;
	if (max_face.x + max_face.width >= img.cols) max_face.width = img.cols - max_face.x - 1;
	if (max_face.y + max_face.height >= img.rows) max_face.height = img.rows - max_face.y - 1;

	// adjust the bounding box
	if (max_face.height > max_face.width) {
		int left = (max_face.height - max_face.width) / 2;
		max_face.y += left;
		max_face.height = max_face.width;
	}
	else {
		int left = (max_face.width - max_face.height) / 2;
		max_face.x += left;
		max_face.width = max_face.height;
	}

	cv::cvtColor(img, curGray, cv::COLOR_BGR2GRAY);

	if (preGray.empty()) {
		preGray = curGray.clone();
		return ERR_IMG;
	}

	if (motionFilter)
	{
		cv::Mat frame_binary;
		cv::absdiff(curGray(max_face), preGray(max_face), frame_binary);
		cv::threshold(frame_binary, frame_binary, 20, 255, cv::THRESH_BINARY); //change fixed value 20 to thr
		static cv::Mat kernel = getStructuringElement(cv::MORPH_RECT, cv::Size(3, 3));
		static cv::Mat kernel_ero = getStructuringElement(cv::MORPH_RECT, cv::Size(5, 5)); //the last variable can be (3,3)
		cv::morphologyEx(frame_binary, frame_binary, cv::MORPH_OPEN, kernel, cv::Point(-1, -1), 1);
		preGray = curGray.clone();
		cv::Scalar img_mean = cv::mean(frame_binary);
		if (img_mean.val[0] > 25) {
			return ERR_IMG;
		}
	}


	float pose_x = fabs(atan2((landmark[3] - landmark[1]), (landmark[2] - landmark[0])) * 180 / M_PI);
	float alpha_1 = atan2((landmark[4] - landmark[0]), (landmark[5] - landmark[1])) * 180 / M_PI;
	float alpha_2 = atan2((landmark[2] - landmark[4]), (landmark[5] - landmark[3])) * 180 / M_PI;
	// compute yaw
	float pose_y = fabs(alpha_1 - alpha_2);
	if (pose_x >= 25 || pose_y >= 25 + pose_x * 2) {
		return ERR_IMG;
	}

	irFace = FaceAlign(curGray, landmark);
	//cv::imshow("IR", irFace);
	if (irFace.rows < 70 || irFace.cols < 70) {
		return ERR_IMG;
	}

	float light, blur;
	GetLightBlurScore(irFace, light, blur);
	if (light < 50 || light > 220 || blur < 10) {
		return ERR_IMG;
	}

	cv::resize(irFace, irFace, cv::Size(128, 128));
	GetLivenessScore(irFace, score);

	if (scores.size() >= uniteFrames) {
		scores.erase(scores.begin());
	}
	scores.push_back(score);

	for (int i = 0; i < scores.size(); i++) {
		if (scores[i] < 0.5f) {
			score = 0.0f;
		}
	}
	if (scores.size() < uniteFrames) {
		return ERR_IMG;
	}
	return ERR_OK;
}

cv::Mat LivenessIRMnn::GetFace(const cv::Mat &src, const float* landmark)
{
	std::vector<cv::Point2f> landmarkPoints;
	for (int l = 0; l < 5; l++) {
		landmarkPoints.push_back(cv::Point2f(landmark[2 * l], landmark[2 * l + 1]));
	}
	cv::Rect tmp = cv::boundingRect(landmarkPoints);
	int xo = tmp.x + tmp.width / 2;
	int yo = tmp.y + tmp.height / 2;
	int L = int(std::max(tmp.width, tmp.height) * 3.0);

	int tmpx, tmpy;
	cv::Rect rect;
	tmpx = std::max(0, xo - L / 2);
	tmpy = std::max(0, yo - L / 2);
	rect = cv::Rect(tmpx, tmpy, std::min(L, src.cols - 1 - tmpx), std::min(L, src.rows - 1 - tmpy));

	float tmp_landmark[10] = { 0.f };
	for (int l = 0; l < 5; l++) {
		tmp_landmark[2 * l] = landmark[2 * l] - rect.x;
		tmp_landmark[2 * l + 1] = landmark[2 * l + 1] - rect.y;
	}
	rect.x = rect.x / 2 * 2;
	rect.y = rect.y / 2 * 2;
	rect.width = rect.width / 2 * 2;
	rect.height = rect.height / 2 * 2;

	cv::Mat face = src(rect).clone();
	cv::Mat ret = FaceAlign(face, tmp_landmark);
	return ret;
}

cv::Mat LivenessIRMnn::FaceAlign(const cv::Mat& frame, const float* landmark, float face_rate)
{
	std::vector<cv::Point> landmarkPoints;
	for (int l = 0; l < 5; l++) {
		landmarkPoints.push_back(cv::Point2f(landmark[2 * l], landmark[2 * l + 1]));
	}

	cv::RotatedRect rotatedRect = cv::minAreaRect(landmarkPoints);
	float tmpLong = rotatedRect.size.width > rotatedRect.size.height ?
		rotatedRect.size.width :
		rotatedRect.size.height;
	rotatedRect.center.y -= tmpLong / 5;
	float ylongSize = tmpLong * 3.0f;
	float xlongSize = tmpLong * 3.0f;
	if ((rotatedRect.center.y - 1) < ylongSize) {
		ylongSize = rotatedRect.center.y - 1;
	}
	if ((rotatedRect.center.x - 1) < xlongSize) {
		xlongSize = rotatedRect.center.x - 1;
	}

	if ((frame.rows - rotatedRect.center.y - 1) < ylongSize) {
		ylongSize = frame.rows - rotatedRect.center.y - 1;
	}
	if ((frame.cols - rotatedRect.center.x - 1) < xlongSize) {
		xlongSize = frame.cols - rotatedRect.center.x - 1;
	}

	if (rotatedRect.center.y - ylongSize < 0 || ylongSize <= 0 || rotatedRect.center.y + ylongSize > frame.rows) {
		//LOGE("error: center y = %f, ylongSize = %f", rotatedRect.center.y, ylongSize);
		return cv::Mat();
	}

	if (rotatedRect.center.x - xlongSize < 0 || xlongSize <= 0 || rotatedRect.center.x + xlongSize > frame.cols) {
		//LOGE("error: center x = %f, xlongSize = %f", rotatedRect.center.x, xlongSize);
		return cv::Mat();
	}

	cv::Mat face = frame(cv::Range(rotatedRect.center.y - ylongSize, rotatedRect.center.y + ylongSize),
		cv::Range(rotatedRect.center.x - xlongSize, rotatedRect.center.x + xlongSize));

	float angle = rotatedRect.angle < -45.0 ? rotatedRect.angle + 90 : rotatedRect.angle;
	cv::Mat r = cv::getRotationMatrix2D(cv::Point(xlongSize, ylongSize), angle, 1.0);
	cv::warpAffine(face, face, r, cv::Size(2.0 * xlongSize, 2.0 * ylongSize));
	float faceLong = tmpLong * face_rate;
	if (faceLong > ylongSize || faceLong > xlongSize) {
		if (ylongSize < tmpLong * 0.85 || xlongSize < tmpLong * 0.85) {
			return cv::Mat();
		}
		faceLong = ylongSize > xlongSize ? xlongSize : ylongSize;
	}
	face = face(cv::Range(ylongSize - faceLong, ylongSize + faceLong),
		cv::Range(xlongSize - faceLong, xlongSize + faceLong));

	return face;
}