/*
 * @Author: xiewenji 527774126@qq.com
 * @Date: 2025-09-03 10:44:31
 * @LastEditors: xiewenji 527774126@qq.com
 * @LastEditTime: 2025-09-06 22:29:55
 * @FilePath: /AI_SO_Test/AIEngineModule/src/AIModel_Impl.cpp
 * @Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
 */

#include "AIModel_Impl.h"
#include "CUDA_DataChange.cuh"
AIModel_Impl::AIModel_Impl()
{
    m_bInitSuccess = false;
    m_GPU_Engine.clear();
    m_modelNodeList.clear();
    m_pNode_input_0 = NULL;
    m_pNode_input_1 = NULL;
    m_pNode_output_0 = NULL;
    m_pNode_output_1 = NULL;
    m_pNode_output_2 = NULL;
    m_DetGPUStream.clear();
    m_nLast_GPUStreamIdx.store(0);
    m_nALLStreamNum = 0;
}
AIModel_Impl::~AIModel_Impl()
{
}
int AIModel_Impl::Init(AIModelRun_Config config)
{
    m_modelRun_Config.Copy(config);
    printf(">>> %s Init Start \n", m_modelRun_Config.strName.c_str());

    if (m_bInitSuccess)
    {
        printf("strName %s Init Success  %d ************ return \n", m_modelRun_Config.strName.c_str(), m_bInitSuccess);
        return 0;
    }
    int re = 0;
    re = ModelRunConfigCheck(m_modelRun_Config);
    if (re != 0)
    {
        printf("strName %s Init Error =%d  run Config error \n", m_modelRun_Config.strName.c_str(), re);
        return re;
    }

    // 不同显卡 单独 载入模型
    if (m_modelRun_Config.gpuconfig.gpu_0)
    {
        re = LoadEngine(0);
        // 加载失败
        if (re != 0)
        {
        }
    }
    if (m_modelRun_Config.gpuconfig.gpu_1)
    {
        re = LoadEngine(1);
        // 加载失败
        if (re != 0)
        {
        }
    }
    int nGPUNum = m_GPU_Engine.size();
    // 模型载入失败
    if (nGPUNum <= 0 || nGPUNum > 2)
    {
        printf(" %s model Init error \n", m_modelRun_Config.strName.c_str());
        return 2;
    }
    for (int i = 0; i < m_modelNodeList.size(); i++)
    {
        if (m_modelNodeList.at(i).type == AI_Buffer_Type_INPUT)
        {
            if (m_pNode_input_0 == NULL)
            {
                m_pNode_input_0 = &m_modelNodeList.at(i);
                input_0.copy(*m_pNode_input_0);
            }
            else
            {
                m_pNode_input_1 = &m_modelNodeList.at(i);
                input_1.copy(*m_pNode_input_1);
            }
        }
        else
        {
            if (m_pNode_output_0 == NULL)
            {
                m_pNode_output_0 = &m_modelNodeList.at(i);
                output_0.copy(*m_pNode_output_0);
            }
            else if (m_pNode_output_1 == NULL)
            {
                m_pNode_output_1 = &m_modelNodeList.at(i);
                output_1.copy(*m_pNode_output_1);
            }
            else
            {
                m_pNode_output_2 = &m_modelNodeList.at(i);
                output_2.copy(*m_pNode_output_2);
            }
        }
    }
    if (m_pNode_input_0 == NULL || m_pNode_output_0 == NULL)
    {
        printf(" %s model Init error \n", m_modelRun_Config.strName.c_str());
        return 3;
    }

    int streamNum = m_GPU_Engine.at(0)->cudaSteams.size();

    // 生产推理用的 参数信息
    if (m_GPU_Engine.size() > 1)
    {
        if (m_GPU_Engine.at(0)->cudaSteams.size() != m_GPU_Engine.at(1)->cudaSteams.size())
        {
            printf(" %s model Init error \n", m_modelRun_Config.strName.c_str());
            return 3;
        }
    }
    if (streamNum <= 0)
    {
        printf(" %s model Init streamNum error \n", m_modelRun_Config.strName.c_str());
        return 4;
    }

    // 显示初始化信息
    if (true)
    {
        printf("Model Name :%s GUP engine Num %ld\n", m_modelRun_Config.strName.c_str(), m_GPU_Engine.size());

        for (const auto &ge : m_GPU_Engine)
        {
            printf("  gpu %d status %d   strem Num %ld \n", ge->nGPUIdx, ge->bsucc, ge->cudaSteams.size());

            for (const auto &psteam : ge->cudaSteams)
            {
                printf("     steam %d \n", psteam->nstreamIdx);
                for (const auto &pcon : psteam->streamConfigList)
                {
                    printf("         type %d [chw] %d %d %d input %p  output %p\n",
                           pcon->nodeConfig.type, pcon->nodeConfig.channel, pcon->nodeConfig.height, pcon->nodeConfig.width,
                           pcon->gpu_buffers, pcon->gpu_ImgData);
                }
            }
        }
    }

    // 把所有gup 的所有stream 都放到一起，方便调用。
    for (int i = 0; i < streamNum; i++)
    {

        for (int igpu = 0; igpu < nGPUNum; igpu++)
        {
            std::shared_ptr<Det_GPU_Stram> temDetStream = std::make_shared<Det_GPU_Stram>();
            temDetStream->nGPUIdx = m_GPU_Engine.at(igpu)->nGPUIdx;
            temDetStream->engine = m_GPU_Engine.at(igpu)->engine;
            temDetStream->cuda_stream = m_GPU_Engine.at(igpu)->cudaSteams.at(i);
            m_DetGPUStream.push_back(temDetStream);
        }
    }
    m_nALLStreamNum = m_DetGPUStream.size();
    printf("========********* Model Name :%s GUP engine Num %ld **************=============\n", m_modelRun_Config.strName.c_str(), m_GPU_Engine.size());
    for (const auto steam : m_DetGPUStream)
    {
        printf(">>> gpu %d stream %d \n",
               steam->nGPUIdx, steam->cuda_stream->nstreamIdx);
    }

    m_bInitSuccess = true;

    // warm-up：首次推理会触发 TensorRT/CUDA 懒初始化（kernel 加载、workspace 分配、首次 enqueueV3 建图等），
    // 提前用一张空图预热，避免首张检测图耗时飙升。
    {
        cv::Mat warmIn = InitMat(m_pNode_input_0->channel, m_pNode_input_0->width, m_pNode_input_0->height);
        cv::Mat warmOut;
        int warmRe = AIDet(warmIn, warmOut);
        if (warmRe != 0)
        {
            printf("%s warm-up error = %d\n", m_modelRun_Config.strName.c_str(), warmRe);
        }
    }

    return 0;
}
cv::Mat AIModel_Impl::InitMat(int channel, int w, int h)
{
    if (w <= 0 || h <= 0 || channel <= 0 || channel > 3)
    {
        return cv::Mat();
    }

    cv::Mat dst;
    if (channel == 1)
    {
        dst = cv::Mat(h, w, CV_8UC1, cv::Scalar(0));
    }
    else
    {
        dst = cv::Mat(h, w, CV_8UC3, cv::Scalar(0, 0, 0));
    }
    return dst;
}
int AIModel_Impl::AIDet(const cv::Mat &inImg, cv::Mat &outimg)
{

    // 1、准备数据
    if (m_pNode_input_0 == NULL || m_pNode_output_0 == NULL)
    {
        printf("AIDet: Node Config INput output  Node error \n");
        return 1;
    }
    // std::this_thread::sleep_for(std::chrono::milliseconds(2)); // Sleep for 100 milliseconds
    uchar *p_indata_0 = (uchar *)inImg.data;
    outimg = InitMat(m_pNode_output_0->channel, m_pNode_output_0->width, m_pNode_output_0->height);
    uchar *p_outdata_1 = (uchar *)outimg.data;
    if (!m_bInitSuccess)
    {
        return 2;
    }

    AI_Det_In_1_Out_1(m_pNode_input_0, m_pNode_output_0, p_indata_0, p_outdata_1);
    return 0;
}

