// Copyright © 2018 - 2019 H2O.AI Inc. All rights reserved.

// WARNING:
//   This API is unstable and under heavy development.
//   Watch https://github.com/h2oai/mojo2/issues/1579 for further information.

#ifndef MOJO_C_API_H_
#define MOJO_C_API_H_

#ifdef __cplusplus
#define EXTERN_C extern "C"
#include <cstdio>

#include "Pipeline.pb.h"
#else
#define EXTERN_C
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#endif

#if defined(_MSC_VER) || defined(_WIN32)
#define MOJO_CAPI_EXPORT EXTERN_C __declspec(dllexport)
#else
#define MOJO_CAPI_EXPORT EXTERN_C
#endif

typedef int8_t mojo_bool;
typedef int32_t mojo_int32;
typedef int64_t mojo_int64;
typedef float mojo_float32;
typedef double mojo_float64;

extern const mojo_bool MOJO_NaN_bool;
extern const mojo_int32 MOJO_NaN_int32;
extern const mojo_int64 MOJO_NaN_int64;
extern const mojo_float32 MOJO_NaN_float32;
extern const mojo_float64 MOJO_NaN_float64;

MOJO_CAPI_EXPORT const char* MOJO_Version();

/// @brief Helper type to unify set of bits defined in MOJO_Transform_Operations
typedef long int MOJO_Transform_Operations_Type;

/// Operations to perform with the pipeline.
typedef enum MOJO_Transform_Operations
#ifdef __cplusplus
    : MOJO_Transform_Operations_Type
#endif
{
  /// standard prediction
  PREDICT = 1 << 0,
  /// prediction interval
  INTERVAL = 1 << 1,
  /// SHAP values
  CONTRIBS_RAW = 1 << 2,
  /// SHAP values mapped to original features
  CONTRIBS_ORIGINAL = 1 << 3,
} MOJO_Transform_Operations;

// mojo data types
typedef enum MOJO_DataType {
  MOJO_UNKNOWN = 0,
  /// [i8] byte-represented boolean, 0=false, 1=true, NA is not defined
  MOJO_BOOL = 1,
  /// [i32] 4 bytes signed integer, NA = INT32_MIN = -2147483648
  MOJO_INT32 = 2,
  /// [i64] 8 bytes signed integer, NA = INT64_MIN = -9223372036854775808LL
  MOJO_INT64 = 3,
  /// [f32] 4 bytes floating point
  MOJO_FLOAT32 = 4,
  /// [f64] 8 bytes floating point
  MOJO_FLOAT64 = 5,
  /// c++ std::string, NA=empty string
  MOJO_STRING = 6,
} MOJO_DataType;

/// @brief mojo frame type
typedef struct MOJO_Frame MOJO_Frame;

///@brief delete the mojo frame and free the memory
MOJO_CAPI_EXPORT void MOJO_DeleteFrame(MOJO_Frame* frame);

///@brief get number of columns in a mojo frame
MOJO_CAPI_EXPORT size_t MOJO_FrameNcol(MOJO_Frame* frame);

///@brief mojo model type
struct MOJO_Model {
#ifdef __cplusplus
  MOJO_Model(MOJO_Transform_Operations_Type supported_ops, const bool is_valid,
             const char* uuid, const char* dai_version, const long time_created,
             const size_t missing_values_count, char** missing_values,
             const size_t feature_count, char** feature_names,
             MOJO_DataType* feature_types, const size_t transformed_names_count,
             char** transformed_names,
             const size_t binomial_problem_labels_count,
             char** binomial_problem_labels,
             const mojo::spec::Problem* problem = nullptr)
      : supported_ops(supported_ops),
        is_valid(is_valid),
        uuid(uuid),
        dai_version(dai_version),
        time_created(time_created),
        missing_values_count(missing_values_count),
        missing_values(missing_values),
        feature_count(feature_count),
        feature_names(feature_names),
        feature_types(feature_types),
        transformed_names_count(transformed_names_count),
        transformed_names(transformed_names),
        binomial_problem_labels_count(binomial_problem_labels_count),
        binomial_problem_labels(binomial_problem_labels),
        problem(problem) {}
#endif
  /// @brief supported_ops - all operations that the model _may_ be able to
  /// support (in call to MOJO_NewPipeline)
  ///                       missing flag means, do not even try
  MOJO_Transform_Operations_Type supported_ops;
  /// @brief is the mojo model is valid
  const bool is_valid;
  /// @brief uuid of the model
  const char* uuid;
  /// @brief dai version
  const char* dai_version;
  /// @brief created time of the model
  const long time_created;
  /// @brief missing values count
  const size_t missing_values_count;
  /// @brief missing values replacement
  char** missing_values;
  /// @brief total features count
  const size_t feature_count;
  /// @brief name of features in the mojo model
  char** feature_names;
  /// @brief type of features in the mojo model
  MOJO_DataType* feature_types;
  /// @brief transformed feature names count
  const size_t transformed_names_count;
  /// @brief transformed feature names
  char** transformed_names;

