#include "stdafx.h"
#include "fuseImg.h"

#include "waveDecompose.h"
#include "waveReconstruct.h"

// 直方图投影
#include "histProject.h"

#include <omp.h>

// 功能：函数 FuseImg() 融合两幅输入图像
bool FuseImg(cv::Mat &dst, const cv::Mat img1, const cv::Mat img2, const string wname, const unsigned int level)
{
	// 输入有效性判断
	if (img1.empty() == true || (img1.type() != CV_8UC1 && img1.type() != CV_8UC3))
	{
		return false;
	}

	if (img2.empty() == true || (img2.type() != CV_8UC1 && img2.type() != CV_8UC3))
	{
		return false;
	}

	if (img1.type() != img2.type()
		|| img1.size != img2.size)
	{
		return false;
	}

	if (wname == "" || level < 1)
	{
		return false;
	}

	// 图像类型
	int imgType = img1.type();
	if (imgType == CV_8UC3)
	{
		// 提取img1的亮度分量
		cv::Mat YCrCb1;
		cv::cvtColor(img1, YCrCb1, cv::COLOR_BGR2YCrCb);
		std::vector<cv::Mat> planes1;
		cv::split(YCrCb1, planes1);

		// 提取img2的亮度分量
		cv::Mat YCrCb2;
		cv::cvtColor(img2, YCrCb2, cv::COLOR_BGR2YCrCb);
		std::vector<cv::Mat> planes2;
		cv::split(YCrCb2, planes2);

		// 融合两幅图像亮度分量
		cv::Mat fusion_Y;
		if (false == FuseImg(fusion_Y, planes1[0], planes2[0], wname, level))
		{
			return false;
		}
		else
		{
			planes1[0] = fusion_Y;

			// 其余两个通道取平均
			planes1[1] = (planes1[1] + planes2[1]) / 2;
			planes1[2] = (planes1[2] + planes2[2]) / 2;

			cv::merge(planes1, YCrCb1);
			cv::cvtColor(YCrCb1, dst, cv::COLOR_YCrCb2BGR);
		}

		return true;
	} 
	else
	{
		// 图像1 小波分解
		std::deque<cv::Mat> coef1;
		std::deque<cv::Vec2i> scf1;

		// 图像2 小波分解
		std::deque<cv::Mat> coef2;
		std::deque<cv::Vec2i> scf2;

		bool ret1 = true;
		bool ret2 = true;

		#pragma omp parallel sections
		{
			#pragma omp section
			{
				ret1 = waveDec2(coef1, scf1, img1, wname, level);
			}

			#pragma omp section
			{
				ret2 = waveDec2(coef2, scf2, img2, wname, level);
			}
		}

		if (ret1 == false || ret2 == false)
		{
			return false;
		}

		// 小波系数融合

		// 1低频:取均值
		coef1[0] = (coef1[0] + coef2[1]) / 2;

		// 2高频：取绝对值较大者 
		// 理论依据：大的小波系数含有较多的能量
		for (unsigned int m  = 1; m < coef1.size(); m++)
		{
			int rows = coef1[m].rows;
			int cols = coef1[m].cols;

			for (int i = 0; i < rows; i++)
			{
				for (int j = 0; j < cols; j++)
				{
					if (abs(coef2[m].at<double>(i, j)) > abs(coef1[m].at<double>(i, j)))
					{
						coef1[m].at<double>(i, j) = coef2[m].at<double>(i, j);
					}
				}
			}
		}

		// 图像重构
		cv::Mat dst_double;
		if (false == waveRec2(dst_double, coef1, scf1, wname, level)) 
		{
			return false;
		}

		// 类型转换
		double minVal = 0;
		double maxVal = 0;
		cv::minMaxIdx(dst_double, &minVal, &maxVal);

		cv::Mat preHP; // 直方图均衡化前的融合图像

		// 压缩 减少信息损失
		dst_double.convertTo(preHP, CV_8UC1, 255.0 / (maxVal - minVal), -minVal / (maxVal - minVal));

		// 自适应直方图投影，提高图像对比度
		SAHP(dst, preHP);

		return true;
	}
}


