#include "SegDetector.h"
#include "DnnDefine.h"

#ifndef _myminmax
#define _myminmax
#define mymin(a,b) (a) < (b) ? (a) : (b)
#define mymax(a,b) (a) > (b) ? (a) : (b)
#endif


CSegDetector::CSegDetector(const char* strModelFile, int iBatchSize, int iInferMode, int iDeviceIndex)
{
	{
		static bool bFirst = true;
		if (bFirst)
		{
			DnnInfer_InitGlobalLog();
			bFirst = false;
		}
	}
	m_bInit = false;
	m_strModelFile = strModelFile;
	m_pInfer = NULL;
	m_iMinArea = 20;
	m_iMinHeight = 5;
	m_bEnableMask = true;
	m_bEnableContourArea = true;
	m_iBatchSize = iBatchSize;
	m_iInferMode = iInferMode;
	m_iDeviceInedx = iDeviceIndex;
}

CSegDetector::~CSegDetector()
{
	UnInit();
}

bool CSegDetector::Init()
{
	bool bRet = false;
	{
		InferenceParameterEx inferParam;
		DnnInfer_GetInferParamFormModelFile(m_strModelFile.c_str(), &inferParam);

		m_pInfer = DnnInfer_Init(m_strModelFile.c_str(), m_iBatchSize, m_iInferMode, m_iDeviceInedx);
		if (m_pInfer != NULL)
		{
			m_bInit = true;
			DnnInfer_GetInputInfo(m_pInfer, &m_iWidth, &m_iHeight, &m_iChannel, &m_iBatchSize, &m_iLabelNum);
			printf("GetInputInfo %d %d %d\n", m_iWidth, m_iHeight, m_iChannel);

			GpuDeviceInfo gpuDeviceInfo;
			DnnInfer_GetDeviceInfo(&gpuDeviceInfo);
			for (int i = 0; i < gpuDeviceInfo.gpu_count; i++)
			{
				std::cout << "显卡名称：" << gpuDeviceInfo.gpu_name[i] << std::endl;
				std::cout << "显存大小：" << gpuDeviceInfo.gpu_memory_size[i] << " MB" << std::endl;
				std::cout << "一个block的共享内存大小：" << gpuDeviceInfo.gpu_sharedMemPerBlock[i] << " KB" << std::endl;
				std::cout << "block最大线程数：" << gpuDeviceInfo.gpu_maxThreadsPerBlock[i] << std::endl;
			}
			//m_iMinArea = 20;
			//m_iMinHeight = 5;
			bRet = true;
		}
	}
	return bRet;
}

bool CSegDetector::UnInit()
{
	if (m_bInit)
	{
		m_bInit = false;
		if (m_pInfer != NULL)
		{
			DnnInfer_Close(m_pInfer);
		}
	}
	return true;
}
bool CSegDetector::Execute(cv::Mat& image)
{
	if (!m_bInit)
		return false;
	m_vecDetectResult.clear();

	m_iHeight = image.rows;
	m_iWidth = image.cols;

	double dFx, dFy, dXoffset, dYoffset;

	ProcessInput(image, m_inputDatum, dFx, dFy, dXoffset, dYoffset);

	std::chrono::steady_clock::time_point startTime = std::chrono::steady_clock::now();
	//Infer(pObject, &input, iLabelArray, fpProbalityArray);

	DnnInfer_Seg_Infer(m_pInfer, &m_inputDatum, &m_outputDatum);

	auto elapsedSeconds = std::chrono::steady_clock::now() - startTime;
	printf("infer cost time %lf\n", std::chrono::duration<double, std::milli>(elapsedSeconds).count());
	int iMaskSize = m_iHeight * m_iWidth;
	cv::Mat mask(m_iHeight, m_iWidth, CV_8UC1, (uchar*)((uchar*)m_outputDatum.data));

	std::vector<DetectorResult> vecDetectResult;

	MaskToDetectResult(mask, m_iMinArea, m_iMinHeight, cv::Size(m_iWidth, m_iHeight), vecDetectResult);

	for (int j = 0; j < vecDetectResult.size(); j++)
	{
		DetectorResult detectorResult = vecDetectResult[j];
		detectorResult.left = mymax((vecDetectResult[j].left - dXoffset) / dFx, 0.0);
		detectorResult.top = mymax((vecDetectResult[j].top - dYoffset) / dFy, 0.0);
		detectorResult.width = mymin(vecDetectResult[j].width / dFx, image.cols - detectorResult.left);
		detectorResult.height = mymin(vecDetectResult[j].height / dFy, image.rows - detectorResult.top);
		if (detectorResult.width > 0 && detectorResult.height > 10)
		{
			if (m_bEnableMask)
			{
				if (!detectorResult.mask.empty())
					cv::resize(detectorResult.mask, detectorResult.mask, cv::Size(detectorResult.width, detectorResult.height), cv::INTER_NEAREST);
			}

			if (m_bEnableContourArea)
			{
				detectorResult.area *= dFx * dFy;
			}

			m_vecDetectResult.emplace_back(detectorResult);
		}
	}
	return true;
}

