// Copyright(c) Microsoft Corporation.All rights reserved.
// Licensed under the MIT License.
//

#define VSI_NPU

/* in and out tensors shapes */
#define IN_W    612
#define IN_H    512
#define IN_D      3
#define IN_NB   (IN_W * IN_H * IN_D)
#define IN_SIZE (IN_NB * sizeof(float))
#define IN_FILE "/home/root/data/inputs/inydet_perm.bin"

#define OUT_W 39
#define OUT_H 32
#define OUT_D  5
#define OUT_NB (OUT_W * OUT_H * OUT_D)
#define OUT_SIZE (OUT_NB * sizeof(float))
#define OUT_FILE "outydet.bin"

const char* model_path = "/home/root/data/models/ydet_model_200213_194043-no_sigmoid.onnx";

#include <assert.h>
#include <onnxruntime_c_api.h>
#include <cmath>
#include <stdlib.h>
#include <stdio.h>
#include <time.h>
#include <vector>

#ifdef VSI_NPU
#include "core/providers/vsi_npu/vsi_npu_provider_factory.h"
#endif

const OrtApi* g_ort = OrtGetApiBase()->GetApi(ORT_API_VERSION);

//-----------------------------------------------------------------------------
float time_diff_secs(struct timespec *start, struct timespec *end) {
//-----------------------------------------------------------------------------
    return (end->tv_sec - start->tv_sec) + 1e-9*(end->tv_nsec - start->tv_nsec);
}

//*****************************************************************************
// helper function to check for status
void CheckStatus(OrtStatus* status) {
    if (status != NULL) {
      const char* msg = g_ort->GetErrorMessage(status);
      fprintf(stderr, "%s\n", msg);
      g_ort->ReleaseStatus(status);
      exit(1);
    }
}

int main(int argc, char* argv[]) {
  //*************************************************************************
  // initialize  enviroment...one enviroment per process
  // enviroment maintains thread pools and other state info
  OrtEnv* env;
  CheckStatus(g_ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env));

  // initialize session options if needed
  OrtSessionOptions* session_options;
  CheckStatus(g_ort->CreateSessionOptions(&session_options));
  g_ort->SetIntraOpNumThreads(session_options, 1);

  // Sets graph optimization level
  g_ort->SetSessionGraphOptimizationLevel(session_options, ORT_ENABLE_BASIC);

  // Optionally add more execution providers via session_options
  // E.g. for CUDA include cuda_provider_factory.h and uncomment the following line:
  // OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions, 0);

#ifdef VSI_NPU
  OrtSessionOptionsAppendExecutionProvider_VsiNpu(session_options, 0);
