// Copyright (c) 2023 Horizon Robotics.All Rights Reserved.
//
// The material in this file is confidential and contains trade secrets
// of Horizon Robotics Inc. This is proprietary information owned by
// Horizon Robotics Inc. No part of this work may be disclosed,
// reproduced, copied, transmitted, or used in any way for any purpose,
// without the express written permission of Horizon Robotics Inc.

// This is a simple program that describes how to run resnet50 classification
// on an image and get its top k results by predict score.
// Should be noted: Only resnet50 is supported here.

#include <algorithm>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <iterator>
#include <map>
#include <queue>
#include <utility>

#include "gflags/gflags.h"
#include "hlog/logging.h"
#include "hobot/dnn/hb_dnn.h"
#include "hobot/hb_ucp.h"
#include "hobot/hb_ucp_sys.h"

using namespace std;
#define EMPTY ""

DEFINE_string(model_file, EMPTY, "model file path");
DEFINE_string(feature_file, EMPTY, "Test feature path");
DEFINE_string(coors_file, EMPTY, "Test coors path");
DEFINE_int32(top_k, 5, "Top k classes, 5 by default");

#define SIGMOID(x) (1.0f / (1.0f + exp(-x)))
#define MOUDULE_NAME "DNN_BASIC_SAMPLE"
#define LOGD(err_msg, ...) HFLOGM_D(MOUDULE_NAME, err_msg, ##__VA_ARGS__)
#define LOGI(err_msg, ...) HFLOGM_I(MOUDULE_NAME, err_msg, ##__VA_ARGS__)
#define LOGE(err_msg, ...) HFLOGM_E(MOUDULE_NAME, err_msg, ##__VA_ARGS__)
#define LOGW(err_msg, ...) HFLOGM_W(MOUDULE_NAME, err_msg, ##__VA_ARGS__)
#define LOGE_AND_RETURN_IF_NULL(ptr, code)  \
  {                                         \
    if (ptr == nullptr) {                   \
      cout << #ptr " is null pointer"; \
      return code;                          \
    }                                       \
  }

#define HB_CHECK_SUCCESS(value, errmsg)             \
  do {                                              \
    /*value can be call of function*/               \
    auto ret_code = value;                          \
    if (ret_code != 0) {                            \
      LOGE("{}, error code: {}", errmsg, ret_code); \
      return ret_code;                              \
    }                                               \
  } while (0);


int prepare_tensor(hbDNNTensor *input_tensor, hbDNNTensor *output_tensor, hbDNNHandle_t dnn_handle);

int32_t read_bin_2_tensor(string &feature_file,
                          string &coors_file,
                          hbDNNTensor *input_tensor);

int loadData(const char *file, void **data, unsigned int *length);

struct Tensor {
    vector<vector<vector<vector<float>>>> data;
    vector<int> size;

    Tensor(vector<int> sz) : size(sz) {
        data.resize(sz[0]);
        for (auto& b : data) {
            b.resize(sz[1]);
            for (auto& c : b) {
                c.resize(sz[2]);
                for (auto& h : c) {
                    h.resize(sz[3]);
                }
            }
        }
    }

    void initFromPointer(const float* ptr) {
        int batch = size[0], cat = size[1], height = size[2], width = size[3];
        for (int b = 0; b < batch; ++b) {
            for (int c = 0; c < cat; ++c) {
                for (int h = 0; h < height; ++h) {
                    for (int w = 0; w < width; ++w) {
                        // Calculate the index in the flat array
                        int index = b * (cat * height * width) + c * (height * width) + h * width + w;
                        data[b][c][h][w] = ptr[index];
                    }
                }
            }
        }
    }

    vector<int> getSize() const {
        return size;
    }
};


class CenterPointBBoxCoder {
public:
    CenterPointBBoxCoder(
        const vector<float>& pc_range,
        int out_size_factor,
        const vector<float>& voxel_size,
        int max_num = 100,
        float score_threshold = -1.0f
    ) : pc_range(pc_range), out_size_factor(out_size_factor), voxel_size(voxel_size),
        max_num(max_num), score_threshold(score_threshold) {}

