#include "CenterFaceMnn.h"



static const unsigned char centerface_model[] = {
#include "algorithm_module/models/centerface_small_mnn.dat"
};

CenterFaceMnn* CenterFaceMnn::m_hInstance;


CenterFaceMnn* CenterFaceMnn::GetInstance()
{
	if (!m_hInstance)
	{
		m_hInstance = new CenterFaceMnn();
	}
	return m_hInstance;
}

CenterFaceMnn::CenterFaceMnn()
{
	m_detector = std::shared_ptr<MNN::Interpreter>(MNN::Interpreter::createFromBuffer(centerface_model, sizeof(centerface_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);

	const float mean_vals[3] = { 127.5f, 127.5f, 127.5f };
	const float norm_vals[3] = { 0.0078431373f, 0.0078431373f, 0.0078431373f };
	::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);

	input_tensor = m_detector->getSessionInput(m_session, NULL);
	hm_tensor = m_detector->getSessionOutput(m_session, "hm");
	wh_tensor = m_detector->getSessionOutput(m_session, "wh");
	reg_tensor = m_detector->getSessionOutput(m_session, "reg");
	lm_tensor = m_detector->getSessionOutput(m_session, "lm");
}

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

int CenterFaceMnn::Detect(const cv::Mat& img, std::vector<FaceInfo>& faces, float scale)
{
	std::lock_guard<std::mutex> lock(m_mt);
	faces.clear();
	if (img.cols * img.rows == 0)
	{
		return -1;
	}

	if (scale <= 0 || scale > 1) {
		scale = img.cols > img.rows ? 160.0f / img.cols : 160.0f / img.rows;
	}
	int resize_w = (int)(img.cols * scale) / 32 * 32;
	int resize_h = (int)(img.rows * scale) / 32 * 32;
	float scale_w = (float)img.cols / (float)resize_w;
	float scale_h = (float)img.rows / (float)resize_h;


	cv::Mat input;
	cv::resize(img, input, cv::Size(resize_w, resize_h));
	m_detector->resizeTensor(input_tensor, 1, 3, resize_h, resize_w);
	m_detector->resizeSession(m_session);
	//prepare data
	std::shared_ptr<MNN::CV::ImageProcess> pretreat(MNN::CV::ImageProcess::create(m_img_config));
	pretreat->convert(input.data, resize_w, resize_h, input.step[0], input_tensor);
	m_detector->runSession(m_session);

	MNN::Tensor tensor_hm(hm_tensor, MNN::Tensor::CAFFE);
	MNN::Tensor tensor_wh(wh_tensor, MNN::Tensor::CAFFE);
	MNN::Tensor tensor_reg(reg_tensor, MNN::Tensor::CAFFE);
	MNN::Tensor tensor_lm(lm_tensor, MNN::Tensor::CAFFE);

	hm_tensor->copyToHostTensor(&tensor_hm);
	wh_tensor->copyToHostTensor(&tensor_wh);
	reg_tensor->copyToHostTensor(&tensor_reg);
	lm_tensor->copyToHostTensor(&tensor_lm);

	std::vector<int> hm_shape = tensor_hm.shape();

	float* heatmap = tensor_hm.host<float>();
	float* wh = tensor_wh.host<float>();
	float* reg = tensor_reg.host<float>();
	float* lm = tensor_lm.host<float>();

	Decode(heatmap, wh, reg, lm, hm_shape[2], hm_shape[3], faces, 0.4, 0.3);
	for (int i = 0; i < faces.size(); i++)
	{
		FaceInfo& tmp = faces[i];
		tmp.x1 *= scale_w;
		tmp.x2 *= scale_w;
		tmp.y1 *= scale_h;
		tmp.y2 *= scale_h;

		for (int j = 0; j < 5; j++)
		{
			tmp.landmarks[j * 2] *= scale_w;
			tmp.landmarks[j * 2 + 1] *= scale_h;
		}
	}

	return faces.size();
}

void CenterFaceMnn::Decode(float* heatmap, float* scale, float* offset, float* landmarks, int h, int w, std::vector<FaceInfo>& faces, float scoreThresh, float nmsThresh)
{
	int fea_h = h;
	int fea_w = w;
	int d_w = w * 4;
	int d_h = h * 4;
	int spacial_size = fea_w*fea_h;


	float *scale1 = (float*)(scale);
	float *scale0 = scale1 + spacial_size;

	float *offset1 = (float*)(offset);
	float *offset0 = offset1 + spacial_size;

	std::vector<int> ids = GetIds(heatmap, fea_h, fea_w, scoreThresh);

	std::vector<FaceInfo> faces_tmp;
	for (int i = 0; i < ids.size() / 2; i++) {
		int id_h = ids[2 * i];
		int id_w = ids[2 * i + 1];
		int index = id_h*fea_w + id_w;

		float s0 = std::exp(scale0[index]) * 4;
		float s1 = std::exp(scale1[index]) * 4;
		//cout << s0 << "," << s1 << endl;
		float o0 = offset0[index];
		float o1 = offset1[index];

		//std::cout << s0 << " " << s1 << " " << o0 << " " << o1 << std::endl;

		float x1 = std::max(0., (id_w + o1 + 0.5) * 4 - s1 / 2);
		float y1 = std::max(0., (id_h + o0 + 0.5) * 4 - s0 / 2);
		float x2 = std::max(0., (id_w + o1 + 0.5) * 4 + s1 / 2);
		float y2 = std::max(0., (id_h + o0 + 0.5) * 4 + s0 / 2);
		x1 = std::min(x1, (float)d_w);
		y1 = std::min(y1, (float)d_h);
		x2 = std::min(x2, (float)d_w);
		y2 = std::min(y2, (float)d_h);

		//std::cout << x1 << " " << y1 << " " << x2 << " " << y2 << std::endl;

		FaceInfo facebox;
		facebox.x1 = x1;
		facebox.y1 = y1;
		facebox.x2 = x2;
		facebox.y2 = y2;
		facebox.score = heatmap[index];

		float box_w = x2 - x1;//s1?
		float box_h = y2 - y1;// s0 ?

		for (int j = 0; j < 5; j++) {
			float *xmap = (float*)landmarks + (2 * j)*spacial_size;
			float *ymap = (float*)landmarks + (2 * j + 1)*spacial_size;
			facebox.landmarks[2 * j] = x1 + xmap[index] * s1;
			facebox.landmarks[2 * j + 1] = y1 + ymap[index] * s0;

		}
		faces_tmp.push_back(facebox);
	}
	NMS(faces_tmp, faces, nmsThresh);
}

void CenterFaceMnn::NMS(std::vector<FaceInfo>& input, std::vector<FaceInfo>& output, float nmsthreshold)
{
	std::sort(input.begin(), input.end(),
		[](const FaceInfo& a, const FaceInfo& b)
	{
		return a.score > b.score;
	});

	int box_num = input.size();

	std::vector<int> merged(box_num, 0);

	for (int i = 0; i < box_num; i++)
	{
		if (merged[i])
			continue;

		output.push_back(input[i]);

		float h0 = input[i].y2 - input[i].y1 + 1;
		float w0 = input[i].x2 - input[i].x1 + 1;

		float area0 = h0 * w0;


		for (int j = i + 1; j < box_num; j++)
		{
			if (merged[j])
				continue;

			float inner_x0 = input[i].x1 > input[j].x1 ? input[i].x1 : input[j].x1;//std::max(input[i].x1, input[j].x1);
			float inner_y0 = input[i].y1 > input[j].y1 ? input[i].y1 : input[j].y1;

			float inner_x1 = input[i].x2 < input[j].x2 ? input[i].x2 : input[j].x2;  //bug fixed ,sorry
			float inner_y1 = input[i].y2 < input[j].y2 ? input[i].y2 : input[j].y2;

			float inner_h = inner_y1 - inner_y0 + 1;
			float inner_w = inner_x1 - inner_x0 + 1;


			if (inner_h <= 0 || inner_w <= 0)
				continue;

			float inner_area = inner_h * inner_w;

			float h1 = input[j].y2 - input[j].y1 + 1;
			float w1 = input[j].x2 - input[j].x1 + 1;

			float area1 = h1 * w1;

			float score;

			score = inner_area / (area0 + area1 - inner_area);

			if (score > nmsthreshold)
				merged[j] = 1;
		}

	}
}

std::vector<int> CenterFaceMnn::GetIds(float *heatmap, int  h, int w, float thresh)
{
	std::vector<int> ids;
	for (int i = 0; i < h; i++) {
		for (int j = 0; j < w; j++) {
			if (heatmap[i*w + j] > thresh) {
				ids.push_back(i);
				ids.push_back(j);
			}
		}
	}
	return ids;
}

const std::string CenterFaceMnn::GetModelVer()
{
	return m_model_ver;
}