#include "LivenessRGB.h"

LivenessRGB* LivenessRGB::m_instance;

static unsigned char liveness_rgb_model[] = {
#include "algorithm_module/models/RGBLiveness_mnn_v1.42.dat"
};


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

LivenessRGB::LivenessRGB()
{
	//char pwd[] = "rushuai.liu_deepcam";
	//static bool decode = false;
	//if (!decode) {
	//	for (int i = 0; i < sizeof(liveness_rgb_model); i++) {
	//		liveness_rgb_model[i] = liveness_rgb_model[i] ^ pwd[i % sizeof(pwd)];
	//	}
	//	decode = true;
	//}

	m_detector = std::shared_ptr<MNN::Interpreter>(MNN::Interpreter::createFromBuffer(liveness_rgb_model, sizeof(liveness_rgb_model)));
	MNN::ScheduleConfig config;
	MNN::BackendConfig backendConfig;
	backendConfig.precision = MNN::BackendConfig::Precision_High;
	backendConfig.power = MNN::BackendConfig::Power_High;
	backendConfig.memory = MNN::BackendConfig::Memory_High;
	config.backendConfig = &backendConfig;
	config.type = MNN_FORWARD_CPU;
	config.numThread = 4;
	m_session = m_detector->createSession(config);
	m_input_tensor = m_detector->getSessionInput(m_session, NULL);
	m_score_tensor = m_detector->getSessionOutput(m_session, "score");


	const float mean_vals[3] = { 123.675f, 116.28f, 103.53f };
	const float norm_vals[3] = { 0.0171247538f, 0.0175070028f, 0.0174291939f };
	::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)2;
	m_img_config.destFormat = (MNN::CV::ImageFormat)1;
	m_img_config.filterType = (MNN::CV::Filter)(1);
	m_img_config.wrap = (MNN::CV::Wrap)(1);

	m_pretreat = std::shared_ptr<MNN::CV::ImageProcess>(MNN::CV::ImageProcess::create(m_img_config));
}

LivenessRGB::~LivenessRGB()
{
	m_detector->releaseSession(m_session);
}

float LivenessRGB::LivenessDetect(const cv::Mat& img, const float* landmark)
{
	std::lock_guard<std::mutex> lock(m_mt);
	cv::Mat input = GetFace(img, landmark);

	if (input.empty())
	{
		return 0;
	}
	//cv::imshow("LivenessRGB", input);
	cv::resize(input, input, cv::Size(128, 128));
	
	m_pretreat->convert(input.data, 128, 128, input.step[0], m_input_tensor);
	m_detector->runSession(m_session);

	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());

	float* out = tensor_score->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]);
	float score = exp(out[1]) / sum;
	return score;
}

cv::Mat LivenessRGB::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 LivenessRGB::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;
}