#include "rknn.h"
#include "postprocess.h"
#include "rknn_api.h"
#include <algorithm>
#include <fstream>
#include <cstdio>
#include <stdlib.h>
#include <string.h>
#include "slog.h"
#include "utils.h"

static inline std::string& ltrim(std::string& s)
{
  s.erase(s.begin(), std::find_if(s.begin(), s.end(), 
                                  [](int c) { return !std::isspace(c); }));
  return s;
}

static inline std::string& rtrim(std::string& s)
{
  s.erase(std::find_if(s.rbegin(), s.rend(), 
                       [](int c) { return !std::isspace(c); }).base(), s.end());
  return s;
}

static inline std::string& trim(std::string& s)
{
  return ltrim(rtrim(s));
}



rknnNet::rknnNet()
{
  int ret = 0;
  memset(&src, 0, sizeof(src));
  memset(&dst, 0, sizeof(dst));
  // memset(&src_rect, 0, sizeof(src_rect));
  // memset(&dst_rect, 0, sizeof(dst_rect));
  // this->model_path = model_path;
}

rknnNet::~rknnNet()
{
  // printf("rknnNet::~rknnNet() calling...\n");

  rknn_envs_free();
}

// unsigned char *rknnNet::load_data(FILE* fp, size_t ofst, size_t sz)
// {
//   unsigned char* data;
//   int            ret;

//   data = NULL;

//   if (NULL == fp) {
//     return NULL;
//   }

//   ret = fseek(fp, ofst, SEEK_SET);
//   if (ret != 0) {
//     printf("blob seek failure.\n");
//     return NULL;
//   }

//   data = (unsigned char*)malloc(sz);
//   if (data == NULL) {
//     printf("buffer malloc failure.\n");
//     return NULL;
//   }
//   ret = fread(data, 1, sz, fp);
//   return data;
// }

// unsigned char *rknnNet::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;
// }

void rknnNet::dump_tensor_attr(rknn_tensor_attr *attr)
{
  printf("\t index = %d, name = %s, n_dims = %d, dims = [%d, %d, %d, %d], n_elems = %d, size = %d, fmt = %s, type = %s, qnt_type = %s, zp = %d, scale = %f\n", 
         attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
         attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type),
         get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
}


