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

using namespace ai_matrix;

TrainStepOneEngine::TrainStepOneEngine() {}

TrainStepOneEngine::~TrainStepOneEngine() {}

APP_ERROR TrainStepOneEngine::Init()
{
    bUseEngine_ = MyUtils::getins()->ChkIsHaveTarget("NUM");
    if (!bUseEngine_)
    {
        LogWarn << "engineId_:" << engineId_ << " not use engine";
        return APP_ERR_OK;
    }

    strPort0_ = engineName_ + "_" + std::to_string(engineId_) + "_0";
    modelConfig_ = MyYaml::GetIns()->GetModelConfig("TrainStepOneEngine");
    iSpaceMinRBXPer_ = MyYaml::GetIns()->GetIntValue("gc_space_minrbx_imgpercent");

    // 读取模型信息
    APP_ERROR ret = ReadModelInfo();
    if (ret != APP_ERR_OK)
    {
        LogError << "Failed to read model info, ret = " << ret;
        return ret;
    }
    ret = InitModel();
    if (ret != APP_ERR_OK)
    {
        LogError << "Failed to read model info, ret = " << ret;
        return ret;
    }

    // 获取几个摄像头识别车号；是否有单独识别平车的摄像头(即定检期摄像头是否识别平车)
    std::map<int, ai_matrix::DataSourceConfig> mapUseDataSouceCfg = MyYaml::GetIns()->GetUseDataSourceConfig();
    for (auto iter = mapUseDataSouceCfg.begin(); iter != mapUseDataSouceCfg.end(); iter++)
    {
        if (iter->second.strTarget.find("NUM") != std::string::npos)
        {
            LogDebug << "DataSource:" << iter->first << " deal NUM";
            mapDataSourceCfg_[iter->first] = iter->second;

            if (iter->second.strTarget.find("CHKDATE") != std::string::npos ||
                iter->second.strTarget.find("FLAT") != std::string::npos)
            {
                LogDebug << "sourceid:" << iter->first << " Target:" << iter->second.strTarget;
                mapSameSideFlatcarSid_[iter->first] = iter->first;
            }
        }
    }

    // 设置同侧识别平车摄像头信息
    for (auto iter = mapDataSourceCfg_.begin(); iter != mapDataSourceCfg_.end(); iter++)
    {
        if (mapSameSideFlatcarSid_.find(iter->first) != mapSameSideFlatcarSid_.end())
        {
            continue;
        }
        for (auto iterFlat = mapSameSideFlatcarSid_.begin(); iterFlat != mapSameSideFlatcarSid_.end(); iterFlat++)
        {
            ai_matrix::DataSourceConfig dataSourceCfg = MyYaml::GetIns()->GetDataSourceConfigById(iterFlat->first);
            if (dataSourceCfg.iLeftFirst == iter->second.iLeftFirst)
            {
                mapSameSideFlatcarSid_[iter->first] = iterFlat->first;
                break;
            }
        }
    }

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

APP_ERROR TrainStepOneEngine::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, modelConfig_.strOmPath);
    if (nRet != 0)
    {
        LogInfo << "YoloV5ClassifyInferenceInit nRet:" << nRet;
        return APP_ERR_COMM_READ_FAIL;
    }
    return APP_ERR_OK;
}

APP_ERROR TrainStepOneEngine::ReadModelInfo()
{
    char szAbsPath[PATH_MAX];
    // Get the absolute path of model file
    if (realpath(modelConfig_.strOmPath.c_str(), szAbsPath) == nullptr)
    {
        LogError << "Failed to get the real path of " << modelConfig_.strOmPath.c_str();
        return APP_ERR_COMM_NO_EXIST;
    }

    // Check the validity of model path
    int iFolderExist = access(szAbsPath, R_OK);
    if (iFolderExist == -1)
    {
        LogError << "ModelPath " << szAbsPath << " doesn't exist or read failed!";
        return APP_ERR_COMM_NO_EXIST;
    }

    //读取模型参数信息文件
    Json::Value jvModelInfo;
    if (!MyUtils::getins()->ReadJsonInfo(jvModelInfo, modelConfig_.strModelInfoPath))
    {
        LogError << "ModelInfoPath:" << modelConfig_.strModelInfoPath << " doesn't exist or read failed!";
        return APP_ERR_COMM_NO_EXIST;
    }

    model_width = jvModelInfo["model_width"].asInt();
    model_height = jvModelInfo["model_height"].asInt();
    clear_num = jvModelInfo["clear"].isArray() ? jvModelInfo["clear"].size() : 0;
    class_num = jvModelInfo["class"].isArray() ? jvModelInfo["class"].size() : 0;
    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;
    score_threshold = modelConfig_.fScoreThreshold;
    nms_threshold = modelConfig_.fNMSTreshold;

    return APP_ERR_OK;
}

