#include "include/version.h"
#include "include/model.h"
#include "include/version.h"
#include "include/context.h"
#include "include/errorcode.h"
#include "include/lite_session.h"

#include <thread>
#include <iostream>
#include <string>
#include <vector>
#include <cstdio>
#include <utility>
#include <tuple>

#include "imagenet_label.inc"

// model size limit: 256 MiB
const int MAX_MODEL_SIZE = 256 * 1024 * 1024;

// dataset
std::string IMAGE_FILE = "/data/val_data_c/%05d.bin";

using namespace mindspore;
using namespace mindspore::lite;
using namespace mindspore::session;

void read_image(int idx, void *tensor_buf, size_t size)
{
  char image[128];
  sprintf(image, IMAGE_FILE.c_str(), idx);
  FILE *fp = fopen(image, "rb");
  fread(tensor_buf, sizeof(char), size, fp);
  fclose(fp);
}

void print_tensor(tensor::MSTensor *t)
{
  float *data_ptr = static_cast<float *>(t->MutableData());
  for (int i = 0; i < t->ElementsNum(); ++i)
  {
    std::cout << data_ptr[i] << ", ";
    if (i % 13 == 12)
    {
      std::cout << std::endl;
    }
  }
}

int arg_max(tensor::MSTensor *t)
{
  float *data_ptr = static_cast<float *>(t->MutableData());
  float max_val = 0.f;
  int max_idx = -1;
  for (int i = 0; i < t->ElementsNum(); ++i)
  {
    if (data_ptr[i] > max_val)
    {
      max_idx = i;
      max_val = data_ptr[i];
    }
  }
  return max_idx;
}

std::tuple<void *, void *, tensor::MSTensor *, tensor::MSTensor *> sessionInOut(mindspore::session::LiteSession *session)
{
  // alloc input mem
  std::vector<tensor::MSTensor *> inputs = session->GetInputs();
  tensor::MSTensor *input = inputs.front();
  void *input_buf = input->MutableData();
  //std::cout << "input tenosr num: " << inputs.size() << std::endl;
  //std::cout << "input tensor[0] shape: ";
  //for (int i : input->shape())
  //{
  //  std::cout << i << " ";
  //}
  //std::cout << std::endl;

  // get output
  std::unordered_map<std::string, tensor::MSTensor *> outputs = session->GetOutputs();
  tensor::MSTensor *output = outputs.begin()->second;
  //std::cout << "output tenosr num: " << outputs.size() << std::endl;
  //std::cout << "output tensor[0] name: " << outputs.begin()->first << ", shape: ";
  void *output_buf = output->MutableData();
  ///for (int i : output->shape())
  ///{
  ///  std::cout << i << " ";
  ///}
  ///std::cout << std::endl;
  return {input_buf, output_buf, input, output};
}

mindspore::session::LiteSession *GenerateSession(mindspore::lite::Model *model)
{
  if (model == nullptr)
  {
    std::cerr << "Read model file failed while running" << std::endl;
    return nullptr;
  }
  Context *context = new (std::nothrow) mindspore::lite::Context;
  if (context == nullptr)
  {
    std::cerr << "New context failed while running" << std::endl;
    return nullptr;
  }

  LiteSession *session = mindspore::session::LiteSession::CreateSession(context);
  delete (context);
  if (session == nullptr)
  {
    std::cerr << "CreateSession failed while running" << std::endl;
    return nullptr;
  }
  int ret = session->CompileGraph(model);
  if (ret != mindspore::lite::RET_OK)
  {
    std::cout << "CompileGraph failed while running" << std::endl;
    delete (session);
    return nullptr;
  }

  return session;
}

int main(int argc, const char *argv[])
{
  if (argc != 5)
  {
    std::cout << "usage: ./classification your_model.ms image_num thread_num imagenetval_path" << std::endl;
    return -1;
  }
  IMAGE_FILE = std::string(argv[4]);
  std::string version = mindspore::lite::Version();
  std::cout << "version: " << version << std::endl;

  // load model
  FILE *fp = fopen(argv[1], "rb");
  char *model_buf = new char[MAX_MODEL_SIZE];
  size_t model_size = fread(model_buf, sizeof(char), MAX_MODEL_SIZE, fp);
  fclose(fp);
  std::cout << "model: " << argv[1] << ", size: " << model_size << " Bytes" << std::endl;

  Model *model = Model::Import(model_buf, model_size);

  int THREAD_NUM = std::atoi(argv[3]);

  std::vector<mindspore::session::LiteSession *> sessions(THREAD_NUM);
  std::vector<void *> inputs(THREAD_NUM);
  std::vector<void *> outputs(THREAD_NUM);
  std::vector<tensor::MSTensor *> inputTensors(THREAD_NUM);
  std::vector<tensor::MSTensor *> outputTensors(THREAD_NUM);
  for (int i = 0; i < THREAD_NUM; ++i)
  {
    sessions[i] = GenerateSession(model);
    std::tuple<void *, void *, tensor::MSTensor *, tensor::MSTensor *> inOut = sessionInOut(sessions[i]);
    inputs[i] = std::get<0>(inOut);
    outputs[i] = std::get<1>(inOut);
    inputTensors[i] = std::get<2>(inOut);
    outputTensors[i] = std::get<3>(inOut);
  }

  model->Free();

  // infer
  int IMAGE_NUM = std::atoi(argv[2]);
  std::vector<int> result(IMAGE_NUM);

  std::vector<std::thread> threads(THREAD_NUM);
  for (int id = 0; id < THREAD_NUM; ++id)
  {
    threads[id] = std::thread([&](int id) {
      for (size_t i = id; i < IMAGE_NUM; i += THREAD_NUM)
      {
        read_image(i, inputs[id], inputTensors[id]->Size());
        int ret = sessions[id]->RunGraph();
        if (ret != RET_OK)
        {
          std::cerr << "Run graph failed." << std::endl;
          return RET_ERROR;
        }
        //print_tensor(output);
        //std::cout << arg_max(output) << std::endl;
        result[i] = arg_max(outputTensors[id]);
      }
    },
                              id);
  }

  for (int i = 0; i < THREAD_NUM; ++i)
  {
    threads[i].join();
  }
  std::cout << "inference finished" << std::endl;
  for (int i = 0; i < THREAD_NUM; ++i)
  {
    delete sessions[i];
  }
  int correct = 0;
  for (int i = 0; i < IMAGE_NUM; ++i)
  {
    if (label[i] == result[i] - 1)
    {
      correct++;
    }
  }
  std::cout << "top1 acc: " << correct / (float)IMAGE_NUM << std::endl;
  return 0;
}