    vector<vector<float>> gatherFeature(const vector<vector<float>>& feat, const vector<vector<int>>& ind) {
        int batch = feat.size();
        int K = ind[0].size();
        vector<vector<float>> gathered(batch, vector<float>(K));

        for (int b = 0; b < batch; ++b) {
            for (int k = 0; k < K; ++k) {
                gathered[b][k] = feat[b][ind[b][k]];
            }
        }

        return gathered;
    }

    void topk(const Tensor& scores, const Tensor& iou, int K,
              vector<vector<float>>& topk_scores,
              vector<vector<int>>& topk_inds,
              vector<vector<int>>& topk_clses,
              vector<vector<float>>& topk_ys,
              vector<vector<float>>& topk_xs) {
        vector<int> sz = scores.getSize();
        int batch = sz[0], cat = sz[1], height = sz[2], width = sz[3];

        topk_scores.resize(batch, vector<float>(K));
        topk_inds.resize(batch, vector<int>(K));
        topk_clses.resize(batch, vector<int>(K));
        topk_ys.resize(batch, vector<float>(K));
        topk_xs.resize(batch, vector<float>(K));

        for (int b = 0; b < batch; ++b) {
            vector<pair<float, int>> score_indices;
            for (int c = 0; c < cat; ++c) {
                for (int h = 0; h < height; ++h) {
                    for (int w = 0; w < width; ++w) {
                        float score = SIGMOID(scores.data[b][c][h][w]);;
                        if (iou.data.size() > 0) {
                            float iou_val = iou.data[b][0][h][w];
                            iou_val = max(0.0f, min(1.0f, iou_val));
                            score *= pow(iou_val, 2);
                            
                        }
                        score_indices.emplace_back(score, c * height * width + h * width + w);
                    }
                }
            }
            sort(score_indices.begin(), score_indices.end(), [](const pair<float, int>& a, const pair<float, int>& b) {
                return a.first > b.first;
            });

            for (int k = 0; k < K; ++k) {
                topk_scores[b][k] = score_indices[k].first;
                topk_inds[b][k] = score_indices[k].second % (height * width);
                topk_clses[b][k] = score_indices[k].second / (height * width);
                topk_ys[b][k] = static_cast<float>(topk_inds[b][k] / width);
                topk_xs[b][k] = static_cast<float>(topk_inds[b][k] % width);
            }
        }
    }

    vector<vector<vector<float>>> transposeAndGatherFeat(const Tensor& feat, const vector<vector<int>>& ind) {
        vector<int> sz = feat.getSize();
        int batch = sz[0], channels = sz[1], height = sz[2], width = sz[3];
        int K = ind[0].size();

        vector<vector<vector<float>>> gathered(batch, vector<vector<float>>(K, vector<float>(channels)));;

        for (int b = 0; b < batch; ++b) {
            for (int k = 0; k < K; ++k) {
                int flat_index = ind[b][k];
                int h = flat_index / width;
                int w = flat_index % width;
                for (int c = 0; c < channels; ++c) {
                    gathered[b][k][c] = feat.data[b][c][h][w];
                }
            }
        }

        return gathered;
    }