bool CSegDetector::Execute(std::vector<cv::Mat>& images)
{
	if (!m_bInit)
		return false;
	m_vecDetectResult.clear();

	m_iHeight = images[0].rows;
	m_iWidth = images[0].cols;

	double dFx, dFy, dXoffset, dYoffset;

	ProcessInputs(images, m_inputDatum, dFx, dFy, dXoffset, dYoffset);

	std::chrono::steady_clock::time_point startTime = std::chrono::steady_clock::now();
	//Infer(pObject, &input, iLabelArray, fpProbalityArray);

	DnnInfer_Seg_Infer(m_pInfer, &m_inputDatum, &m_outputDatum);

	for (int i = 0; i < m_outputDatum.shape.N; i++)
	{
		auto elapsedSeconds = std::chrono::steady_clock::now() - startTime;
		printf("infer cost time %lf\n", std::chrono::duration<double, std::milli>(elapsedSeconds).count());
		int iMaskSize = m_iHeight * m_iWidth;
		cv::Mat mask(m_iHeight, m_iWidth, CV_8UC1, (uchar*)((uchar*)m_outputDatum.data + i * iMaskSize));

		std::vector<DetectorResult> vecDetectResult;

		MaskToDetectResult(mask, m_iMinArea, m_iMinHeight, cv::Size(m_iWidth, m_iHeight), vecDetectResult);

		for (int j = 0; j < vecDetectResult.size(); j++)
		{
			DetectorResult detectorResult = vecDetectResult[j];
			detectorResult.left = mymax((vecDetectResult[j].left - dXoffset) / dFx, 0.0);
			detectorResult.top = mymax((vecDetectResult[j].top - dYoffset) / dFy, 0.0);
			detectorResult.width = mymin(vecDetectResult[j].width / dFx, m_iWidth - detectorResult.left);
			detectorResult.height = mymin(vecDetectResult[j].height / dFy, m_iHeight - detectorResult.top);
			if (detectorResult.width > 0 && detectorResult.height > 10)
			{
				if (m_bEnableMask)
				{
					if (!detectorResult.mask.empty())
						cv::resize(detectorResult.mask, detectorResult.mask, cv::Size(detectorResult.width, detectorResult.height), cv::INTER_NEAREST);
				}

				if (m_bEnableContourArea)
				{
					detectorResult.area *= dFx * dFy;
				}

				m_vecDetectResult.emplace_back(detectorResult);
			}
		}
	}

	return true;
}