int rknnNet::rknn_envs_init(const char* model_path, const char *label_path, int n)
{
  int ret = 0;    
  /* 加载rknn文件，创建网络 */
  int model_data_size = 0;
  model_data = load_model_data(model_path, &model_data_size);
  int num = load_labels(label_path);
  ret = rknn_init(&ctx, model_data, model_data_size, 0, NULL);
  if (ret < 0) {
    printf("rknn_init error ret=%d\n", ret);
    return -1;
  }

  #if 0
  // RK3568 no support set core mask
  // set rknn core mask
  rknn_core_mask core_mask;
  if (n == 0) {
    core_mask = RKNN_NPU_CORE_0;
  }
  else if (n == 1) {
    core_mask = RKNN_NPU_CORE_1;
  }
  else {
    core_mask = RKNN_NPU_CORE_2;
  }
  ret = rknn_set_core_mask(ctx, core_mask);
  if (ret < 0) {
    printf("rknn set core mask error ret = %d\n", ret);
    return -1;
  }
  #endif

  // RKNN query SDK version
  rknn_sdk_version version;
  ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &version, sizeof(rknn_sdk_version));
  if (ret < 0) {
    printf("rknn_query RKNN_QUERY_SDK_VERSION error ret=%d\n", ret);
    return -1;
  }

  printf("RKNN SDK version: %s driver version: %s\n", version.api_version, version.drv_version);

  // ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
  // if (ret < 0) {
  //   printf("rknn_init RKNN_QUERY_IN_OUT_NUM error ret=%d\n", ret);
  //   return -1;
  // }

  // RKNN query input/output parameters
  ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(rknn_input_output_num));
  if (ret < 0) {
    printf("rknn query RKNN_QUERY_IN_OUT_NUM error ret = %d\n", ret);
    return -1;
  }
  printf("Model Input num: %d\n", io_num.n_input);

  input_attrs = (rknn_tensor_attr *)malloc(io_num.n_input * sizeof(rknn_tensor_attr));
  memset(input_attrs, 0, io_num.n_input * sizeof(rknn_tensor_attr));
  for (size_t 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) {
      printf("rknn_query RKNN_QUERY_INPUT_ATTR error ret = %d\n", ret);
      return -1;
    }
    dump_tensor_attr(&(input_attrs[i]));
  }

  printf("Model Output num: %d\n", io_num.n_output);
  output_attrs = (rknn_tensor_attr *)malloc(io_num.n_output * sizeof(rknn_tensor_attr));
  memset(output_attrs, 0, io_num.n_output * sizeof(rknn_tensor_attr));
  for (size_t 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) {
      printf("rknn_query RKNN_QUERY_OUTPUT_ATTR error ret = %d\n", ret);
      return -1;
    }
    dump_tensor_attr(&(output_attrs[i]));
  }

  //模型的输入通道数、宽、高  
  model_channel = 3;
  model_width   = 0;
  model_height  = 0;

  if (input_attrs[0].fmt == RKNN_TENSOR_NCHW) {
    printf("Model is NCHW input format!\n");
    model_channel = input_attrs[0].dims[1];
    model_width   = input_attrs[0].dims[2];
    model_height  = input_attrs[0].dims[3];
  } 
  else {
    printf("Model is NHWC input format!\n");
    model_width   = input_attrs[0].dims[1];
    model_height  = input_attrs[0].dims[2];
    model_channel = input_attrs[0].dims[3];
  }
  // printf("Model input channel = %d, width = %d, height = %d\n", model_channel, model_width, model_height);

  return 0;
}



int rknnNet::inference(cv::Mat& orig_img, detect_result_group_t *detect_result_group)
{
  int ret = 0;
  // void* resize_buf = nullptr;
  // rknn_tensor_attr input_attrs[io_num.n_input];  //存放输入参数


  //先对传来的图像进行处理
  cv::Mat img;   //用于NPU推理的图像
  // cv::Mat tImg;  //用于图像的转化
  if (orig_img.empty()) {
    SLOG_INFO(" input image empty!!!\n");
    return -1;
  }
  cv::cvtColor(orig_img, img, cv::COLOR_BGR2RGB);
  // cv::resize(tImg, img, cv::Size(640, 640), 0, 0, cv::INTER_NEAREST);    //模型的输入图像分辨率为[640 × 640], 原始图像需要缩放一次

  int img_width  = img.cols;
  int img_height = img.rows;


  // 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) {
  //     printf("rknn_init error ret=%d\n", ret);
  //     return -1;
  //   }
  // }

  // rknn_tensor_attr output_attrs[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));
  // }

  //模型的输入通道数、宽、高  
  int channel = model_channel;
  int width   = model_width;
  int height  = model_height;

  // if (input_attrs[0].fmt == RKNN_TENSOR_NCHW) {  //输入的通道数、宽、高
  //   channel = input_attrs[0].dims[1];
  //   width   = input_attrs[0].dims[2];
  //   height  = input_attrs[0].dims[3];
  // } else {
  //   width   = input_attrs[0].dims[1];
  //   height  = input_attrs[0].dims[2];
  //   channel = input_attrs[0].dims[3];
  // }

  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;


  // 指定目标大小和预处理方式,默认采用LetterBox预处理
  BOX_RECT pads;
  memset(&pads, 0, sizeof(BOX_RECT));
  cv::Mat resize_image(height, width, CV_8UC3);

  // 计算绽放比例
  float scale_w = (float)width / img_width;
  float scale_h = (float)height / img_height;

  if (img_width != width || img_height != height) {
    if (0) {
      // 直接采用RGA加速缩放
      preprocessing(img, resize_image);
    }
    else {
      // 采用OpenCV LetterBox缩放
      float min_scale = std::min(scale_w, scale_h);
      scale_w = min_scale;
      scale_h = min_scale;
      preprocessing(img, resize_image, pads, min_scale);
    }

    inputs[0].buf = resize_image.data;
  }
  else {
    inputs[0].buf = img.data;
  }

  // gettimeofday(&start_time, NULL);      //开始时间
  size_t start_time = _get_ms();
  rknn_inputs_set(ctx, io_num.n_input, inputs);  //设置NPU的输入，

  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);  //使用模型推理
  ret = rknn_outputs_get(ctx, io_num.n_output, outputs, NULL);  //获得模型结果
  // gettimeofday(&stop_time, NULL);        //结束时间
  // size_t stop_time = _get_ms();
  // std::cout << "Once run sue " << (_get_ms() - start_time) / 1000. << " s." << std::endl;


  // 后处理
  // float scale_w = (float)width / img_width;
  // float scale_h = (float)height / img_height;
  // scale_w = (float)width / orig_img.cols;
  // scale_h = (float)height / orig_img.rows;

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

  //将模型推理的结果存放到detect_result_group
  // post_process((int8_t*)outputs[0].buf, (int8_t*)outputs[1].buf, (int8_t*)outputs[2].buf, height, width,
  //              _conf_threshold, _nms_threshold, pads, scale_w, scale_h, out_zps, out_scales, detect_result_group);  
  // 
  postprocessing((int8_t *)outputs[0].buf, (int8_t *)outputs[1].buf, (int8_t *)outputs[2].buf, 
                 pads, scale_w, scale_h, out_zps, out_scales, detect_result_group);

  ret = rknn_outputs_release(ctx, io_num.n_output, outputs);

  return 0;
}


