#include "PL_BlockGrouping.h"
#include "MaterialBuffer.h"
#include "logger.h"
#include "MediaHelper.h"
#include "GraphicHelper.h"

#include <string.h> // for memcpy
#include <opencv2/core/mat.hpp>
#include <opencv2/imgproc.hpp>
#include <set>

#include <opencv/cv.h>
#include <opencv/cxcore.h>

struct MBFrameWrapper
{
	MB_Frame frame;
	int frameRefCount;

	MBFrameWrapper(bool _copyData, const MB_Frame _frame) :
		frame(_frame), frameRefCount(_copyData ? 0 : -1)
	{
		if (_copyData)
		{
			uint8_t* newBuffer = new uint8_t[_frame.buffSize];
			frame.buffer = newBuffer;
			memcpy(newBuffer, _frame.buffer, _frame.buffSize);
		}
	}

   void reference()
   {
	   if (frameRefCount >= 0)
		   frameRefCount++;
	   else
		   frameRefCount--;
   }

	void release()
	{
		if (frameRefCount == 0)
		{
			delete[] (uint8_t*)frame.buffer;
			frame.reset();
		}
		else if (frameRefCount == -1)
		{
			frame.reset();
		}
		
		if (frameRefCount < -1)
		{
			frameRefCount++;
		}
		else if (frameRefCount > 0)
		{
			frameRefCount--;
		}
	}
	
	~MBFrameWrapper()
	{
		release();
	}
};

struct PLBG_Block
{
	const uint16_t blockID;

	RectWrapper rectWrapper;
	MBFrameWrapper* mbfWrapper;
	uint8_t* croppedData;
	size_t croppedDataSize;
	
	PLBG_Block() : blockID(++_blockID), rectWrapper(), mbfWrapper(nullptr), croppedData(nullptr), croppedDataSize(0)
	{ }
	
