#include "VPTProcess.h"

#include <stdlib.h>
#include <cuda_runtime.h>
#include "ErrorInfo.h"

#include "../utils/logger.hpp"

#define THRESHOLD 0.5

#define USE_YOLOV5


// #include "../../model/vptModeTrt/ga_vpt_init_net.h"
// #include "../../model/vptModeTrt/ga_vpt_predict_net.h"
// //#include "../../model/hidemodel_caffe_1108/ga_vpt_init_net_caffe2.h"
// //#include "../../model/hidemodel_caffe_1108/ga_vpt_predict_net_caffe2.h"

//model trt onnx 20210715
#include "../../model/vptModeTrt/ga_vpt_det_yolo_640x640.h" // debug by zsh
#include "../../model/vptModeTrt/ga_trt_yolo_vpt_calibrator.h" // debug by zsh

// // fpn
// #include "vpt_fpn_plugin_factory.h"
// #include "../../model/vptModeTrt/ga_trt_fpn_vpt_calibrator.h"


struct DetectResultInfo {
	vector< vector <float>> det_result;
	int task_id;
	int width;
	int height;
	unsigned long long ts;
};



static long long get_cur_time_ms(){
	chrono::time_point<chrono::system_clock, chrono::milliseconds> tpMicro
		= chrono::time_point_cast<chrono::milliseconds>(chrono::system_clock::now());

	return tpMicro.time_since_epoch().count();
}

VPTProcess::VPTProcess(/* args */)
{
	m_taskTrackerMap.clear();
}

VPTProcess::~VPTProcess()
{
}

void VPTProcess::config_fpn_param(ctools_init_params& param, int batch_size) {
	// param.model_type_ = MODEL_FPN;

	// param.weight_file_ = NULL;
	// param.net_file_ = NULL;

	// param.data_process_str_ =
	// 	//"CopyData_CPU2GPU_U8;"
	// 	"TypeConvert_U8_F32;"
	// 	"ResizeMaxPad_F32_F32,test_size,720,test_max_size,1280,max_height,736,max_width,1280,"
	// 	"submean_b,103.94,submean_g,116.78,submean_r,123.68,"
	// 	"variance_rev_b,0.017,variance_rev_g,0.017,variance_rev_r,0.017;"
	// 	"NHWC2NCHW_F32"
	// 	;
	// param.need_im_info_ = 1; // true 

	// if (param.engine_type_ == ENGINE_MCAFFE2)
	// {
	// 	/*param.weight_array_ = (unsigned char*)ga_vpt_init_net_caffe2;
	// 	param.weight_array_len_ = ga_vpt_init_net_len_caffe2;
	// 	param.net_array_ = (unsigned char*)ga_vpt_predict_net_caffe2;
	// 	param.net_array_len_ = ga_vpt_predict_net_len_caffe2;*/
	// }
	// else if (param.engine_type_ == ENGINE_TENSORRT)
	// {
	// 	param.weight_array_ = (uint8_t*)ga_vpt_init_net;
	// 	param.weight_array_len_ = ga_vpt_init_net_len;
	// 	param.net_array_ = (uint8_t*)ga_vpt_predict_net;
	// 	param.net_array_len_ = ga_vpt_predict_net_len;

	// 	memset(param.tensorrt_param_str_, 0, sizeof(param.tensorrt_param_str_));

	// 	std::string g_data_mode = "FP32";
	// 	bool g_is_create_calibrator = false;

	// 	sprintf(param.tensorrt_param_str_, "max_batchsize %d,"
	// 		"data_mode %s,"
	// 		"is_create_calibrator %d,"
	// 		"input_names data im_info,"
	// 		"output_names cls_prob bbox_pred_final rois_count_each_img",
	// 		batch_size, g_data_mode.c_str(), g_is_create_calibrator);

	// 	param.tensorrt_calibrator_file_ = NULL;// "trt_fpn_vpt_calibrator";
	// 	param.tensorrt_calibrator_array_len_ = ga_trt_fpn_vpt_calibrator_len;// "trt_fpn_vpt_calibrator";
	// 	param.tensorrt_calibrator_array_ = (unsigned char*)ga_trt_fpn_vpt_calibrator;// "trt_fpn_vpt_calibrator";

	// 	param.tensorrt_plugin_factory_ptr_ = &(tensorrt_plugin_factory);
	// }
}

