/**
  ******************************************************************************
  * @file    stai_main_app.c
  * @author  AST Embedded Analytics Research Platform
  * @date    2024-05-23T10:34:45+0200
  ******************************************************************************
  * Copyright (c) 2024 STMicroelectronics.
  * All rights reserved.
  *
  * This software is licensed under terms that can be found in the LICENSE file
  * in the root directory of this software component.
  * If no LICENSE file comes with this software, it is provided AS-IS.
  ******************************************************************************
  */
#include <inttypes.h>
#include <stdio.h>
#include <string.h>

#include "stai.h"             /* include ST Edge AI macros, types and data structures */
#include "network_inputs.h"   /* where input1 buffer is defined */
#include "network.h"          /* include ST Edge AI generated network model */


#define LOG_PRINT(fmt, ...) \
  { printf(fmt, ##__VA_ARGS__); fflush(stdout); }


int main(int argc, char* argv[])
{
  stai_return_code return_code = STAI_SUCCESS;

  LOG_PRINT("name:%s\n", STAI_NETWORK_MODEL_NAME)

  LOG_PRINT("n_inputs:%d\n", STAI_NETWORK_IN_NUM)
  LOG_PRINT("n_outputs:%d\n", STAI_NETWORK_OUT_NUM)

  LOG_PRINT("activations:%d\n", STAI_NETWORK_ACTIVATIONS_SIZE_BYTES)
  LOG_PRINT("weights:%d\n", STAI_NETWORK_WEIGHTS_SIZE_BYTES)
  LOG_PRINT("runtime_name:STM.AI\n")

  /*  !New!: Declare and allocate memory for private network context instance */
  STAI_NETWORK_CONTEXT_DECLARE(network, STAI_NETWORK_CONTEXT_SIZE)

  /* !New!: Runtime initialization */
  return_code = stai_runtime_init();

  /*  Initialize network context  */
  return_code = stai_network_init(network);
  if (return_code != STAI_SUCCESS) {
    LOG_PRINT("  ## Test Failed executing stai init: 0x%x.\n\n", return_code)
    return -1;
  }

  /*  Declare activations buffer pointers array  */
  stai_ptr activation_buffers[STAI_NETWORK_ACTIVATIONS_NUM] = {0};

  /*  Allocate and set activation buffer #1  */
  STAI_ALIGNED(STAI_NETWORK_ACTIVATION_1_ALIGNMENT)
  uint8_t activation1[STAI_NETWORK_ACTIVATION_1_SIZE] = {0};
  activation_buffers[0] = (stai_ptr)(activation1);

  /*  Allocate and set activation buffer #2  */
  STAI_ALIGNED(STAI_NETWORK_ACTIVATION_2_ALIGNMENT)
  uint8_t activation2[STAI_NETWORK_ACTIVATION_2_SIZE] = {0};
  activation_buffers[1] = (stai_ptr)(activation2);

  /*  !New!: Set network activations buffers  */
  return_code = stai_network_set_activations(network, activation_buffers, STAI_NETWORK_ACTIVATIONS_NUM);
  if (return_code != STAI_SUCCESS) {
    LOG_PRINT("  ## Test Failed executing stai set activations: 0x%x.\n\n", return_code)
    return -1;
  }

#if 0
  /* NOTE: This step is no more required now in ST Edge AI Client APIs since weights buffers are generated and bind 
     directly to the C model */
  /*  !New!: Declare activations buffer pointers array  */
  stai_ptr weight_buffers[STAI_NETWORK_WEIGHTS_NUM] = {0};
  stai_size n_weights = 0;
  return_code = stai_network_get_weights(network, weight_buffers, &n_weights);
  if ((return_code == STAI_SUCCESS) && (n_weights==STAI_NETWORK_WEIGHTS_NUM)) {
    return_code = stai_network_set_weights(network, weight_buffers, n_weights);
  } else {
    LOG_PRINT("  ## Test Failed executing stai set weights: 0x%x.\n\n", return_code)
    return -1;
  }
#endif

  /* Inputs Buffer Setup */
  /* C-Table declaring inputs buffer pointers array and set inputs addresses */
  stai_ptr input_buffers[STAI_NETWORK_IN_NUM] = {
    (stai_ptr)input1 /* defined in network_inputs.h */
  };

  /*  !New!: Set network inputs buffers  */
  return_code = stai_network_set_inputs(network, input_buffers, STAI_NETWORK_IN_NUM);
  if (return_code != STAI_SUCCESS) {
    LOG_PRINT("  ## Test Failed executing stai set inputs: 0x%x.\n\n", return_code)
    return -1;
  }

  /* Outputs Buffers Setup */
  /*  Allocate and declare output buffer #1  */
  STAI_ALIGNED(STAI_NETWORK_OUT_1_ALIGNMENT)
  float output1[STAI_NETWORK_OUT_1_SIZE];
  
  /*  Declare outputs buffer pointers array and set outputs addresses */
  stai_ptr output_buffers[STAI_NETWORK_OUT_NUM] = {
    (stai_ptr)output1
  };

  /*  Set network outputs buffers  */
  return_code = stai_network_set_outputs(network, output_buffers, STAI_NETWORK_OUT_NUM);
  if (return_code != STAI_SUCCESS) {
    LOG_PRINT("  ## Test Failed executing stai set outputs: 0x%x.\n\n", return_code)
    return -1;
  }

  stai_network_info info;
  return_code = stai_network_get_info(network, &info);
  if (return_code != STAI_SUCCESS) {
    LOG_PRINT("  ## Test Failed executing stai get network info: 0x%x.\n\n", return_code)
    return -1;
  }

  LOG_PRINT("* Runtime version   : %d.%d.%d\n",
    info.runtime_version.major, info.runtime_version.minor, info.runtime_version.micro)
  LOG_PRINT("* Tool version      : %d.%d.%d\n",
    info.tool_version.major, info.tool_version.minor, info.tool_version.micro)
  LOG_PRINT("* APIs version      : %d.%d.%d\n",
    info.api_version.major, info.api_version.minor, info.api_version.micro)
  LOG_PRINT("* Network nodes     : %d\n", info.n_nodes)
  LOG_PRINT("* Network macc      : %d\n", info.n_macc)
  LOG_PRINT("* Network inputs    : %d\n", info.n_inputs)
  LOG_PRINT("* Network outputs   : %d\n", info.n_outputs)

  /*  Execute network model inference on sample test (synchronous mode)  */
  LOG_PRINT("Starting inference\n")

  /* The run API now supports both sync and async modes */
  return_code = stai_network_run(network, STAI_MODE_SYNC);

  LOG_PRINT("Completed inference\n")
  if (return_code != STAI_SUCCESS) {
    LOG_PRINT("  ## Test Failed executing stai network run: 0x%x.\n\n", return_code)
    return -1;
  }

  LOG_PRINT("__START_OUTPUT1 __\n")
  for(int32_t o = 0; o < STAI_NETWORK_OUT_1_SIZE; o++) {
    const float value = ((float*)output1)[o];
    if (o != 0 && o % 10 == 0) {
      LOG_PRINT("\n")
    }
    LOG_PRINT("%f " ", value);
  }
  LOG_PRINT("\n__END_OUTPUT1 __\n")

  /*  Network de-initialization  */
  return_code = stai_network_deinit(network);

  /* Runtime de-initialization */
  return_code = stai_runtime_deinit();

  return (return_code == STAI_SUCCESS) ? 0 : -1;
}