#include "ImageColorBalance.h"


ImageColorBalance::ImageColorBalance()
{
	//set the refresh frequency 
	m_nCurrentRefreshFrequency = 0;
}

void ImageColorBalance::Init(unsigned char* pImage, int nPanoWidth, int nPanoHeight, int nPanoPitch, int nCameraWidth, int nCameraHeight, int nCameraPitch,int nRefreshFrequency)
{
	m_pPanoImageBuffer = pImage;
	m_nPanoWidth = nPanoWidth;
	m_nPanoHeight = nPanoHeight;
	m_nPanoPitch = nPanoPitch;
	m_nCameraWidth = nCameraWidth;
	m_nCameraHeight = nCameraHeight;
	m_nCameraPitch = nCameraPitch;

	m_nRefreshFrequency = nRefreshFrequency;
}


void ImageColorBalance::FusionImageByBorderColor(unsigned char* pPano, cv::Rect InferRc, unsigned char* pCvTarget, cv::Rect TargetOverlapRc)
{
	if (TargetOverlapRc.width < 2)
		return;
	
	std::vector<bool> bAddOrSub;
	//according to pano image pointer to build a cv mat to store the pano image
	cv::Mat PanoImage = cv::Mat(this->m_nPanoHeight, this->m_nPanoWidth, CV_8UC3, pPano);
	//get the fusion part in pano image
	cv::Mat CvInferMat = PanoImage(InferRc);
	
	if(TargetOverlapRc.width==0 || TargetOverlapRc.height==0)
		return;

	//get the target image to cv mat
	cv::Mat CvTargetMat = cv::Mat(this->m_nCameraHeight, this->m_nCameraWidth, CV_8UC3, pCvTarget)(TargetOverlapRc);
	std::vector<cv::Mat> Mask;


	//GetGradientMask(CvInferMat, CvTargetMat, Mask, bAddOrSub);

	if (m_nCurrentRefreshFrequency == 0)
	{
		m_nCurrentRefreshFrequency = m_nRefreshFrequency;
		
	}
	else
	{
		m_nCurrentRefreshFrequency--;
	}

	std::vector<cv::Mat> InferChannels;
	cv::split(CvInferMat, InferChannels);
	
	for (int i = 0; i < 3; i++)
	{
		if (bAddOrSub[i])
			InferChannels[i] = InferChannels[i] + Mask[i];
		else
			InferChannels[i] = InferChannels[i] - Mask[i];
	}
	cv::merge(InferChannels, CvInferMat);
}



void ImageColorBalance::GetGradientMask(cv::Rect Infer, cv::Rect Target, cv::Mat& MaskTarget, cv::Mat& MaskInfer, std::vector<bool>& bAddOrSub)
{
	bAddOrSub.resize(3);

	//according to intersect rect rgb avg val to get the light base in this image

	cv::Mat WeigthMapTarget = cv::Mat(cv::Size(Infer.width, Infer.height), CV_8U, cv::Scalar(255));
	cv::Mat WeightMapInfer = cv::Mat(cv::Size(Infer.width, Infer.height), CV_8U, cv::Scalar(255));
	cv::Mat DistanceMaskTarget, DistanceMaskInfer;

	//set a column 0 in Mask
	WeigthMapTarget(cv::Rect(Infer.width - 1, 0, 1, Infer.height)).setTo(cv::Scalar(0));
	WeightMapInfer(cv::Rect(0, 0, 1, Infer.height)).setTo(cv::Scalar(0));
	DistanceMaskTarget.create(WeigthMapTarget.size(), WeigthMapTarget.type());
	DistanceMaskInfer.create(WeightMapInfer.size(), WeightMapInfer.type());
	//calculate a Weightmap in DistanceMask
	distanceATS_L1_8u(WeigthMapTarget.data, DistanceMaskTarget.cols, DistanceMaskTarget.rows, DistanceMaskTarget.step, DistanceMaskTarget.data);
	distanceATS_L1_8u(WeightMapInfer.data, DistanceMaskInfer.cols, DistanceMaskInfer.rows, DistanceMaskInfer.step, DistanceMaskInfer.data);

	std::shared_ptr<float> pfLightBaseDif(new float[3] {0});
	cv::Mat vTargetChannels, vInferChannels;
	DistanceMaskTarget.convertTo(DistanceMaskTarget, CV_32F);
	DistanceMaskInfer.convertTo(DistanceMaskInfer, CV_32F);

	//step is according to inferRc`s width
	float fStep = abs(1.f / Infer.width);

	vTargetChannels = DistanceMaskTarget * fStep;
	vInferChannels = DistanceMaskInfer * fStep;

	MaskTarget = (vTargetChannels);
	MaskInfer = (vInferChannels);

}


void ImageColorBalance::GetColorBalanceImage(unsigned char* pFusionImage, int& nWidth, int& nHeight, int& nPitch)
{
	pFusionImage = this->m_pPanoImageBuffer;
	nWidth = this->m_nPanoWidth;
	nHeight = this->m_nPanoHeight;
	nPitch = this->m_nPanoPitch;
}