void VPTProcess::config_yolo_param(ctools_init_params& param, int batch_size) {

	param.model_type_ = MODEL_YOLOV5; 
	param.weight_file_ = NULL;
	param.net_file_ = NULL;

	param.data_process_str_ =
		//"CopyData_CPU2GPU_U8;"
		"TypeConvert_U8_F32;"
		"ResizeMaxMidPad_F32_F32,test_size,640,test_max_size,640,max_height,640,max_width,640,"
		"submean_b,0,submean_g,0,submean_r,0,"
		"variance_rev_b,0.00392,variance_rev_g,0.00392,variance_rev_r,0.00392;"
		"BGR2RGB_F32_F32;"
		"NHWC2NCHW_F32"
		;
	param.need_im_info_ = 0;

	if (param.engine_type_ == ENGINE_MCAFFE2)
	{
		/*param.weight_array_ = (unsigned char*)ga_vpt_init_net_caffe2;
		param.weight_array_len_ = ga_vpt_init_net_len_caffe2;
		param.net_array_ = (unsigned char*)ga_vpt_predict_net_caffe2;
		param.net_array_len_ = ga_vpt_predict_net_len_caffe2;*/
	}
	else if (param.engine_type_ == ENGINE_TENSORRT)
	{
		param.net_array_ = (uint8_t*)ga_vpt_det_yolo_640x640;
		param.net_array_len_ = ga_vpt_det_yolo_640x640_len; //debug by zsh

		memset(param.tensorrt_param_str_, 0, sizeof(param.tensorrt_param_str_));

		std::string g_data_mode = "FP32";
		bool g_is_create_calibrator = false;

		int g_is_onnx_model = 1;
		sprintf(param.tensorrt_param_str_, "max_batchsize %d,"
			"data_mode %s,"
			"is_create_calibrator %d,"
			"is_onnx_model %d,"
			"input_names images,"
			"output_names output",
			batch_size, g_data_mode.c_str(), g_is_create_calibrator, g_is_onnx_model);

		param.tensorrt_calibrator_array_len_ = ga_trt_yolo_vpt_calibrator_len;
		param.tensorrt_calibrator_array_ = (unsigned char*)ga_trt_yolo_vpt_calibrator;

	}
}

int VPTProcess::init(int gpuid, int max_batch_size)
{
	ctools_init_params param;
	param.thres_ = 0.5;
	param.log_level_ = 0;
	param.device_type_ = DEVICE_GPU;
	param.device_id_ = gpuid;
	param.engine_type_ = ENGINE_TENSORRT;
	param.trt_serialize_file_ = "./serialize_file/VPT";

	m_max_batch_size = max_batch_size;

#ifdef USE_YOLOV5
	config_yolo_param(param, m_max_batch_size);
#else
	config_fpn_param(param, m_max_batch_size);
#endif

	int flag = ctools_init(&det_handle, &param);
	if (SUCCESS != flag) {
		LOG_ERROR("VPTProcess init failed!");
	}

	return flag;
}

void VPTProcess::check_VPT_Result(VPT_Result & vResult) {
	int index = 0;
	for (int i = 0; i < vResult.objCount; i++) {
		if ((vResult.obj[i].right - vResult.obj[i].left) > 10 && (vResult.obj[i].bottom - vResult.obj[i].top) > 10)
		{
			if (index == i) {
				index++;
				continue;
			}
			vResult.obj[index].left = vResult.obj[i].left;
			vResult.obj[index].top = vResult.obj[i].top;
			vResult.obj[index].right = vResult.obj[i].right;
			vResult.obj[index].bottom = vResult.obj[i].bottom;
			vResult.obj[index].center_x = vResult.obj[i].center_x;
			vResult.obj[index].center_y = vResult.obj[i].center_y;
			vResult.obj[index].index = vResult.obj[i].index;
			vResult.obj[index].id = vResult.obj[i].id;
			vResult.obj[index].num = vResult.obj[i].num;
			vResult.obj[index].confidence = vResult.obj[i].confidence;
			vResult.obj[index].snap_flag = vResult.obj[i].snap_flag;
			index++;
		}
	}

	vResult.objCount = index;
}

