#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; // debug by zsh

	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;// debug by zsh

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

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

	bool isUseDet = true;
	int channels = 3;

	long long t1 = get_cur_time_ms();

	ctools_result *detresult;
	int res_status = ctools_process(det_handle, batch_img, batchsize, &detresult);

#ifdef LOG_INFO2
	long long t2 =  get_cur_time_ms();
	cout << "ctools_process time_using = "<< t2 - t1 << endl;
#endif

	vector <vector< vector <float>>> detectResult(batchsize);	//sort

	for (int b = 0; b < batchsize; 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];

			int imgid = b;
			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[imgid].push_back(obj);
			}
		}
	}

	int resIndex = startBatch;
	int detectIndex = 0;

	for (int i = startBatch; i < taskTrackers.size(); 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[detectIndex], result[resIndex].obj, deleteObjectID[detectIndex]);
			result[resIndex].objCount = objCount;
			vector<vector<float>>().swap(detectResult[detectIndex]);
			detectResult[detectIndex].clear();

			task_tracker.lastFrameIndex = vec_frameIndex[detectIndex];
			continue;
		}

		int update_times = vec_frameIndex[detectIndex] - task_tracker.lastFrameIndex - 1;
		if (update_times < 0)
		{
			cout << "FrameIndex error !! lastFrameIndex= "<< task_tracker.lastFrameIndex <<" cur_frameindex = " << vec_frameIndex[detectIndex] << 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;
			VPT_Result unresult;
			VPT_ObjInfo obj[MAX_OBJ_COUNT];
			unresult.objCount =task_tracker.tracker.update(task_tracker.ratioWidth, task_tracker.ratioHeight, false, detectResult[detectIndex], unresult.obj, deleteObjectID[detectIndex]);
			check_VPT_Result(unresult);
			unUsedResult[resIndex].push_back(unresult);
		}

		//有检测框输入的跟踪
		int objCount = task_tracker.tracker.update(task_tracker.ratioWidth, task_tracker.ratioHeight, true, detectResult[detectIndex], result[resIndex].obj, deleteObjectID[detectIndex]);
		result[resIndex].objCount = objCount;
		vector<vector<float>>().swap(detectResult[detectIndex]);
		detectResult[detectIndex].clear();

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

		check_VPT_Result(result[resIndex]);

		resIndex++;
		detectIndex++;
		if (resIndex == startBatch + batchsize)
			break;
	}

#ifdef LOG_INFO2
	long long t3 =  get_cur_time_ms();
	cout << "track time_using = "<< t3 - t2 << endl;
#endif

	vector <vector< vector <float>>>().swap(detectResult);

	return SUCCESS;
}

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