//
// Created by nvidia on 24-9-12.
//

#ifndef TRAIN_COMMONSTRUCT_H
#define TRAIN_COMMONSTRUCT_H

//模型参数
typedef struct _YoloV5ModelParam{
    unsigned int uiClassNum;
    unsigned int uiDetSize;
    float fScoreThreshold;
    float fNmsThreshold;
}YoloV5ModelParam;

typedef struct _YoloV5ClearityModelParam{
    unsigned int uiClassNum;
    unsigned int uiClearNum;
    unsigned int uiDetSize;
    float fScoreThreshold;
    float fNmsThreshold;
}YoloV5ClearityModelParam;

typedef struct _ModelCommonInfo{
    unsigned int uiModelWidth;
    unsigned int uiModelHeight;
    unsigned int uiInputSize;
    unsigned int uiOutputSize;
    unsigned int uiChannel;
    unsigned int uiBatchSize;
    std::string strInputBlobName;
    std::string strOutputBlobName;
}ModelCommonInfo;

typedef struct _YoloV5ModelInfo{
    YoloV5ModelParam yolov5ModelParam;
    ModelCommonInfo modelCommonInfo;
}YoloV5ModelInfo;

typedef struct _YoloV5ClearityModelInfo{
    YoloV5ClearityModelParam yolov5ClearityModelParam;
    ModelCommonInfo modelCommonInfo;
}YoloV5ClearityModelInfo;

struct ModelConfig
{
    std::string strAnchorPath;  //锚框文件路径
    std::string strClassPath;   //类别名称集合
    std::string strOmPath;      //om文件路径
    int iModelWidth;            //模型输入宽度
    int iModelHeight;           //模型输入高度
    float fScoreThreshold;      //得分阈值
    float fNMSTreshold;         //nms阈值
    int iClassNum;              //类别个数
    int iLineNum;               //字段代号
};

typedef struct
{
    float fLTX = 0; //左上X坐标
    float fLTY = 0; //左上y坐标
    float fRBX = 0; //右下y坐标
    float fRBY = 0; //右下y坐标
}Location;


// 定义日期结构体
typedef struct Date
{
    int year;
    int month;
    int day;
} Date;

//数据源基本信息
typedef struct
{
    uint32_t iWidth = 0;
    uint32_t iHeight = 0;
    int iRate = 0;  //帧率
} DataSourceInfo;

// 原始帧信息
typedef struct {
    // 帧画面
    std::shared_ptr<void> pData = nullptr;
    // 帧图片
    cv::Mat cvImage;
    // Size of memory, bytes
    uint32_t iSize = 0;
    // 帧数据时间戳
    uint64_t i64TimeStamp = 0;
} SourceFrameData;

//识别处理数据
typedef struct
{
    // 数据来源标识
    int iDataSource;
    // 数据源信息
    DataSourceInfo dataSourceInfo;
    // 原始帧数据
    SourceFrameData sourceFrameData;
    // 解码参数
    AVCodecParameters *pCodecParameters_ = nullptr;
    // 识别结束标识
    bool bIsEnd = false;
} ProcessData;

typedef struct
{
    // 帧数据时间戳
    uint64_t i64TimeStamp = 0;
    // Size of memory, bytes
    uint32_t iSize = 0;
    // Width of image
    uint32_t iWidth = 0;
    // Height of image
    uint32_t iHeight = 0;
    // 帧率
    int iRate = 0;
    // 原始数据
    std::shared_ptr<void> pData = nullptr;
    AVCodecID avCodecID;

} VFrameData;

// 识别区域坐标
typedef struct {
    // 字段代号（仅第二步使用）
    int iLine = -1;
    // 大小框的类别
    int iClassId = -1;
    // 框的名称
    std::string strClassName;
    // 分数
    float fScore = 0;
    // 清晰度
    float fClear = 0;
    // 左上X坐标 Left Top x
    float fLTX = 0;
    // 左上y坐标 Left Top y
    float fLTY = 0;
    // 右下y坐标 Right Bottom x
    float fRBX = 0;
    // 右下y坐标 Right Bottom y
    float fRBY = 0;
} VCoordinate;

typedef struct {
    // 模型地址
    std::string strModelPath;
    // 图片地址
    std::string strImagePath;
    // 模型类型[0:检测 detect 1:分类 classify]
    uint8_t iModel = 0;
    // 模型版本["v5","v8"]
    std::string strModelVersion;
    // 类别模板地址
    std::string strClassModelPath;
} VYoloTestInfo;

//动态检测结果
typedef struct
{
    // 数据来源标识
    int iDataSource;
    // 文件存储目录
    std::string strFilePath;
    // 文件名字
    std::string strFileName;
    // 帧号
    uint32_t iFrameId = 0;
    // 帧数据时间戳
    uint64_t i64TimeStamp = 0;
    // 结束识别标识
    bool bIsEnd = false;
    // 过车日期(格式：YYYY-MM-DD)
    std::string strDetectDate;
    // 过车时间(格式：HH-MM-SS)
    std::string strDetectTime;
    // 图片
    cv::Mat cvImage;
} VDetectInfo;