int rknnNet::preprocessing(const cv::Mat &image, cv::Mat &resize_image) {
  // rga_buffer_t src;
  // rga_buffer_t dst;
  im_rect src_rect;
  im_rect dst_rect;
  memset(&src, 0, sizeof(rga_buffer_t));
  memset(&dst, 0, sizeof(rga_buffer_t));
  memset(&src_rect, 0, sizeof(src_rect));
  memset(&dst_rect, 0, sizeof(dst_rect));

  size_t img_width = image.cols;
  size_t img_height = image.rows;

  // printf("Input image (%dx%d) \n", image.cols, image.rows);
  if (image.type() != CV_8UC3) {
    printf("source image type is %d\n", image.type());
    return -1;
  }

  size_t target_width = model_width;
  size_t target_height = model_height;
  src = wrapbuffer_virtualaddr((void *)image.data, img_width, img_height, RK_FORMAT_RGB_888);
  dst = wrapbuffer_virtualaddr((void *)resize_image.data, target_width, target_height, RK_FORMAT_RGB_888);
  int ret = imcheck(src, dst, src_rect, dst_rect);
  if (IM_STATUS_NOERROR != ret) {
    printf("RGA check error %s !!!\n", imStrError((IM_STATUS)ret));
    return -1;
  }

  IM_STATUS STATUS = imresize(src, dst);

  return 0;
}

int rknnNet::preprocessing(const cv::Mat &image, cv::Mat &resize_image,
                           BOX_RECT &pads, const float scale,
                           const cv::Scalar &pad_color) {
  size_t target_width = model_width;
  size_t target_height = model_height;

  // printf("target width = %d, height = %d\n", target_width, target_height);

  // 调整图像大小
  cv::Mat resized;
  cv::resize(image, resized, cv::Size(), scale, scale);
  // printf("orig image widht = %d, height = %d\n", image.cols, image.rows);
  // printf("resized image widht = %d, height = %d\n", resized.cols, resized.rows);

  // 计算填充大小
  int pad_width = target_width - resized.cols;
  int pad_height = target_height - resized.rows;
  // printf("padding width = %d, height = %d\n", pad_width, pad_height);

  pads.left = pad_width / 2;
  pads.right = pad_width - pads.left;
  pads.top = pad_height / 2;
  pads.bottom = pad_height - pads.top;
  // printf("padding left = %d, right = %d, top = %d, bottom = %d\n", pads.left, pads.right, pads.top, pads.bottom);

  // 在图像周围填充
  cv::copyMakeBorder(resized, resize_image, pads.top, pads.bottom, pads.left, pads.right, cv::BORDER_CONSTANT, pad_color);
  // printf("resized image widht = %d, height = %d\n", resize_image.cols, resize_image.rows);

  return 0;
}

