#include "face_det_ai_engine.h"

#include "../common/logger.hpp"
#include "../ai_platform/mvpt_process_assist.h"
#include "../ai_platform/macro_definition.h"

#include "opencv2/opencv.hpp"
#include "face_detect.h"

face_det_ai_engine::face_det_ai_engine(){
	m_max_batchsize = 10;
}

face_det_ai_engine::~face_det_ai_engine(){
	fd_release(&handle);
	clear();

	if(m_algorthim_ctx){
		aclrtSetDevice(m_devId);
        aclrtDestroyContext(m_algorthim_ctx);
    }
}

int face_det_ai_engine::init_ai_engine(const facedet_ai_engine_param &ai_param/*, person_det_algorthim_cache * cache*/){

	task_param_manager_ = task_param_manager::getInstance();

	m_devId = ai_param.sdk_param.devId;

	ACL_CALL(aclrtSetDevice(m_devId), ACL_ERROR_NONE, -1);
    ACL_CALL(aclrtCreateContext(&m_algorthim_ctx, m_devId), ACL_ERROR_NONE, -1);

	return fd_init(&handle, ai_param.sdk_param);
}

int face_det_ai_engine::ai_engine_process_batch(std::vector<std::string> &task_ids, sy_img *image_data_array, std::vector<onelevel_det_result> &result , std::vector<std::vector<int>> &deleteObjectID){

	map<string, map<algo_type, task_param_manager::algo_param_type_t_*>> && algor_param = task_param_manager_->get_task_other_params();

	const int total_batchsize = task_ids.size();
	if (total_batchsize <= 0){
		return 0;
	}

	aclrtSetDevice(m_devId);
	int ret = aclrtSetCurrentContext(m_algorthim_ctx);
	if(ACL_ERROR_NONE != ret){
		return 0;
	}

	fd_result *fd_result_ = new fd_result[total_batchsize];
	for (int i = 0; i < total_batchsize; ++i)
		fd_result_[i].info = new fd_info[50];

	do{
		int stride = m_max_batchsize;
		int steps = (total_batchsize + stride - 1) / stride;
		
		bool bError = false;
		for (int c = 0; c < steps; ++c) {
			int offset = c * m_max_batchsize;
			const int batchsize = (c == steps - 1) ? (total_batchsize - offset) : stride;
			
			int ret = fd_detect_batch(handle, image_data_array + offset, SY_FORMAT_BGR888, batchsize, fd_result_ + offset);
			if(ret < 0){
				LOG_ERROR(" fd_detect_batch error!!! image_size: {} model_batch_size: {}, step: [{}/{}] offset: {} batchsize: {}", total_batchsize, m_max_batchsize, c, steps, offset, batchsize);
				bError = true;
				break;
			}
		}

		if(bError){
			break;
		}

		// 属性检测使用人脸检测的原图，不需要切图
		int cur_index = 0;
		int img_index = 0;

		vector <vector< vector <float>>> detectResult(total_batchsize);  // sort
		auto task_id_iter = task_ids.cbegin();
		for (int c = 0; c < total_batchsize; ++c)
		{
			task_param_manager::algo_param_type_t_* cur_task_params = algor_param[*task_id_iter++][algorithm_type_t::FACE_SNAPSHOT];
			for (int i = 0; i < fd_result_[c].count; ++i)
			{
				if (!snapshot_legal_inarea(cur_task_params->basic_param->algor_valid_rect,
					fd_result_[c].info[i].face_position.left_, fd_result_[c].info[i].face_position.top_,
					fd_result_[c].info[i].face_position.left_ + fd_result_[c].info[i].face_position.width_,
					fd_result_[c].info[i].face_position.top_ + fd_result_[c].info[i].face_position.height_)
					|| fd_result_[c].info[i].face_pos_score < ((algor_config_param_snapshot *)cur_task_params->algor_param)->threshold)
				{
					continue;
				}

				vector <float> obj;
				obj.push_back(fd_result_[c].info[i].face_position.left_);
				obj.push_back(fd_result_[c].info[i].face_position.top_);
				obj.push_back(fd_result_[c].info[i].face_position.left_ + fd_result_[c].info[i].face_position.width_); //right
				obj.push_back(fd_result_[c].info[i].face_position.top_ + fd_result_[c].info[i].face_position.height_); //bottom
				obj.push_back(fd_result_[c].info[i].score);
				obj.push_back(1);		//统一index值为1
				//存入关键点信息
				for(int j = 0; j < FACIALFEAPOINTSIZE; ++j)
				{
					obj.push_back(fd_result_[c].info[i].facial_fea_point[j].x_);
					obj.push_back(fd_result_[c].info[i].facial_fea_point[j].y_);
				}

				//-added by zsh 添加姿态角信息------------------------------
				obj.push_back(fd_result_[c].info[i].roll);
				obj.push_back(fd_result_[c].info[i].yaw);
				obj.push_back(fd_result_[c].info[i].pitch);
				// cout << fabs(fd_result_[c].info[i].roll) << " " << fabs(fd_result_[c].info[i].yaw) << " " << fabs(fd_result_[c].info[i].pitch) << endl;
				//--------------------------------------------------------
				detectResult[c].push_back(obj);

	#if 0
				if (fd_result_[img_index].count > 1)
				{
					//选择居中且靠上的人脸作为唯一的结果
					float min_dis = numeric_limits<float>::max();
					int min_index = 0;
					float person_center_x = (float)(cur_persondet_result[c]->obj[i].right - cur_persondet_result[c]->obj[i].left) / 2.0;
					float person_center_y = (float)(cur_persondet_result[c]->obj[i].bottom - cur_persondet_result[c]->obj[i].top) / 6.0;
					for (int j = 0; j < fd_result_[img_index].count; ++j)
					{
						float cx = (float)fd_result_[img_index].info[j].face_position.left_ + (float)(fd_result_[img_index].info[j].face_position.width_) / 2.0;
						float cy = (float)fd_result_[img_index].info[j].face_position.top_ + (float)(fd_result_[img_index].info[j].face_position.height_) / 2.0;
						float dis = (person_center_x - cx) * (person_center_x - cx) + (person_center_y - cy) * (person_center_y - cy);
						if (dis < min_dis)
						{
							min_dis = dis;
							min_index = j;
						}
					}
					//姿态角控制
					if (fabs(fd_result_[img_index].info[min_index].roll) < pose_thresld[c] && fabs(fd_result_[img_index].info[min_index].yaw) < pose_thresld[c] && fabs(fd_result_[img_index].info[min_index].pitch) < pose_thresld[c])
					{
						cur_res.count = 1;
						cur_res.info = new fd_info[1];
						memcpy(&cur_res.info[0], &fd_result_[img_index].info[min_index], sizeof(fd_info));
					}
					else
					{
						cur_res.info = new fd_info[1];
						cur_res.count = 0;
					}

				}
				else if (fd_result_[img_index].count == 1 && fabs(fd_result_[img_index].info[0].roll) < pose_thresld[c] && fabs(fd_result_[img_index].info[0].yaw) < pose_thresld[c] && fabs(fd_result_[img_index].info[0].pitch) < pose_thresld[c]) //姿态角控制
				{
					cur_res.count = 1;
					cur_res.info = new fd_info[1];
					memcpy(&cur_res.info[0], &fd_result_[img_index].info[0], sizeof(fd_info));
				}
				else
				{
					cur_res.info = new fd_info[1];
					cur_res.count = 0;
				}
				_fd_result[vec_ids[c]].push_back(cur_res);

				for (int j = 0; j < cur_res.count; ++j)
				{
					++cur_index;
				}
	#endif

			}
		}

		//跟踪
		for (size_t real_index = 0; real_index < total_batchsize; real_index++) {
			string task_id = task_ids[real_index];

			bool isUseDet = true;
			vector<int> delete_ids;
			const float maxLen = std::sqrt(image_data_array[real_index].w_ * image_data_array[real_index].w_ + image_data_array[real_index].h_ * image_data_array[real_index].h_); //-modified by zsh 220719
			for (int j = 0; j < task_trackers[task_id].fusion_interval; ++j) {
				if (j == 0) {
					int objCount = task_trackers[task_id].tracker.update_v2(isUseDet, /*save lk = */true, /*center_dist = */true, maxLen, detectResult[real_index], result[real_index].obj, deleteObjectID[real_index]);
					result[real_index].obj_count = objCount;
					vector<vector<float>>().swap(detectResult[real_index]);
					detectResult[real_index].clear();
					isUseDet = false;
				} else {
					onelevel_det_result unresult;
					unresult.obj_count = task_trackers[task_id].tracker.update_v2(isUseDet, true, true, maxLen, detectResult[real_index], unresult.obj, deleteObjectID[real_index]);
				}
			}

			++real_index;
		}

		vector <vector< vector <float>>>().swap(detectResult);  // free memory.

		ret = total_batchsize;
	}while(0);

	if (fd_result_) {
		for (int i = 0; i < total_batchsize; ++i) {
			delete[] fd_result_[i].info;
			fd_result_[i].info = nullptr;
		}
		delete[] fd_result_;
		fd_result_ = nullptr;
	}

	return ret;
}

void face_det_ai_engine::clear()
{
	for (auto it = _fd_result.begin(); it != _fd_result.end();)
	{
		for (auto &fd : it->second)
		{
			delete[] fd.info;
			fd.info = nullptr;
		}
		_fd_result.erase(it++);
	}
}

void face_det_ai_engine::add_tracker(std::string task_id, int fusion_interval){
	LOG_INFO("face: tracker add task {}", task_id.c_str());
	task_tracker t;
	t.task_id = task_id;
	t.tracker.FusionInterval = fusion_interval; // 221117byzsh
	t.fusion_interval = fusion_interval;
	task_trackers.insert(std::make_pair(task_id, t));
}


void face_det_ai_engine::finish_task(std::string task_id)
{
	LOG_INFO("face: tracker finish task {}", task_id.c_str());
	auto iter = task_trackers.find(task_id);
	if (iter != task_trackers.end())
	{
		task_trackers.erase(task_id);
	}
}