int& ImageColorBalance::GetRefreshFrequency()
{
	return this->m_nCurrentRefreshFrequency;
}

const int ImageColorBalance::GetStdRefreshFrequency()
{
	return this->m_nRefreshFrequency;
}

int ImageColorBalance::FusionImageByAutoLightBalance(unsigned char* pPano, cv::Rect InferRc, unsigned char* pCvTarget, cv::Rect TargetOverlapRc)
{
	//将pPano全景数据转为mat
	cv::Mat PanoImage = cv::Mat(this->m_nPanoHeight, this->m_nPanoWidth, CV_8UC3, pPano);
	
	//从全景图像中获取需要色彩融合的部分
	cv::Mat ColorBalanceROI = PanoImage(InferRc);

	//if (m_InferWeightMask.empty())
	//	GetGradientMask();

	return 0;
}

cv::Mat ImageColorBalance::stretchImage(cv::Mat src)
{
	//stretch image to 0~1

	int row = src.rows;
	int col = src.cols;
	cv::Mat dst(row, col, CV_64FC1);
	double MaxValue = 0;
	double MinValue = 256.0;
	//find  maxvalue and minvalue in image
	for (int i = 0; i < row; i++) {
		for (int j = 0; j < col; j++) {
			MaxValue = std::max(MaxValue, src.at<double>(i, j));
			MinValue = std::min(MinValue, src.at<double>(i, j));
		}
	}
	//stretch image
	for (int i = 0; i < row; i++) {
		for (int j = 0; j < col; j++) {
			dst.at<double>(i, j) = (1.0 * src.at<double>(i, j) - MinValue) / (MaxValue - MinValue);
			//protect 
			if (dst.at<double>(i, j) > 1.0) {
				dst.at<double>(i, j) = 1.0;
			}
			else if (dst.at<double>(i, j) < 0) {
				dst.at<double>(i, j) = 0;
			}
		}
	}
	return dst;
}

cv::Mat ImageColorBalance::getPara(int radius)
{
	int size = radius * 2 + 1;
	cv::Mat dst(size, size, CV_64FC1);
	for (int i = -radius; i <= radius; i++) {
		for (int j = -radius; j <= radius; j++) {
			//center point is 0
			if (i == 0 && j == 0) {
				dst.at<double>(i + radius, j + radius) = 0;
			}
			else {
				//distance to center point
				dst.at<double>(i + radius, j + radius) = 1.0 / sqrt(i * i + j * j);
			}
		}
	}
	//get the sum of all value
	double sum = 0;
	for (int i = 0; i < size; i++) {
		for (int j = 0; j < size; j++) {
			sum += dst.at<double>(i, j);
		}
	}
	//normalize
	for (int i = 0; i < size; i++) {
		for (int j = 0; j < size; j++) {
			dst.at<double>(i, j) = dst.at<double>(i, j) / sum;
		}
	}
	return dst;
}

cv::Mat ImageColorBalance::NormalACE(cv::Mat src, int ratio, int radius)
{
	cv::Mat para = getPara(radius);
	int row = src.rows;
	int col = src.cols;
	int size = 2 * radius + 1;
	cv::Mat Z(row + 2 * radius, col + 2 * radius, CV_64FC1);
	//padding
	for (int i = 0; i < Z.rows; i++) {
		for (int j = 0; j < Z.cols; j++) {
			if ((i - radius >= 0) && (i - radius < row) && (j - radius >= 0) && (j - radius < col)) {
				//in the range of src,copy the src in z
				Z.at<double>(i, j) = src.at<double>(i - radius, j - radius);
			}
			else {
				Z.at<double>(i, j) = 0;
			}
		}
	}
	//init the dst
	cv::Mat dst(row, col, CV_64FC1);
	for (int i = 0; i < row; i++) {
		for (int j = 0; j < col; j++) {
			dst.at<double>(i, j) = 0.f;
		}
	}
	//convolution
	for (int i = 0; i < size; i++) {
		for (int j = 0; j < size; j++) {
			if (para.at<double>(i, j) == 0) continue;
			for (int x = 0; x < row; x++) {
				for (int y = 0; y < col; y++) {
					double sub = src.at<double>(x, y) - Z.at<double>(x + i, y + j);
					double tmp = sub * ratio;
					if (tmp > 1.0) tmp = 1.0;
					if (tmp < -1.0) tmp = -1.0;
					dst.at<double>(x, y) += tmp * para.at<double>(i, j);
				}
			}
		}
	}
	return dst;
}

