// mnist float 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"

// Compiled module embedded here to avoid file IO:
#include "samples/float_model_embedding/mnist_bytecode_module_dylib_c.h"
#include "samples/float_model_embedding/mnist_input_c.h"

const MlModel kModel = {
    .num_input = 1,
    .num_input_dim = {4},
    .input_shape = {{1, 28, 28, 1}},
    .input_length = {28 * 28 * 1},
    .input_size_bytes = {sizeof(float)},
    .num_output = 1,
    .output_length = {10},
    .output_size_bytes = sizeof(float),
    .hal_element_type = IREE_HAL_ELEMENT_TYPE_FLOAT_32,
    .entry_func = "module.predict",
    .model_name = "mnist",
};

const iree_const_byte_span_t load_bytecode_module_data() {
  const struct iree_file_toc_t *module_file_toc =
      samples_float_model_embedding_mnist_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 = mnist_input_c_create();
  memcpy(*buffer, input_file_toc->data, input_file_toc->size);
  return result;
}

iree_status_t check_output_data(const MlModel *model,
                                iree_hal_buffer_mapping_t *mapped_memory,
                                int index_output) {
  iree_status_t result = iree_ok_status();
  if (index_output > model->num_output ||
      mapped_memory->contents.data_length / model->output_size_bytes !=
          model->output_length[index_output]) {
    result = iree_make_status(IREE_STATUS_UNKNOWN, "output length mismatches");
  }
  LOG_INFO("Output #%d data length: %d", index_output,
           mapped_memory->contents.data_length / model->output_size_bytes);
  // find the label index with best prediction
  float best_out = 0.0;
  int best_idx = -1;
  for (int i = 0; i < model->output_length[index_output]; ++i) {
    float out = ((float *)mapped_memory->contents.data)[i];
    if (out > best_out) {
      best_out = out;
      best_idx = i;
    }
  }
  LOG_INFO("Digit recognition result is: digit: %d", best_idx);
  return result;
}