APP_ERROR TrainStepOneEngine::DeInit()
{
    if (!bUseEngine_)
    {
        LogWarn << "engineId_:" << engineId_ << " not use engine";
        return APP_ERR_OK;
    }

    yolov5model.YoloV5ClearityInferenceDeinit();
    LogInfo << "TrainStepOneEngine DeInit ok";
    return APP_ERR_OK;
}

/**
* push数据到队列，队列满时则休眠一段时间再push
* inParam : const std::string strPort                           push的端口
          : const std::shared_ptr<ProcessData> &pProcessData    push的数据
* outParam: N/A
* return  : N/A
*/
void TrainStepOneEngine::PushData(const std::string &strPort, const std::shared_ptr<ProcessData> &pProcessData)
{
    while (true)
    {
        int iRet = outputQueMap_[strPort]->push(std::static_pointer_cast<void>(pProcessData));
        if (iRet != 0)
        {
            LogDebug << "sourceid:" << pProcessData->iDataSource << " frameid:" << pProcessData->iFrameId << " push fail iRet:" << iRet;
            if (iRet == 2)
            {
                usleep(10000); // 10ms
                continue;
            }
        }
        break;
    }
}

/**
* 过滤无效信息
* inParam : std::vector<stDetection> &vecRet            :识别结果数据
          : std::shared_ptr<ProcessData> pProcessData   :帧信息数据
* outParam: N/A
* return  : N/A
*/
void TrainStepOneEngine::FilterInvalidInfo(std::vector<stDetection> &vecRet, std::shared_ptr<ProcessData> &pProcessData)
{
    ai_matrix::DataSourceConfig dataSourceCfg = mapDataSourceCfg_[pProcessData->iDataSource];
    std::vector<stDetection> vecSpaceInfo;
    for (auto it = vecRet.begin(); it != vecRet.end();)
    {
        LogDebug << "frameId:" << pProcessData->iFrameId
                 << " bigclassid:" << it->class_id << " ltx:" << it->bbox[0] << " lty:" << it->bbox[1]
                 << " rbx:" << it->bbox[2] << " rby:" << it->bbox[3];
        // 根据配置文件中  设置的识别范围，过滤掉无效数据
        if (!(it->bbox[0] >= dataSourceCfg.fIdentifyAreasLTX &&
              it->bbox[1] >= dataSourceCfg.fIdentifyAreasLTY &&
              it->bbox[2] <= dataSourceCfg.fIdentifyAreasRBX &&
              it->bbox[3] <= dataSourceCfg.fIdentifyAreasRBY))
        {
            LogDebug << "frameId:" << pProcessData->iFrameId
                    << " bigclassid:" << it->class_id << " 超出识别区域-识别区域：("
                    << dataSourceCfg.fIdentifyAreasLTX << "," << dataSourceCfg.fIdentifyAreasLTY << "),("
                    << dataSourceCfg.fIdentifyAreasRBX << "," << dataSourceCfg.fIdentifyAreasRBY << ")";
            it = vecRet.erase(it);
            continue;
        }

        // 如果设置了不识别车头，则去掉车头所有大框
        if (!MyYaml::GetIns()->GetBoolValue("gc_train_heard_detect") && it->class_id == TRAIN_HEAD)
        {
            LogDebug << "frameId:" << pProcessData->iFrameId << " 过滤掉车头编号";
            it = vecRet.erase(it);
            continue;
        }

        // 过滤车头部分的非车头编号大框
        if(pProcessData->nMonitorState == MONITOR_MODEL_TRAIN_HEAD )
        {
            if(it->class_id  != TRAIN_HEAD)
            {
                LogDebug << " 帧号:" << pProcessData->iFrameId
                        << " 大类:" << it->class_id << " 识别于车头位置，无效！";
                it = vecRet.erase(it);
                continue;
            }
        }

        // 去除车尾的间隔信息
        if (pProcessData->nMonitorState == MONITOR_MODEL_TRAIN_TAIL
        && ((it->class_id >= 9 && it->class_id <= 17 && it->class_id != 15) || it->class_id == 18))
        {
            LogDebug << " frameId:" << pProcessData->iFrameId
                     << " bigclassid:" << it->class_id
                     <<" 识别于车尾部分,无效！";
            it = vecRet.erase(it);
            continue;
        }

        // 去除无车状态下的大框信息
        if (pProcessData->nMonitorState == MONITOR_MODEL_NO_TRAIN)
        {
            LogDebug << " frameId:" << pProcessData->iFrameId
                     << " bigclassid:" << it->class_id
                     <<" 识别于无车状态,无效！";
            it = vecRet.erase(it);
            continue;
        }

        // 按大框高度剔除远股道识别的信息
        int iClassHeight = it->bbox[3] - it->bbox[1];
        if (dataSourceCfg.mapClassMinH.find(it->class_id) != dataSourceCfg.mapClassMinH.end() &&
            iClassHeight < dataSourceCfg.mapClassMinH[it->class_id])
        {
            LogDebug << " frameId:" << pProcessData->iFrameId
                     << " bigclassid:" << it->class_id << " iClassHeight:" << iClassHeight
                     << " minH:" << dataSourceCfg.mapClassMinH[it->class_id] << " 过滤疑似远股道识别";
            it = vecRet.erase(it);
            continue;
        }

        //如果有平车摄像头，则平车摄像头只识别大框类别1，6。
        auto iterFlat = mapSameSideFlatcarSid_.find(pProcessData->iDataSource);
        if (iterFlat != mapSameSideFlatcarSid_.end() && pProcessData->iDataSource == iterFlat->second)
        {
            if (it->class_id != 1 && it->class_id != 6)
            {
                LogDebug << " frameId:" << pProcessData->iFrameId << " flat camera only deal 1 or 6";
                it = vecRet.erase(it);
                continue;
            }
        }

        //剔除高度大于宽的车号大框
        if (((it->class_id >= 2 && it->class_id <= 6) || it->class_id == J_TRAIN_NUM || it->class_id == W_TRAIN_NUM) &&
            (it->bbox[3] - it->bbox[1]) > (it->bbox[2] - it->bbox[0]))
        {
            LogWarn << " frameId:" << pProcessData->iFrameId
                    << " bigclassid:" << it->class_id << " 过滤 高度大于宽度的车号";
            it = vecRet.erase(it);
            continue;
        }
        if (((it->class_id >= 2 && it->class_id <= 6) || it->class_id == J_TRAIN_NUM || it->class_id == W_TRAIN_NUM) &&
                (it->bbox[3] - it->bbox[1]) < MyYaml::GetIns()->GetIntValue("gc_num_frame_height"))
        {
            LogWarn << "疑似误识别到远股道车号，帧号:" << pProcessData->iFrameId
                    << "大框高度:" << (it->bbox[3] - it->bbox[1]);
            it = vecRet.erase(it);
            continue;
        }

        if ((it->class_id == 1 || it->class_id == TRAIN_PRO)
        && (it->bbox[3] - it->bbox[1]) < MyYaml::GetIns()->GetIntValue("gc_pro_frame_height")) {
            LogWarn << "疑似误识别到远股道属性，帧号:" << pProcessData->iFrameId
                    << "大框高度:" << (it->bbox[3] - it->bbox[1]);
            it = vecRet.erase(it);
            continue;
        }

        if ((it->class_id == 9 || it->class_id == C_TRAIN_SPACE)
            && (it->bbox[2] - it->bbox[0]) > MyYaml::GetIns()->GetIntValue("gc_c_space_frame_width")) {
            LogWarn << "疑似误识别到敞车间隔，帧号:" << pProcessData->iFrameId
                    << "大框宽度:" << (it->bbox[2] - it->bbox[0]);
            it = vecRet.erase(it);
            continue;
        }

        //补连塔的相机比较近，间隔基本在画面底部，因此当间隔比较靠画面上时过滤掉。
        if ((it->class_id >= 9 && it->class_id <= 17 && it->class_id != 15) || it->class_id == U_TRAIN_SPACE)
        {
            if (it->bbox[3] < (pProcessData->iHeight * iSpaceMinRBXPer_ / 100))
            {
                LogWarn << " frameId:" << pProcessData->iFrameId
                        << " bigclassid:" << it->class_id << " 过滤间隔过于靠下的间隔信息 fRBY:" << it->bbox[3];
                it = vecRet.erase(it);
                continue;
            }
            vecSpaceInfo.emplace_back(*it);
        }

        ++it;
    }

    //主摄像头1帧如果只识别2个大框，如果非平车的车号和属性场景，则必有间隔框
    if (pProcessData->iDataSource == 0 && vecRet.size() == 2)
    {
        int iHeight0 = vecRet[0].bbox[1] / 2 + vecRet[0].bbox[3] / 2;
        int iHeight1 = vecRet[1].bbox[1] / 2 + vecRet[1].bbox[3] / 2;
        int iCenterY = pProcessData->iHeight / 2;
        if (iHeight0 < iCenterY && iHeight1 < iCenterY) //非平车
        {
            if (!((vecRet[0].class_id >= 9 && vecRet[0].class_id <= 17 && vecRet[0].class_id != 15) || vecRet[0].class_id == U_TRAIN_SPACE) && 
                !((vecRet[1].class_id >= 9 && vecRet[1].class_id <= 17 && vecRet[1].class_id != 15) || vecRet[1].class_id == U_TRAIN_SPACE))
            {
                LogDebug << " frameId:" << pProcessData->iFrameId << " no space";
                vecRet.clear();
            }
        }
    }

    // //K车间隙大，导致识别对面股道信息的过滤。 获取间隔最宽的大框，跟该大框X轴有重合的过滤掉。
    // if (vecSpaceInfo.size() > 0 && vecRet.size() > 1)
    // {
    //     std::sort(vecRet.begin(), vecRet.end(), [](stDetection &a, stDetection &b){return (a.bbox[2]-a.bbox[0]) > (b.bbox[2]-b.bbox[0]);});
    //     auto iterMaxWSpace = vecSpaceInfo.begin();
    //     if (iterMaxWSpace->class_id == K_TRAIN_SPACE)
    //     {
    //         for (auto it = vecRet.begin(); it != vecRet.end();)
    //         {
    //             if ((it->class_id <= 6 || it->class_id == W_TRAIN_NUM) && !(it->bbox[2] <= iterMaxWSpace->bbox[0] || it->bbox[0] >= iterMaxWSpace->bbox[2]))
    //             {
    //                 LogError << "sourceid:" << pProcessData->iDataSource << " frameId:" << pProcessData->iFrameId
    //                          << " classid:" << it->class_id << " overlap x";
    //                 it = vecRet.erase(it);
    //                 continue;
    //             }
    //             ++it;
    //         }
    //     }
    // }
}

