#include "FaceCache.h"
#include "PipeLine.h"
#include "MaterialBuffer.h"
#include "PL_SensetimeFaceTrackMultiTrd.h"
#include <logger.h>
#include <Logger/src/logger.hpp>

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

#include <vector>

#include <PbFaceList.pb.h>

#include <libyuv.h>
#include <MediaHelper.h>

//#define YUV420_TO_RGB888 1
#define YUV420_TO_RGB565 1
//#define YUV420_TO_ABGR8888 1
//#define NV12_TO_ABGR8888 1

struct FcPmBreackerContext
{
	uint8_t frameYUV[1920*1080*4];
	size_t frameYUVSize;
    MB_Frame::MBFType frameYUVType;
	
	uint8_t frameRGB[1920*1088*4];
	size_t frameRGBSize;

    size_t width;
    size_t height;
	
	st_ff_vect_t faceFeatures;
	
	bool dataAvailable;
	
	FcPmBreackerContext() :
            frameYUV(), frameYUVSize(0), frameYUVType(MB_Frame::MBFT__FIRST), frameRGB(), frameRGBSize(0),
            faceFeatures(), width(0), height(0), dataAvailable(false)
	{
	}
	
	void reset()
	{
		frameYUVSize = 0;
        frameRGBSize = 0;
        width = 0;
        height = 0;
		dataAvailable = false;
		faceFeatures.clear();
	}
	
	bool convertYUV420ToRGB888()
	{
        int src_height = height;
        int src_width = width;
        const uint8* src_y = (const uint8*)(frameYUV);
        const uint8* src_u = (const uint8*)(src_y + (src_height * src_width));
        const uint8* src_v = (const uint8*)(src_u + (src_height * src_width / 4));

        libyuv::I420ToRGB24(src_y, src_width,
                           src_u, MH_SUBSAMPLE1(src_width, 2),
                           src_v, MH_SUBSAMPLE1(src_width, 2),
                           frameRGB, 3 * src_width,
                           src_width, src_height);
        frameRGBSize = src_height * src_width * 3;

        //{
        //    static size_t f = 0;
        //    char fname[50];
        //    sprintf(fname, "/sdcard/face-%u.rgb", ++f);
        //    FILE *pFile = fopen(fname, "wb");
        //    fwrite(frameRGB, 1, frameRGBSize, pFile);
        //    fclose(pFile);
        //    if (f > 20)exit(0);
        //}

        return true;
	}

    bool convertYUV420ToRGB565()
    {
        int src_height = height;
        int src_width = width;
        const uint8* src_y = (const uint8*)(frameYUV);
        const uint8* src_u = (const uint8*)(src_y + (src_height * src_width));
        const uint8* src_v = (const uint8*)(src_u + (src_height * src_width / 4));

        if (frameYUVType == MB_Frame::MBFT_YUV420)
        {
            libyuv::I420ToRGB565(src_y, src_width,
                                 src_u, MH_SUBSAMPLE1(src_width, 2),
                                 src_v, MH_SUBSAMPLE1(src_width, 2),
                                 frameRGB, 2 * src_width,
                                 src_width, src_height);
        }
        else if (frameYUVType == MB_Frame::MBFT_NV12)
        {
            libyuv::NV12ToRGB565(src_y, src_width,
                                 src_u, src_width,
                                 frameRGB, 2 * src_width,
                                 src_width, src_height);

            //{
            //    static size_t f = 0;
            //    char fname[50];
            //    sprintf(fname, "/sdcard/face-%u.rgb", ++f);
            //    FILE *pFile = fopen(fname, "wb");
            //    fwrite(frameRGB, 1, frameRGBSize, pFile);
            //    fclose(pFile);
            //    if (f > 20)exit(0);
            //}
        }
        frameRGBSize = src_height * src_width * 2;
        return true;
    }

