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

using namespace ai_matrix;

MoveEngine::MoveEngine() {}

MoveEngine::~MoveEngine() {}

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

    this->modelConfig_ = Config::getins()->getModelByMoveConfig();
    this->identifyConfig_ = Config::getins()->getIdentifyConfig();
    this->baseConfig_ = Config::getins()->getBaseConfig();
    this->dataSourceConfig_ = Config::getins()->getDataSourceConfig();

    if (identifyConfig_.bNeedMoveDetectFlag)
    {
        // 读取模型信息
        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;
        }
    }

    InitParam();

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

APP_ERROR MoveEngine::InitModel()
{
    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 = yolov8model.YoloV8InferenceInit(&modelinfo, strModelName, modelConfig_.strModelPath);
    if (nRet != 0)
    {
        LogInfo << "YoloV5ClassifyInferenceInit nRet:" << nRet;
        return APP_ERR_COMM_READ_FAIL;
    }
    return APP_ERR_OK;
}

APP_ERROR MoveEngine::ReadModelInfo()
{
    int iFolderExist = access(modelConfig_.strModelPath.c_str(), R_OK);
    if (iFolderExist == -1)
    {
        LogError << "ModelPath " << modelConfig_.strModelPath << " doesn't exist or read failed!";
        return APP_ERR_COMM_NO_EXIST;
    }
    return APP_ERR_OK;
}

APP_ERROR MoveEngine::DeInit()
{
    if (identifyConfig_.bNeedMoveDetectFlag)
    {
        yolov8model.YoloV8InferenceDeinit();
    }
    LogInfo << "MoveEngine DeInit ok";
    return APP_ERR_OK;
}


/**
* 参数初始化(列车结束时需调用)
* inParam : N/A
* outParam: N/A
* return  : N/A
*/
void MoveEngine::InitParam()
{
    this->iStepInter_ = 0;
    for (auto & map_it : this->mapMoveDataNO_)
    {
        map_it.second = 1;
    }
    this->strTrainTime_ = "";
    this->strTrainName_ = "";
    this->strTrainDate_ = "";
    for (bool & value : this->bHaveTrain_)
    {
        value = false;
    }
}

bool MoveEngine::getIsHaveTrain()
{
    for (const bool & type : this->bHaveTrain_)
    {
        if (type) return true;
    }
//    for (int i = 0; i < 2; ++i)
//    {
//        if (g_bHaveTrain[i].load()) return true;
//    }
    return false;
}
bool MoveEngine::isNewTrain()
{
    bool bIsNew = true;
    for (auto & map_it : this->mapMoveDataNO_)
    {
        if (map_it.second != 1) bIsNew = false;
    }

    return bIsNew;
}