// 功能：函数 FuseImgRGBplusIR() 融合两幅输入图像
bool FuseImgRGBplusIR(cv::Mat &dst, const cv::Mat CCDimg, const cv::Mat IRimg, const string wname, const unsigned int level)
{
	// 输入有效性判断
	if (CCDimg.empty() == true ||  CCDimg.type() != CV_8UC3)
	{
		return false;
	}

	if (IRimg.empty() == true || (IRimg.type() != CV_8UC1 && IRimg.type() != CV_8UC3))
	{
		return false;
	}

	if (CCDimg.size != IRimg.size)
	{
		return false;
	}

	if (wname == "" || level < 1)
	{
		return false;
	}

	// 根据类型区分
	if (IRimg.type() == CV_8UC3)
	{
		// 提取IRimg亮度分量
		cv::Mat IR;
		std::vector<cv::Mat> planesIR;
		cv::split(IRimg, planesIR);
		planesIR[0].copyTo(IR);

		// 提取CCDimg亮度分量
		cv::Mat YCrCb;
		cv::cvtColor(CCDimg, YCrCb, cv::COLOR_BGR2YCrCb);
		std::vector<cv::Mat> planesCCD;
		cv::split(YCrCb, planesCCD);

		// 融合两幅图像亮度分量
		cv::Mat dst_Y;
		if (!FuseImg(dst_Y, planesCCD[0], IR, wname, level))
		{
			return false;
		}

		// 重构图像
		dst_Y.copyTo(planesCCD[0]);
		cv::merge(planesCCD, YCrCb);
		cv::cvtColor(YCrCb, dst, cv::COLOR_YCrCb2BGR);
	}
	else
	{
		// 提取CCDimg亮度分量
		cv::Mat YCrCb;
		cv::cvtColor(CCDimg, YCrCb, cv::COLOR_BGR2YCrCb);
		std::vector<cv::Mat> planesCCD;
		cv::split(YCrCb, planesCCD);

		// 融合两幅图像亮度分量
		cv::Mat dst_Y;
		if (false == FuseImg(dst_Y, planesCCD[0], IRimg, wname, level))
		{
			return false;
		}

		// 重构图像
		dst_Y.copyTo(planesCCD[0]);
		cv::merge(planesCCD, YCrCb);
		cv::cvtColor(YCrCb, dst, cv::COLOR_YCrCb2BGR);
	}

	return true;
}

// 功能：函数 FastFuseImgCCDplusIR() 快速融合两幅输入图像
bool FastFuseImgRGBplusIR(cv::Mat &dst, const cv::Mat CCDimg, const cv::Mat IRimg)
{
	// 输入有效性判断
	if (CCDimg.empty() == true ||  CCDimg.type() != CV_8UC3)
	{
		return false;
	}

	if (IRimg.empty() == true || (IRimg.type() != CV_8UC1 && IRimg.type() != CV_8UC3))
	{
		return false;
	}

	if (CCDimg.size != IRimg.size)
	{
		return false;
	}

	// 单通道红外图像
	cv::Mat IR;

	// 根据类型区分
	if (IRimg.type() == CV_8UC3)
	{
		// 提取IRimg亮度分量
		std::vector<cv::Mat> planesIR;
		cv::split(IRimg, planesIR);
		planesIR[0].copyTo(IR);
	}
	else
	{
		IR = IRimg;
	}

	// 提取CCDimg亮度分量
	cv::Mat YCrCb;
	cv::cvtColor(CCDimg, YCrCb, cv::COLOR_BGR2YCrCb);
	std::vector<cv::Mat> planesCCD;
	cv::split(YCrCb, planesCCD);

	// 融合两幅图像亮度分量
	cv::Mat dst_Y = 0.5 * planesCCD[0] + 0.5 * IR;

	// 增强亮度分量
	cv::Mat enchanceDstY;
	SAHP(enchanceDstY, dst_Y);

	// 重构图像
	planesCCD[0] = enchanceDstY;
	cv::merge(planesCCD, YCrCb);
	cv::cvtColor(YCrCb, dst, cv::COLOR_YCrCb2BGR);

	return true;
}