// Person_detection quant model
// MlModel struct initialization to include model I/O info.
// Bytecode loading, input/output processes.

#include <springbok.h>

#include "iree/base/api.h"
#include "iree/hal/api.h"
#include "samples/util/util.h"
#include "person_detection.h"

// Compiled module embedded here to avoid file IO:
#include "samples/quant_model/person_detection_bytecode_module_dylib_c.h"
#include "samples/quant_model/person_detection_quant_input_c.h"

const MlModel kModel = {
    .num_input = 1,
    .num_input_dim = {4},
    .input_shape = {{1, 96, 96, 1}},
    .input_length = {96 * 96 * 1},
    .input_size_bytes = {sizeof(int8_t)},
    .num_output = 1,
    .output_length = {2},
    .output_size_bytes = sizeof(int8_t),
    .hal_element_type = IREE_HAL_ELEMENT_TYPE_SINT_8,
    .entry_func = "module.main",
    .model_name = "person_detection_quant",
};

PersonDetectionOutput detection;

const iree_const_byte_span_t load_bytecode_module_data() {
  const struct iree_file_toc_t *module_file_toc =
      samples_quant_model_person_detection_bytecode_module_dylib_create();
  return iree_make_const_byte_span(module_file_toc->data,
                                   module_file_toc->size);
}

iree_status_t load_input_data(const MlModel *model, void **buffer) {
  iree_status_t result = alloc_input_buffer(model, buffer);
  const struct iree_file_toc_t *input_file_toc =
      person_detection_quant_input_c_create();
  memcpy(*buffer, input_file_toc->data, input_file_toc->size);
  return result;
}

iree_status_t process_output(const MlModel *model,
                              iree_hal_buffer_mapping_t *buffers,
                              MlOutput *output) {
  iree_status_t result = iree_ok_status();
  int8_t *data = (int8_t *)buffers[0].contents.data;
  detection.non_person_score = data[0];
  detection.person_score = data[1];

  LOG_INFO("Output: Non-person Score: %d; Person Score: %d",
            detection.non_person_score,
            detection.person_score);
  output->result = &detection;
  output->len = sizeof(detection);

  return result;
}