#include <Inc/FaceMask.h>

FaceMask* FaceMask::m_instance = nullptr;

unsigned char face_mask_model[] = {
#include "algorithm_module/models/deepcam_mask_v1.1_mnn.dat"
};


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


FaceMask::FaceMask()
{
	const float mean_vals[3] = { 127.5f, 127.5f, 127.5f };
	const float norm_vals[3] = { 0.0078431373, 0.0078431373, 0.0078431373 };
	const unsigned char pw[] = "rushuai.liu_deepcam";

	static bool decode = false;
	if (!decode) {
		for (int i = 0; i < sizeof(face_mask_model); i++) {
			face_mask_model[i] = face_mask_model[i] ^ pw[i % sizeof(pw)];
		}
		decode = true;
	}

	m_detector = std::shared_ptr<MNN::Interpreter>(MNN::Interpreter::createFromBuffer(face_mask_model, sizeof(face_mask_model)));


	
	MNN::ScheduleConfig config;
	config.type = MNN_FORWARD_CPU;
	config.numThread = 4;
	MNN::BackendConfig backendConfig;
	backendConfig.precision = MNN::BackendConfig::Precision_High;
	backendConfig.power = MNN::BackendConfig::Power_High;
	config.backendConfig = &backendConfig;

	m_sess_mask = m_detector->createSession(config);
	m_input_tensor = m_detector->getSessionInput(m_sess_mask, NULL);
	m_score_tensor = m_detector->getSessionOutput(m_sess_mask, "fc2");
	m_detector->resizeTensor(m_input_tensor, 1, 3, 112, 112);
	m_detector->resizeSession(m_sess_mask);

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

}

FaceMask::~FaceMask()
{
	m_detector->releaseSession(m_sess_mask);
}

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

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

	if (input.empty())
	{
		return ERR_IMG;
	}
	
	cv::resize(input, input, cv::Size(112, 112));
	cv::imshow("FaceMask", input);

	std::shared_ptr<MNN::CV::ImageProcess> pretreat(MNN::CV::ImageProcess::create(m_img_config));
	pretreat->convert(input.data, 112, 112, input.step[0], m_input_tensor);
	m_detector->runSession(m_sess_mask);

	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->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[0]) / sum;
	return ERR_OK;
}