/**
* 使用单device处理
* inParam : std::shared_ptr<ProcessData> pProcessData
* outParam: N/A
* return  : N/A
*/
void MoveEngine::SingleDeviceProcess(std::shared_ptr<ProcessData> pProcessData, int iType)
{
    if(this->isNewTrain() || strTrainDate_.empty())
    {
        this->strTrainDate_ = TimeUtil::getins()->getDate();
        this->strTrainTime_ = TimeUtil::getins()->getTime();
        this->strTrainName_ = StringUtil::getins()->replace_all_distinct(strTrainTime_, ":", "-");

        //创建该该列车的数据存储目录(防止后续多线程创建报错)
        std::string strTrainPath = baseConfig_.strDebugResultPath + "/"
                + strTrainDate_ + "/"
                + this->strTrainName_ + "/jpg_0";
        FileUtil::getins()->createDirPath(strTrainPath);
        //创建该该列车的数据存储目录(防止后续多线程创建报错)
        std::string strTrainPath2 = baseConfig_.strDebugResultPath + "/"
                                   + strTrainDate_ + "/"
                                   + this->strTrainName_ + "/jpg_1";
        FileUtil::getins()->createDirPath(strTrainPath2);
        std::string strBestImgPath = baseConfig_.strResultPath + "/"
                + strTrainDate_ + "/"
                + this->strTrainName_;
        FileUtil::getins()->createDirPath(strBestImgPath);
    }

//    if (g_come_direction != 0 && (g_come_direction != this->dataSourceConfig_.iDirection && this->dataSourceConfig_.iDirection))
//    {
//        pProcessData->bIsTrainEnd = true;
//    }

//    if (pProcessData->bIsTrainEnd && this->isNewTrain())
//    {
//        return;
//    }

    std::string strFilePath = baseConfig_.strDebugResultPath + "/"
            + strTrainDate_ + "/"
            + strTrainName_ + "/jpg_" + std::to_string(pProcessData->iDataSource);

    // 通知存储来车检测结果
    std::shared_ptr<VMoveInfo> pVMoveInfo = std::make_shared<VMoveInfo>();
    pVMoveInfo->iDataSource = pProcessData->iDataSource;
    pVMoveInfo->strFilePath = strFilePath;
    pVMoveInfo->strFileName = std::to_string(this->mapMoveDataNO_[pProcessData->iDataSource]) + ".json";
    pVMoveInfo->iFrameId = this->mapMoveDataNO_[pProcessData->iDataSource]; //当前帧号
    pVMoveInfo->i64TimeStamp = pProcessData->sourceFrameData.i64TimeStamp;
    pVMoveInfo->bIsEnd = pProcessData->bIsTrainEnd;
    pVMoveInfo->strTrainDate = strTrainDate_;
    pVMoveInfo->strTrainTime = strTrainTime_;
    pVMoveInfo->strTrainName = strTrainName_;
    pVMoveInfo->iTrainStage = iType;
    outputQueMap_[strPort1_]->push(std::static_pointer_cast<void>(pVMoveInfo), true);

    // push端口,存图
    std::shared_ptr<SaveImgData> pSaveImgData = std::make_shared<SaveImgData>();
    pSaveImgData->iDataSource = pProcessData->iDataSource;
    pSaveImgData->strFilePath = strFilePath;
    pSaveImgData->strFileName = std::to_string(this->mapMoveDataNO_[pProcessData->iDataSource]) + ".jpg";
    pSaveImgData->strTrainStage = this->intTrainStage_2_str(iType);
    pSaveImgData->frameData.pData = pProcessData->sourceFrameData.pData;
    pSaveImgData->frameData.iSize = pProcessData->sourceFrameData.iSize;
    pSaveImgData->frameData.iWidth = pProcessData->dataSourceInfo.iWidth;
    pSaveImgData->frameData.iHeight = pProcessData->dataSourceInfo.iHeight;
    pSaveImgData->frameData.i64TimeStamp = pProcessData->sourceFrameData.i64TimeStamp;
    pSaveImgData->bIsEnd = pProcessData->bIsTrainEnd;
    outputQueMap_[strPort0_]->push(std::static_pointer_cast<void>(pSaveImgData), true);

    if (pProcessData->bIsTrainEnd)
    {
        std::shared_ptr<SaveImgData> pSaveImgData2 = std::make_shared<SaveImgData>();
        pSaveImgData2->iDataSource = pProcessData->iDataSource == 0 ? 1: 0;
        pSaveImgData2->strFileName = std::to_string(this->mapMoveDataNO_[pSaveImgData2->iDataSource]-1) + ".jpg";
        pSaveImgData2->strFilePath = baseConfig_.strDebugResultPath + "/"
                                    + strTrainDate_ + "/"
                                    + strTrainName_ + "/jpg_" + std::to_string(pSaveImgData2->iDataSource);
        pSaveImgData2->bIsEnd = pProcessData->bIsTrainEnd;
        outputQueMap_[strPort0_]->push(std::static_pointer_cast<void>(pSaveImgData2), true);

        LogInfo << "--->  火车结束信号  数据源：" << pSaveImgData2->iDataSource << "  帧：" << this->mapMoveDataNO_[pSaveImgData2->iDataSource];
        LogInfo << "--->  火车结束信号  数据源：" << pProcessData->iDataSource << "  帧：" << this->mapMoveDataNO_[pProcessData->iDataSource];
        this->InitParam();
    }
}

std::string MoveEngine::intTrainStage_2_str(int iTrainStage)
{
    switch (iTrainStage) {
        case MONITOR_MODEL_TRAIN_HEAD:
            return "train_head";
        case MONITOR_MODEL_HEAD_FIRST:
            return "head_and_carriage";
        case MONITOR_MODEL_NO_TRAIN:
            return "no_train";
        case MONITOR_MODEL_TRAIN_TAIL:
            return "train_tail";
        case MONITOR_MODEL_TRAIN_BODY:
            return "carriage";
        default:
            return "";
    }
}