int AIModel_Impl::AIDet(const cv::Mat &inImg, cv::Mat &outimg0, cv::Mat &outimg1)
{
    return 0;
}

int AIModel_Impl::AIClass(const cv::Mat &inImg, float *fmaxScore)
{
    // 1、准备数据
    if (m_pNode_input_0 == NULL || m_pNode_output_0 == NULL)
    {
        printf("AIDet: Node Config INput output  Node error \n");
        return -1;
    }
    if (!m_bInitSuccess)
    {
        return -2;
    }
    uchar *p_indata_0 = (uchar *)inImg.data;

    int class_dix = 0;
    class_dix = AI_Det_In_1_Out_1_class(p_indata_0, fmaxScore);

    return class_dix;
}

// 检测模型的运行参数是否合法，是否正常
int AIModel_Impl::ModelRunConfigCheck(AIModelRun_Config &runConfig)
{
    // GPU 数量判断
    if (runConfig.gpuconfig.GetNum() == 0)
    {
        return 1;
    }
    if (runConfig.strPath == "")
    {
        printf(" %s model path is error \n", runConfig.strName.c_str());
        return 2;
    }
    if (runConfig.Stream_num <= 0)
    {
        runConfig.Stream_num = 1;
    }
    if (runConfig.Stream_num > 4)
    {
        runConfig.Stream_num = 4;
    }

    return 0;
}