	void crop_data()
	{
		if (croppedData != nullptr || mbfWrapper == nullptr)
			return;

		PLGH_Rect& rect(rectWrapper.rect);
		if (mbfWrapper->frame.type == MB_Frame::MBFT_YUV420)
		{
			int src_width = mbfWrapper->frame.width;
			int src_height = mbfWrapper->frame.height;
			uint8_t* src_y = (uint8_t*)(mbfWrapper->frame.buffer);
			uint8_t* src_u = (uint8_t*)(src_y + (src_height * src_width));
			uint8_t* src_v = (uint8_t*)(src_u + (src_height * src_width / 4));

			cv::Mat matY(cv::Size(src_width, src_height), CV_8UC1, src_y);
			cv::Mat roiMatY(matY, cv::Rect(rect.leftTop.X, rect.leftTop.Y, rect.width(), rect.height()));
			cv::Mat cloneRoiMatY(roiMatY.clone());

			cv::Mat matU(cv::Size(MH_SUBSAMPLE1(src_width, 2), MH_SUBSAMPLE1(src_height, 2)), CV_8UC1, src_u);
			cv::Mat roiMatU(matU, cv::Rect(MH_SUBSAMPLE1(rect.leftTop.X, 2), MH_SUBSAMPLE1(rect.leftTop.Y, 2), MH_SUBSAMPLE1(rect.width(), 2), MH_SUBSAMPLE1(rect.height(), 2)));
			cv::Mat cloneRoiMatU(roiMatU.clone());

			cv::Mat matV(cv::Size(MH_SUBSAMPLE1(src_width, 2), MH_SUBSAMPLE1(src_height, 2)), CV_8UC1, src_v);
			cv::Mat roiMatV(matV, cv::Rect(MH_SUBSAMPLE1(rect.leftTop.X, 2), MH_SUBSAMPLE1(rect.leftTop.Y, 2), MH_SUBSAMPLE1(rect.width(), 2), MH_SUBSAMPLE1(rect.height(), 2)));
			cv::Mat cloneRoiMatV(roiMatV.clone());

			const size_t yPlanarSize = rect.width() * rect.height();
			croppedDataSize = yPlanarSize * 1.5;
			croppedData = new uint8_t[croppedDataSize];

			uint8_t* dst_y = (uint8_t*)croppedData;
			uint8_t* dst_u = (uint8_t*)(dst_y + yPlanarSize);
			uint8_t* dst_v = (uint8_t*)(dst_u + yPlanarSize / 4);

			memcpy(dst_y, cloneRoiMatY.ptr(), yPlanarSize);
			memcpy(dst_u, cloneRoiMatU.ptr(), yPlanarSize / 4);
			memcpy(dst_v, cloneRoiMatV.ptr(), yPlanarSize / 4);

			//{
			//	static size_t f = 0;
			//	char fname[50];
			//	sprintf(fname, "/sdcard/face-%u-w%u-h%u.yuv420", ++f, rect.width(), rect.height());
			//	FILE *pFile = fopen(fname, "wb");
			//	fwrite(croppedData, 1, croppedDataSize, pFile);
			//	fclose(pFile);
			//	if (f > 10)exit(0);
			//}
		}
		else if (mbfWrapper->frame.type == MB_Frame::MBFT_NV12)
		{
			int src_width = mbfWrapper->frame.width;
			int src_height = mbfWrapper->frame.height;
			uint8_t* src_y = (uint8_t*)(mbfWrapper->frame.buffer);
			uint8_t* src_uv = (uint8_t*)(src_y + (src_height * src_width));
			
			cv::Mat matY(cv::Size(src_width, src_height), CV_8UC1, src_y);
			cv::Mat roiMatY(matY, cv::Rect(rect.leftTop.X, rect.leftTop.Y, rect.width(), rect.height()));
			cv::Mat cloneRoiMatY(roiMatY.clone());
			
			cv::Mat matUV(cv::Size(MH_SUBSAMPLE1(src_width, 2), MH_SUBSAMPLE1(src_height, 2)), CV_16UC1, src_uv);
			cv::Mat roiMatUV(matUV, cv::Rect(MH_SUBSAMPLE1(rect.leftTop.X, 2), MH_SUBSAMPLE1(rect.leftTop.Y, 2), MH_SUBSAMPLE1(rect.width(), 2), MH_SUBSAMPLE1(rect.height(), 2)));
			//cv::Mat roiMatUV(matUV, cv::Rect(rect.leftTop.X, rect.leftTop.Y, rect.width(), rect.height()));
			cv::Mat cloneRoiMatUV(roiMatUV.clone());
			
			const size_t yPlanarSize = rect.width() * rect.height();
			croppedDataSize = yPlanarSize * 1.5;
			croppedData = new uint8_t[croppedDataSize];//#todo optimize static memory
			
			uint8_t* dst_y = (uint8_t*)croppedData;
			uint8_t* dst_uv = (uint8_t*)(dst_y + yPlanarSize);
			
			memcpy(dst_y, cloneRoiMatY.ptr(), yPlanarSize);
			memcpy(dst_uv, cloneRoiMatUV.ptr(), yPlanarSize * 0.5);
		}
		else if (mbfWrapper->frame.type == MB_Frame::MBFT_RGB888)
		{
			LOG_WARN << "frame type not supported" << LOG_ENDL;
		}
		else
		{
			LOG_WARN << "frame type not supported" << LOG_ENDL;
		}
	}
	
	void release_crop()
	{
		delete[] (uint8_t*)croppedData;
		croppedDataSize = 0;
	}
	
	~PLBG_Block()
	{
	}

private:
	static uint16_t _blockID;
};

uint16_t PLBG_Block::_blockID = 0;

struct PLBG_Group
{
	const uint16_t groupID;
	
	std::vector<PLBG_Block> blocks; // sequence of paid blocks
	int paid_rect_interval;
	int paid_frame_count;
	int _current_frame_accepted_blocks;
	