int rknnNet::draw_results(cv::Mat& orig_img, detect_result_group_t *detect_result_group)
{

  char text[256];
  //printf("count: %d\n", detect_result_group.count);
  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);
    printf("name: %s @ size:(%d %d %d %d) %f\n", det_result->name, det_result->box.left, det_result->box.top,
           det_result->box.right, det_result->box.bottom, det_result->prop);
    int x1 = det_result->box.left;
    int y1 = det_result->box.top;
    int x2 = det_result->box.right;
    int y2 = det_result->box.bottom;

    cv::rectangle(orig_img, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(255, 0, 0, 255), 3);
    cv::putText(orig_img, text, cv::Point(x1, y1 + 12), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0));
  }

  cv::imshow("demo", orig_img);
  cv::waitKey(1); 

  return 0;
}

void rknnNet::rknn_envs_free()
{
  // release
  if (ctx)
    rknn_destroy(ctx);

  if (model_data)
    free(model_data);

  if (input_attrs)
    free(input_attrs);
  if (output_attrs)
    free(output_attrs);
}



int rknnNet::postprocessing(int8_t *input0, int8_t *input1, int8_t *input2,
                            BOX_RECT pads, float scale_w, float scale_h,
                            std::vector<int32_t> &qnt_zps,
                            std::vector<float> &qnt_scales,
                            detect_result_group_t *group) {
  const int anchor0[6] = {10, 13, 16, 30, 33, 23};
  const int anchor1[6] = {30, 61, 62, 45, 59, 119};
  const int anchor2[6] = {116, 90, 156, 198, 373, 326};

  int model_in_w = model_width;
  int model_in_h = model_height;
  float conf_threshold = _conf_threshold;
  float nms_threshold = _nms_threshold;
  memset(group, 0, sizeof(detect_result_group_t));

  std::vector<float> filterBoxes;
  std::vector<float> objProbs;
  std::vector<int> classId;

  int class_num = classNames.size();
  // printf("class num: %d\n", class_num);

  // stride 8
  int stride0 = 8;
  int grid_h0 = model_in_h / stride0;
  int grid_w0 = model_in_w / stride0;
  int validCount0 = 0;
  validCount0 = process(input0, (int *)anchor0, grid_h0, grid_w0, class_num,/* model_in_h, model_in_w, */ 
                        stride0, filterBoxes, objProbs, classId,
                        conf_threshold, qnt_zps[0], qnt_scales[0]);
  // printf("validCount0: %d\n", validCount0);

  // stride 16
  int stride1 = 16;
  int grid_h1 = model_in_h / stride1;
  int grid_w1 = model_in_w / stride1;
  int validCount1 = 0;
  validCount1 = process(input1, (int *)anchor1, grid_h1, grid_w1, class_num,/* model_in_h, model_in_w, */ 
                        stride1, filterBoxes, objProbs, classId,
                        conf_threshold, qnt_zps[1], qnt_scales[1]);
  // printf("validCount1: %d\n", validCount1);

  // stride 32
  int stride2 = 32;
  int grid_h2 = model_in_h / stride2;
  int grid_w2 = model_in_w / stride2;
  int validCount2 = 0;
  validCount2 = process(input2, (int *)anchor2, grid_h2, grid_w2, class_num,/* model_in_h, model_in_w, */ 
                        stride2, filterBoxes, objProbs, classId,
                        conf_threshold, qnt_zps[2], qnt_scales[2]);
  // printf("validCount2: %d\n", validCount2);

  int validCount = validCount0 + validCount1 + validCount2;
  // no object detect
  if (validCount <= 0) {
    return 0;
  }

  std::vector<int> indexArray;
  for (int i = 0; i < validCount; ++i) {
    indexArray.push_back(i);
  }

  quick_sort_indice_inverse(objProbs, 0, validCount - 1, indexArray);

  std::set<int> class_set(std::begin(classId), std::end(classId));

  for (auto c : class_set) {
    nms(validCount, filterBoxes, classId, indexArray, c, nms_threshold);
  }

  int last_count = 0;
  group->count = 0;
  // box valid detect targe
  for (int i = 0; i < validCount; ++i) {
    if (indexArray[i] == -1 || last_count >= OBJ_NUMB_MAX_SIZE) {
      continue;
    }
    int n = indexArray[i];

    float x1 = filterBoxes[n * 4 + 0] - pads.left;
    float y1 = filterBoxes[n * 4 + 1] - pads.top;
    float x2 = x1 + filterBoxes[n * 4 + 2];
    float y2 = y1 + filterBoxes[n * 4 + 3];
    int id = classId[n];
    float obj_conf = objProbs[i];

    group->results[last_count].id = id;
    group->results[last_count].box.left = (int)(clamp(x1, 0, model_in_w) / scale_w);
    group->results[last_count].box.top = (int)(clamp(y1, 0, model_in_h) / scale_h);
    group->results[last_count].box.right = (int)(clamp(x2, 0, model_in_w) / scale_w);
    group->results[last_count].box.bottom = (int)(clamp(y2, 0, model_in_h) / scale_h);
    group->results[last_count].prop = obj_conf;
    const char *label = classNames[id].c_str();
    strncpy(group->results[last_count].name, label, OBJ_NAME_MAX_SIZE);

    // printf("result %2d: (%4d, %4d, %4d, %4d), %s\n", i,
    //        group->results[last_count].box.left,
    //        group->results[last_count].box.top,
    //        group->results[last_count].box.right,
    //        group->results[last_count].box.bottom, label);
    last_count++;
  }
  // printf("[%s] detect count: %d\n", __func__, last_count);
  group->count = last_count;

  return 0;
}



