#include "PL_DlibFaceTrack.h"
#include "MaterialBuffer.h"
#include "logger.h"

#include <opencv2/opencv.hpp>
#include <opencv2/highgui/highgui.hpp>

#include <dlib/opencv.h>
#include <dlib/image_processing/frontal_face_detector.h>
#include <dlib/image_processing/render_face_detections.h>
#include <dlib/image_processing.h>
#include <dlib/gui_widgets.h>

#include <libyuv.h>

#define DLIB_POSE_MODEL_PATH "/opt/dlib/models/shape_predictor_68_face_landmarks.dat"

struct PL_DlibFaceTrack_Internal
{
	//uint8_t buffer[1920*1080*4];
	//size_t buffSize;
	//size_t buffSizeMax;
	MB_Frame lastFrame;
	PL_DlibFaceTrack_Config config;

	bool payError;
	
	dlib::frontal_face_detector detector; // #todo reset
	dlib::shape_predictor pose_model;

	PL_DlibFaceTrack_Internal() : 
		//buffSize(0), buffSizeMax(sizeof(buffer)), 
		lastFrame(), config(), payError(true)
	{
	}
	
	~PL_DlibFaceTrack_Internal()
	{
	}
	
	void reset()
	{
		//buffSize = 0;
		payError = true;
		
		MB_Frame _lastFrame;
		lastFrame = _lastFrame;
		PL_DlibFaceTrack_Config _config;
		config = _config;
	}
};

PipeLineElem* create_PL_DlibFaceTrack()
{
	return new PL_DlibFaceTrack;
}

PL_DlibFaceTrack::PL_DlibFaceTrack() : internal(new PL_DlibFaceTrack_Internal)
{
}

PL_DlibFaceTrack::~PL_DlibFaceTrack()
{
	delete (PL_DlibFaceTrack_Internal*)internal;
	internal= nullptr;
}

bool PL_DlibFaceTrack::init(void* args)
{
	PL_DlibFaceTrack_Internal* in = (PL_DlibFaceTrack_Internal*)internal;
	in->reset();
	
	PL_DlibFaceTrack_Config* config = (PL_DlibFaceTrack_Config*)args;
	if (config != nullptr)
		in->config = *config;
	
#ifdef __AVX__
	LOG_DEBUG << "AVX on" << std::endl;
	#ifdef DLIB_HAVE_SSE2
		LOG_DEBUG << "DLIB_HAVE_SSE2 on" << std::endl;
	#endif
	#ifdef DLIB_HAVE_SSE3
		LOG_DEBUG << "DLIB_HAVE_SSE3 on" << std::endl;
	#endif
	#ifdef DLIB_HAVE_SSE41
		LOG_DEBUG << "DLIB_HAVE_SSE41 on" << std::endl;
	#endif
	#ifdef DLIB_HAVE_AVX
		LOG_DEBUG << "DLIB_HAVE_AVX on" << std::endl;
	#endif
#endif
	
	in->detector = dlib::get_frontal_face_detector();
	dlib::deserialize(DLIB_POSE_MODEL_PATH) >> in->pose_model;
	
	return true;
}

void PL_DlibFaceTrack::finit()
{
	PL_DlibFaceTrack_Internal* in = (PL_DlibFaceTrack_Internal*)internal;

}

uint64_t time_msec()
{
	timeval tv;
	gettimeofday(&tv, nullptr);
	
	return (tv.tv_sec * 1000 * 1000 + tv.tv_usec) / 1000;
}