// 载入模型。。。。
int AIModel_Impl::LoadEngine(int ngpuIdx)
{
    std::shared_ptr<GPU_Engine> gpuEngine = std::make_shared<GPU_Engine>();
    gpuEngine->engine = std::make_shared<Engine>(ngpuIdx);
    gpuEngine->nGPUIdx = ngpuIdx;
    if (!gpuEngine->engine)
    {
        printf("Engine Create  error\n");
        return -1;
    }
    // 1、加载模型
    bool reload = gpuEngine->engine->loadFromFile(m_modelRun_Config.strPath);

    if (!reload)
    {
        printf("GPU %d  Load error  %s\n", ngpuIdx, m_modelRun_Config.strPath.c_str());
        printf("Load model  error\n");
        return -2;
    }
    printf("%s  GPU %d  Load succ  %s\n", m_modelRun_Config.strName.c_str(), ngpuIdx, m_modelRun_Config.strPath.c_str());

    // 2、 解析 egengine的 相关信息
    int re = GetEngineInfo1(gpuEngine);
    if (re != 0)
    {
        return re;
    }
    // 没有输入输出 节点。
    if (m_modelNodeList.size() <= 0 || m_modelNodeList.size() > MAX_MODEL_NODE_NUM)
    {
        printf("Load model  input output error\n");
        return -3;
    }

    cudaSetDevice(ngpuIdx);
    // 3、 申请显存
    for (int i = 0; i < m_modelRun_Config.Stream_num; i++)
    {
        std::shared_ptr<Cuda_Stream_Config> cudasteam = std::make_shared<Cuda_Stream_Config>();
        cudasteam->nstreamIdx = i;

        if (cudaStreamCreate(&cudasteam->stream) != cudaSuccess)
        {
            printf("Failed to create CUDA stream\n");
            return -4;
        }
        cudasteam->context = std::unique_ptr<nvinfer1::IExecutionContext>(
            gpuEngine->engine->engine_->createExecutionContext());
        // printf("----- stream %d \n", i);
        // 初始化 数据 内存
        for (int j = 0; j < m_modelNodeList.size(); j++)
        {
            std::shared_ptr<Stream_Node_Config> steamNodeconfig = std::make_shared<Stream_Node_Config>();
            steamNodeconfig->nodeConfig.copy(m_modelNodeList.at(j));

            cudaMalloc(&steamNodeconfig->gpu_buffers, steamNodeconfig->nodeConfig.floatsize);
            cudaMalloc(&steamNodeconfig->gpu_ImgData, steamNodeconfig->nodeConfig.ucharsize);
            // 分类 的模型
            if (m_modelRun_Config.IsClass)
            {
                if (steamNodeconfig->cpu_floatData)
                {
                    delete steamNodeconfig->cpu_floatData;
                    steamNodeconfig->cpu_floatData = NULL;
                }

                steamNodeconfig->cpu_floatData = new float[steamNodeconfig->nodeConfig.datalength];
            }

            printf("-----******GPU %d  stream %d cuda malloc %d inOrout %d [chw] %d*%d*%d \n",
                   ngpuIdx, i, j, steamNodeconfig->nodeConfig.type, steamNodeconfig->nodeConfig.channel,
                   steamNodeconfig->nodeConfig.height, steamNodeconfig->nodeConfig.width);

            if (m_modelNodeList.at(j).type == AI_Buffer_Type_INPUT)
            {
                if (cudasteam->pstreamNode_input_0 == nullptr)
                {
                    cudasteam->pstreamNode_input_0 = steamNodeconfig;
                }
                else
                {
                    cudasteam->pstreamNode_input_1 = steamNodeconfig;
                }
            }
            else
            {
                if (cudasteam->pstreamNode_output_0 == nullptr)
                {
                    cudasteam->pstreamNode_output_0 = steamNodeconfig;
                }
                else if (cudasteam->pstreamNode_output_1 == nullptr)
                {
                    cudasteam->pstreamNode_output_1 = steamNodeconfig;
                }
                else
                {
                    cudasteam->pstreamNode_output_2 = steamNodeconfig;
                }
            }
            cudasteam->streamConfigList.push_back(steamNodeconfig);
        }
        gpuEngine->cudaSteams.push_back(cudasteam);
    }
    gpuEngine->bsucc = true;

    m_GPU_Engine.push_back(gpuEngine);

    return 0;
}