unsigned char *rknnNet::load_model_data(const char *file_path, int *size)
{
  // printf("%s model path: %s\n", __func__, file_path);
  FILE *fp = fopen(file_path, "rb");
  if (NULL == fp) {
    printf("open model %s file failed.\n", file_path);
  }

  fseek(fp, 0, SEEK_END);
  int sz = ftell(fp);
  fseek(fp, 0, SEEK_SET);
  // printf("%s model %s size : %d\n", __func__, file_path, sz);

  unsigned char *data = (unsigned char *)malloc(sz);
  if (data == NULL) {
    printf("malloc model data buffer failed.\n");
  }

  int ret = fread(data, 1, sz, fp);
  // printf("%s fread result: %d\n", __func__, ret);
  if (ret < 0) {
    printf("read model data failed.\n");
    free(data);
  }
  // printf("read model size: %d - %d\n", ret, sz);
  fclose(fp);
  
  *size = sz;

  return data;
}

int rknnNet::load_labels(const char *label_path)
{
  // printf("load label file: %s\n", label_path);
  std::ifstream infile(label_path);
  std::string line;

  while (getline(infile, line)) {
    trim(line);
    if (line.size() > 0) {
      // label_info tmp;
      // std::string id = line.substr(0, line.find(':'));
      // std::string name = line.substr(line.find(':')+1);
      // std::cout << "ID: " << id << ", Name: " << name << std::endl;
      // tmp.id = std::atoi(id.c_str());
      // tmp.name = name;
      classNames.push_back(line);
    }
  }

  // printf("labels count: %ld\n", classNames.size());
  for (size_t i = 0; i < classNames.size(); i++) {
    // std::cout << "ID: " << i << ", Name: " << classNames[i] << std::endl;
  }
  infile.close();

  return classNames.size();
}