	PLBG_Group() : groupID(++_groupID), blocks(), paid_rect_interval(0), paid_frame_count(0), _current_frame_accepted_blocks(0)
	{ }
	
private:
	static uint16_t _groupID;
};

uint16_t PLBG_Group::_groupID = 0;

struct PL_BlockGrouping_Internal
{
	PL_BlockGrouping_Config config;

	std::set<MBFrameWrapper*> mbfws;
	
	std::list<PLBG_Group> groups;
	std::list<PLBG_Group> groupsMature;
	
	bool payError;
	
	plbg_output_vec_t outputs;

	PL_BlockGrouping_Internal() : 
		config(), mbfws(),
		groups(), groupsMature(), 
		payError(true), 
		outputs()
	{
	}
	
	~PL_BlockGrouping_Internal()
	{
		reset();
	}
	
	void reset()
	{
		PL_BlockGrouping_Config _config;
		config = _config;
		
		//for (std::set<uint8_t*>::iterator iter = buffers.begin(); iter != buffers.end(); iter++)
		//	delete[] *iter;

		mbfws.clear(); //#todo delete
		
		groups.clear();
		groupsMature.clear(); //#todo delete
		
		payError = true;
		
		outputs.clear(); //#todo delete croppedData
	}
};

PipeLineElem* create_PL_BlockGrouping()
{
	return new PL_BlockGrouping;
}

PL_BlockGrouping::PL_BlockGrouping() : internal(new PL_BlockGrouping_Internal)
{
}

PL_BlockGrouping::~PL_BlockGrouping()
{
	delete (PL_BlockGrouping_Internal*)internal;
	internal= nullptr;
}

bool PL_BlockGrouping::init(void* args)
{
	PL_BlockGrouping_Internal* in = (PL_BlockGrouping_Internal*)internal;
	in->reset();
	
	if (args != nullptr)
	{
		PL_BlockGrouping_Config* config = (PL_BlockGrouping_Config*)args;
		in->config = *config;
	}

	return true;
}

void PL_BlockGrouping::finit()
{
	PL_BlockGrouping_Internal* in = (PL_BlockGrouping_Internal*)internal;
	
}

bool plbg_pay_breaker(const PipeMaterial* pm, void* args)
{
	*(PipeMaterial**)args = const_cast<PipeMaterial*>(pm);
	return false;
}

bool test_group_accept_rect(PL_BlockGrouping_Internal* in, const PLBG_Group& group, const PLGH_Rect& rect)
{
	const PLGH_Rect& lastRect(group.blocks.back().rectWrapper.rect);
	
	// test block center diff
	PLGH_Point lastBlockCenter(lastRect.center());
	PLGH_Point testBlockCenter(rect.center());
	if (std::abs(lastBlockCenter.X - testBlockCenter.X) > in->config.block_center_x_diff || 
		std::abs(lastBlockCenter.Y - testBlockCenter.Y > in->config.block_center_y_diff))
		return false;
		
	// test block area
	if (std::abs(lastRect.area() - rect.area()) > in->config.block_area_diff)
		return false;
	
	// test frame continuity
	if (group.paid_rect_interval > in->config.continuity_max)
		return false;
	
	// test blocks count alreadly accepted this frame
	if (group._current_frame_accepted_blocks >= in->config.accept_blocks_max_per_frame)
		return false;
	
	return true;
}

