#include "PostProcessEngine.h"

using namespace std;
using namespace ai_matrix;

const std::vector<std::string> class_names={
    "person", "bicycle", "car", "motorcycle", "airplane",
    "bus", "train", "truck", "boat", "traffic light", "fire hydrant",
    "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse",
    "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
    "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis",
    "snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
    "skateboard", "surfboard", "tennis racket", "bottle", "wine glass",
    "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich",
    "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake",
    "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv",
    "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave",
    "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
    "scissors", "teddy bear", "hair drier", "toothbrush"
};


PostProcessEngine::PostProcessEngine() {}
PostProcessEngine::~PostProcessEngine() {}

APP_ERROR PostProcessEngine::Init()
{
    strPort0_ = engineName_ + "_" + std::to_string(engineId_) + "_0";

    width_ = IMAGE_WIDTH, height_ = IMAGE_HEIGHT;

    LogInfo << "engineId_:" << engineId_ << " ImagePreprocessEngine Init ok";
    return APP_ERR_OK;
}

APP_ERROR PostProcessEngine::DeInit()
{   
    LogInfo << "engineId_:" << engineId_ << " ImagePreprocessEngine DeInit ok";
    return APP_ERR_OK;
}


APP_ERROR PostProcessEngine::Process()
{
    uint64_t u64image_num = 0;
    std::string result;
    
    uint64_t u64count_num = 0;
    int iRet = APP_ERR_OK;

    while (!isStop_)
    {
        //后处理
        std::shared_ptr<void> pVoidData0 = nullptr;
        inputQueMap_[strPort0_]->pop(pVoidData0);
        if (nullptr == pVoidData0)
        {
            usleep(1*1000); //n ms
            continue;
        }

        // std::cout<<"Enter PostProcessEngine Thread "<<++u64count_num<<" Times!"<<std::endl;
        std::shared_ptr<InferenceData> pInferenceModelData = std::static_pointer_cast<InferenceData>(pVoidData0);

        std::vector<Yolo::Detection> res;     
        nms(res, static_cast<float *>(pInferenceModelData->pData.get()), CONF_THRESH, NMS_THRESH);  //进行非极大值抑制后res中保留着结果
        std::cout<<"this frame find obj num is:"<<res.size()<<std::endl;
        //根据非极大值抑制的结果标注相关信息(画框,文字信息等)
        //res.size()为每张图片上的识别到的对象数目
        cv::Mat RGBImage(height_, width_, CV_8UC3, pInferenceModelData->pSrcData.get());   //RGB源数据
        
        // for (size_t j = 0; j < res.size(); j++) {
        //     cv::Rect r = get_rect(RGBImage, res[j].bbox);
        //     cv::rectangle(RGBImage, r, cv::Scalar(0x27, 0xC1, 0x36), 2);
        //     cv::putText(RGBImage, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
        // }
        // result = "postprocess_result_"+to_string(++u64image_num)+".jpg";
        // cv::imwrite(result, RGBImage);   //Jpeg库与NVIDIA底层库有冲突,使用硬件Jpeg编码,使用OpenCV也不行,会报段错误!!!

        for (size_t j = 0; j < res.size(); j++) {
            uint8_t r, g, b;
            tie(r, g, b) = random_color(res[j].class_id);
            std::string put_text = class_names.at(res[j].class_id) + ": " + std::to_string(res[j].conf);
            cv::Rect rect = get_rect(RGBImage, res[j].bbox);
            cv::rectangle(RGBImage, rect, cv::Scalar(r, g, b), 5);
            // cv::putText(RGBImage, put_text, cv::Point(rect.x, rect.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(r, g, b), 2);

            int width = cv::getTextSize(put_text, 0, 1, 2, nullptr).width + 10;
            cv::rectangle(RGBImage, cv::Point(rect.x-3, rect.y-33), cv::Point(rect.x + width, rect.y), cv::Scalar(r, g, b), -1);
            cv::putText(RGBImage, put_text, cv::Point(rect.x, rect.y-5), 0, 1, cv::Scalar::all(0), 2, 16);
        }
        
        //构造画框后的RGB数据
        void* pPostProcessRGBBuffer = nullptr;
		unsigned int pPostProcessRGBBuffer_Size = pInferenceModelData->iSrcSize;
		pPostProcessRGBBuffer = new uint8_t[pPostProcessRGBBuffer_Size];
        memcpy(pPostProcessRGBBuffer, pInferenceModelData->pSrcData.get(), pPostProcessRGBBuffer_Size);

		std::shared_ptr<FrameData> pPostProcessRGBFrameData = std::make_shared<FrameData>();
        
        pPostProcessRGBFrameData->iDataSource = engineId_;
        pPostProcessRGBFrameData->iSize = pPostProcessRGBBuffer_Size;
        pPostProcessRGBFrameData->pData.reset(pPostProcessRGBBuffer, [](void* data){if(data) {delete[] data; data = nullptr;}}); //智能指针管理内存
        // pPostProcessRGBFrameData->pData.reset(pPostProcessRGBBuffer, Deleter); //智能指针管理内存
        pPostProcessRGBFrameData->i64TimeStamp =  pInferenceModelData->i64TimeStamp;


        #if 1
        //构造画框后的RGB数据压入下一引擎
        iRet = outputQueMap_[strPort0_]->push(std::static_pointer_cast<void>(pPostProcessRGBFrameData));
        if (iRet != APP_ERR_OK){
			LogError << "push the postprocess rgb frame data failed...";
            // std::cerr<<"push the postprocess rgb frame data failed..."<<std::endl;
		}else{
            // std::cout<<"push the postprocess rgb frame data success!"<<std::endl;
        }
        #endif
    }
}