#include <iostream>
#include <fstream>
#include <cmath>
#include <kpu.h>

/**
 * @file kpu_inference.ino
 * @brief Arduino/KPU run program: loads a model, runs inference on input.bin.
 */

// File paths for KPU model and input binary data
std::string kmodel_path = "test.kmodel";
std::string input_bin   = "input.bin";

void setup() {
  // put your setup code here, to run once:
  std::ifstream ifs_input(input_bin, std::ios::binary);
  ifs_input.seekg(0, ifs_input.end);
  size_t len_input = ifs_input.tellg();
  char* input_bin_data = new char[len_input];
  ifs_input.seekg(0, ifs_input.beg);
  ifs_input.read(input_bin_data, len_input);
  ifs_input.close();

  /** Initialize Kmodel */
  KPU kpu(kmodel_path);
  int input_size  = kpu.get_input_size();   /**< Number of input tensors */
  int output_size = kpu.get_output_size();  /**< Number of output tensors */
  printf("input size: %d\n", input_size);
  printf("output size: %d\n", output_size);

  /** Print input tensor shapes */
  for (int i = 0; i < input_size; i++) {
      printf("input %d shape: ", i);
      dims_t shape_i = kpu.get_input_shape(i);
      for (size_t j = 0; j < shape_i.size(); j++) {
          printf("%zu ", shape_i[j]);
      }
      printf("\n");
  }

  /** Print output tensor shapes */
  for (int i = 0; i < output_size; i++) {
      printf("output %d shape: ", i);
      dims_t shape_i = kpu.get_output_shape(i);
      for (size_t j = 0; j < shape_i.size(); j++) {
          printf("%zu ", shape_i[j]);
      }
      printf("\n");
  }

  /** Get input type and size information */
  typecode_t input_typecode = kpu.get_input_typecode(0);
  printf("input typecode:%d\n", input_typecode);
  dims_t input_shape = kpu.get_input_shape(0);
  size_t input_data_size = kpu.get_input_data_size(0);
  printf("input data size: %d\n", input_data_size);
  size_t input_data_bytes = kpu.get_input_data_bytes(0);
  printf("input data bytes: %d\n", input_data_bytes);

  /** Create runtime tensor and copy input data */
  runtime_tensor input_tensor = host_runtime_tensor::create(
      input_typecode,
      input_shape,
      { (gsl::byte *)input_bin_data, input_data_bytes },
      true,
      hrt::pool_shared
  ).expect("cannot create input tensor");
  hrt::sync(input_tensor, sync_op_t::sync_write_back, true).unwrap();

  /** Set input tensor and run KPU inference */
  kpu.set_input_tensor(0, input_tensor);
  kpu.run();

  /** Get output type and size information */
  typecode_t output_typecode = kpu.get_output_typecode(0);
  printf("output typecode:%d\n", output_typecode);
  dims_t output_shape = kpu.get_output_shape(0);
  size_t output_data_size = kpu.get_output_data_size(0);
  printf("output data size: %d\n", output_data_size);
  size_t output_data_bytes = kpu.get_output_data_bytes(0);
  printf("output data bytes: %d\n", output_data_bytes);

  /** Retrieve KPU output data */
  std::vector<char> output_kmodel_data = kpu.get_output_data(0);
}

void loop() {
  // put your main code here, to run repeatedly:

}
