#include "ContainerStep2InferenceEngine.h"
#include <opencv2/opencv.hpp>
#include "myqueue.h"

using namespace ai_matrix;

ContainerStep2InferenceEngine::ContainerStep2InferenceEngine() {}

ContainerStep2InferenceEngine::~ContainerStep2InferenceEngine() {}

APP_ERROR ContainerStep2InferenceEngine::Init()
{
    strPort0_ = engineName_ + "_" + std::to_string(engineId_) + "_0";
    this->modelConfig_ = Config::getins()->getModelByContainerStep2Config();

    //读取模型信息
    int iFolderExist = access(modelConfig_.strModelPath.c_str(), R_OK);
    if (iFolderExist == -1)
    {
        LogError << "模型：" << modelConfig_.strModelPath << " 不存在!";
        return false;
    }
    class_num = this->modelConfig_.vecClass.size();
    score_threshold = this->modelConfig_.fScoreThreshold;
    input_size = GET_INPUT_SIZE(model_width , model_height);
    output_size = GET_OUTPUT_SIZE(model_width, model_height, clear_num, class_num);
    det_size = clear_num + class_num + 5;

    int ret = initModel();
    if (ret != APP_ERR_OK)
    {
        LogError << "Failed to read model info, ret = " << ret;
        return ret;
    }

    LogInfo << "ContainerStep2InferenceEngine Init ok";
    return APP_ERR_OK;
}

APP_ERROR ContainerStep2InferenceEngine::initModel()
{

    modelinfo.yolov5ClearityModelParam.uiClassNum = class_num;
    modelinfo.yolov5ClearityModelParam.uiClearNum = clear_num;
    modelinfo.yolov5ClearityModelParam.uiDetSize = det_size;
    modelinfo.yolov5ClearityModelParam.fScoreThreshold = score_threshold;
    modelinfo.yolov5ClearityModelParam.fNmsThreshold = nms_threshold;

    modelinfo.modelCommonInfo.uiModelWidth = model_width;
    modelinfo.modelCommonInfo.uiModelHeight = model_height;
    modelinfo.modelCommonInfo.uiInputSize = input_size;
    modelinfo.modelCommonInfo.uiOutputSize = output_size;
    modelinfo.modelCommonInfo.uiChannel = INPUT_CHANNEL;
    modelinfo.modelCommonInfo.uiBatchSize = batch_size;
    modelinfo.modelCommonInfo.strInputBlobName = INPUT_BLOB_NAME;
    modelinfo.modelCommonInfo.strOutputBlobName = OUTPUT_BLOB_NAME;

    string strModelName;

    int nRet = yolov5model.YoloV5ClearityInferenceInit(&modelinfo,
                                                       strModelName,
                                                       this->modelConfig_.strModelPath);
    if (nRet != 0)
    {
        LogInfo << "YoloV5ClassifyInferenceInit nRet:" << nRet;
        return APP_ERR_COMM_READ_FAIL;
    }
    return APP_ERR_OK;
}

APP_ERROR ContainerStep2InferenceEngine::DeInit()
{
    yolov5model.YoloV5ClearityInferenceDeinit();
    LogInfo << "ContainerStep2InferenceEngine DeInit ok";
    return APP_ERR_OK;
}

void ContainerStep2InferenceEngine::resetLocation(SingleData &singleData, SingleData &step1SingleData, float fResizeRatio)
{
    singleData.fLTX = singleData.fLTX * fResizeRatio + step1SingleData.fLTX;
    singleData.fLTY = singleData.fLTY * fResizeRatio + step1SingleData.fLTY;
    singleData.fRBX = singleData.fRBX * fResizeRatio + step1SingleData.fLTX;
    singleData.fRBY = singleData.fRBY * fResizeRatio + step1SingleData.fLTY;

    singleData.fLTX = (singleData.fLTX < IMAGE_WIDTH) ? singleData.fLTX : IMAGE_WIDTH;
    singleData.fLTY = (singleData.fLTY < IMAGE_HEIGHT) ? singleData.fLTY : IMAGE_HEIGHT;
    singleData.fRBX = (singleData.fRBX < IMAGE_WIDTH) ? singleData.fRBX : IMAGE_WIDTH;
    singleData.fRBY = (singleData.fRBY < IMAGE_HEIGHT) ? singleData.fRBY : IMAGE_HEIGHT;
}

