/*!
 *  \file elcorenn.h
 *  \brief Main ElcoreNN functions for models loading and inference.
 *  \copyright 2023-2024 RnD Center ELVEES, JSC. All rights reserved.
 */

#ifndef ELCORENN_H_
#define ELCORENN_H_

#include <cstdint>

/*!
 * Maximum number of tensor dimensions in ElcoreNN framework.
 */
const int MAX_NDIM_VALUE = 6;

/*!
 * Backend ID.
 */
typedef unsigned int ENNBackendId;

/*!
 * Model ID.
 */
typedef unsigned int ENNModelId;

/*!
 * \brief Describes data type
 * should be equal to ElcoreNN data types
 */
enum class ENNDataType {
  FLOAT32 = 0,
  FLOAT16 = 1,
  INT32 = 2,
  UINT32 = 3,
  UINT8 = 4,
  INT8 = 5
};

/*!
 * \brief Heap size allocated by DSP
 */
enum class ENNHeapSize {
  Size_64MB,
  Size_128MB,
  Size_256MB,
  Size_512MB,
  Size_1GB,
  Size_2GB,
  Size_3GB
};

/*!
 * \brief Initializes internal DSPs resources.
 *        DSPs are not statically assigned and
 *        are chosen dynamically by DSP scheduler when executing models.
 *        Execution is paralleled on all available DSPs: each DSP computes
 *        its own data portion divided by batch.
 * \param[in] dsp_heap Heap size allocated by each DSP.
 * \return Backend ID
 */
ENNBackendId InitBackend(ENNHeapSize dsp_heap = ENNHeapSize::Size_512MB);

/*!
 * \brief Initializes internal DSPs resources.
 *        DSPs are not statically assigned and
 *        are chosen dynamically by DSP scheduler when executing models.
 *        Execution is paralleled on device_count DSPs: each DSP computes
 *        its own data portion divided by batch.
 * \param[in] devices_count Number of selected devices.
 * \param[in] dsp_heap Heap size allocated by each DSP.
 * \return Backend ID
 */
ENNBackendId InitBackend(uint32_t devices_count,
                         ENNHeapSize dsp_heap = ENNHeapSize::Size_512MB);

/*!
 * \brief Initializes internal DSPs resources.
 *        Statically assign DSPs ids to execute models.
 * \param[in] devices_count Number of selected devices.
 * \param[in] devices Indices of DSP cores.
 * \param[in] dsp_heap Heap size allocated by each DSP.
 * \return Backend ID
 */
ENNBackendId InitBackend(uint32_t devices_count, uint32_t* devices,
                         ENNHeapSize dsp_heap = ENNHeapSize::Size_512MB);

/*!
 * \brief Release internal backend resources and all backend models.
 * \param[in] backend_id Backend ID
 */
void ReleaseBackend(ENNBackendId backend_id);

/*!
 * \brief Releases internal DSPs resources.
 */
void ReleaseDevice();

/*!
 * \brief Loads model from files.
 * \param[in] model_json The description of the model saved in json format
 * \param[in] model_weights The binary file of model's weights
 * \param[in] optimization Data type for model optimization (default float16)
 * \param[in] backend_id Backend used to load model
 * \return Model ID
 */
ENNModelId LoadModel(const char* model_json, const char* model_weights,
                     ENNDataType optimization = ENNDataType::FLOAT16,
                     ENNBackendId backend_id = 0);

/*!
 * \brief Get number of model inputs.
 * \param[in] model_id Model ID
 * \return Number of model inputs
 */
int GetInputsNumber(ENNModelId model_id);

/*!
 * \brief Get number of model outputs.
 * \param[in] model_id Model ID
 * \return Number of model outputs
 */
int GetOutputsNumber(ENNModelId model_id);

/*!
 * \brief Get input name from a model by input layer index.
 * \param[in] model_id Model ID
 * \param[in] layer_idx Input layer index in the model
 * \return Input layer name specified by index
 */
const char* GetInputName(ENNModelId model_id, uint32_t layer_idx);

/*!
 * \brief Get input shape from a model by input layer index. First position
 * in array is dimensions number of the layer.
 * \param[in] model_id Model ID
 * \param[in] layer_idx Input layer index for the model
 * \param[out] shape Array to save model layer shape (It should be at least
 * [MAX_NDIM_VALUE + 1] values. Arrangement of array values is as follows:
 * First position - number of input tensor dimensions;
 * All the following values - dimensions of each tensor axis)
 */
void GetInputShape(ENNModelId model_id, uint32_t layer_idx, uint32_t* shape);

/*!
 * \brief Get output name from a model by output layer index.
 * \param[in] model_id Model ID
 * \param[in] layer_idx Output layer index in the model
 * \return Output layer name specified by index
 */
const char* GetOutputName(ENNModelId model_id, uint32_t layer_idx);

/*!
 * \brief Get output shape from a model by output layer index. First position
 * in array is dimensions number of the layer.
 * \param[in] model_id Model ID
 * \param[in] layer_idx Output layer index for a model
 * \param[out] shape Array to save model layer shape (It should be at least
 * [MAX_NDIM_VALUE + 1] values. Arrangement of array values is as follows:
 * First position - number of output tensor dimensions;
 * All the following values - dimensions of each tensor axis)
 */
void GetOutputShape(ENNModelId model_id, uint32_t layer_idx, uint32_t* shape);

/*!
 * \brief Runs model inference (input user pointers array, output user pointers
 * array).
 * \param[in] model_id Model ID
 * \param[in] input_data Array of pointers for each input data arrays
 * \param[in] input_data_type Array of data types for each input data arrays
 * \param[out] output_data Array of pointers for each output data arrays
 * \param[in] batch_size Batch size
 */
void InvokeModel(ENNModelId model_id, void** input_data,
                 ENNDataType* input_data_type, float** output_data,
                 uint32_t batch_size);

/*!
 * \brief Runs model inference (input pointers array, output user pointers
 * array).
 * \param[in] model_id Model ID
 * \param[in] input_data Array of pointers for each input data arrays
 * \param[out] output_data Array of pointers for each output data arrays
 * \param[in] batch_size Batch size
 */
void InvokeModel(ENNModelId model_id, float** input_data, float** output_data,
                 uint32_t batch_size);

/*!
 * \brief Runs model inference (input dmabuf, output user pointers array).
 * \param[in] model_id Model ID
 * \param[in] input_dmabuf_fd_array Input data as array of int's dmabuf file
 * descriptors
 * \param[out] output_data Array of pointers for each output data arrays
 * \param[in] batch_size Batch size dimension in input data
 */
void InvokeModel(ENNModelId model_id, int* input_dmabuf_fd_array,
                 float** output_data, uint32_t batch_size);

/*!
 * \brief Runs model inference (input dmabuf, output dmabuf).
 * \param[in] model_id Model ID
 * \param[in] input_dmabuf_fd_array Input data as array of int's dmabuf file
 * descriptors
 * \param[out] output_dmabuf_fd_array Output data as array of int's dmabuf file
 * descriptors
 * \param[in] batch_size Batch size dimension in input data
 */
void InvokeModel(ENNModelId model_id, int* input_dmabuf_fd_array,
                 int* output_dmabuf_fd_array, uint32_t batch_size);

/*!
 * \brief Release model
 * \param[in] model_id Model ID
 */
void ReleaseModel(ENNModelId model_id);

/*!
 * \brief Saves model time statistic to a CSV file.
 * \param[in] model_id Model ID
 * \param[in] file_path Path to CSV file
 */
void SaveModelStatisticToCSV(ENNModelId model_id, const char* file_path);

#endif // ELCORENN_H_