APP_ERROR MoveEngine::Process()
{
    int iRet = APP_ERR_OK;
    while (!isStop_)
    {
        //pop端口0
        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);

//        int iType = MONITOR_MODEL_INIT_STATE;

        //1. 无需动态检测
        if (!this->identifyConfig_.bNeedMoveDetectFlag)
        {
            SingleDeviceProcess(pProcessData, MONITOR_MODEL_INIT_STATE);

            this->mapMoveDataNO_[pProcessData->iDataSource]++;

            if (pProcessData->bIsTrainEnd)
            {
                InitParam();
            }
            continue;
        }

        // 进行推理

        if (!pProcessData->sourceFrameData.pData || !pProcessData->sourceFrameData.iSize)
        {
            continue;
        }

        if (this->getIsHaveTrain() && !this->mapMoveDataNO_[pProcessData->iDataSource] % 2)
        {
            SingleDeviceProcess(pProcessData, iType[pProcessData->iDataSource]);

            this->mapMoveDataNO_[pProcessData->iDataSource]++;
            continue;
        }

        cv::Mat img; // BGR
        cv::Mat img_gray(pProcessData->dataSourceInfo.iHeight, pProcessData->dataSourceInfo.iWidth, CV_8UC1);
        cvtColor(cv::Mat(pProcessData->dataSourceInfo.iHeight, pProcessData->dataSourceInfo.iWidth, CV_8UC3, static_cast<uint8_t*>(pProcessData->sourceFrameData.pData.get())), img_gray, COLOR_BGR2GRAY);
        cvtColor(img_gray, img, COLOR_GRAY2BGR);

        auto start = std::chrono::system_clock::now(); // 计时开始
        float fReturnVal[STEP0_OUTPUT_ARRAY];
        memset(fReturnVal, 0x00, sizeof(fReturnVal));

        yolov8model.YoloV8InferenceModelGetType(img, fReturnVal, STEP0_OUTPUT_ARRAY * sizeof(float));
        // exit(0);

        float fScore = 0.0f;
        for(int n = 0; n < STEP0_OUTPUT_ARRAY; n++){
            if(fReturnVal[n] > fScore){
                fScore = fReturnVal[n];
                iType[pProcessData->iDataSource] = n;
            }
        }
//            LogDebug <<"模型得分  车头:"<< fReturnVal[0]<<"  无车:"<< fReturnVal[1]<<"  车尾:"<< fReturnVal[2]<<"  有车:"<< fReturnVal[3];
//            LogInfo<<"来车当前状态:"<< (nType == 0 ? "有车头" : (nType == 1 ? "无车"));
        switch (iType[pProcessData->iDataSource]) {
            case MONITOR_MODEL_INIT_STATE:
                LogDebug << "数据源：" << pProcessData->iDataSource << " 来车检测初始状态";
                break;
            case MONITOR_MODEL_TRAIN_HEAD:
                LogDebug << "数据源：" << pProcessData->iDataSource << " 来车状态：有车头";
                break;
            case MONITOR_MODEL_HEAD_FIRST:
                LogDebug << "数据源：" << pProcessData->iDataSource << " 来车状态：车头与车身之间";
            case MONITOR_MODEL_NO_TRAIN:
//                    LogDebug << "来车状态：无车";
                break;
            case MONITOR_MODEL_TRAIN_TAIL:
                LogDebug << "数据源：" << pProcessData->iDataSource << " 来车状态：车尾";
                break;
            case MONITOR_MODEL_TRAIN_BODY:
                LogDebug << "数据源：" << pProcessData->iDataSource << " 来车状态：车身";
                break;
            default:
                LogWarn << "数据源：" << pProcessData->iDataSource << " 来车状态：未定义状态";
                break;
        }

        if (this->identifyConfig_.bTrainHeardDetect)
        {
            this->bHaveTrain_[pProcessData->iDataSource] = (iType[pProcessData->iDataSource] == MONITOR_MODEL_TRAIN_BODY
                              || iType[pProcessData->iDataSource] == MONITOR_MODEL_TRAIN_HEAD
                              || iType[pProcessData->iDataSource] == MONITOR_MODEL_HEAD_FIRST
                              || iType[pProcessData->iDataSource] == MONITOR_MODEL_TRAIN_TAIL);
//            g_bHaveTrain[pProcessData->iDataSource].store(iType == MONITOR_MODEL_TRAIN_BODY
//                                                           || iType == MONITOR_MODEL_TRAIN_HEAD
//                                                           || iType == MONITOR_MODEL_HEAD_FIRST
//                                                           || iType == MONITOR_MODEL_TRAIN_TAIL);
        }
        else
        {
            this->bHaveTrain_[pProcessData->iDataSource] = (iType[pProcessData->iDataSource] == MONITOR_MODEL_TRAIN_BODY
                              || iType[pProcessData->iDataSource] == MONITOR_MODEL_TRAIN_TAIL
                              || iType[pProcessData->iDataSource] == MONITOR_MODEL_HEAD_FIRST);
//            g_bHaveTrain[pProcessData->iDataSource].store(iType == MONITOR_MODEL_TRAIN_BODY
//                              || iType == MONITOR_MODEL_TRAIN_TAIL
//                              || iType == MONITOR_MODEL_HEAD_FIRST);
        }

        auto end = std::chrono::system_clock::now();
        uint64_t i64Time = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();

        if (getIsHaveTrain())
        {
            iHasTrainNum_ = iHasTrainNum_ > 20 ? iHasTrainNum_ : iHasTrainNum_ + 1;
//            if (iHasTrainNum_ > 0) LogInfo << "数据源：" << pProcessData->iDataSource << "当前有车, 计数：" << iHasTrainNum_;
        }
        else
        {
            iHasTrainNum_ = iHasTrainNum_ == 0 ? iHasTrainNum_ : iHasTrainNum_ - 1;
            if (iHasTrainNum_ == 0) LogInfo << "数据源：" << pProcessData->iDataSource << " ----- 当前无车 -----";
        }

        if (iHasTrainNum_ > 0)
            //有车开始识别
        {
            iStepInter_ = 1;

        }
        else
            //无车停止识别
        {
            if (iStepInter_ == 1)
            {
                iStepInter_ = 2;
            }
        }
        //有车识别处理
        if (iStepInter_ != 0)
        {
            pProcessData->bIsTrainEnd = (iStepInter_ == 2); //动态检测无车，设置列车结束标识

            SingleDeviceProcess(pProcessData, iType[pProcessData->iDataSource]);

            this->mapMoveDataNO_[pProcessData->iDataSource]++;

            if (iStepInter_ == 2)
            {
                InitParam();
            }
        }
    }
    return APP_ERR_OK;
}