#include"belt.hpp"

#define checkRuntime(op)  __check_cuda_runtimeBelt((op), #op, __FILE__, __LINE__)

bool __check_cuda_runtimeBelt(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\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);


BELT::~BELT() {

	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 BELT::initConfig(const char* enginefile, int* leftarea, int* rightarea, int* beltarea, float bigcoal_threshold, float beltdeviation_threshold, float confThreshold) {

	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;

	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;

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


	point left_p, right_p, belt_p;
	for (int i = 0; i < 8; i += 2) {

		left_p.x = leftarea[i];
		left_p.y = leftarea[i + 1];
		leftarea_.push_back(left_p);
	}

	for (int j = 0; j < 8; j += 2) {

		right_p.x = rightarea[j];
		right_p.y = rightarea[j + 1];
		rightarea_.push_back(right_p);
	}

	for (int q = 0; q < 8; q += 2) {

		belt_p.x = beltarea[q];
		belt_p.y = beltarea[q + 1];
		beltarea_.push_back(belt_p);
	}

	//Ƥ������
	belt_top_cy = (beltarea[1] + beltarea[3]) / 2;
	belt_bottom_cy = (beltarea[5] + beltarea[7]) / 2;

	belt_top_length = abs(beltarea[2] - beltarea[0]);
	belt_bottom_length = abs(beltarea[4] - beltarea[6]);

	k = (1 / belt_top_length - 1 / belt_bottom_length) / (belt_top_cy - belt_bottom_cy);
	if (abs(k) < 0.000001) {
		k += (abs(k) / k * 0.000001);
	}

	confThreshold_ = confThreshold;
	threshold = bigcoal_threshold;
	threshold_ = beltdeviation_threshold;

	return true;

}

int BELT::detect(cv::Mat& image, std::vector<Box>& boxes) {

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

		cv::Rect box;
		box.x = cx;
		box.y = cy;
		box.width = w1;
		box.height = h1;

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

	}


	if (bboxes.size() == 0) {

		return 1;
	}


	point box_p;
	float left_roller_num = 0, right_roller_num = 0;//���������
	float left_roller_length = 0, right_roller_length = 0;//�����й���������ռƤ��������ȵı���
	struct Box box;

	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 id = classIds[idx];

		int image_base_cx = static_cast<int>(affine.d2i[0] * x + affine.d2i[2]);
		int image_base_cy = 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;

		box_p.x = image_base_cx;
		box_p.y = image_base_cy;

		float mmpx = k * (image_base_cy - belt_top_cy) + 1 / belt_top_length;

		//�������й��Ҳ���Ƥ������
		if ((id != 1) && (!panduan(beltarea_, box_p))) {
			continue;
		}

		//�ж�λ��Ƥ�������ú���Ƿ��ú��
		float width_ratio = mmpx * image_base_width;
		float height_ratio = mmpx * image_base_height;
		float diameter = std::max(width_ratio, height_ratio);
		if (id == 0) {
			if (diameter < threshold) {
				continue;
			}

		}

		if (id == 1) {
			//λ����������й�
			if (panduan(leftarea_, box_p)) {
				left_roller_num += 1;
				left_roller_length += (mmpx * image_base_width);

			}
			//λ����������й�
			else if (panduan(rightarea_, box_p)) {
				right_roller_num += 1;
				right_roller_length += (mmpx * image_base_width);
			}

			else {
				continue;
			}
		}

		box.class_id = id;
		box.x = image_base_cx;
		box.y = image_base_cy;
		box.width = image_base_width;
		box.height = image_base_height;
		box.score = scores[idx];
		boxes.push_back(box);
	}

	if (left_roller_num == 0) {
		left_roller_num = 1;
	}

	if (right_roller_num == 0) {
		right_roller_num = 1;
	}

	//printf("�������=%f,��=%f;�������=%f,��=%f\n", left_roller_num,left_roller_length / left_roller_num, right_roller_num,right_roller_length / right_roller_num);

	if (abs(left_roller_length / left_roller_num - right_roller_length / right_roller_num) > threshold_) {

		return 0;//��ƫ
	}

	else {
		return 1;//����ƫ
	}
}