  /// @brief binomial problem labels count
  const size_t binomial_problem_labels_count;
  /// @brief binomial problem labels
  char** binomial_problem_labels;
/// @brief problem details as a protobuf message
#ifdef __cplusplus
  // TODO: add c support
  const mojo::spec::Problem* problem = nullptr;
#endif
};

///@brief load a new mojo model from mojo file
///@param filename name of the `.mojo` file
///@param tf_lib_prefix only used with TensorFlow, pass "." if not sure
///@return pointer to model instance or NULL if the file is corrupt or doesn't
/// exist
MOJO_CAPI_EXPORT struct MOJO_Model* MOJO_NewModel(const char* filename,
                                                  const char* tf_lib_prefix);

/// @brief Describes model's metadata for defined operations
struct MOJO_Pipeline {
#ifdef __cplusplus
  MOJO_Pipeline(const MOJO_Model* model,
                const MOJO_Transform_Operations_Type operations,
                const size_t output_count, const char** output_names,
                MOJO_DataType* output_types,
                MOJO_Transform_Operations_Type* output_ops)
      : model(model),
        operations(operations),
        output_count(output_count),
        output_names(output_names),
        output_types(output_types),
        output_ops(output_ops) {}
#endif
  /// @brief the mojo model
  const struct MOJO_Model* model;
  /// @brief operations
  const MOJO_Transform_Operations_Type operations;
  /// @brief output columns count
  const size_t output_count;
  /// @brief names of outpus in the mojo model
  const char** output_names;
  /// @brief types of outputs in the mojo model
  MOJO_DataType* output_types;
  /// @brief operations effectively causing presence of column in output
  MOJO_Transform_Operations_Type* output_ops;
};

/// @brief create pipeline for the model and operations
/// @param model
/// @param operations
/// @return new pipeline, or null if model doesn't support one of requested
/// operations
MOJO_CAPI_EXPORT struct MOJO_Pipeline* MOJO_NewPipeline(
    struct MOJO_Model* model, MOJO_Transform_Operations_Type operations);

/// @brief release memory of the pipeline
/// @param pipeline

///@brief Deletes the pipeline. Must be called *before* calling
///`MOJO_DeleteModel`.
MOJO_CAPI_EXPORT void MOJO_DeletePipeline(struct MOJO_Pipeline* pipeline);

/// @brief create empty dataframe for the pipeline
/// @param pipeline
/// @param nrow
/// @return instance of a new frame
MOJO_CAPI_EXPORT MOJO_Frame* MOJO_Pipeline_NewFrame(
    struct MOJO_Pipeline* pipeline, size_t nrow);

/// @brief perform transform operation for the pipeline
/// @param pipeline
/// @param frame
/// @param count if nonzero, indicates the minimal number of records to be
/// transformed;
///              if zero, all rows are transformed. Needed for batching; use 0
///              if not sure.
/// @param debug
MOJO_CAPI_EXPORT void MOJO_Transform(struct MOJO_Pipeline* pipeline,
                                     MOJO_Frame* frame, size_t count,
                                     bool debug);

// delete the mojo model and free the memory
MOJO_CAPI_EXPORT void MOJO_DeleteModel(struct MOJO_Model* model);

/// @brief extract the data from frame's input column
/// @param frame
/// @param index input column index
MOJO_CAPI_EXPORT void* MOJO_Input_Data(const struct MOJO_Pipeline* pipeline,
                                       MOJO_Frame* frame, size_t index);

/// @brief extract the data from frame's output column
/// @param frame
/// @param index output column index
MOJO_CAPI_EXPORT void* MOJO_Output_Data(const struct MOJO_Pipeline* pipeline,
                                        MOJO_Frame* frame, size_t index);

/// @brief Set string value in a column. Unlike with primitive types, STRING
/// cannot be accessed directly.
/// @param data data pointer returned from MOJO_Output_Data
/// @param index index
/// @param value value
MOJO_CAPI_EXPORT void MOJO_Column_Write_Str(void* data, size_t index,
                                            const char* value);

/// @brief Get a string value in a column. Unlike with primitive types, STRING
/// cannot be accessed directly.
/// @param data data pointer returned from MOJO_Output_Data
/// @param index index
/// @return result as const char*
MOJO_CAPI_EXPORT const char* MOJO_Column_Read_Str(void* data, size_t index);

#endif  // MOJO_C_API_H_