int AIModel_Impl::GetEngineInfo1(std::shared_ptr<GPU_Engine> &pgpuengine)
{
    int nInputNum = pgpuengine->engine->getNbBindings();

    m_modelNodeList.erase(m_modelNodeList.begin(), m_modelNodeList.end());
    m_modelNodeList.clear();

    for (int i = 0; i < nInputNum; i++)
    {
        Node_Config node_config;

        std::string name = pgpuengine->engine->getBindingName(i);
        node_config.name = name;
        Dims dims = pgpuengine->engine->getBindingDims(i);

        if (pgpuengine->engine->bindingIsInput(i))
        {
            std::cout << "Input: " << name << "\n";
            node_config.type = AI_Buffer_Type_INPUT;
        }
        else
        {
            std::cout << "Output: " << name << "\n";
            node_config.type = AI_Buffer_Type_OUTPUT;
        }
        std::cout << "dims.nbDims: " << dims.nbDims << "\n";
        if (dims.nbDims == 4)
        {

            switch (m_modelRun_Config.inputType)
            {
            case Input_CHW:
                node_config.channel = dims.d[1];
                node_config.height = dims.d[2];
                node_config.width = dims.d[3];
                break;
            case Input_HWC:

                node_config.height = dims.d[1];
                node_config.width = dims.d[2];
                node_config.channel = dims.d[3];
                break;
            default:
                node_config.channel = dims.d[1];
                node_config.height = dims.d[2];
                node_config.width = dims.d[3];
                break;
            }
        }
        else if (dims.nbDims == 3)
        {
            switch (m_modelRun_Config.inputType)
            {
            case Input_CHW:
                node_config.channel = dims.d[0];
                node_config.height = dims.d[1];
                node_config.width = dims.d[2];
                break;
            case Input_HWC:

                node_config.height = dims.d[0];
                node_config.width = dims.d[1];
                node_config.channel = dims.d[2];
                break;
            default:
                node_config.channel = dims.d[0];
                node_config.height = dims.d[1];
                node_config.width = dims.d[2];
                break;
            }
        }
        else if (dims.nbDims == 2)
        {
            node_config.channel = 1;
            node_config.height = dims.d[0];
            node_config.width = dims.d[1];
        }
        else if (dims.nbDims == 1)
        {
            node_config.channel = 1;
            node_config.height = dims.d[0];
            node_config.width = 1;
        }
        else
        {
            printf("engine  C H W is  error\n");
            return -2;
        }
        node_config.CalDataSize();

        m_modelNodeList.push_back(node_config);
        printf("dim: ");
        for (int di = 0; di < dims.nbDims; di++)
        {
            printf(" %ld ", dims.d[di]);
        }
        printf(" \n");
        printf("C H W [%d %d %d]\n", node_config.channel, node_config.height, node_config.width);
    }

    return 0;
}