    bool convertYUV420ToABGR8888()
    {
        int src_height = height;
        int src_width = width;
        const uint8* src_y = (const uint8*)(frameYUV);
        const uint8* src_u = (const uint8*)(src_y + (src_height * src_width));
        const uint8* src_v = (const uint8*)(src_u + (src_height * src_width / 4));

        libyuv::I420ToABGR(src_y, src_width, // android ARGB_8888 is ABGR
                             src_u, MH_SUBSAMPLE1(src_width, 2),
                             src_v, MH_SUBSAMPLE1(src_width, 2),
                             frameRGB, 4 * src_width,
                             src_width, src_height);
        frameRGBSize = src_height * src_width * 4;

        //{
        //    static size_t f = 0;
        //    char fname[50];
        //    sprintf(fname, "/sdcard/face-%u.argb", ++f);
        //    FILE *pFile = fopen(fname, "wb");
        //    fwrite(frameRGB, 1, frameRGBSize, pFile);
        //    fclose(pFile);
        //    if (f > 10)exit(0);
        //}

        return true;
    }
};

bool fc_pm_breaker_ptr(const PipeMaterial* pm, void* args)
{
	FcPmBreackerContext* ctx = (FcPmBreackerContext*)args;
	
	const st_ff_vect_t& faceFeatures(*(const st_ff_vect_t*)(pm->buffer));
	
	ctx->faceFeatures = faceFeatures;
	
	ctx->dataAvailable &= true;
	return true;
}

bool fc_pm_breaker_ptr_count(const PipeMaterial* pm, void* args)
{
    size_t* count = (size_t*)args;

    const st_ff_vect_t& faceFeatures(*(const st_ff_vect_t*)(pm->buffer));
    *count = faceFeatures.size();
    return false;
}

bool fc_pm_breaker_frame(const PipeMaterial* pm, void* args)
{
	FcPmBreackerContext* ctx = (FcPmBreackerContext*)args;
	
	const MB_Frame* lastFrame((const MB_Frame*)(pm->buffer));
	if (lastFrame->type != MB_Frame::MBFT_YUV420 && lastFrame->type != MB_Frame::MBFT_NV12)
	{
		ctx->dataAvailable &= false;
		return false;
	}
	
	if (lastFrame->buffSize == 0)
	{
		ctx->dataAvailable &= false;
		return false;
	}

    ctx->frameYUVSize = lastFrame->buffSize;
    ctx->frameYUVType = lastFrame->type;
    ctx->width = lastFrame->width;
    ctx->height = lastFrame->height;
	memcpy(ctx->frameYUV, lastFrame->buffer, ctx->frameYUVSize);
	
	ctx->dataAvailable &= true;
	return true;
}

NativeImgIdx::operator std::string() const
{
    char buf[256];
    sprintf(buf, "offset=%d, size=%d, type=%d, width=%d, height=%d", offset, size, type, width, height);
    return std::string(buf);
}

#ifdef USE_ST_SDK
FaceCache::FaceCache() : _ctx(new FcPmBreackerContext)
{
}

FaceCache::~FaceCache()
{
	delete (FcPmBreackerContext*)_ctx;
}
#endif

// returns count of face
int FaceCache::cachePm(const PipeMaterial& pm)
{
	if (_ctx == nullptr)
		_ctx = new FcPmBreackerContext;
	
	FcPmBreackerContext& ctx(*(FcPmBreackerContext*)_ctx);
	ctx.reset();
	ctx.dataAvailable = true;
	pm.breake(PipeMaterial::PMT_PTR, MB_Frame::MBFT__FIRST, fc_pm_breaker_ptr, _ctx);
	pm.breake(PipeMaterial::PMT_FRAME, MB_Frame::MBFT__FIRST, fc_pm_breaker_frame, _ctx);
	
	if (ctx.dataAvailable)
		return ctx.faceFeatures.size();
	else
		return 0;
}

// returns count of face
size_t FaceCache::getFaceCount(const PipeMaterial& pm) const
{
    size_t count = 0;
    pm.breake(PipeMaterial::PMT_PTR, MB_Frame::MBFT__FIRST, fc_pm_breaker_ptr_count, &count);

    return count;
}

