#ifndef __BERT_H__
#define __BERT_H__

#include <tokenization.h>

#include <cstdint>
#include <migraphx/program.hpp>
#include <string>

namespace migraphxSamples {
typedef enum _ErrorCode
{
    SUCCESS = 0,
    MODEL_NOT_EXIST,
    CONFIG_FILE_NOT_EXIST,
    FAIL_TO_LOAD_MODEL,
    FAIL_TO_OPEN_CONFIG_FILE,
} ErrorCode;

typedef struct _Sort_st
{
    int index;
    float value;
} Sort_st;

typedef struct _ResultOfPredictions
{
    int start_index;
    int end_index;
    float start_predictionvalue;
    float end_predictionvalue;
} ResultOfPredictions;

class Bert
{
    public:
    Bert();

    ~Bert();

    ErrorCode Initialize();

    ErrorCode Inference(const std::vector<std::vector<long unsigned int>>& input_ids,
                        const std::vector<std::vector<long unsigned int>>& input_masks,
                        const std::vector<std::vector<long unsigned int>>& segment_ids,
                        std::vector<float>& start_position,
                        std::vector<float>& end_position);

    ErrorCode Preprocessing(cuBERT::FullTokenizer tokenizer,
                            int batch_size,
                            int max_seq_length,
                            const char* text,
                            char* question,
                            std::vector<std::vector<long unsigned int>>& input_ids,
                            std::vector<std::vector<long unsigned int>>& input_masks,
                            std::vector<std::vector<long unsigned int>>& segment_ids);

    ErrorCode Postprocessing(int n_best_size,
                             int max_answer_length,
                             const std::vector<float>& start_position,
                             const std::vector<float>& end_position,
                             std::string& answer);

    private:
    std::vector<std::string> tokens_text;
    std::vector<std::string> tokens_question;

    migraphx::program net;
    std::string inputName1;
    std::string inputName2;
    std::string inputName3;
    std::string inputName4;
    migraphx::shape inputShape1;
    migraphx::shape inputShape2;
    migraphx::shape inputShape3;
    migraphx::shape inputShape4;
};

} // namespace migraphxSamples

#endif