int AIModel_Impl::AI_Det_In_1_Out_1(Node_Config *pConfig_in, Node_Config *pConfig_out, const unsigned char *p_indata_0, unsigned char *p_outdata_1)
{
    // printf("=== s1 ");
    std::shared_ptr<Det_GPU_Stram> pdetStream;
    if (GetStream(pdetStream) != 0 || !pdetStream)
    {
        printf("AI_Det_In_1_Out_1: GetStream error \n");
        return 3;
    }
    //  printf("=== s2 ");
    std::lock_guard<std::mutex> lock(pdetStream->AI_mutex);
    // printf("  ss g %d s %d --  ", pdetStream->nGPUIdx, pdetStream->cuda_stream->nstreamIdx);
    // // 设置 显卡ID
    cudaSetDevice(pdetStream->nGPUIdx);

    cudaMemcpyAsync(pdetStream->cuda_stream->pstreamNode_input_0->gpu_ImgData,
                    p_indata_0, pConfig_in->ucharsize, cudaMemcpyHostToDevice, pdetStream->cuda_stream->stream);
    // // //  在显存中 图片数据从 uchar 转到 float
    Cuda_ucharToFloat_stream((unsigned char *)pdetStream->cuda_stream->pstreamNode_input_0->gpu_ImgData,
                             (float *)pdetStream->cuda_stream->pstreamNode_input_0->gpu_buffers, pConfig_in->datalength, pdetStream->cuda_stream->stream);

    pdetStream->cuda_stream->context->setTensorAddress(pConfig_in->name.c_str(), pdetStream->cuda_stream->pstreamNode_input_0->gpu_buffers);
    pdetStream->cuda_stream->context->setTensorAddress(pConfig_out->name.c_str(), pdetStream->cuda_stream->pstreamNode_output_0->gpu_buffers);

    pdetStream->cuda_stream->context->enqueueV3(pdetStream->cuda_stream->stream);

    Cuda_FloatTouchar_stream((float *)pdetStream->cuda_stream->pstreamNode_output_0->gpu_buffers,
                             (unsigned char *)pdetStream->cuda_stream->pstreamNode_output_0->gpu_ImgData, pConfig_out->datalength, pdetStream->cuda_stream->stream);
    cudaMemcpyAsync(p_outdata_1, pdetStream->cuda_stream->pstreamNode_output_0->gpu_ImgData,
                    pConfig_out->ucharsize, cudaMemcpyDeviceToHost, pdetStream->cuda_stream->stream);
    cudaStreamSynchronize(pdetStream->cuda_stream->stream);
    // printf("**ee g %d s %d **  \n", pdetStream->nGPUIdx, pdetStream->cuda_stream->nstreamIdx);
    return 0;
}