bool FaceCache::getFaceListPb(uint8_t* buffer, size_t& buffMaxSize)
{
    FcPmBreackerContext& ctx(*(FcPmBreackerContext*)_ctx);
    if (!ctx.dataAvailable)
        return false;
	
	PbFaceList pbFaceList;
    pbFaceList.set_magic(pbFaceList.magic());
    pbFaceList.set_image_count(ctx.faceFeatures.size());
	pbFaceList.set_src_width(ctx.width);
	pbFaceList.set_src_height(ctx.height);
	
    cv::Mat yMat(cv::Size(ctx.width, ctx.height), CV_8UC1, ctx.frameYUV);

    for (int i = 0; i < ctx.faceFeatures.size(); i++)
    {
        if (! ctx.faceFeatures[i].test_face_in_cone(35.0f, 35.0f, 35.0f))
            continue;

        const PLGH_Rect& faceRect(ctx.faceFeatures[i].rect);
        cv::Mat roiMat1(yMat, cv::Rect(faceRect.leftTop.X, faceRect.leftTop.Y, faceRect.width(), faceRect.height()));
		cv::Mat roiMat(roiMat1.clone()); // #todo copy data should be avoid!!!!
		size_t s = roiMat.total() * roiMat.elemSize();

		PbFaceList_FaceListImage& pbFaceListImage(*(pbFaceList.add_images()));
		pbFaceListImage.set_idx(i);
		pbFaceListImage.set_size(s);
		pbFaceListImage.set_type(PbFaceList_FaceListImage_ImageType_MBFT_Y8);
		pbFaceListImage.set_width(roiMat.cols);
		pbFaceListImage.set_height(roiMat.rows);
		pbFaceListImage.set_top_left_x(faceRect.leftTop.X);
		pbFaceListImage.set_top_left_y(faceRect.leftTop.Y);
		
		pbFaceListImage.add_img(roiMat.ptr(), s);

        //{
        //    static size_t f = 0;
        //    char fname[50];
        //    sprintf(fname, "/sdcard/face-%u.yuv", ++f);
        //    FILE *pFile = fopen(fname, "wb");
        //    fwrite(roiMat.ptr(), 1, s, pFile);
        //    fclose(pFile);
        //    if (f > 20)exit(0);
        //}

        //{
        //    static size_t f = 0;
        //    char fname[50];
        //    sprintf(fname, "/sdcard/face-%u.yuv", ++f);
        //    FILE *pFile = fopen(fname, "wb");
        //    const std::string* img(*pbFaceListImage.img().data());
        //    fwrite(img->data(), 1, img->size(), pFile);
        //    fclose(pFile);
        //    if (f > 20)exit(0);
        //}
    }

	size_t s = pbFaceList.ByteSize();
	buffMaxSize = std::min(s, buffMaxSize);
	pbFaceList.SerializeToArray(buffer, buffMaxSize);

    //{
    //    static size_t f = 0;
    //    char fname[50];
    //    sprintf(fname, "/sdcard/facelist-%u.pb", ++f);
    //    FILE *pFile = fopen(fname, "wb");
    //    fwrite(buffer, 1, buffMaxSize, pFile);
    //    fclose(pFile);
    //    if (f > 10)exit(0);
    //}

    return buffMaxSize > 0;
}

#ifdef YUV420_TO_RGB888