APP_ERROR ContainerStep2InferenceEngine::Process()
{
    int iRet = APP_ERR_OK;

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

        std::shared_ptr<InferenceResultData> pInferenceResultData = std::static_pointer_cast<InferenceResultData>(pVoidData0);

        std::shared_ptr<VStep2OutputData> pVStep2OutputData = std::make_shared<VStep2OutputData>();
//        VStep2OutputData vStep2OutputData;
        pVStep2OutputData->iFrameId = pInferenceResultData->iFrameId;
        pVStep2OutputData->iDataSource = pInferenceResultData->iDataSource;
        pVStep2OutputData->bIsEnd = pInferenceResultData->bIsEnd;
        pVStep2OutputData->strDetectDate = pInferenceResultData->strDetectDate;
        pVStep2OutputData->strDetectTime = pInferenceResultData->strDetectTime;
        pVStep2OutputData->cvImage = pInferenceResultData->cvImage.clone();
//        *pVStep2OutputData = vStep2OutputData;

        if (pInferenceResultData->bIsEnd)
        {
            outputQueMap_[strPort0_]->push(std::static_pointer_cast<void>(pVStep2OutputData), true);
            continue;
        }

        if (pInferenceResultData->cvImage.empty())
        {
//            LogWarn << " 帧：" << pVStep2OutputData->iFrameId << " 数据源：" << pVStep2OutputData->iDataSource << " 所传图像为空";
            outputQueMap_[strPort0_]->push(std::static_pointer_cast<void>(pVStep2OutputData), true);
            usleep(1000); //1ms
            continue;
        }

//        LogWarn << "-- 0  -->" << pInferenceResultData->vecSingleData.size();
        for (auto & single : pInferenceResultData->vecSingleData)
        {
            Step2ResultData step2ResultData;
            step2ResultData.fLTX = single.fLTX;
            step2ResultData.fLTY = single.fLTY;
            step2ResultData.fRBX = single.fRBX;
            step2ResultData.fRBY = single.fRBY;
            step2ResultData.iClassId = single.iClassId;
            step2ResultData.fScore = single.fScore;
            step2ResultData.iTrainIndex = single.iTrainIndex;

            if (single.iTargetType == CORNER)
            {
                pVStep2OutputData->vecCornerResultData.emplace_back(step2ResultData);
                continue;
            };
        }

//        LogInfo << " -V--- " << pInferenceResultData->singleData.fScore;
        if (pInferenceResultData->singleData.fScore > 0.0f)
        {
            pVStep2OutputData->step2ResultData.fLTX = pInferenceResultData->singleData.fLTX;
            pVStep2OutputData->step2ResultData.fLTY = pInferenceResultData->singleData.fLTY;
            pVStep2OutputData->step2ResultData.fRBX = pInferenceResultData->singleData.fRBX;
            pVStep2OutputData->step2ResultData.fRBY = pInferenceResultData->singleData.fRBY;
            pVStep2OutputData->step2ResultData.iClassId = pInferenceResultData->singleData.iClassId;
            pVStep2OutputData->step2ResultData.fScore = pInferenceResultData->singleData.fScore;
            pVStep2OutputData->step2ResultData.iTrainIndex = pInferenceResultData->singleData.iTrainIndex;

            cv::Rect rect(cv::Point(pInferenceResultData->singleData.fLTX, pInferenceResultData->singleData.fLTY),
                          cv::Point(pInferenceResultData->singleData.fRBX, pInferenceResultData->singleData.fRBY));
            cv::Mat image = pInferenceResultData->cvImage(rect).clone();

            //进行推理
            std::vector<stDetection> vecInferenceResult;
            auto start = std::chrono::system_clock::now(); // 计时开始
            yolov5model.YoloV5ClearityInferenceModel(image, vecInferenceResult, 2);
            auto end = std::chrono::system_clock::now();

            for (const auto & j : vecInferenceResult)
            {
                if (j.class_conf < STEP2_SCORE_THRESH)
                {
                    continue;
                }
                SingleData singledata;
                singledata.iLine = j.clear_id;
                singledata.iClassId = j.class_id;
                singledata.fScore = j.class_conf;
                singledata.fLTX = j.bbox[0];
                singledata.fLTY = j.bbox[1];
                singledata.fRBX = j.bbox[2];
                singledata.fRBY = j.bbox[3];

                this->resetLocation(singledata, pInferenceResultData->singleData);
                pVStep2OutputData->step2ResultData.vecSingleData.emplace_back(singledata);

//                LogDebug << "frameId:" << pVStep2OutputData->iFrameId
//                         << " --iClassId:" << singledata.iClassId
//                         << " iLine:" << singledata.iLine
//                         << " score=" << singledata.fScore
//                         << " ["
//                         << singledata.fLTX << "," << singledata.fLTY
//                         << "],["
//                         << singledata.fRBX << "," << singledata.fRBY
//                         << "]";
            }
        }

        outputQueMap_[strPort0_]->push(std::static_pointer_cast<void>(pVStep2OutputData), true);
    }
    return APP_ERR_OK;
}