bool PL_BlockGrouping::pay(const PipeMaterial& pm)
{
	PL_BlockGrouping_Internal* in = (PL_BlockGrouping_Internal*)internal;

	//#todo
	//if (in->mbfs.size() + 1 >= PLBG_MAX_CACHE_FRAMES)
	//{
	//	LOG_WARN << "cached frame exceed " << PLBG_MAX_CACHE_FRAMES << LOG_ENDL;
	//	return false;
	//}


	PipeMaterial pmPtr;
	{
		PipeMaterial* _pmPtr = nullptr;
		pm.breake(PipeMaterial::PMT_PTR, MB_Frame::MBFT__FIRST, plbg_pay_breaker, &_pmPtr);
		if (_pmPtr == nullptr)
		{
			LOG_WARN << "PMT_PM_LIST need PMT_PTR" << LOG_ENDL;
			return false;
		}
		pmPtr = *_pmPtr;
	}

	MB_Frame mbfFrame;
	{
		PipeMaterial* _pmFrame = nullptr;

		if (_pmFrame == nullptr)
			pm.breake(PipeMaterial::PMT_FRAME, MB_Frame::MBFT_YUV420, plbg_pay_breaker, &_pmFrame);
		if (_pmFrame == nullptr)
			pm.breake(PipeMaterial::PMT_FRAME, MB_Frame::MBFT_NV12, plbg_pay_breaker, &_pmFrame);

		if (_pmFrame == nullptr)
		{
			LOG_WARN << "PMT_PM_LIST need PMT_FRAME of (MBFT_YUV420 or MBFT_NV12)" << LOG_ENDL;
			return false;
		}

		mbfFrame = *(MB_Frame*)_pmFrame->buffer;
	}

	// cache frame
	MBFrameWrapper* mbfWrapper = new MBFrameWrapper(in->config.copyData, mbfFrame);

	// get blocks
	std::list<RectWrapper> rectws;
	in->config.get_rect_func(pmPtr, mbfFrame, rectws);

	// classify blocks to exists group
	for (std::list<PLBG_Group>::iterator iterGrp = in->groups.begin(); iterGrp != in->groups.end(); ++iterGrp)
	{
		PLBG_Group& currGroup(*iterGrp);
		currGroup.paid_rect_interval++;
		currGroup.paid_frame_count++;

		for (std::list<RectWrapper>::iterator iterRectw = rectws.begin(); iterRectw != rectws.end(); )
		{
			const RectWrapper& currRectw(*iterRectw);

			if (!test_group_accept_rect(in, currGroup, currRectw.rect))
			{
				++iterRectw;
				continue;
			}
			
			// accept this block into group
			PLBG_Block currBlock;
			currBlock.rectWrapper = currRectw;
			currBlock.mbfWrapper = mbfWrapper;
			mbfWrapper->reference();
			currBlock.crop_data();
			currGroup.blocks.push_back(currBlock);
			currGroup.paid_rect_interval = 0;
			currGroup._current_frame_accepted_blocks++;

			iterRectw = rectws.erase(iterRectw);
		}
	}

	//#todo
	// we test remain groups asume remain rectws
	//if (in->groups.size() > PLBG_MAX_GROUPS - rectws.size())
	//{
	//	if (mbfWrapper->frameRefCount == 0)
	//	{
	//		in->mbfs.erase(mbfWrapper);
	//		delete mbfWrapper;
	//		mbfWrapper = nullptr;
	//	}
//
	//	LOG_WARN << "Exceed max groups " << PLBG_MAX_GROUPS << LOG_ENDL;
	//	return false;
	//}
	
	// grouping left blocks
	std::list<PLBG_Group>::iterator iterNewGrp = in->groups.end();
	if (!rectws.empty())
	{
		// add the first rect (not classified) into new group
		PLBG_Block currBlock;
		currBlock.rectWrapper = *rectws.begin();
		currBlock.mbfWrapper = mbfWrapper;
		mbfWrapper->reference();
		currBlock.crop_data();
		rectws.erase(rectws.begin());
		
		PLBG_Group currGroup;
		currGroup.blocks.push_back(currBlock);
		currGroup.paid_rect_interval = 0;
		currGroup._current_frame_accepted_blocks++;
		iterNewGrp = in->groups.insert(in->groups.end(), currGroup);
	}
	
	for (std::list<RectWrapper>::iterator iterRectw = rectws.begin(); iterRectw != rectws.end(); ++iterRectw)
	{
		const RectWrapper& currRectw(*iterRectw);
		
		std::list<PLBG_Group>::iterator iterGrp = iterNewGrp;
		for ( ; iterGrp != in->groups.end(); ++iterGrp)
		{
			if (test_group_accept_rect(in, *iterGrp, iterRectw->rect))
				break;
		}
		
		// add rect
		PLBG_Block currBlock;
		currBlock.rectWrapper = *iterRectw;
		currBlock.mbfWrapper = mbfWrapper;
		mbfWrapper->reference();
		currBlock.crop_data();
		if (iterGrp == in->groups.end())
		{
			PLBG_Group currGroup;
			iterGrp = in->groups.insert(in->groups.end(), currGroup);
		}
		
		iterGrp->blocks.push_back(currBlock);
		iterGrp->paid_rect_interval = 0;
		iterGrp->_current_frame_accepted_blocks++;
	}
	
	// reset
	for (std::list<PLBG_Group>::iterator iterGrp = in->groups.begin(); iterGrp != in->groups.end(); ++iterGrp)
		iterGrp->_current_frame_accepted_blocks = 0;

	if (mbfWrapper->frameRefCount != -1 && mbfWrapper->frameRefCount != 0)
		in->mbfws.insert(mbfWrapper);
	else
		delete mbfWrapper;

	return true;
}