int AIModel_Impl::GetStream(std::shared_ptr<Det_GPU_Stram> &pdetStream)
{
    if (m_nALLStreamNum <= 0 || m_DetGPUStream.empty())
    {
        return 1;
    }

    // 多线程并发时原子轮询分配 stream：
    int sidx = m_nLast_GPUStreamIdx.load();
    while (true)
    {
        int next = sidx + 1;
        if (next >= m_nALLStreamNum)
        {
            next = 0;
        }
        if (m_nLast_GPUStreamIdx.compare_exchange_weak(sidx, next))
        {
            sidx = next;
            break;
        }
        // CAS 失败时 sidx 已被更新为当前实际值，重新计算 next
    }

    pdetStream = m_DetGPUStream.at(sidx);
    return 0;
}
int AIModel_Impl::AI_Det_In_1_Out_1_class(unsigned char *p_indata_0, float *fmaxScore)
{
    std::shared_ptr<Det_GPU_Stram> pdetStream;
    if (GetStream(pdetStream) != 0 || !pdetStream)
    {
        printf("AI_Det_In_1_Out_1_class: GetStream error \n");
        *fmaxScore = 0.0f;
        return -3;
    }
    //  printf("=== s2 ");
    std::lock_guard<std::mutex> lock(pdetStream->AI_mutex);
    // printf("  ss g %d s %d --  ", pdetStream->nGPUIdx, pdetStream->cuda_stream->nstreamIdx);
    // // 设置 显卡ID
    std::shared_ptr<Stream_Node_Config> p_input_0 = pdetStream->cuda_stream->pstreamNode_input_0;
    std::shared_ptr<Stream_Node_Config> p_output_0 = pdetStream->cuda_stream->pstreamNode_output_0;
    cudaSetDevice(pdetStream->nGPUIdx);

    cudaMemcpyAsync(p_input_0->gpu_ImgData, p_indata_0,
                    p_input_0->nodeConfig.ucharsize, cudaMemcpyHostToDevice, pdetStream->cuda_stream->stream);
    // // //  在显存中 图片数据从 uchar 转到 float
    Cuda_ucharToFloat_stream((unsigned char *)p_input_0->gpu_ImgData,
                             (float *)p_input_0->gpu_buffers, p_input_0->nodeConfig.datalength, pdetStream->cuda_stream->stream);

    pdetStream->cuda_stream->context->setTensorAddress(p_input_0->nodeConfig.name.c_str(), p_input_0->gpu_buffers);
    pdetStream->cuda_stream->context->setTensorAddress(p_output_0->nodeConfig.name.c_str(), p_output_0->gpu_buffers);

    pdetStream->cuda_stream->context->enqueueV3(pdetStream->cuda_stream->stream);

    cudaMemcpyAsync(p_output_0->cpu_floatData, p_output_0->gpu_buffers,
                    p_output_0->nodeConfig.floatsize, cudaMemcpyDeviceToHost, pdetStream->cuda_stream->stream);
    cudaStreamSynchronize(pdetStream->cuda_stream->stream);

    int label = F2softmaxId(p_output_0->cpu_floatData, p_output_0->nodeConfig.datalength, fmaxScore);
    return label;
}
int AIModel_Impl::F2softmaxId(float *data, int class_num, float *fmaxScore)
{
    int cls_num = class_num;
    float total = 0;
    for (int i = 0; i < cls_num; i++)
    {
        // std::cout<<"data["<<i<<"]="<<data[i]<<std::endl;
        total += exp(data[i]);
    }
    // std::cout<<"total="<<total<<std::endl;
    // std::vector<float> output;//这样才可以用push_back-但相比于=更耗时
    std::vector<float> output(cls_num);
    for (int i = 0; i < output.size(); i++)
    {
        float outi = (float)exp(data[i]) / (float)total;
        output[i] = outi;
        // output.push_back(outi);
        // std::cout<<"output["<<i<<"]="<<output[i]<<std::endl;
    }

    int max_id = 0;
    float max_score = 0;
    for (int i = 0; i < output.size(); i++)
    {
        // std::cout << "output[" << i << "]=" << output[i] << std::endl;
        if (output[i] > max_score)
        {
            max_id = i;
            max_score = output[i];
        }
    }
    *fmaxScore = max_score;
    // std::cout << "F2softmaxId()----------------------------------max_score=" << max_score << std::endl; // 1023dyy
    // std::cout << "F2softmaxId()----------------------------------max_id=" << max_id << std::endl;       // 1023dyy
    return max_id;
}