cv::Mat ImageColorBalance::FastACE(cv::Mat src, int ratio, int radius)
{
	int row = src.rows;
	int col = src.cols;
	//if the image is too small,just return 0.5
	if (std::min(row, col) <= 2) {
		cv::Mat dst(row, col, CV_64FC1);
		for (int i = 0; i < row; i++) {
			for (int j = 0; j < col; j++) {
				dst.at<double>(i, j) = 0.5;
			}
		}
		return dst;
	}
	//resize the image to half
	cv::Mat Rs((row + 1) / 2, (col + 1) / 2, CV_64FC1);
	cv::resize(src, Rs, cv::Size((col + 1) / 2, (row + 1) / 2));
	//recursive
	//convolution and resize
	cv::Mat Rf = FastACE(Rs, ratio, radius);
	cv::resize(Rf, Rf, cv::Size(col, row));
	cv::resize(Rs, Rs, cv::Size(col, row));
	cv::Mat dst(row, col, CV_64FC1);
	cv::Mat dst1 = NormalACE(src, ratio, radius);
	cv::Mat dst2 = NormalACE(Rs, ratio, radius);
	for (int i = 0; i < row; i++) {
		for (int j = 0; j < col; j++) {
			dst.at<double>(i, j) = Rf.at<double>(i, j) + dst1.at<double>(i, j) - dst2.at<double>(i, j);
		}
	}
	return dst;
}

cv::Mat ImageColorBalance::getACE(cv::Mat src, int ratio, int radius)
{
	int row = src.rows;
	int col = src.cols;
	std::vector <cv::Mat> v;
	split(src, v);
	v[0].convertTo(v[0], CV_64FC1);
	v[1].convertTo(v[1], CV_64FC1);
	v[2].convertTo(v[2], CV_64FC1);
	cv::Mat src1(row, col, CV_64FC1);
	cv::Mat src2(row, col, CV_64FC1);
	cv::Mat src3(row, col, CV_64FC1);

	for (int i = 0; i < row; i++) {
		for (int j = 0; j < col; j++) {
			src1.at<double>(i, j) = 1.0 * src.at<cv::Vec3b>(i, j)[0] / 255.0;
			src2.at<double>(i, j) = 1.0 * src.at<cv::Vec3b>(i, j)[1] / 255.0;
			src3.at<double>(i, j) = 1.0 * src.at<cv::Vec3b>(i, j)[2] / 255.0;
		}
	}
	src1 = stretchImage(FastACE(src1, ratio, radius));
	src2 = stretchImage(FastACE(src2, ratio, radius));
	src3 = stretchImage(FastACE(src3, ratio, radius));

	//
	cv::Mat dst1(row, col, CV_8UC1);
	cv::Mat dst2(row, col, CV_8UC1);
	cv::Mat dst3(row, col, CV_8UC1);
	for (int i = 0; i < row; i++) {
		for (int j = 0; j < col; j++) {
			dst1.at<uchar>(i, j) = (int)(src1.at<double>(i, j) * 255);
			if (dst1.at<uchar>(i, j) > 255) dst1.at<uchar>(i, j) = 255;
			else if (dst1.at<uchar>(i, j) < 0) dst1.at<uchar>(i, j) = 0;
			dst2.at<uchar>(i, j) = (int)(src2.at<double>(i, j) * 255);
			if (dst2.at<uchar>(i, j) > 255) dst2.at<uchar>(i, j) = 255;
			else if (dst2.at<uchar>(i, j) < 0) dst2.at<uchar>(i, j) = 0;
			dst3.at<uchar>(i, j) = (int)(src3.at<double>(i, j) * 255);
			if (dst3.at<uchar>(i, j) > 255) dst3.at<uchar>(i, j) = 255;
			else if (dst3.at<uchar>(i, j) < 0) dst3.at<uchar>(i, j) = 0;
		}
	}
	std::vector <cv::Mat> out;
	out.push_back(dst1);
	out.push_back(dst2);
	out.push_back(dst3);
	cv::Mat dst;
	merge(out, dst);
	return dst;
}

void ImageColorBalance::InterectRectDiff(cv::Mat& Infer, cv::Mat& Target, float& fInferToTarget_mul,  float& fInferToTarget_sub)
{
	cv::Scalar MeanInfer;
	cv::Scalar StdInfer;
	//获取infer图像的均值
	cv::meanStdDev(Infer, MeanInfer, StdInfer);

	cv::Scalar MeanTarget;
	cv::Scalar StdTarget;
	//获取target图像的均值
	cv::meanStdDev(Target, MeanTarget, StdTarget);

	float fMeanInfer = MeanInfer.val[0] + MeanInfer.val[1] + MeanInfer.val[2];
	fMeanInfer = fMeanInfer / 3.0f;
	float fMeanTarget = MeanTarget.val[0] + MeanTarget.val[1] + MeanTarget.val[2];
	fMeanTarget = fMeanTarget / 3.0f;

	fInferToTarget_sub = fMeanInfer;
	fInferToTarget_mul = fMeanTarget / fMeanInfer;
}

void ImageColorBalance::FlatParamInChain(std::vector<float>& vInParams, std::vector<float>& vOutParams)
{
	//获取均值
	float Mean = 0.0f;
	for (int i = 0; i < int(vInParams.size()); i++)
	{
		Mean += vInParams[i];
	}
	Mean = Mean / int(vInParams.size());

	

}

void ImageColorBalance::IntersectRectChange(cv::Mat& InferImg, cv::Mat& TargetImg, cv::Rect InferRc, cv::Rect TargetRc, float fInferToTarget_mul)
{
	InferImg =  InferImg * fInferToTarget_mul;
}