bool CSegDetector::ExecuteV2(cv::Mat& image)
{
	if (!m_bInit)
		return false;
	m_vecDetectResult.clear();

	m_iHeight = image.rows;
	m_iWidth = image.cols;

	double dFx, dFy, dXoffset, dYoffset;

	ProcessInput(image, m_inputDatum, dFx, dFy, dXoffset, dYoffset);

	std::chrono::steady_clock::time_point startTime = std::chrono::steady_clock::now();
	//Infer(pObject, &input, iLabelArray, fpProbalityArray);

	DnnInfer_Seg_InferV2(m_pInfer, &m_inputDatum, &m_outputDatum, &m_outputProbDatum);

	auto elapsedSeconds = std::chrono::steady_clock::now() - startTime;
	printf("infer cost time %lf\n", std::chrono::duration<double, std::milli>(elapsedSeconds).count());
	int iMaskSize = m_iHeight * m_iWidth;

	cv::Mat mask(m_iHeight, m_iWidth, CV_8UC1, (uchar*)((uchar*)m_outputDatum.data));
	cv::Mat maxProbMask(m_iHeight, m_iWidth, CV_32FC1, ((float*)m_outputProbDatum.data));
	cv::imshow("maxProb", maxProbMask);
	cv::waitKey(0);

	std::vector<DetectorResult> vecDetectResult;

	//ProbToDetectResult(m_outputProbDatum, m_iMinArea, m_iMinHeight, cv::Size(m_iWidth, m_iHeight), vecDetectResult);

	MaskToDetectResult(mask, m_iMinArea, m_iMinHeight, cv::Size(m_iWidth, m_iHeight), vecDetectResult);

	for (int j = 0; j < vecDetectResult.size(); j++)
	{
		DetectorResult detectorResult = vecDetectResult[j];
		detectorResult.left = mymax((vecDetectResult[j].left - dXoffset) / dFx, 0.0);
		detectorResult.top = mymax((vecDetectResult[j].top - dYoffset) / dFy, 0.0);
		detectorResult.width = mymin(vecDetectResult[j].width / dFx, image.cols - detectorResult.left);
		detectorResult.height = mymin(vecDetectResult[j].height / dFy, image.rows - detectorResult.top);
		if (detectorResult.width > 0 && detectorResult.height > 10)
		{
			if (m_bEnableMask)
			{
				if (!detectorResult.mask.empty())
				{
					cv::resize(detectorResult.mask, detectorResult.mask, cv::Size(detectorResult.width, detectorResult.height), cv::INTER_NEAREST);
				}
			}

			if (m_bEnableContourArea)
			{
				detectorResult.area *= dFx * dFy;
			}

			m_vecDetectResult.emplace_back(detectorResult);
		}
	}
	return true;
}

bool CSegDetector::ExecuteV3(cv::Mat& image)
{
	if (!m_bInit)
		return false;
	m_vecDetectResult.clear();

	m_iHeight = image.rows;
	m_iWidth = image.cols;

	double dFx, dFy, dXoffset, dYoffset;

	ProcessInput(image, m_inputDatum, dFx, dFy, dXoffset, dYoffset);

	std::chrono::steady_clock::time_point startTime = std::chrono::steady_clock::now();
	//Infer(pObject, &input, iLabelArray, fpProbalityArray);

	SegResult segResult;
	memset(&segResult, 0, sizeof(segResult));
	DnnInfer_Seg_InferV3(m_pInfer, &m_inputDatum, &segResult);

	auto elapsedSeconds = std::chrono::steady_clock::now() - startTime;
	printf("infer cost time %lf\n", std::chrono::duration<double, std::milli>(elapsedSeconds).count());
	int iMaskSize = m_iHeight * m_iWidth;

	cv::Mat mask(segResult.inferShape.H, segResult.inferShape.W, CV_8UC1, (segResult.pLabel));
	cv::Mat maxProbMask(segResult.inferShape.H, segResult.inferShape.W, CV_32FC1, (segResult.pProb + 1 * segResult.inferShape.H * segResult.inferShape.W));
	cv::imshow("maxProb", maxProbMask);
	cv::waitKey(0);

	std::vector<DetectorResult> vecDetectResult;

	dFx = segResult.dScale;
	dFy = segResult.dScale;
	dXoffset = segResult.offsetX;
	dYoffset = segResult.offsetY;

	//ProbToDetectResult(m_outputProbDatum, m_iMinArea, m_iMinHeight, cv::Size(m_iWidth, m_iHeight), vecDetectResult);

	MaskToDetectResult(mask, m_iMinArea, m_iMinHeight, cv::Size(segResult.inferShape.W, segResult.inferShape.H), vecDetectResult);

	for (int j = 0; j < vecDetectResult.size(); j++)
	{
		DetectorResult detectorResult = vecDetectResult[j];
		detectorResult.left = mymax((vecDetectResult[j].left - dXoffset) / dFx, 0.0);
		detectorResult.top = mymax((vecDetectResult[j].top - dYoffset) / dFy, 0.0);
		detectorResult.width = mymin(vecDetectResult[j].width / dFx, image.cols - detectorResult.left);
		detectorResult.height = mymin(vecDetectResult[j].height / dFy, image.rows - detectorResult.top);
		if (detectorResult.width > 0 && detectorResult.height > 10)
		{
			if (m_bEnableMask)
			{
				if (!detectorResult.mask.empty())
				{
					cv::resize(detectorResult.mask, detectorResult.mask, cv::Size(detectorResult.width, detectorResult.height), cv::INTER_NEAREST);
				}
			}

			if (m_bEnableContourArea)
			{
				detectorResult.area *= dFx * dFy;
			}

			m_vecDetectResult.emplace_back(detectorResult);
		}
	}
	return true;
}