#endif

  //*************************************************************************
  // create session and load model into memory
  // using squeezenet version 1.3
  // URL = https://github.com/onnx/models/tree/master/squeezenet
  OrtSession* session;

  printf("Using Onnxruntime C API\n");
  CheckStatus(g_ort->CreateSession(env, model_path, session_options, &session));

  //*************************************************************************
  // print model input layer (node names, types, shape etc.)
  size_t num_input_nodes;
  OrtStatus* status;
  OrtAllocator* allocator;
  CheckStatus(g_ort->GetAllocatorWithDefaultOptions(&allocator));

  // print number of model input nodes
  status = g_ort->SessionGetInputCount(session, &num_input_nodes);
  std::vector<const char*> input_node_names(num_input_nodes);
  std::vector<int64_t> input_node_dims;

  printf("Number of inputs = %zu\n", num_input_nodes);
  assert(num_input_nodes == 1);

  // iterate over all input nodes
  for (size_t i = 0; i < num_input_nodes; i++) {
    // print input node names
    char* input_name;
    status = g_ort->SessionGetInputName(session, i, allocator, &input_name);
    printf("Input %zu : name=%s\n", i, input_name);
    input_node_names[i] = input_name;

    // print input node types
    OrtTypeInfo* typeinfo;
    status = g_ort->SessionGetInputTypeInfo(session, i, &typeinfo);
    const OrtTensorTypeAndShapeInfo* tensor_info;
	CheckStatus(g_ort->CastTypeInfoToTensorInfo(typeinfo, &tensor_info));
    ONNXTensorElementDataType type;
	CheckStatus(g_ort->GetTensorElementType(tensor_info, &type));
    printf("Input %zu : type=%d\n", i, type);

    // print input shapes/dims
    size_t num_dims;
	CheckStatus(g_ort->GetDimensionsCount(tensor_info, &num_dims));
    printf("Input %zu : num_dims=%zu\n", i, num_dims);
    input_node_dims.resize(num_dims);
	g_ort->GetDimensions(tensor_info, (int64_t*)input_node_dims.data(), num_dims);
    for (size_t j = 0; j < num_dims; j++)
      printf("Input %zu : dim %zu=%jd\n", i, j, input_node_dims[j]);

	g_ort->ReleaseTypeInfo(typeinfo);
  }

  //*************************************************************************
  // Similar operations to get output node information.
  // Use OrtSessionGetOutputCount(), OrtSessionGetOutputName()
  // OrtSessionGetOutputTypeInfo() as shown above.

  //*************************************************************************
  // Score the model using sample data, and inspect values

  size_t input_tensor_size = IN_NB;

  std::vector<float> input_tensor_values(input_tensor_size);
  std::vector<const char*> output_node_names = {"134"};

  // Set input data
  FILE* pFILEin = fopen(IN_FILE, "r");
  if (pFILEin == 0) {
      printf("Failed opening %s\n",IN_FILE);
      return -1;
  }
  fread(input_tensor_values.data(), sizeof(float), IN_NB, pFILEin);
  fclose(pFILEin);

  // create input tensor object from data values
  OrtMemoryInfo* memory_info;
  CheckStatus(g_ort->CreateCpuMemoryInfo(OrtArenaAllocator, OrtMemTypeDefault, &memory_info));
  OrtValue* input_tensor = NULL;
  CheckStatus(g_ort->CreateTensorWithDataAsOrtValue(memory_info, input_tensor_values.data(), IN_SIZE, input_node_dims.data(), 4, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, &input_tensor));
  int is_tensor;
  CheckStatus(g_ort->IsTensor(input_tensor, &is_tensor));
  assert(is_tensor);
  g_ort->ReleaseMemoryInfo(memory_info);

  //--------------------- Inference ---------------------------
  struct timespec end_p, start_p;
  clock_gettime(CLOCK_REALTIME, &start_p);

  OrtValue* output_tensor = NULL;
  CheckStatus(g_ort->Run(session, NULL, input_node_names.data(), (const OrtValue* const*)&input_tensor, 1, output_node_names.data(), 1, &output_tensor));
  clock_gettime(CLOCK_REALTIME, &end_p);
  printf("durée = %f ms\n", time_diff_secs(&start_p, &end_p) * 1000.0);

  CheckStatus(g_ort->IsTensor(output_tensor, &is_tensor));
  assert(is_tensor);

  //--------------------- Get output tensor and write it ---------------------------
  float* floatarr;
  CheckStatus(g_ort->GetTensorMutableData(output_tensor, (void**)&floatarr));
  FILE* pFILEout = fopen(OUT_FILE, "w");
  if (pFILEout == 0) {
      printf("Failed opening %s\n", OUT_FILE);
      return -1;
  }
  fwrite(floatarr, sizeof(float), OUT_NB, pFILEout);
  fclose(pFILEout);

  g_ort->ReleaseValue(output_tensor);
  g_ort->ReleaseValue(input_tensor);
  g_ort->ReleaseSession(session);       // End-up stuck inside this call sometimes with VsiNpu
  g_ort->ReleaseSessionOptions(session_options);
  g_ort->ReleaseEnv(env);
  return 0;
}