bool FaceCache::getFaceListImage(std::vector<NativeImgIdx>& imgIdxes, uint8_t* buffImg, size_t& buffImgMaxSize)
{
    FcPmBreackerContext& ctx(*(FcPmBreackerContext*)_ctx);
    if (!ctx.dataAvailable)
        return false;

    if (ctx.frameRGBSize == 0)
    {
        //#todo should optimize not convert the whole image
        if (! ctx.convertYUV420ToRGB888())
            return false;
    }

    imgIdxes.clear();
	size_t offset = 0;

    cv::Mat rgbMat(cv::Size(ctx.width, ctx.height), CV_8UC3, ctx.frameRGB);
    for (st_ff_vect_t::const_iterator ffiter = ctx.faceFeatures.begin(); ffiter != ctx.faceFeatures.end(); ++ffiter)
    {
        const FaceRect& faceRect(ffiter->rect);
        cv::Mat roiMat1(rgbMat, cv::Rect(faceRect.leftTop.x, faceRect.leftTop.y, faceRect.width(), faceRect.height()));
        cv::Mat roiMat(roiMat1.clone()); // #todo copy data should be avoid!!!!

        NativeImgIdx imgidx;
        imgidx.offset = offset;
        imgidx.size = roiMat.total() * roiMat.elemSize();
        imgidx.type = MB_Frame::MBFT_RGB888;
        imgidx.width = roiMat.cols;
        imgidx.height = roiMat.rows;

        const size_t newSize = imgidx.offset + imgidx.size;

		//isContinuous
		if (newSize > buffImgMaxSize)
        {
            LOG_ERROR << "FaceCache::getFaceListImage buffImgMaxSize truncated" << LOG_ENDL;
            return false;
        }


        uint8_t* pbuf = buffImg + imgidx.offset;
        memcpy(pbuf, roiMat.ptr(), imgidx.size);

        //uint8_t* pbuf = buffImg + imgidx.offset;
        //cv::Mat roiMat1(roiMat.cols, roiMat.rows, CV_8UC3, pbuf);
        //pbuf[0]=123;
        //roiMat.copyTo(roiMat1);
        //memcpy(buffImg + imgidx.offset, roiMat.ptr(), imgidx.size);

        //uint8_t* pbuf = buffImg + imgidx.offset;
        //for(size_t r = 0; r < roiMat.rows; r++)
        //{
        //    const uint8_t* pcol = roiMat.ptr<uint8_t>(r);
        //    memcpy(pbuf, pcol, roiMat.cols);
        //    pbuf += roiMat.cols;
        //}

        //{
        //    static size_t f = 0;
        //    char fname[50];
        //    sprintf(fname, "/sdcard/face-%u.rgb", ++f);
        //    FILE *pFile = fopen(fname, "wb");
        //    fwrite(pbuf, 1, imgidx.size, pFile);
        //    fclose(pFile);
        //    if (f > 20)exit(0);
        //}

        imgIdxes.push_back(imgidx);
        offset = newSize;
    }

    buffImgMaxSize = offset;
    return true;
}

#elif YUV420_TO_RGB565

bool FaceCache::getFaceListImage(std::vector<NativeImgIdx>& imgIdxes, uint8_t* buffImg, size_t& buffImgMaxSize)
{
    FcPmBreackerContext& ctx(*(FcPmBreackerContext*)_ctx);
    if (!ctx.dataAvailable)
        return false;

    if (ctx.frameRGBSize == 0)
    {
        //#todo should optimize not convert the whole image
        if (! ctx.convertYUV420ToRGB565())
            return false;
    }

    imgIdxes.clear();
    size_t offset = 0;

    cv::Mat rgbMat(cv::Size(ctx.width, ctx.height), CV_16UC1, ctx.frameRGB);
    for (st_ff_vect_t::const_iterator ffiter = ctx.faceFeatures.begin(); ffiter != ctx.faceFeatures.end(); ++ffiter)
    {
        if (! ffiter->test_face_in_cone(35.0f, 35.0f, 35.0f))
            continue;

        const PLGH_Rect& faceRect(ffiter->rect);
        cv::Mat roiMat1(rgbMat, cv::Rect(faceRect.leftTop.X, faceRect.leftTop.Y, faceRect.width(), faceRect.height()));
        cv::Mat roiMat(roiMat1.clone()); // #todo copy data should be avoid!!!!

        NativeImgIdx imgidx;
        imgidx.offset = offset;
        imgidx.size = roiMat.total() * roiMat.elemSize();
        imgidx.type = MB_Frame::MBFT_RGB565;
        imgidx.width = roiMat.cols;
        imgidx.height = roiMat.rows;

        const size_t newSize = imgidx.offset + imgidx.size;

        //isContinuous
        if (newSize > buffImgMaxSize)
        {
            LOG_ERROR << "FaceCache::getFaceListImage buffImgMaxSize truncated" << LOG_ENDL;
            return false;
        }

        /*
        uint8_t* pbuf = buffImg + imgidx.offset;
        uint8_t* psrc = ctx.frameRGB + (ctx.width * faceRect.leftTop.y + faceRect.leftTop.x) * 2;
        for (int row = 0; row < faceRect.height(); row++)
        {
            memcpy(pbuf, psrc, imgidx.width * 2);
            pbuf += imgidx.width * 2;
            psrc += ctx.width * 2;
        }
         */

        uint8_t* pbuf = buffImg + imgidx.offset;
        memcpy(pbuf, roiMat.ptr(), imgidx.size);

        //{
        //    static size_t f = 0;
        //    char fname[50];
        //    sprintf(fname, "/sdcard/face-%u-w%d-h%d.rgb565", ++f, imgidx.width, imgidx.height);
        //    FILE *pFile = fopen(fname, "wb");
        //    fwrite(pbuf, 1, imgidx.size, pFile);
        //    fclose(pFile);
        //    if (f > 20)exit(0);
        //}

        imgIdxes.push_back(imgidx);
        offset = newSize;
    }

    buffImgMaxSize = offset;
    return true;
}

