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

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

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

bool RKNNManager::initialize(const std::string &model_path)
{
    model_data = load_model(model_path.c_str(), &model_data_size);
    if (!model_data)
        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 < 0)
    {
        std::cerr << "rknn_query error! ret=" << ret << std::endl;
        return false;
    }

    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;
        }
    }

    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 < 0)
            return 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 RKNNManager::infer(int index,const cv::Mat &input_image, cv::Mat &output_image)
{
    rknn_input inputs[1];
    memset(inputs, 0, sizeof(inputs));
    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;

    if (input_image.cols != width || input_image.rows != height)
    {
        resize_buf = malloc(height * width * channel);
        memset(resize_buf, 0x00, height * width * channel);

        rga_buffer_t src = wrapbuffer_virtualaddr((void *)input_image.data, input_image.cols, input_image.rows, RK_FORMAT_RGB_888);
        rga_buffer_t dst = wrapbuffer_virtualaddr((void *)resize_buf, width, height, RK_FORMAT_RGB_888);
        im_rect src_rect, dst_rect;
        IM_STATUS status = imresize(src, dst);
        if (status != IM_STATUS_SUCCESS)
            return false;

        inputs[0].buf = resize_buf;
    }
    else
    {
        inputs[0].buf = (void *)input_image.data;
    }

    int ret = rknn_inputs_set(ctx, io_num.n_input, inputs);
    if (ret < 0)
        return false;

    rknn_output outputs[io_num.n_output];
    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;

    float scale_w = (float)width / input_image.cols;
    float scale_h = (float)height / input_image.rows;

    detect_result_group_t detect_result_group;
    std::vector<float> out_scales;
    std::vector<int32_t> out_zps;
    for (int i = 0; i < io_num.n_output; ++i)
    {
        out_scales.push_back(output_attrs[i].scale);
        out_zps.push_back(output_attrs[i].zp);
    }

    post_process((int8_t *)outputs[0].buf, (int8_t *)outputs[1].buf, (int8_t *)outputs[2].buf, height, width,
                 BOX_THRESH, NMS_THRESH, scale_w, scale_h, out_zps, out_scales, &detect_result_group);

    output_image = input_image.clone();
    char text[256];
    for (int i = 0; i < detect_result_group.count; i++)
    {
        detect_result_t *det_result = &(detect_result_group.results[i]);
        //sprintf(text, "%s %.1f%%", det_result->name, det_result->prop * 100);
        //std::cout << index << " Class: " << det_result->name << ", Confidence: " << det_result->prop * 100 << "%" << std::endl;
        
        // int x1 = det_result->box.left;
        // int y1 = det_result->box.top;
        // int x2 = det_result->box.right;
        // int y2 = det_result->box.bottom;
        // //std::cout << "Box: (" << x1 << ", " << y1 << ") - (" << x2 << ", " << y2 << ")" << std::endl;
        // rectangle(output_image, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(255, 0, 0, 255), 3);
        // putText(output_image, text, cv::Point(x1, y1 + 12), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0));
    }

    rknn_outputs_release(ctx, io_num.n_output, outputs);
    return true;
}

void RKNNManager::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;
}