#include"yolov5.hpp"
#define checkRuntime(op)  __check_cuda_runtime_yolov5((op), #op, __FILE__, __LINE__)

bool __check_cuda_runtime_yolov5(cudaError_t code, const char* op, const char* file, int line) {
	if (code != cudaSuccess) {
		const char* err_name = cudaGetErrorName(code);
		const char* err_message = cudaGetErrorString(code);
		printf("runtime error %s:%d  %s failed. \n  code = %s, message = %s\n", file, line, op, err_name, err_message);
		return false;
	}
	return true;
}

inline const char* severity_string(nvinfer1::ILogger::Severity t) {
	switch (t) {
	case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: return "internal_error";
	case nvinfer1::ILogger::Severity::kERROR:   return "error";
	case nvinfer1::ILogger::Severity::kWARNING: return "warning";
	case nvinfer1::ILogger::Severity::kINFO:    return "info";
	case nvinfer1::ILogger::Severity::kVERBOSE: return "verbose";
	default: return "unknow";
	}
}

class Logger : public ILogger
{
	void log(Severity severity, const char* msg)  noexcept
	{
		// suppress info-level messages
		if (severity != Severity::kINFO)
			printf("%s: %s zzj\n", severity_string(severity), msg);
	}
};

void warp_affine_bilinear(uint8_t* src, int src_line_size, int src_width, int src_height,
	float* dst, int dst_width, int dst_height, uint8_t fill_value, AffineMatrix matrix, cudaStream_t stream);

HAT::~HAT() {

	// ͬ同步结束，释放资源
	checkRuntime(cudaStreamSynchronize(stream));
	if (stream != nullptr)
	{
		checkRuntime(cudaStreamDestroy(stream));
	}

	if (!context) {
		context->destroy();
	}
	if (!engine) {
		engine->destroy();
	}
	if (!runtime) {
		runtime->destroy();
	}

	if (!output_data_host) {
		delete[] output_data_host;
	}

	if (!input_data_host) {
		delete[] input_data_host;
	}

	if (!input_data_device) {
		checkRuntime(cudaFree(input_data_device));
	}
}

bool HAT::initConfig(const char* enginefile, double confThreshold, double nmsThreshold) {

	//const std::string trtfile = enginefile;
	std::ifstream file(enginefile, std::ios::binary);
	char* trtModelStream = NULL;
	int size = 0;
	if (file.good()) {
		file.seekg(0, file.end);
		size = file.tellg();
		file.seekg(0, file.beg);
		trtModelStream = new char[size];
		assert(trtModelStream);
		file.read(trtModelStream, size);
		file.close();
	}
	else {
		return false;
	}

	// 初始化几个对象
	Logger logger;
	this->runtime = createInferRuntime(logger);

	assert(this->runtime != nullptr);
	this->engine = runtime->deserializeCudaEngine(trtModelStream, size);
	assert(this->engine != nullptr);
	this->context = engine->createExecutionContext();
	assert(this->context != nullptr);
	delete[] trtModelStream;

	auto input_dims = engine->getBindingDimensions(0);
	input_batch = input_dims.d[0];
	input_channel = input_dims.d[1];
	input_height = input_dims.d[2];
	input_width = input_dims.d[3];
	input_numel = input_batch * input_channel * input_height * input_width;

	//printf("input size=%d*%d*%d*%d\n", input_batch, input_channel, input_height,input_width);
	input_data_host = new float[input_numel * sizeof(float)]{ 0 };

	auto output_dims = engine->getBindingDimensions(1);
	output_numbox = output_dims.d[1];
	output_numprob = output_dims.d[2];
	num_classes = output_numprob - 5;
	output_numel = input_batch * output_numbox * output_numprob;

	//printf("output size=%d*%d*%d\n", input_batch, output_numbox, output_numprob);

	output_data_host = new float[sizeof(float) * output_numel]{ 0 };

	checkRuntime(cudaMalloc(&input_data_device, input_numel * sizeof(float)));
	checkRuntime(cudaMalloc(&output_data_device, output_numel * sizeof(float)));

	checkRuntime(cudaStreamCreate(&stream));

	confThreshold_ = confThreshold;
	nmsThreshold_ = nmsThreshold;

	return true;

}

void HAT::detect(cv::Mat& image, std::vector<Box>& boxes) 
{
	try
	{
		int width = image.cols;
		int height = image.rows;
		int channels = image.channels();
		int src_size = width * height * channels;
		uint8_t* psrc_device = nullptr;
		checkRuntime(cudaMalloc(&psrc_device, src_size));
		checkRuntime(cudaMemcpyAsync(psrc_device, image.data, src_size, cudaMemcpyHostToDevice, stream));

		AffineMatrix affine;
		affine.compute(width, height, input_width, input_height);
		warp_affine_bilinear(psrc_device, width * 3, width, height, input_data_device, input_width, input_height, 114, affine, stream);

		float* bindings[] = { input_data_device,output_data_device };
		context->enqueueV2((void**)bindings, stream, nullptr);

		checkRuntime(cudaMemcpyAsync(output_data_host, output_data_device, sizeof(float) * output_numel, cudaMemcpyDeviceToHost, stream));
		checkRuntime(cudaFree(psrc_device));

		std::vector<cv::Rect> bboxes;
		std::vector<int> classIds;
		std::vector<float> scores;
		for (int i = 0; i < output_numbox; ++i) {
			float* ptr = output_data_host + i * output_numprob;
			float objness = ptr[4];
			if (objness < confThreshold_) {
				continue;
			}
			float* pclass = ptr + 5;
			int label = std::max_element(pclass, pclass + num_classes) - pclass;
			float prob = pclass[label] * objness;
			if (prob < confThreshold_)
				continue;

			float cx = ptr[0];
			float cy = ptr[1];
			float w1 = ptr[2];
			float h1 = ptr[3];
			float x1 = cx - w1 / 2;
			float y1 = cy - h1 / 2;

			cv::Rect box;
			box.x = x1;
			box.y = y1;
			box.width = w1;
			box.height = h1;

			bboxes.emplace_back(box);
			classIds.emplace_back(label);
			scores.emplace_back(prob);
		}

		if (bboxes.size() <= 0) {
			return;
		}

		std::vector<int> indexes;
		cv::dnn::NMSBoxes(bboxes, scores, confThreshold_, nmsThreshold_, indexes);
		//有框
		for (size_t i = 0; i < indexes.size(); i++) {
			int idx = indexes[i];
			auto& ibox = bboxes[idx];
			float x = ibox.x;
			float y = ibox.y;
			float w = ibox.width;
			float h = ibox.height;

			// 对应原图左上角点、宽、高
			int image_base_x = static_cast<int>(affine.d2i[0] * x + affine.d2i[2]);
			int image_base_y = static_cast<int>(affine.d2i[0] * y + affine.d2i[5]);
			int image_base_width = static_cast<int>(affine.d2i[0] * w);
			int image_base_height = static_cast<int>(affine.d2i[0] * h);
			if ((image_base_width <= 0) || (image_base_height <= 0))
				continue;

			struct Box box;
			box.x = image_base_x;
			box.y = image_base_y;
			box.width = image_base_width;
			box.height = image_base_height;
			box.score = scores[idx];
			box.class_id = classIds[idx];
			boxes.push_back(box);
		}
	}
	catch (const std::exception& ex)
	{
		std::string errorMessage = "yolov5-";
		errorMessage += ex.what();
		throw std::runtime_error(errorMessage);
	}	
}