vector<VPTProcessResult> VPTProcess::process(vector<DataInfo> vec_data) {

	vector<VPTProcessResult> vec_result;

	if(nullptr == det_handle){
		return vec_result;
	}

	long long t1 = get_cur_time_ms();

	int batchsize = vec_data.size();

	vector<DetectResultInfo> vec_detectResult;

	int cycle_time = batchsize / m_max_batch_size;
	cycle_time = (batchsize % m_max_batch_size) == 0 ?  cycle_time : (cycle_time + 1) ;
	
	for (int i = 0; i < cycle_time; i++) {
		int start_index = i * m_max_batch_size;
		int end_index = start_index + m_max_batch_size;
		if(end_index >= batchsize) {
			end_index = batchsize;
		}

		vector<sy_img> vec_img;
		vector<int> vec_task_id;
		vector<unsigned long long> vec_ts;
		for (int j = start_index; j < end_index; j++) {
			vec_img.push_back(vec_data[j].img);
			vec_task_id.push_back(vec_data[j].task_id);
			vec_ts.push_back(vec_data[j].frameIndex);
		}
		
		ctools_result *detresult;
		int res_status = ctools_process(det_handle, vec_img.data(), vec_img.size(), &detresult);

		for (size_t b = 0; b < vec_img.size(); b++) {
			ctools_result &cur_result = detresult[b];

			DetectResultInfo result_info;
			result_info.task_id = vec_task_id[b];
			result_info.ts = vec_ts[b];
			result_info.width = vec_img[b].w_;
			result_info.height = vec_img[b].h_;
			for (int c = 0; c < cur_result.obj_count_ && c < MAX_OBJ_COUNT; c++)
			{
				float x1 = cur_result.obj_results_[c].data_[2];
				float y1 = cur_result.obj_results_[c].data_[3];
				float x2 = cur_result.obj_results_[c].data_[4];
				float y2 = cur_result.obj_results_[c].data_[5];

				float class_id = cur_result.obj_results_[c].data_[0];
				float score = cur_result.obj_results_[c].data_[1];

				if (score >= THRESHOLD)
				{
					vector <float> obj;

					obj.push_back(x1);
					obj.push_back(y1);
					obj.push_back(x2);
					obj.push_back(y2);
					obj.push_back(score);
					obj.push_back(class_id);
					// detectResult[real_index].push_back(obj);
					result_info.det_result.push_back(obj);
				}
			}

			vec_detectResult.push_back(result_info);
		}
	}

	for (int i = 0; i < vec_detectResult.size(); i++) {
		DetectResultInfo& det_result_info = vec_detectResult[i];
		TaskTracker& task_tracker = m_taskTrackerMap[det_result_info.task_id];

		// TaskTracker& task_tracker = tools->taskTrackers[i];
		if (!task_tracker.tracker.GetState()) {
			continue;
		}

		int widht = det_result_info.width;
		int height = det_result_info.height;

		VPTProcessResult oneResult;
		oneResult.task_id = det_result_info.task_id;
		if (task_tracker.lastFrameIndex > 0) {
			// 非第一帧
			int update_times = det_result_info.ts - task_tracker.lastFrameIndex - 1;
			if (update_times < 0) {
				cout << "FrameIndex error !! lastFrameIndex= " << task_tracker.lastFrameIndex << " cur_frameindex = " << det_result_info.ts << endl;
			}

			cout << "update_times " << update_times << endl;

			for (int j = 0; j < update_times; j++) { // 无检测框跟踪
				VPT_Result unresult;
				unresult.objCount = task_tracker.tracker.update(widht, height, false, task_tracker.lastDetectResult, unresult.obj, task_tracker.lastDeleteObjectID);
				check_VPT_Result(unresult);
				oneResult.vecUnUsedResult.push_back(unresult);
			}
		}
		oneResult.vptResult.objCount = task_tracker.tracker.update(widht, height, true, det_result_info.det_result, oneResult.vptResult.obj, oneResult.vecDeleteObj);

		check_VPT_Result(oneResult.vptResult);

		task_tracker.lastDetectResult = det_result_info.det_result;
		task_tracker.lastDeleteObjectID = oneResult.vecDeleteObj;
		task_tracker.lastFrameIndex = det_result_info.ts;

		oneResult.ts = det_result_info.ts;

		vec_result.push_back(oneResult);
	}

	return vec_result;
}

void VPTProcess::release() {
	if (det_handle) {
		ctools_release(&det_handle);
		det_handle = NULL;
	}

	m_taskTrackerMap.clear();
}

void VPTProcess::AddTaskTracker(const int taskID)
{
	TaskTracker t;
	t.TaskID = taskID;
	t.lastFrameIndex = 0;
	t.tracker.setYOLOv5(true);  // YOLOv5 要设为true, fpn 要设为false

	m_taskTrackerMap[taskID] = t;
}

void VPTProcess::FinishTaskTracker(const int taskID)
{
	if (m_taskTrackerMap.find(taskID) == m_taskTrackerMap.end()) {
		return;
	}
	
	m_taskTrackerMap.erase(taskID);
}

void VPTProcess::PauseTaskTracker(const int taskID)
{
	if (m_taskTrackerMap.find(taskID) == m_taskTrackerMap.end()) {
		return;
	}

	TaskTracker& t = m_taskTrackerMap[taskID];
	t.tracker.Pause();
}

void VPTProcess::RestartTaskTraker(const int taskID)
{
	if (m_taskTrackerMap.find(taskID) == m_taskTrackerMap.end()) {
		return;
	}

	TaskTracker& t = m_taskTrackerMap[taskID];
	t.tracker.ReSet();
}

void VPTProcess::DrawTracker(const int taskID, cv::Mat *img)
{
	if (m_taskTrackerMap.find(taskID) == m_taskTrackerMap.end()) {
		return;
	}

	TaskTracker& t = m_taskTrackerMap[taskID];
	t.tracker.addTracker(img);
}