#include <omp.h>
#include <MobileFaceFeatureMnn.h>


MobileFaceFeatureMnn* MobileFaceFeatureMnn::m_hInstance = nullptr;


#define FEATURE_IMG_SIZE 112

extern unsigned char mobile_face_model0[];
extern unsigned char mobile_face_model1[];
extern unsigned char mobile_face_model2[];

extern int mobile_face_model0_len;
extern int mobile_face_model1_len;
extern int mobile_face_model2_len;


MobileFaceFeatureMnn* MobileFaceFeatureMnn::GetInstance(){
	if (m_hInstance == nullptr){
		m_hInstance = new MobileFaceFeatureMnn;
	}
	return m_hInstance;
}

MobileFaceFeatureMnn::MobileFaceFeatureMnn(){
	const unsigned char pw[] = "~12321@AAA";
	std::vector<unsigned char> model;
	model.insert(model.end(), mobile_face_model0, mobile_face_model0 + mobile_face_model0_len);
	model.insert(model.end(), mobile_face_model1, mobile_face_model1 + mobile_face_model1_len);
	model.insert(model.end(), mobile_face_model2, mobile_face_model2 + mobile_face_model2_len);

	int len = mobile_face_model0_len + mobile_face_model1_len + mobile_face_model2_len;

	for (int i = 0; i < len; i++) {
		model[i] = model[i] ^ pw[i % sizeof(pw)];
	}

	//m_detector = std::shared_ptr<MNN::Interpreter>(MNN::Interpreter::createFromFile(R"(I:\DeepCamPro\DeepCamFace\Release\MobileFaceNet_v6.6.4.mnn)"));
	m_detector = std::shared_ptr<MNN::Interpreter>(MNN::Interpreter::createFromBuffer(model.data(), len));

	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_mobileface = m_detector->createSession(config);
	m_input_tensor = m_detector->getSessionInput(m_sess_mobileface, NULL);
	m_feature_tensor = m_detector->getSessionOutput(m_sess_mobileface, "feature");
	m_detector->resizeTensor(m_input_tensor, 1, 3, FEATURE_IMG_SIZE, FEATURE_IMG_SIZE);
	m_detector->resizeSession(m_sess_mobileface);

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

MobileFaceFeatureMnn::~MobileFaceFeatureMnn(){
	m_detector->releaseModel();
}

ret_t MobileFaceFeatureMnn::GetFaceFeature(const cv::Mat& img, const float* landmarks, float* feature){
	std::lock_guard<std::mutex> lock(m_mt);
	cv::Mat alignFace = FaceAlign(img, landmarks);
	//cv::imshow("alignFace", alignFace);
	//cv::waitKey(1);
	if (alignFace.cols != FEATURE_IMG_SIZE || alignFace.rows != FEATURE_IMG_SIZE){
		cv::resize(alignFace, alignFace, cv::Size(FEATURE_IMG_SIZE, FEATURE_IMG_SIZE));
	}

	std::shared_ptr<MNN::CV::ImageProcess> pretreat(MNN::CV::ImageProcess::create(m_img_config));
	pretreat->convert(alignFace.data, FEATURE_IMG_SIZE, FEATURE_IMG_SIZE, alignFace.step[0], m_input_tensor);

	m_detector->runSession(m_sess_mobileface);

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

	//std::vector<int> feature_shape = tensor_feature->shape();
	
	//float* face_feature = tensor_feature->host<float>();

	memcpy(feature, tensor_feature->host<float>(), 512 * sizeof(float));

	//feature.insert(feature.end(), face_feature, face_feature + 512);
	//for (int i = 0; i < 512; i++) {
	//	feature.push_back(face_feature[i]);
	//}
	return RET_OK;
}

cv::Mat MobileFaceFeatureMnn::estimateTrans(std::vector<cv::Point2f>& srcLmks, std::vector<cv::Point2f>& dstLmks)
{
	cv::Mat A(2 * srcLmks.size(), 4, CV_32FC1);
	cv::Mat B(2 * srcLmks.size(), 1, CV_32FC1);
	for (int i = 0; i < srcLmks.size(); i++)
	{
		A.at<float>(2 * i, 0) = srcLmks[i].x;
		A.at<float>(2 * i, 1) = -srcLmks[i].y;
		A.at<float>(2 * i, 2) = 1;
		A.at<float>(2 * i, 3) = 0;
		B.at<float>(2 * i, 0) = dstLmks[i].x;

		A.at<float>(2 * i + 1, 0) = srcLmks[i].y;
		A.at<float>(2 * i + 1, 1) = srcLmks[i].x;
		A.at<float>(2 * i + 1, 2) = 0;
		A.at<float>(2 * i + 1, 3) = 1;
		B.at<float>(2 * i + 1, 0) = dstLmks[i].y;
	}

	cv::Mat C, X;
	C = (A.t() * A);
	X = C.inv() * A.t() * B;
	cv::Mat transformMatrix(2, 3, CV_32FC1);
	transformMatrix.at<float>(0, 0) = X.at<float>(0, 0);
	transformMatrix.at<float>(0, 1) = -X.at<float>(1, 0);
	transformMatrix.at<float>(0, 2) = X.at<float>(2, 0);
	transformMatrix.at<float>(1, 0) = X.at<float>(1, 0);
	transformMatrix.at<float>(1, 1) = X.at<float>(0, 0);
	transformMatrix.at<float>(1, 2) = X.at<float>(3, 0);
	return transformMatrix;
}

cv::Mat MobileFaceFeatureMnn::FaceAlign(const cv::Mat& inputImage, const float* landmarks)
{
	std::vector<cv::Point2f> templateLmks;
	std::vector<cv::Point> salientLmks;
	for (int l = 0; l < 5; l++) {
		salientLmks.push_back(cv::Point(landmarks[l * 2], landmarks[l * 2 + 1]));
	}
	static double const template_points[] = { 30.2946, 65.5318, 48.0252, 33.5493, 62.7299, 51.6963, 51.5014, 71.7366, 92.3655, 92.2041 };
	static double shift_x = 8.0;
	static size_t n = sizeof(template_points) / (2 * sizeof(template_points[0]));
	for (int i = 0; i < n; i++)
	{
		templateLmks.push_back(cv::Point2f(template_points[i] + shift_x, template_points[i + n]));
	}

	cv::Mat normalizedImage;
	cv::Mat noFace(0, 0, CV_8UC1);

	cv::Point2f leftEye;
	cv::Point2f rightEye;

	leftEye = salientLmks.at(0);
	rightEye = salientLmks.at(1);

	cv::Point2f nose;
	cv::Point2f leftMouth;
	cv::Point2f rightMouth;

	nose = salientLmks.at(2);
	leftMouth = salientLmks.at(3);
	rightMouth = salientLmks.at(4);

	std::vector<cv::Point2f> srcLmks;
	srcLmks.push_back(leftEye);
	srcLmks.push_back(rightEye);
	srcLmks.push_back(nose);
	srcLmks.push_back(leftMouth);
	srcLmks.push_back(rightMouth);

	cv::Mat matTransform = estimateTrans(srcLmks, templateLmks);

	if (matTransform.empty())
	{
		normalizedImage = noFace;
	}
	else {
		cv::warpAffine(inputImage, normalizedImage, matTransform, cv::Size(FEATURE_IMG_SIZE, FEATURE_IMG_SIZE));
	}
	return normalizedImage;
}

ret_t MobileFaceFeatureMnn::MobileFaceFeatureCompare(const float* feature1, const float* feature2, float& fSimilarity) {
    fSimilarity = 0.f;
    
    if (feature1 == nullptr || feature2 == nullptr) {
        return RET_INVALID_FEATURE;
    }
    float tmp0 = 0, tmp1 = 0, tmp2 = 0;
    
    #pragma omp parallel for reduction(+:tmp0,tmp1,tmp2)
    for (int i = 0; i < 512; i++) {
        tmp0 += feature1[i] * feature2[i];
        tmp1 += feature1[i] * feature1[i];
        tmp2 += feature2[i] * feature2[i];
    }

    float score = (float(tmp0 / (sqrt(tmp1) * sqrt(tmp2))));

    // 拉高分数
    float max_score = log10f(1.01);
    float min_score = log10f(0.02);
    float ret, temp;
    if (score <= 0) {
        temp = fabs(score);
        if (temp > 0.3) {
            ret = 0.0001;
        }
        else {
            ret = temp;
        }
    }
    else {
        temp = log10f(score + 0.01);
        fSimilarity = fabsf(temp - min_score) / (max_score - min_score);
    }
    return RET_OK;
}