float scoring_canny_focus(const PLBG_Block& block)
{
	const int w = block.rectWrapper.rect.width();
	const int h = block.rectWrapper.rect.height();
	if (block.croppedData == nullptr || block.croppedDataSize < w * h)
		return 0.0f;

	cv::Mat yMat(cv::Size(w, h), CV_8UC1, block.croppedData);
	cv::Mat edges;
	cv::GaussianBlur(yMat, edges, cv::Size(5, 5), 1.5, 1.5);
	cv::Canny(edges, edges, 50, 100, 3);
	//cv::Sobel(yMat, edges, CV_8UC1, 1, 0, 3);

	//memcpy(block.croppedData, edges.data, w * h);

	size_t sum = 0;
	std::for_each(edges.begin<uint8_t>(), edges.end<uint8_t>(), [&](uint8_t v)
	{
		sum += (v != 0);
	});

	float focusRate = (float)sum / (w * h);
	return focusRate;
}

float scoring_histogram_focus(const PLBG_Block& block)
{
	return 1.0f;
}

bool PL_BlockGrouping::gain(PipeMaterial& pm)
{
	PL_BlockGrouping_Internal* in = (PL_BlockGrouping_Internal*)internal;

	// clean last mature groups and outputs
	for (std::list<PLBG_Group>::iterator iterGrp = in->groupsMature.begin(); iterGrp != in->groupsMature.end(); ++iterGrp)
	{
		for (std::vector<PLBG_Block>::iterator iterBlk = iterGrp->blocks.begin(); iterBlk != iterGrp->blocks.end(); ++iterBlk)
		{
			iterBlk->mbfWrapper->release();
			iterBlk->release_crop();
			if (iterBlk->rectWrapper.user_data_deleter != nullptr)
				iterBlk->rectWrapper.user_data_deleter(iterBlk->rectWrapper);
		}
	}
	
	in->groupsMature.clear();
	
	for (std::set<MBFrameWrapper*>::iterator iter = in->mbfws.begin(); iter != in->mbfws.end(); )
	{
		MBFrameWrapper* mbfw = *iter;
		if (mbfw->frameRefCount == 0 || mbfw->frameRefCount == -1)
		{
			delete mbfw;
			iter = in->mbfws.erase(iter);
		}
		else
			++iter;
	}
	
	in->outputs.clear();
	
	// find "mature" groups
	for (std::list<PLBG_Group>::iterator iterGrp = in->groups.begin(); iterGrp != in->groups.end(); )
	{
		bool grpMature = false;
		PLBG_Group& currGroup(*iterGrp);

		if (currGroup.blocks.size() >= in->config.group_size_max)
			grpMature = true;
		
		if (currGroup.paid_frame_count >= in->config.frame_count_max)
			grpMature = true;
		
		if (grpMature && (currGroup.blocks.size() < in->config.group_size_min))
		{
			//drop group
			for (std::vector<PLBG_Block>::iterator iterBlk = currGroup.blocks.begin(); iterBlk != currGroup.blocks.end(); ++iterBlk)
			{
				iterBlk->mbfWrapper->release();
				iterBlk->release_crop();
			}
			
			iterGrp = in->groups.erase(iterGrp);
			continue;
		}
		
		if (grpMature)
		{
			// generate output group
			in->groupsMature.push_back(currGroup);
			iterGrp = in->groups.erase(iterGrp);
			continue;
		}

		++iterGrp;
	}
	
	// scoring blocks
	const bool calc_canny = !MH_F_ZEQ(in->config.canny_focus);
	const bool calc_histogram = !MH_F_ZEQ(in->config.histogram_uniformy);
	const bool calc_user_score_1 = !MH_F_ZEQ(in->config.user_score_1_rate);
	const float user_score_2_rate = 1.0f - in->config.canny_focus - in->config.histogram_uniformy - in->config.user_score_1_rate;
	const bool calc_user_score_2 = (!MH_F_ZEQ(user_score_2_rate) && in->config.user_score_2_func != nullptr);
	
	for (std::list<PLBG_Group>::iterator iterGrp = in->groupsMature.begin(); iterGrp != in->groupsMature.end(); ++iterGrp)
	{
		std::vector<PLBG_Block>::iterator iterBestBlk = iterGrp->blocks.end();
		float bestScore = 0.0f;
		
		for (std::vector<PLBG_Block>::iterator iterBlk = iterGrp->blocks.begin(); iterBlk != iterGrp->blocks.end(); ++iterBlk)
		{
			float totalScore = 0.0f;
			
			if (calc_canny)
			{
				float cannyScore = scoring_canny_focus(*iterBlk) * 10.0f;
				totalScore += in->config.canny_focus * cannyScore;
			}
			if (calc_histogram)
			{
				float histScore = scoring_histogram_focus(*iterBlk);
				totalScore += in->config.histogram_uniformy * histScore;
			}
			if (calc_user_score_1)
			{
				totalScore += in->config.user_score_1_rate * iterBlk->rectWrapper.user_score_1;
			}
			if (calc_user_score_2)
			{
				float userScore2 = in->config.user_score_2_func(&(iterBlk->mbfWrapper->frame), iterBlk->rectWrapper.rect, iterBlk->croppedData);
				totalScore += user_score_2_rate * userScore2;
			}
			
			if (MH_F_GT(totalScore, bestScore))
			{
				iterBestBlk = iterBlk;
				bestScore = totalScore;
			}
		}
		
		if (iterBestBlk == iterGrp->blocks.end() || MH_F_ZEQ(bestScore))
			continue;
		
		PLBG_Output output;
		output.groupID = iterGrp->groupID;
		output.blockID = iterBestBlk->blockID;
		output.rectInOriginFrame = iterBestBlk->rectWrapper;
		output.originframe = &(iterBestBlk->mbfWrapper->frame);
		output.score = bestScore;
		output.croppedData = iterBestBlk->croppedData;
		output.croppedDataSize = iterBestBlk->croppedDataSize;

		in->outputs.push_back(output);
	}
	
	pm.type = PipeMaterial::PMT_PTR;
	pm.buffer = &(in->outputs);
	pm.buffSize = 0;
	pm.former = this;
	pm.deleter = nullptr;
	pm.args = nullptr;

	return true;
}