    vector<vector<vector<float>>> decode(
        const Tensor& heat,
        const Tensor& rot,
        const Tensor& hei,
        const Tensor& dim,
        const Tensor& iou,
        const Tensor& reg = Tensor({}),
        int task_id = -1
    ) {
        vector<int> sz = heat.getSize();
        int batch = sz[0], cat = sz[1], height = sz[2], width = sz[3];

        vector<vector<float>> topk_scores;
        vector<vector<int>> topk_inds, topk_clses;
        vector<vector<float>> topk_ys, topk_xs;

        topk(heat, iou, max_num, topk_scores, topk_inds, topk_clses, topk_ys, topk_xs);

        vector<vector<vector<float>>> reg_gathered;
        if (reg.data.size() > 0) {
            reg_gathered = transposeAndGatherFeat(reg, topk_inds);
            for (int b = 0; b < batch; ++b) {
                for (int k = 0; k < max_num; ++k) {
                    topk_xs[b][k] += reg_gathered[b][k][0];
                    topk_ys[b][k] += reg_gathered[b][k][1];
                }
            }
        } else {
            for (int b = 0; b < batch; ++b) {
                for (int k = 0; k < max_num; ++k) {
                    topk_xs[b][k] += 0.5f;
                    topk_ys[b][k] += 0.5f;
                }
            }
        }

        vector<vector<vector<float>>> rot_gathered = transposeAndGatherFeat(rot, topk_inds);
        vector<vector<vector<float>>> hei_gathered = transposeAndGatherFeat(hei, topk_inds);
        vector<vector<vector<float>>> dim_gathered = transposeAndGatherFeat(dim, topk_inds);
        vector<vector<vector<float>>> final_box_preds(batch, vector<vector<float>>(max_num, vector<float>(9)));

        for (int b = 0; b < batch; ++b) {
            for (int k = 0; k < max_num; ++k) {
                float rot = atan2(rot_gathered[b][k][0], rot_gathered[b][k][1]);
                final_box_preds[b][k][0] = topk_xs[b][k] * out_size_factor * voxel_size[0] + pc_range[0];
                final_box_preds[b][k][1] = topk_ys[b][k] * out_size_factor * voxel_size[1] + pc_range[1];
                final_box_preds[b][k][2] = hei_gathered[b][k][0];
                final_box_preds[b][k][3] = dim_gathered[b][k][0];
                final_box_preds[b][k][4] = dim_gathered[b][k][1];
                final_box_preds[b][k][5] = dim_gathered[b][k][2];
                final_box_preds[b][k][6] = rot;
                final_box_preds[b][k][7] = topk_scores[b][k];
                final_box_preds[b][k][8] = topk_clses[b][k];

            }
        }

        return final_box_preds;
    }

private:
    vector<float> pc_range;
    int out_size_factor;
    vector<float> voxel_size;
    int max_num;
    float score_threshold;
};


/**
 * Step1: get model handle
 * Step2: prepare input and output tensor
 * Step3: set input data to input tensor
 * Step4: run inference
 * Step5: do postprocess with output data
 * Step6: release resources
 */