/**
* 设置大框类型
* inParam : PostSubData &postSubData   :推理结果
* outParam: PostSubData &postSubData   :推理结果
* return  : N/A
*/
void TrainStepOneEngine::SetTargetType(PostSubData &postSubData)
{
    if (postSubData.iBigClassId == TRAIN_HEAD)
    {
        postSubData.iTargetType = HEAD;
    }
    else if (postSubData.iBigClassId == TRAIN_PRO)
    {
        postSubData.iTargetType = PRO;
    }
    else if ((postSubData.iBigClassId >= 2 && postSubData.iBigClassId <= 6) ||
             postSubData.iBigClassId == J_TRAIN_NUM ||
             postSubData.iBigClassId == W_TRAIN_NUM)
    {
        postSubData.iTargetType = NUM;
    }
    else if (postSubData.iBigClassId >= 9 && postSubData.iBigClassId <= 17 && postSubData.iBigClassId != 15)
    {
        postSubData.iTargetType = TRAINSPACE;
    }
    else if (postSubData.iBigClassId == U_TRAIN_SPACE)
    {
        postSubData.iTargetType = SPACE;
    }
}

APP_ERROR TrainStepOneEngine::Process()
{
    if (!bUseEngine_)
    {
        LogWarn << "engineId_:" << engineId_ << " not use engine";
        return APP_ERR_OK;
    }
    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<ProcessData> pProcessData = std::static_pointer_cast<ProcessData>(pVoidData0);
        //组织输出数据
        std::shared_ptr<PostData> pPostData = std::make_shared<PostData>();
        pPostData->iModelType = MODELTYPE_NUM;
        pPostData->nMonitorState = pProcessData->nMonitorState; //来车检测的四个分类

		//获取图片
        if (pProcessData->iStatus == TRAINSTATUS_RUN || pProcessData->bIsEnd)
        {
            if (pProcessData->pData != nullptr && pProcessData->iSize != 0)
            {
                cv::Mat img(pProcessData->iHeight, pProcessData->iWidth, CV_8UC3, static_cast<uint8_t *>(pProcessData->pData.get()));   //BGR

                //进行推理
                std::vector<stDetection> res;
                //auto start = std::chrono::system_clock::now();  //计时开始
                yolov5model.YoloV5ClearityInferenceModel(img, res);
                //auto end = std::chrono::system_clock::now();
                //LogInfo << "nopr1 inference time: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms";

                //过滤无效信息
                FilterInvalidInfo(res, pProcessData);

                MyUtils::getins()->GetMaxScoreResult(res);
                //整理推理结果
                //根据非极大值抑制的结果标注相关信息(画框,文字信息等)
                //res.size()为每张图片上的识别到的对象数目
                for (size_t j = 0; j < res.size(); j++)
                {
                    /*
                    [0:车头; 1:属性; 2:煤炭漏斗车(兖矿自备,枣矿自备); 3:敞车; 4:棚车; 5:罐车; 6:平车
                    7:定检期; 8:牲畜车; 9:敞车间隔; 10:自备车间隔; 11:平车间隔; 12:罐车间隔; 13:棚车车间隔;
                    14:牲畜车间隔; 15:毒品车; 16: 毒品车间隔; 17:混合车厢间隔; 18:连接轴通用间隔]
                    */
                    if (res[j].class_id < 0 || res[j].class_id > 18)
                    {
                        continue;
                    }
                    SingleData singledata;
                    // singledata.iLine = -1;
                    singledata.iClassId = res[j].class_id;
                    singledata.fScore = res[j].class_conf;
                    // singledata.iAnchorId = -1;
                    singledata.fLTX = res[j].bbox[0];
                    singledata.fLTY = res[j].bbox[1];
                    singledata.fRBX = res[j].bbox[2];
                    singledata.fRBY = res[j].bbox[3];
                    singledata.fClear = res[j].clear_id;

                    PostSubData postSubData;
                    postSubData.iBigClassId = res[j].class_id;
                    postSubData.vecSingleData.emplace_back(singledata);
                    postSubData.step1Location.fLTX = res[j].bbox[0];
                    postSubData.step1Location.fLTY = res[j].bbox[1];
                    postSubData.step1Location.fRBX = res[j].bbox[2];
                    postSubData.step1Location.fRBY = res[j].bbox[3];
                    SetTargetType(postSubData);
                    pPostData->vecPostSubData.emplace_back(postSubData);

//                    LogDebug << "数据源:" << pProcessData->iDataSource << " 帧:" << pProcessData->iFrameId
//                             << " --iClassId:" << singledata.iClassId << " iLine:" << singledata.iLine << " confidence=" << singledata.fScore
//                             << " lx=" << singledata.fLTX << " ly=" << singledata.fLTY << " rx=" << singledata.fRBX << " ry=" << singledata.fRBY
//                             << " clear:" << singledata.fClear;
                }
            }
        }

        //及时释放内存
        if (pProcessData->pData != nullptr)
        {
            pProcessData->pData = nullptr;
            pProcessData->iSize = 0;
        }

        //push端口0，第1步推理
        pProcessData->pVoidData = std::static_pointer_cast<void>(pPostData);
        iRet = outputQueMap_[strPort0_]->push(std::static_pointer_cast<void>(pProcessData), true);
//        PushData(strPort0_, pProcessData);
    }
    return APP_ERR_OK;
}