bool CSegDetector::GetResults(std::vector<DetectorResult>& vecDetectorResults)
{
	vecDetectorResults = m_vecDetectResult;
	return true;
}

void CSegDetector::EnableMask(bool bEnableMask)
{
	m_bEnableMask = bEnableMask;
}

void CSegDetector::EnableContourArea(bool bEnableContourArea)
{
	m_bEnableContourArea = bEnableContourArea;
}

void CSegDetector::SetMinAreaAndHeight(int minArea, int minHeight)
{
	m_iMinArea = minArea;
	m_iMinHeight = minHeight;
}

bool CSegDetector::ProcessInput(cv::Mat& cvInput, Datum& input, double& dFx, double& dFy, double& dXoffset, double& dYoffset)
{

	cv::Mat resizeImage;
	int iModelWidth = m_iWidth;
	int iModelHeight = m_iHeight;
	float fResizeScale = 1.0;
	int padoffsetX = 0;
	int padoffsetY = 0;
	dFx = 1.0;
	dFy = 1.0;
	dXoffset = 0;
	dYoffset = 0;

	Shape shape(1, cvInput.channels(), cvInput.rows, cvInput.cols);
	input.Reshape(shape);
	cv::Size input_size = cvInput.size();
	const int channels = cvInput.channels();
	input.SetDataItemSize(1);
	uchar* datum_data = (uchar*)input.Getdata();

	unsigned int offsetN = shape.C * shape.H * shape.W;

	if (cvInput.isContinuous())
	{
		memcpy(datum_data, cvInput.data, offsetN);
	}

	return true;

	if (false && (cvInput.rows != iModelHeight || cvInput.cols != iModelWidth))
	{
		int iheight = cvInput.rows;
		int iwidth = cvInput.cols;

		double dHeightResizeScale = 1.0 * iModelHeight / iheight;
		double dWidthResizeScale = 1.0 * iModelWidth / iwidth;

		if (dHeightResizeScale < dWidthResizeScale)
		{
			fResizeScale = dHeightResizeScale;
		}
		else
		{
			fResizeScale = dWidthResizeScale;
		}
		cv::resize(cvInput, resizeImage, cv::Size(), fResizeScale, fResizeScale);
	}
	else
	{
		resizeImage = cvInput;
	}
	if (m_iChannel == 1)
	{
		if (resizeImage.channels() == 3)
		{
			cv::cvtColor(resizeImage, resizeImage, cv::COLOR_BGR2GRAY);
		}
		cv::Mat paddMat = cv::Mat::zeros(cv::Size(iModelWidth, iModelHeight), CV_8UC1);

		padoffsetX = (iModelWidth - resizeImage.cols) / 2;
		padoffsetY = (iModelHeight - resizeImage.rows) / 2;
		resizeImage.copyTo(paddMat(cv::Rect(padoffsetX, padoffsetY, resizeImage.cols, resizeImage.rows)));
		//cvProcessedImage = paddMat;

		Shape shape(1, paddMat.channels(), paddMat.rows, paddMat.cols);
		input.Reshape(shape);
		cv::Size input_size = paddMat.size();
		const int channels = paddMat.channels();
		input.SetDataItemSize(1);
		uchar* datum_data = (uchar*)input.Getdata();

		unsigned int offsetN = shape.C * shape.H * shape.W;

		if (paddMat.isContinuous())
		{
			memcpy(datum_data, paddMat.data, offsetN);
		}
	}
	else
	{
		if (resizeImage.channels() == 1)
		{
			cv::cvtColor(resizeImage, resizeImage, cv::COLOR_GRAY2BGR);
		}

		cv::Mat paddMat = cv::Mat::zeros(cv::Size(iModelWidth, iModelHeight), CV_8UC3);

		padoffsetX = (iModelWidth - resizeImage.cols) / 2;
		padoffsetY = (iModelHeight - resizeImage.rows) / 2;
		resizeImage.copyTo(paddMat(cv::Rect(padoffsetX, padoffsetY, resizeImage.cols, resizeImage.rows)));

		Shape shape(1, paddMat.channels(), paddMat.rows, paddMat.cols);
		input.Reshape(shape);
		cv::Size input_size = paddMat.size();
		const int channels = paddMat.channels();
		input.SetDataItemSize(1);
		uchar* datum_data = (uchar*)input.Getdata();

		unsigned int offsetN = shape.C * shape.H * shape.W;

		if (paddMat.isContinuous())
		{
			memcpy(datum_data, paddMat.data, offsetN);
		}
	}

	dFx = fResizeScale;
	dFy = fResizeScale;
	dXoffset = padoffsetX;
	dYoffset = padoffsetY;


	return true;
}