// 存图数据
typedef struct
{
    // 车厢位置状态
    std::string strTrainStage;
    // 行车方向
    std::string strDirection;
    // 车厢运动状态
    std::string strTrainStatus;
    // 图片
    cv::Mat cvImage;
    // 文件存储目录
    std::string strFilePath;
    // 文件名字
    std::string strFileName;
    // 图片需要标记的框坐标
    std::vector<VCoordinate> vecCoordinate;
    // 火车结束标志
    bool bIsEnd;
} SaveImgData;

// 模型具体结果
typedef struct
{
    // 目标类别 (0:车号; 1:属性; 2:车头; 3:定检期; 4:集装箱; 5:公共间隔; 6:车厢大间隔)
    int iTargetType = -1;
    // 字段代号（仅第二步使用）
    int iLine = -1;
    // 大小框的类别
    int iClassId = -1;
    // 分数
    float fScore = 0;
    // 左上X坐标
    float fLTX = 0;
    // 左上y坐标
    float fLTY = 0;
    // 右下y坐标
    float fRBX = 0;
    // 右下y坐标
    float fRBY = 0;
    // 清晰度
    float fClear = 0;
    // 所属车厢(第二步用)
    int iTrainIndex = 0;
} SingleData;

//后处理数据
typedef struct
{
    // 数据来源标识
    int iDataSource;
    // 帧号
    uint32_t iFrameId = 0;
    // 过车日期(格式：YYYY-MM-DD)
    std::string strDetectDate;
    // 过车时间(格式：HH-MM-SS)
    std::string strDetectTime;
    // 火车结束标志
    bool bIsEnd;
    // 当前大框所属车厢
    int iCarXH = 0;
    // 模型具体结果
    SingleData singleData;
    // 箱角识别结果
    std::vector<SingleData> vecSingleData;
    // 图片
    cv::Mat cvImage;
} InferenceResultData;

// 用于方向计算
typedef struct
{
    // 帧号
    uint32_t iFrameId = 0;
    // 大框id (给第二步用 1-属性 23456-编号)
    int iBigClassId = -1;
    // 第一步识别目标中心点X坐标
    float fCenterX;
    // 第一步识别目标X坐标宽度
    float fTargetWidth;
} CalculateInfo;

//转换具体内容
typedef struct
{
    // 字段代号
    int iLine = -1;
    // 是否符合基本数据校验
    bool bIsChkFlag = false;
    // 转换后的每个字符内容
    std::vector<std::string> vecValue;
    // 每个字符对应的分数
    std::vector<float> vecScore;
    // 本框的识别结果
    std::string strTmpResult;
} TransInfo;

typedef  struct
{
    std::string strDetectDate;
    std::string strDetectTime;
    // 帧序号
    int iFrameId;
    // 图片
    cv::Mat cvImage;
    // 大框集合
    std::vector<SingleData> vecSingleData;
    // 是否最后一节
    bool bIsEnd = false;
} VStep2InputData;

typedef  struct
{
    // 所属集装箱
    int iTrainIndex = 0;
    // 大小框的类别
    int iClassId = -1;
    // 分数
    float fScore = 0;
    // 左上X坐标
    float fLTX = 0;
    // 左上y坐标
    float fLTY = 0;
    // 右下y坐标
    float fRBX = 0;
    // 右下y坐标
    float fRBY = 0;
    // 子框得分总和
    float fSubScoreSum = 0;
    // 子框信息集合
    std::vector<SingleData> vecSingleData;
    // 箱号识别转换结果
    TransInfo transInfo;
} Step2ResultData;

typedef  struct
{
    // 数据来源标识
    int iDataSource;
    std::string strDetectDate;
    std::string strDetectTime;
    // 帧序号
    int iFrameId;
    // 箱号大框以及对应小框的集合
    Step2ResultData step2ResultData;
    // 箱角大框
    std::vector<Step2ResultData> vecCornerResultData;
    // 图片
    cv::Mat cvImage;
    // 是否最后一节
    bool bIsEnd = false;
} VStep2OutputData;


typedef  struct
{
    // 识别日期
    std::string strDetectDate;
    // 识别时间
    std::string strDetectTime;
    // 数据来源标识
    int iDataSource;
    // 帧序号
    int iFrameId;
    // 所属集装箱
    int iContainerIndex = 0;
    // 是否符合基本数据校验
    bool bIsChkFlag = false;
    // 本框的识别编号结果
    std::string strNumResult;
    // 集装箱号得分合
    float fSumScore = 0.0f;
    // 本框的识别箱型结果
    std::string strTypeResult;
    // 大框中心点X坐标
    float fX = 0;
    // 大框中心点y坐标
    float fY = 0;
    // 是否识别到顶部两箱子间隔
    bool bHaveTwoContainer = false;
    // 结束识别标志
    bool bIsEnd = false;
} VSelectBestData;

typedef struct
{
    // 识别日期
    std::string strDetectDate;
    // 识别时间
    std::string strDetectTime;
    // 集装箱结果集合
    std::vector<std::string> vecContainerNO;
    // 图片列表
    std::vector<std::string> vecImage;

} DetectResultData;


// 存图数据
typedef struct
{
    // 数据来源标识
    int iDataSource;
    std::string strDetectDate;
    std::string strDetectTime;
    // 帧序号
    int iFrameId;
    // 箱号大框以及对应小框的集合
    Step2ResultData step2ResultData;
    // 箱角大框
    std::vector<Step2ResultData> vecCornerResultData;
    // 图片
    cv::Mat cvImage;
} SaveDebugImgData;

#endif //TRAIN_COMMONSTRUCT_H