#include "beforeProcess.h"
#include "PPYOLOE.hpp"
#include "RgaUtils.h"
#include "im2d.h"
#include "postprocess.h"
#include <dlfcn.h>
#include <sys/time.h>
#include <exception>
#include <iostream>

PPYOLOE::PPYOLOE() : ctx(0), model_data(nullptr), model_data_size(0), resize_buf(nullptr) {}

PPYOLOE::~PPYOLOE()
{
    release();
}

bool PPYOLOE::initialize(const std::string &model_path)
{
    static int index_flag = 0;
    this->index = index_flag++;

    model_data = load_model(model_path.c_str(), &model_data_size);
    if (!model_data)
    {
        std::cerr << "read model failed!" << std::endl;
        return false;
    }
    int ret = rknn_init(&ctx, model_data, model_data_size, 0, NULL);
    if (ret < 0)
    {
        std::cerr << "rknn_init error! ret=" << ret << std::endl;
        return false;
    }
    // 获取模型输入输出数量
    ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret != RKNN_SUCC)
    {
        std::cerr << "rknn_query error! ret=" << ret << std::endl;
        return false;
    }
    // 获取模型输入属性
    input_attrs = new rknn_tensor_attr[io_num.n_input];
    memset(input_attrs, 0, sizeof(input_attrs));
    for (int i = 0; i < io_num.n_input; i++)
    {
        input_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret < 0)
        {
            std::cerr << "rknn_query error! ret=" << ret << std::endl;
            return false;
        }
    }
    // 获取模型输出属性
    output_attrs = new rknn_tensor_attr[io_num.n_output];
    memset(output_attrs, 0, sizeof(output_attrs));
    for (int i = 0; i < io_num.n_output; i++)
    {
        output_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret < RKNN_SUCC)
        {
            return false;
        }
    }

    if (output_attrs[0].qnt_type == RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC && output_attrs[0].type == RKNN_TENSOR_INT8)
    {
        is_quant = true;
    }
    else
    {
        is_quant = false;
    }

    if (input_attrs[0].fmt == RKNN_TENSOR_NCHW)
    {
        channel = input_attrs[0].dims[1];
        height = input_attrs[0].dims[2];
        width = input_attrs[0].dims[3];
    }
    else
    {
        height = input_attrs[0].dims[1];
        width = input_attrs[0].dims[2];
        channel = input_attrs[0].dims[3];
    }

    return true;
}

/// @brief 用于给输入图像进行推理
/// @param input_image 输入图像的格式为BGR
/// @param output_image 直接将结果绘制到输出图像
/// @return
bool PPYOLOE::infer(int index, unsigned char *input_data, int input_width, int intput_height, DataPackage *result)
{
    rknn_input inputs[io_num.n_input];
    rknn_output outputs[io_num.n_output];
    memset(inputs, 0, sizeof(inputs));
    memset(outputs, 0, sizeof(outputs));

    const float nms_threshold = 0.45;     // 默认的NMS阈值
    const float box_conf_threshold = 0.3; // 默认的置信度阈值

    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].size = width * height * channel;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].pass_through = 0;
    inputs[0].buf = (void *)input_data;

    int ret = rknn_inputs_set(ctx, io_num.n_input, inputs);
    if (ret < 0)
    {
        return false;
    }
    memset(outputs, 0, sizeof(outputs));
    for (int i = 0; i < io_num.n_output; i++)
    {
        outputs[i].want_float = 0;
    }

    ret = rknn_run(ctx, NULL);
    if (ret < 0)
    {
        return false;
    }
    ret = rknn_outputs_get(ctx, io_num.n_output, outputs, NULL);
    if (ret < 0)
    {
        return false;
    }

    post_process(this, outputs, box_conf_threshold, nms_threshold, &result->Result, 1);
    rknn_outputs_release(ctx, io_num.n_output, outputs);

    return true;
}

void PPYOLOE::release()
{
    if (ctx)
    {
        rknn_destroy(ctx);
        ctx = 0;
    }
    if (model_data)
    {
        free(model_data);
        model_data = nullptr;
    }
    if (resize_buf)
    {
        free(resize_buf);
        resize_buf = nullptr;
    }
}

unsigned char *load_model(const char *filename, int *model_size)
{
    FILE *fp;
    unsigned char *data;

    fp = fopen(filename, "rb");
    if (NULL == fp)
    {
        printf("Open file %s failed.\n", filename);
        return NULL;
    }

    fseek(fp, 0, SEEK_END);
    int size = ftell(fp);

    data = load_data(fp, 0, size);

    fclose(fp);

    *model_size = size;
    return data;
}