bool CSegDetector::ProcessInputs(std::vector<cv::Mat>& cvInput, Datum& input, double& dFx, double& dFy, double& dXoffset, double& dYoffset)
{
	cv::Mat resizeImage;
	int iModelWidth = m_iWidth;
	int iModelHeight = m_iHeight;
	float fResizeScale = 1.0;
	int padoffsetX = 0;
	int padoffsetY = 0;
	dFx = 1.0;
	dFy = 1.0;

	Shape shape(cvInput.size(), cvInput[0].channels(), cvInput[0].rows, cvInput[0].cols);
	input.SetDataItemSize(1);
	input.Reshape(shape);

	uchar* datum_data = (uchar*)input.Getdata();

	for (int i = 0; i < cvInput.size(); i++)
	{
		unsigned int offsetN = shape.C * shape.H * shape.W;

		if (cvInput[0].isContinuous())
		{
			memcpy(datum_data + i * offsetN, cvInput[0].data, offsetN);
		}
	}


	return true;
}

bool CSegDetector::ProcessInput(cv::Mat& cvInputImage, cv::Mat& cvProcessedImage, double& dFx, double& dFy, double& dXoffset, double& dYoffset)
{
	cv::Mat resizeImage;
	int iModelWidth = m_iWidth;
	int iModelHeight = m_iHeight;
	float fResizeScale = 1.0;
	int padoffsetX = 0;
	int padoffsetY = 0;

	if (true || cvInputImage.rows != iModelHeight || cvInputImage.cols != iModelWidth)
	{
		int iheight = cvInputImage.rows;
		int iwidth = cvInputImage.cols;

		double dHeightResizeScale = 1.0 * iModelHeight / iheight;
		double dWidthResizeScale = 1.0 * iModelWidth / iwidth;

		if (dHeightResizeScale < dWidthResizeScale)
		{
			fResizeScale = dHeightResizeScale;
		}
		else
		{
			fResizeScale = dWidthResizeScale;
		}
		cv::resize(cvInputImage, resizeImage, cv::Size(), fResizeScale, fResizeScale);
	}
	else
	{
		resizeImage = cvInputImage;
	}
	if (m_iChannel == 1)
	{
		if (resizeImage.channels() == 3)
		{
			cv::cvtColor(resizeImage, resizeImage, cv::COLOR_BGR2GRAY);
		}
		cv::Mat paddMat = cv::Mat::zeros(cv::Size(iModelWidth, iModelHeight), CV_8UC1);

		padoffsetX = (iModelWidth - resizeImage.cols) / 2;
		padoffsetY = (iModelHeight - resizeImage.rows) / 2;
		resizeImage.copyTo(paddMat(cv::Rect(padoffsetX, padoffsetY, resizeImage.cols, resizeImage.rows)));
		cvProcessedImage = paddMat;
	}
	else
	{
		cv::Mat paddMat = cv::Mat::zeros(cv::Size(iModelWidth, iModelHeight), CV_8UC3);

		padoffsetX = (iModelWidth - resizeImage.cols) / 2;
		padoffsetY = (iModelHeight - resizeImage.rows) / 2;
		resizeImage.copyTo(paddMat(cv::Rect(padoffsetX, padoffsetY, resizeImage.cols, resizeImage.rows)));

		cvProcessedImage = paddMat;
	}

	dFx = fResizeScale;
	dFy = fResizeScale;
	dXoffset = padoffsetX;
	dYoffset = padoffsetY;


	return true;
}

