#include "VPTProcess.h"

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

#include "../FFNvDecoder/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"



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 */)
{
}

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;
}

int VPTProcess::process(sy_img * batch_img, int batchsize, vector<unsigned long long> vec_frameIndex, vector<VPT_Result>& result, vector<vector<int>>& deleteObjectID, vector<vector<VPT_Result>>& unUsedResult)
{
	if(nullptr == det_handle){
		return FAILED;
	}

	long long t1 = get_cur_time_ms();

	vector <vector< vector <float>>> detectResult(batchsize);
	int real_index = 0;

	int cycle_time = batchsize / m_max_batch_size;
	cycle_time = (batchsize % m_max_batch_size) == 0 ?  cycle_time : (cycle_time + 1) ;
	vector<sy_img> vec_img;
	for (int i = 0; i < cycle_time; i++) {
		vec_img.clear();

		int start_index = i * m_max_batch_size;
		int end_index = start_index + m_max_batch_size;
		if(end_index >= batchsize) {
			end_index = batchsize;
		}
		for (int j = start_index; j < end_index; j++) {
			vec_img.push_back(batch_img[j]);
		}
		
		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];
			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);
				}
			}
			real_index++;
		}
	}
	
	for (int i = 0; i < batchsize; i++){
		TaskTracker& task_tracker = taskTrackers[i];
		if (!task_tracker.tracker.GetState()) {
			continue;
		}

		if (task_tracker.lastFrameIndex <= 0) {
			// 第一帧，只做带检测框的跟踪
			int objCount = task_tracker.tracker.update(task_tracker.ratioWidth, task_tracker.ratioHeight, true, detectResult[i], result[i].obj, deleteObjectID[i]);
			result[i].objCount = objCount;
		} else {
			int update_times = vec_frameIndex[i] - task_tracker.lastFrameIndex - 1;
			if (update_times < 0) {
				cout << "FrameIndex error !! lastFrameIndex= " << task_tracker.lastFrameIndex << " cur_frameindex = " << vec_frameIndex[i] << endl;
			}
			// cout << "lastFrameIndex= " << task_tracker.lastFrameIndex << " cur_frameindex = " << vec_frameIndex[detectIndex] << endl;
			// cout << "update_times = " << update_times << endl;

			for (int j = 0; j < update_times; j++) { // 无检测框跟踪
				//cout << "taskTrackers size: " << task_tracker.size() << "  batchsize:" << detectResult.size() << "  update_times: "<< update_times << endl;
				task_tracker.lastUnUsedResult.objCount = task_tracker.tracker.update(task_tracker.ratioWidth, task_tracker.ratioHeight, false, task_tracker.lastDetectResult, task_tracker.lastUnUsedResult.obj, task_tracker.lastDeleteObjectID);
			}

			//有检测框输入的跟踪
			int objCount = task_tracker.tracker.update(task_tracker.ratioWidth, task_tracker.ratioHeight, true, detectResult[i], result[i].obj, deleteObjectID[i]);
			result[i].objCount = objCount;
		}

		check_VPT_Result(result[i]);

		task_tracker.lastDetectResult = detectResult[i];
		task_tracker.lastUnUsedResult.objCount = result[i].objCount;
		memcpy(task_tracker.lastUnUsedResult.obj, result[i].obj, sizeof(VPT_ObjInfo) * MAX_OBJ_COUNT);
		task_tracker.lastDeleteObjectID = deleteObjectID[i];

		// 记录帧序号
		task_tracker.lastFrameIndex = vec_frameIndex[i];
	}
}

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

	vector<TaskTracker>().swap(taskTrackers);
}

void VPTProcess::AddTaskTracker(const int taskID, const double rWidth, const double rHeight)
{
	TaskTracker t;
	t.TaskID = taskID;
	t.ratioWidth = rWidth;
	t.ratioHeight = rHeight;
	t.lastFrameIndex = 0;
	t.tracker.setYOLOv5(true);  // YOLOv5 要设为true, fpn 要设为false
	taskTrackers.push_back(t);
}

void VPTProcess::FinishTaskTracker(const int taskID)
{
	for (int i = 0; i < taskTrackers.size(); i++)
	{
		if (taskTrackers[i].TaskID == taskID)
		{
			taskTrackers.erase(taskTrackers.begin() + i);
			break;
		}
	}
}

void VPTProcess::PauseTaskTracker(const int taskID)
{
	for (int i = 0; i < taskTrackers.size(); i++)
	{
		if (taskTrackers[i].TaskID == taskID)
		{
			taskTrackers[i].tracker.Pause();
			break;
		}
	}
}

void VPTProcess::RestartTaskTraker(const int taskID)
{
	for (int i = 0; i < taskTrackers.size(); i++)
	{
		if (taskTrackers[i].TaskID == taskID)
		{
			taskTrackers[i].tracker.ReSet();
			break;
		}
	}
}

void VPTProcess::DrawTracker(const int taskID, cv::Mat *img)
{
	for (int i = 0; i < taskTrackers.size(); i++)
	{
		if (taskTrackers[i].TaskID == taskID)
		{
			taskTrackers[i].tracker.addTracker(img);
			break;
		}
	}
}