#ifndef __BERT_H__
#define __BERT_H__

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

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;

};

}

#endif