int main(int argc, char **argv) {
  // Parsing command line arguments
  gflags::SetUsageMessage(argv[0]);
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  cout << gflags::GetArgv() << endl;
  
  // Init logging
  hobot::hlog::HobotLog::Instance()->SetLogLevel(
      "DNN_BASIC_SAMPLE", hobot::hlog::LogLevel::log_info);

  hbDNNPackedHandle_t packed_dnn_handle;
  hbDNNHandle_t dnn_handle;
  const char **model_name_list;
  auto modelFileName = FLAGS_model_file.c_str();

  int model_count = 0;
  // Step1: get model handle
  {
    HB_CHECK_SUCCESS(
        hbDNNInitializeFromFiles(&packed_dnn_handle, &modelFileName, 1),
        "hbDNNInitializeFromFiles failed");
    HB_CHECK_SUCCESS(hbDNNGetModelNameList(&model_name_list, &model_count,
                                           packed_dnn_handle),
                     "hbDNNGetModelNameList failed");
    HB_CHECK_SUCCESS(
        hbDNNGetModelHandle(&dnn_handle, packed_dnn_handle, model_name_list[0]),
        "hbDNNGetModelHandle failed");
  }

  vector<hbDNNTensor> input_tensors;
  vector<hbDNNTensor> output_tensors;
  int input_count = 0;
  int output_count = 0;
  // Step2: prepare input and output tensor
  {
    HB_CHECK_SUCCESS(hbDNNGetInputCount(&input_count, dnn_handle),
                     "hbDNNGetInputCount failed");
    HB_CHECK_SUCCESS(hbDNNGetOutputCount(&output_count, dnn_handle),
                     "hbDNNGetOutputCount failed");
    input_tensors.resize(input_count);
    output_tensors.resize(output_count);
    prepare_tensor(input_tensors.data(), output_tensors.data(), dnn_handle);
  }
  
  // Step3: set input data to input tensor
  {
    // read a single picture for input_tensor[0], for multi_input model, you
    // should set other input data according to model input properties.
    HB_CHECK_SUCCESS(
        read_bin_2_tensor(FLAGS_feature_file, FLAGS_coors_file, input_tensors.data()),
        "read_bin_2_tensor failed");

    LOGI("read point to tensor success");
  }

  hbUCPTaskHandle_t task_handle{nullptr};
  hbDNNTensor *output = output_tensors.data();
  // Step4: run inference
  {
    // generate task handle
    HB_CHECK_SUCCESS(
        hbDNNInferV2(&task_handle, output, input_tensors.data(), dnn_handle),
        "hbDNNInferV2 failed");

    // submit task
    hbUCPSchedParam ctrl_param;
    HB_UCP_INITIALIZE_SCHED_PARAM(&ctrl_param);
    ctrl_param.backend = HB_UCP_BPU_CORE_ANY;
    HB_CHECK_SUCCESS(hbUCPSubmitTask(task_handle, &ctrl_param),
                     "hbUCPSubmitTask failed");

    // wait task done
    HB_CHECK_SUCCESS(hbUCPWaitTaskDone(task_handle, 0),
                     "hbUCPWaitTaskDone failed");
  }

  // Step5: do postprocess with output data
  for (int i = 0; i < output_count; i++) {
      hbUCPMemFlush(&output_tensors[i].sysMem, HB_SYS_MEM_CACHE_INVALIDATE);
  }
  
  auto quanti_type{output->properties.quantiType};
  if (quanti_type != hbDNNQuantiType::NONE) {
    LOGE("quanti_type is not NONE, and the output needs to be dequantized!");
  }

  auto reg =     reinterpret_cast<float *>(output_tensors[0].sysMem.virAddr),
       height =  reinterpret_cast<float *>(output_tensors[1].sysMem.virAddr),
       dim =     reinterpret_cast<float *>(output_tensors[2].sysMem.virAddr),
       rot =     reinterpret_cast<float *>(output_tensors[3].sysMem.virAddr),
       iou =     reinterpret_cast<float *>(output_tensors[4].sysMem.virAddr),
       heatmap = reinterpret_cast<float *>(output_tensors[5].sysMem.virAddr);


  ofstream ofs;  
  ofs.open("heatmap.txt",ios::out );
  for ( int i=0; i<49152; i++) {    //加了个for循环去将第一列写成全部Q，第二列写成123...
    ofs << heatmap[i] << endl;  //endl用于换行
  } 

  ofs.close();

  vector<float> pc_range = {0, -39.68, -3, 69.12, 39.68, 1};
  int out_size_factor = 4;
  vector<float> voxel_size = {0.135, 0.155, 4};
  int max_num = 100;
  float score_threshold = 0.1f;
  CenterPointBBoxCoder coder(pc_range, out_size_factor, voxel_size, max_num, score_threshold);
  int H = 128, W = 128;
  Tensor _reg({1, 2, H, W});
  Tensor _hei({1, 1, H, W});
  Tensor _dim({1, 3, H, W});
  Tensor _rot({1, 2, H, W});
  Tensor _iou({1, 1, H, W});
  Tensor _heat({1, 3, H, W});
  
  _reg.initFromPointer(reg);
  _hei.initFromPointer(height);
  _dim.initFromPointer(dim);
  _rot.initFromPointer(rot);
  _iou.initFromPointer(iou);
  _heat.initFromPointer(heatmap);
  //   Decode bounding boxes
  auto predictions = coder.decode(_heat, _rot, _hei, _dim, _iou, _reg);
  // Print predictions
  for (int b = 0; b < predictions.size(); ++b) {
      cout << "Batch " << b << ":\n";
      for (int k = 0; k < predictions[b].size(); ++k) {
          cout << "  Box " << k << ": ";
          for (float val : predictions[b][k]) {
              cout << val << " ";
          }
          cout << "\n";
      }
  }
  // Step6: release resources
  {
    // release task handle
    HB_CHECK_SUCCESS(hbUCPReleaseTask(task_handle), "hbUCPReleaseTask failed");
    // free input mem
    for (int i = 0; i < input_count; i++) {
      HB_CHECK_SUCCESS(hbUCPFree(&(input_tensors[i].sysMem)),
                       "hbUCPFree failed");
    }
    // free output mem
    for (int i = 0; i < output_count; i++) {
      HB_CHECK_SUCCESS(hbUCPFree(&(output_tensors[i].sysMem)),
                       "hbUCPFree failed");
    }
    // release model
    HB_CHECK_SUCCESS(hbDNNRelease(packed_dnn_handle), "hbDNNRelease failed");
  }

  return 0;
}