bool CSegDetector::MaskToDetectResult(const cv::Mat& mask, float min_area, float min_height,
	const cv::Size& image_size, std::vector<DetectorResult>& vecDetectResult) {
	double min_val;
	double max_val;
	cv::minMaxLoc(mask, &min_val, &max_val);
	int max_bbox_idx = static_cast<int>(max_val);
	cv::Mat resized_mask;
	if (mask.size() != image_size)
	{
		cv::resize(mask, resized_mask, image_size, 0, 0, cv::INTER_NEAREST);
	}
	else
	{
		resized_mask = mask.clone();
	}
	for (int i = 1; i <= max_bbox_idx; i++) {
		cv::Mat bbox_mask = resized_mask == i;
		std::vector<std::vector<cv::Point>> contours;
		cv::findContours(bbox_mask, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
		if (contours.empty())
			continue;
		for (int j = 0; j < contours.size(); j++)
		{
			cv::Rect rect = cv::boundingRect(contours[j]);
			if (std::min(rect.size().width, rect.size().height) < min_height)
				continue;
			if (rect.size().area() < min_area)
				continue;
			DetectorResult detectResult;
			detectResult.left = rect.x;
			detectResult.top = rect.y;
			detectResult.width = rect.width;
			detectResult.height = rect.height;
			detectResult.label = i;
			detectResult.score = 1.0;
			if (m_bEnableContourArea)
			{
				detectResult.area = cv::contourArea(contours[j]);
			}
			else
			{
				detectResult.area = detectResult.width * detectResult.height;
			}
			//detectResult.area = cv::contourArea(contours[j]);
			sprintf_s(detectResult.name, "%d", detectResult.label);
			if (m_bEnableMask)
			{
				detectResult.hasMask = true;
				detectResult.mask = mask(rect).clone();
			}
			else
			{
				detectResult.hasMask = false;
			}

			vecDetectResult.emplace_back(detectResult);
		}
	}
	return true;
}

bool CSegDetector::ProbToDetectResult(Datum& probDatum, float min_area, float min_height, const cv::Size& image_size, std::vector<DetectorResult>& vecDetectResult)
{
	int score_map_size = probDatum.shape.C * probDatum.shape.H * probDatum.shape.W;
	const float* score_map_data =
		reinterpret_cast<const float*>(probDatum.data);
	int num_map_pixels = probDatum.shape.H * probDatum.shape.W;

	for (int i = 0; i < probDatum.shape.N; ++i) {
		const float* current_start_ptr = score_map_data + i * score_map_size;
		cv::Mat ori_score_mat(probDatum.shape.C,
			probDatum.shape.H * probDatum.shape.W,
			CV_32FC1, const_cast<float*>(current_start_ptr));
		ori_score_mat = ori_score_mat.t();
		cv::Mat score_mat(probDatum.shape.H, probDatum.shape.W, CV_32FC1);
		cv::Mat label_mat(probDatum.shape.H, probDatum.shape.W, CV_8UC1);
		for (int j = 0; j < ori_score_mat.rows; ++j) {
			double max_value;
			cv::Point max_id;
			minMaxLoc(ori_score_mat.row(j), 0, &max_value, 0, &max_id);
			if (max_value > 0.5)
			{
				score_mat.at<float>(j) = max_value;
				label_mat.at<uchar>(j) = max_id.x;
			}
			else
			{
				score_mat.at<float>(j) = max_value;
				label_mat.at<uchar>(j) = 0;
			}
		}
		MaskToDetectResult(label_mat, min_area, min_height, image_size, vecDetectResult);
	}
	return true;
}