/*
 * @Author: your name
 * @Date: 2022-04-20 15:49:50
 * @LastEditTime: 2022-09-23 21:51:58
 * @LastEditors: sueRimn
 * @Description: 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
 * @FilePath: /ZCXD_MonitorPlatform/src/CoreLogicModule/include/CamDeal.h
 */
#ifndef AICheck_H_
#define AICheck_H_

// #define USE_TERNSORRT10_BIGMODEL

#define CUDA_API_PER_THREAD_DEFAULT_STREAM 1

#include <vector>
#include <thread>
#include <mutex>
#ifdef USE_TERNSORRT10_BIGMODEL

#include "NvInfer.h"
#include "cuda_runtime_api.h"
using namespace nvinfer1;
#else
#include "argsParser.h"
#include "buffers.h"
#include "common.h"
#include "logger.h"
#include "NvCaffeParser.h"
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#endif
using namespace std;

#define MAX_AI_BUFFER_SIZE 5

enum AIBufferType_
{
    AIBufferType_NULL, // default
    AIBufferType_IN,   // 输入
    AIBufferType_OUT,  // 输出
};
struct AIBuffer
{
    int ntype;           // 数据类型 AIBufferType_
    int ndatalength;     // 数据长度
    int ndataSize;       // 数据字节大小  ndatalength*sizeof(float)
    int nuchardataSize;  // 数据字节大小  ndatalength*sizeof(char)
    std::string strName; // 名称
    AIBuffer()
    {
        ntype = AIBufferType_NULL;
        ndatalength = 0;
        ndataSize = 0;
        nuchardataSize = 0;
        strName = "";
    }
    void copy(AIBuffer tem)
    {
        this->ntype = tem.ntype;
        this->ndatalength = tem.ndatalength;
        this->ndataSize = tem.ndataSize;
        this->nuchardataSize = tem.nuchardataSize;
        this->strName = tem.strName;
    }
    void print(std::string str = "")
    {
        printf("ntype %d strName %s ndatalength %d ndataSize %d nuchardataSize %d\n", ntype, strName.c_str(), ndatalength, ndataSize, nuchardataSize);
    }
};

struct AIInitConfig
{
    int nGpuIdx;
    AIBuffer bufferList[MAX_AI_BUFFER_SIZE];
    std::string engine_file_path;
    AIInitConfig()
    {
        nGpuIdx = -1;
        engine_file_path = "";
    }
    void copy(AIInitConfig tem)
    {
        this->nGpuIdx = tem.nGpuIdx;
        this->engine_file_path = tem.engine_file_path;
        for (int i = 0; i < MAX_AI_BUFFER_SIZE; i++)
        {
            this->bufferList[i].copy(tem.bufferList[i]);
        }
    }
    bool Checking()
    {
        if (nGpuIdx < 0 || nGpuIdx > 4)
        {
            return false;
        }
        if (engine_file_path.empty())
        {
            return false;
        }

        // 第一个不是 输入
        if (bufferList[0].ntype != AIBufferType_IN)
        {
            return false;
        }
        // 第二个是 空
        if (bufferList[1].ntype == AIBufferType_NULL)
        {
            return false;
        }

        bool bhaveout = false;
        int npretype = AIBufferType_IN;
        for (int i = 1; i < MAX_AI_BUFFER_SIZE; i++)
        {
            // 前序是in
            if (npretype == AIBufferType_IN)
            {
                if (bufferList[i].ntype == AIBufferType_IN)
                {
                    continue;
                }
                else if (bufferList[i].ntype == AIBufferType_OUT)
                {
                    npretype = AIBufferType_OUT;
                    bhaveout = true;
                }
                else
                {
                    return false;
                }
            }
            else if (npretype == AIBufferType_OUT)
            {
                if (bufferList[i].ntype == AIBufferType_IN)
                {
                    return false;
                }
                else if (bufferList[i].ntype == AIBufferType_OUT)
                {
                    continue;
                }
                else
                {
                    npretype = AIBufferType_NULL;
                }
            }
            else
            {
                if (bufferList[i].ntype == AIBufferType_IN)
                {
                    return false;
                }
                else if (bufferList[i].ntype == AIBufferType_OUT)
                {
                    return false;
                }
                else
                {
                    npretype = AIBufferType_NULL;
                }
            }
        }

        if (!bhaveout)
        {
            return false;
        }

        return true;
    }
    void CalSize(int mulSzie)
    {
        for (int i = 0; i < MAX_AI_BUFFER_SIZE; i++)
        {
            if (bufferList[i].ntype == AIBufferType_NULL)
            {
                continue;
            }
            bufferList[i].ndataSize = bufferList[i].ndatalength * mulSzie;
            bufferList[i].nuchardataSize = bufferList[i].ndatalength * sizeof(unsigned char);
        }
    }
};
class AI_defect
{
public:
    AI_defect();
    ~AI_defect();
    int model_init(AIInitConfig config);

    int model_Cuda_AI_In_1_Out_1(unsigned char *p_indata_0, unsigned char *p_outdata_1);
    int model_Cuda_AI_In_1_Out_1_float(unsigned char *p_indata_0, float *p_outdata_1);

private:
    void destroy();

private:
    IRuntime *runtime;
    ICudaEngine *engine;
    IExecutionContext *context;
    void *buffers[MAX_AI_BUFFER_SIZE];
    void *ImgData[MAX_AI_BUFFER_SIZE]; // uchar
    const int BATCH_SIZE = 1;

    bool m_bInitialized;

    std::mutex g_mutex;
    AIInitConfig m_config;
    float* floatData[MAX_AI_BUFFER_SIZE];

    int kk = 0;
};

#endif