#elif YUV420_TO_ABGR8888

bool FaceCache::getFaceListImage(std::vector<NativeImgIdx>& imgIdxes, uint8_t* buffImg, size_t& buffImgMaxSize)
{
    FcPmBreackerContext& ctx(*(FcPmBreackerContext*)_ctx);
    if (!ctx.dataAvailable)
        return false;

    if (ctx.frameRGBSize == 0)
    {
        //#todo should optimize not convert the whole image
        if (! ctx.convertYUV420ToABGR8888())
            return false;
    }

    imgIdxes.clear();
    size_t offset = 0;

    cv::Mat rgbMat(cv::Size(ctx.width, ctx.height), CV_8UC4, ctx.frameRGB);
    for (st_ff_vect_t::const_iterator ffiter = ctx.faceFeatures.begin(); ffiter != ctx.faceFeatures.end(); ++ffiter)
    {
        const FaceRect& faceRect(ffiter->rect);

        cv::Mat roiMat1(rgbMat, cv::Rect(faceRect.leftTop.x, faceRect.leftTop.y, faceRect.rightBottom.x - faceRect.leftTop.x, faceRect.rightBottom.y - faceRect.leftTop.y));
        cv::Mat roiMat(roiMat1.clone()); // #todo copy data should be avoid!!!!

        NativeImgIdx imgidx;
        imgidx.offset = offset;
        imgidx.size = roiMat.total() * roiMat.elemSize();
        imgidx.type = MB_Frame::MBFT_ABGR8888;
        imgidx.width = roiMat.cols;
        imgidx.height = roiMat.rows;

        const size_t newSize = imgidx.offset + imgidx.size;

        //isContinuous
        if (newSize > buffImgMaxSize)
        {
            LOG_ERROR << "FaceCache::getFaceListImage buffImgMaxSize truncated" << LOG_ENDL;
            return false;
        }

        uint8_t* pbuf = buffImg + imgidx.offset;
        memcpy(pbuf, roiMat.ptr(), imgidx.size);

        //{
        //    static size_t f = 0;
        //    char fname[50];
        //    sprintf(fname, "/sdcard/face-%u-w%d-h%d.argb", ++f, imgidx.width, imgidx.height);
        //    FILE *pFile = fopen(fname, "wb");
        //    fwrite(pbuf, 1, imgidx.size, pFile);
        //    fclose(pFile);
        //    if (f > 20)exit(0);
        //}

        imgIdxes.push_back(imgidx);
        offset = newSize;
    }

    buffImgMaxSize = offset;
    return true;
}

#endif