int doFaceTrack(PL_DlibFaceTrack_Internal* in, 
				uint8_t* buffer, size_t width, size_t height, MB_Frame::MBFType pixFmt)
{
	uint64_t tbegin = time_msec();

	//uint8_t* pBuffer = new uint8_t[width * height * 3];
	
	//#define SUBSAMPLE(v, a) ((((v) + (a) - 1)) / (a))

	//libyuv::I420ToRGB24(buffer, width, 
	//				   buffer + width * height, SUBSAMPLE(width, 2), 
	//				   buffer + width * height + width * height / 4, SUBSAMPLE(width, 2), 
	//				   pBuffer, 3 * width, 
	//				   width, height);

	//#test
	//static size_t f=0;
	//char fname[50];
	//sprintf(fname, "%u.rgb", ++f);
	//FILE * pFile = fopen (fname, "wb");
	//fwrite (pBuffer , sizeof(char), width * height * 4, pFile);
	//fclose(pFile);

	cv::Mat yMat(cv::Size(width,height), CV_8UC1, buffer);
	//cv::Mat rgbMat(cv::Size(width,height), CV_8UC3, pBuffer);

	//dlib::cv_image<dlib::bgr_pixel> rgbImg(rgbMat);
	dlib::cv_image<unsigned char> yImg(yMat);
	
	dlib::array2d<unsigned char> downImg;
	// Call pyr N times should have the same result, but one time may be faster
	
	switch(in->config.pyramid_down_layers)
	{
	case 2:
	{
		// best for 800x600 image on a 2-core Intel(R) Xeon(R) CPU E3-1220 v5 @ 3.00GHz in VMware
		// face distance from camera is about 50cm
		dlib::pyramid_down<2> pyr;
		if (in->config.pyramid_down_n >= 1)
		{
			pyr(yImg, downImg);
			for(uint8_t i = 0; i < in->config.pyramid_down_n - 1; i++)
				pyr(downImg, downImg);
		}
	}
	break;
	case 3:
	{
		dlib::pyramid_down<3> pyr;
		pyr(yImg, downImg);
		if (in->config.pyramid_down_n >= 1)
		{
			pyr(yImg, downImg);
			for(uint8_t i = 0; i < in->config.pyramid_down_n - 1; i++)
				pyr(downImg, downImg);
		}
	}
	case 4:
	{
		dlib::pyramid_down<4> pyr;
		pyr(yImg, downImg);
		if (in->config.pyramid_down_n >= 1)
		{
			pyr(yImg, downImg);
			for(uint8_t i = 0; i < in->config.pyramid_down_n - 1; i++)
				pyr(downImg, downImg);
		}
	}
	break;
	case 0:
	default:
		// do nothing
	break;
	}

	LOGP(DEBUG, "downImg L=%u, n=%u, nr=%u, nc=%u", 
		in->config.pyramid_down_layers, in->config.pyramid_down_n, downImg.nr(), downImg.nc());

	//LOGP(DEBUG, "doFaceTrack time1=%llu (ms)", time_msec() - tbegin);

	uint8_t factPosFactor = 0;
	if (downImg.nr() > 0 && downImg.nc() > 0)
		factPosFactor = width / downImg.nc();
	
	// Detect faces 
	std::vector<dlib::rectangle> faces;
	if (factPosFactor > 0)
		faces = in->detector(downImg);
	else
		faces = in->detector(yImg);
	
	//LOGP(DEBUG, "doFaceTrack time2=%llu (ms)", time_msec() - tbegin);

	// Find the pose of each face.
	std::vector<dlib::full_object_detection> shapes;
	for (unsigned long i = 0; i < faces.size(); ++i)
	{
		dlib::full_object_detection posePoint;
		if (factPosFactor > 0)
		{
			posePoint = in->pose_model(downImg, faces[i]);
			
			for (size_t i = 0; i < 68; i++)
				posePoint.part(i) = dlib::point(posePoint.part(i).x() * factPosFactor, posePoint.part(i).y() * factPosFactor);
		}
		else
		{
			posePoint = in->pose_model(yImg, faces[i]);
		}
		
		shapes.push_back(posePoint);

		if (factPosFactor > 0)
		{
			faces[i].set_left(faces[i].left() * factPosFactor);
			faces[i].set_top(faces[i].top() * factPosFactor);
			faces[i].set_right(faces[i].right() * factPosFactor);
			faces[i].set_bottom(faces[i].bottom() * factPosFactor);
		}

		LOGP(DEBUG, "face: %d, poseLTRB: [%d, %d, %d, %d]", i, 
			faces[i].left(), faces[i].top(), faces[i].right(), faces[i].bottom());

		cv::rectangle(yMat, 
			cv::Point2i(faces[i].left(), faces[i].top()), 
			cv::Point2i(faces[i].right(), faces[i].bottom()), CV_RGB(128, 128, 128), 2);
			
		for (int i = 0; i < shapes.size(); i++)
		{  
			for (int j = 0; j < 68; j++)
				cv::circle(yMat, cvPoint(shapes[i].part(j).x(), shapes[i].part(j).y()), 2, cvScalar(255, 0, 0), -1);
		}  
	}
	
	LOGP(DEBUG, "doFaceTrack time3=%llu (ms)", time_msec() - tbegin);

	return faces.size();
}

bool PL_DlibFaceTrack::pay(const PipeMaterial& pm)
{
	PL_DlibFaceTrack_Internal* in = (PL_DlibFaceTrack_Internal*)internal;

	if (pm.type != PipeMaterial::PMT_FRAME)
	{
		LOG_ERROR << "Only support PMT_FRAME" << std::endl;
		return false;
	}
	
	if (pm.buffer == nullptr)
		return false;
	
	MB_Frame* frame = (MB_Frame*)pm.buffer;
	if (frame->type != MB_Frame::MBFT_YUV420)
	{
		LOG_ERROR << "Only support MBFT_YUV420" << std::endl;
		return false;
	}

	int face_count = doFaceTrack(
						in, (uint8_t*)frame->buffer, frame->width, frame->height, MB_Frame::MBFT_YUV420);
	if (face_count < 0)
	{
		in->payError = true;
		return false;
	}
	else
		in->payError = false;
	
	//in->buffer readly
	in->lastFrame.type = MB_Frame::MBFT_YUV420;
	in->lastFrame.buffer = frame->buffer;//#todo should copy
	in->lastFrame.buffSize = frame->buffSize;
	in->lastFrame.width = frame->width;
	in->lastFrame.height = frame->height;
	in->lastFrame.pts = frame->pts;
	gettimeofday(&(in->lastFrame.pts),NULL);

	return true;
}

bool PL_DlibFaceTrack::gain(PipeMaterial& pm)
{
	PL_DlibFaceTrack_Internal* in = (PL_DlibFaceTrack_Internal*)internal;

	if (!in->payError)
	{
		pm.type = PipeMaterial::PMT_FRAME;
		pm.buffer = &(in->lastFrame);
		pm.buffSize = 0;
		pm.former = this;
	}
	pm.former = this;
	return !in->payError;
}