#include"clasi.h"

#define checkRuntime3(op)  __check_cuda_runtime3((op), #op, __FILE__, __LINE__)

bool __check_cuda_runtime3(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_string3(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 Logger3 : public ILogger
{
	void log(Severity severity, const char* msg)  noexcept
	{
		// suppress info-level messages
		if (severity != Severity::kINFO)
			printf("%s: %s\n", severity_string3(severity), msg);
	}
};



bool clasification::initConfig(const char* modelpath) {
	std::ifstream file(modelpath, 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;
	}

	// ��ʼ����������
	Logger3 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;

	image_area = input_width * input_height;

	printf("input size=%d*%d*%d*%d\n", input_batch, input_channel, input_height, input_width);
	auto output_dims = engine->getBindingDimensions(1);
	output_classs = output_dims.d[0];
	num_classes = output_dims.d[1];
	printf("output size=%d*%d\n", output_classs, num_classes);

	output_numel = output_classs * num_classes;

	checkRuntime3(cudaMallocHost(&input_data_host, input_numel * sizeof(float)));
	checkRuntime3(cudaMalloc(&input_data_device, input_numel * sizeof(float)));

	checkRuntime3(cudaMallocHost(&output_data_host, sizeof(float) * output_numel));
	checkRuntime3(cudaMalloc(&output_data_device, sizeof(float) * output_numel));


	checkRuntime3(cudaStreamCreate(&stream));


	return true;


}


int clasification::detect(cv::Mat& image) {
	cv::Mat blob;
	cv::resize(image,blob, cv::Size(input_width, input_height));
	unsigned char* pimage = blob.data;
	float* host_b = input_data_host + image_area * 0;
	float* host_g = input_data_host + image_area * 1;
	float* host_r = input_data_host + image_area * 2;
	for (int i = 0; i < image_area; ++i, pimage += 3) {
		*host_r++ = (pimage[0] / 255.0f - 0.406) / 0.225;
		*host_g++ = (pimage[1] / 255.0f - 0.456) / 0.224;
		*host_b++ = (pimage[2] / 255.0f - 0.485) / 0.229;
	}

	checkRuntime3(cudaMemcpyAsync(input_data_device, input_data_host, input_numel * sizeof(float), cudaMemcpyHostToDevice, stream));
	float* bindings[] = { input_data_device, output_data_device };
	bool success = context->enqueueV2((void**)bindings, stream, nullptr);
	checkRuntime3(cudaMemcpyAsync(output_data_host, output_data_device, sizeof(float) * output_numel, cudaMemcpyDeviceToHost, stream));
	checkRuntime3(cudaStreamSynchronize(stream));

	float* prob = output_data_host;
	int predict_label = std::max_element(prob, prob + num_classes) - prob;
	//float confidence = prob[predict_label]; 
	//printf("predict: confidence = %f, label = %d\n", confidence, predict_label);

	return predict_label;


}


clasification::~clasification() {

	// ͬ���������ͷ���Դ

	checkRuntime3(cudaStreamSynchronize(stream));
	//if (stream != nullptr)
	//{
		checkRuntime3(cudaStreamDestroy(stream));
	//}

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

	if (!output_data_host) {
		delete[] output_data_host;

	}

	if (!output_data_device) {
		checkRuntime3(cudaFree(output_data_device));
	}

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


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

}