int loadData(const char *file, void **data, unsigned int *length)
{
  fstream dataFile(file, ifstream::in);

  if (!dataFile.is_open())
  {
    cout << "Can't open files: "<< file<<endl;
    return -1;
  }

  //get length of file:
  unsigned int len = 0;
  dataFile.seekg (0, dataFile.end);
  len = dataFile.tellg();
  dataFile.seekg (0, dataFile.beg);

  //allocate memory:
  char *buffer = new char[len];
  if(buffer==NULL) {
    cout << "Can't malloc buffer."<<endl;
    dataFile.close();
    exit(-1);
  }

  //read data as a block:
  dataFile.read(buffer, len);
  dataFile.close();

  *data = (void*)buffer;
  *length = len;
  return 0;  
}

int prepare_tensor(hbDNNTensor *input_tensor, hbDNNTensor *output_tensor,
                   hbDNNHandle_t dnn_handle) {
  int input_count = 0;
  int output_count = 0;
  hbDNNGetInputCount(&input_count, dnn_handle);
  hbDNNGetOutputCount(&output_count, dnn_handle);

  /** Tips:
   * For input memory size:
   * *   input_memSize = input[i].properties.alignedByteSize
   * For output memory size:
   * *   output_memSize = output[i].properties.alignedByteSize
   */
  hbDNNTensor *input = input_tensor;
  for (int i = 0; i < input_count; i++) {
    HB_CHECK_SUCCESS(
        hbDNNGetInputTensorProperties(&input[i].properties, dnn_handle, i),
        "hbDNNGetInputTensorProperties failed");
    int input_memSize = input[i].properties.alignedByteSize;
    HB_CHECK_SUCCESS(hbUCPMallocCached(&input[i].sysMem, input_memSize, 0),
                     "hbUCPMallocCached failed");

    // Show how to get input name
    const char *input_name;
    HB_CHECK_SUCCESS(hbDNNGetInputName(&input_name, dnn_handle, i),
                     "hbDNNGetInputName failed");
    LOGI("input[{}] name is {}", i, input_name);
  }

  hbDNNTensor *output = output_tensor;
  for (int i = 0; i < output_count; i++) {
    HB_CHECK_SUCCESS(
        hbDNNGetOutputTensorProperties(&output[i].properties, dnn_handle, i),
        "hbDNNGetOutputTensorProperties failed");
    int output_memSize = output[i].properties.alignedByteSize;
    HB_CHECK_SUCCESS(hbUCPMallocCached(&output[i].sysMem, output_memSize, 0),
                     "hbUCPMallocCached failed");

    // Show how to get output name
    const char *output_name;
    HB_CHECK_SUCCESS(hbDNNGetOutputName(&output_name, dnn_handle, i),
                     "hbDNNGetOutputName failed");
    LOGI("output[{}] name is {}", i, output_name);
  }
  return 0;
}

int32_t read_bin_2_tensor(string &feature_file,
                          string &coors_file,
                          hbDNNTensor *input_tensor) {

  unsigned int f_length = 0, c_length = 0;
  void *feature = NULL, *coors = NULL;
  shared_ptr<float> buffer_f((float *)coors, default_delete<float[]>());
  loadData(feature_file.data(), &feature, &f_length);
  buffer_f.reset((float *)feature);
  float* p_feature = (float*)buffer_f.get();

  if (f_length==0) {
    LOGE("feature file not exist!");
    return -1;
  }

  shared_ptr<int32_t> buffer_c((int32_t *)coors, default_delete<int32_t[]>());
  loadData(coors_file.data(), &coors, &c_length);
  buffer_c.reset((int32_t *)coors);
  int32_t* p_coors = (int32_t*)buffer_c.get();

  if (c_length==0) {
    LOGE("coors file not exist!");
    return -1;
  }

  // hbDNNTensor *input = input_tensor;
  auto &properties = input_tensor->properties;  
  float *f_data{reinterpret_cast<float *>(input_tensor[0].sysMem.virAddr)};
  int32_t *c_data{reinterpret_cast<int32_t *>(input_tensor[1].sysMem.virAddr)};

  memcpy(f_data, feature, f_length / sizeof(float));
  memcpy(c_data, coors, c_length / sizeof(int32_t));
  hbUCPMemFlush(&input_tensor->sysMem, HB_SYS_MEM_CACHE_CLEAN);
  return 0;
}