#include <fstream>
#include <sstream>
#include <migraphx/onnx.hpp>
#include <migraphx/gpu/target.hpp>
#include <migraphx/gpu/hip.hpp>
#include <migraphx/generate.hpp>
#include <migraphx/quantization.hpp>
#include <migraphx/reshape2.hpp>
#include <CommonUtility.h>
#include <Filesystem.h>
#include <SimpleLog.h>
#include <algorithm>
#include <string>
#include <vector>
#include <stdexcept>
#include <GPT2.h>
#include <tokenization.h>

namespace migraphxSamples
{

GPT2::GPT2():logFile(NULL)
{

}

GPT2::~GPT2()
{
    configurationFile.release();
}

ErrorCode GPT2::Initialize(InitializationParameterOfNLP initParamOfNLPGPT2)
{
    // 初始化(获取日志文件,加载配置文件等)
    ErrorCode errorCode=DoCommonInitialization(initParamOfNLPGPT2);
    if(errorCode!=SUCCESS)
    {
        LOG_ERROR(logFile,"fail to DoCommonInitialization\n");
        return errorCode;
    }
    LOG_INFO(logFile,"succeed to DoCommonInitialization\n");

    // 获取配置文件参数
    FileNode netNode = configurationFile["GPT2"];
    std::string modelPath=initializationParameter.parentPath+(std::string)netNode["ModelPath"];

    // 设置最大输入shape
    migraphx::onnx_options onnx_options;
    onnx_options.map_input_dims["input"]={1,1000};

    // 加载模型
    if(Exists(modelPath)==false)
    {
        LOG_ERROR(logFile,"%s not exist!\n",modelPath.c_str());
        return MODEL_NOT_EXIST;
    }
    net = migraphx::parse_onnx(modelPath, onnx_options);        
    LOG_INFO(logFile,"succeed to load model: %s\n",GetFileName(modelPath).c_str());

    // 获取模型输入属性
    std::pair<std::string, migraphx::shape> inputAttribute=*(net.get_parameter_shapes().begin());
    inputName=inputAttribute.first;
    inputShape=inputAttribute.second;

    // 设置模型为GPU模式
    migraphx::target gpuTarget = migraphx::gpu::target{};

    // 编译模型
    migraphx::compile_options options;
    options.device_id=0;                          // 设置GPU设备，默认为0号设备
    options.offload_copy=true;                    // 设置offload_copy
    net.compile(gpuTarget,options);
    LOG_INFO(logFile,"succeed to compile model: %s\n",GetFileName(modelPath).c_str());                       

    return SUCCESS;
}

ErrorCode GPT2::DoCommonInitialization(InitializationParameterOfNLP initParamOfNLPGPT2)
{
    initializationParameter = initParamOfNLPGPT2;

    // 获取日志文件
    logFile=LogManager::GetInstance()->GetLogFile(initializationParameter.logName);

    // 加载配置文件
    std::string configFilePath=initializationParameter.configFilePath;
    if(!Exists(configFilePath))     
    {
        LOG_ERROR(logFile, "no configuration file!\n");
        return CONFIG_FILE_NOT_EXIST;
    }
    if(!configurationFile.open(configFilePath, FileStorage::READ))
    {
        LOG_ERROR(logFile, "fail to open configuration file\n");
        return FAIL_TO_OPEN_CONFIG_FILE;
    }
    LOG_INFO(logFile, "succeed to open configuration file\n");

    // 修改父路径
    std::string &parentPath = initializationParameter.parentPath;
    if (!parentPath.empty())
    {
        if(!IsPathSeparator(parentPath[parentPath.size() - 1]))
        {
           parentPath+=PATH_SEPARATOR;
        }
    }
    return SUCCESS;
}

static bool CompareM(Predictions a, Predictions b)
{
	return a.predictionvalue > b.predictionvalue;
}

long unsigned int GPT2::Inference(const std::vector<long unsigned int> &input_id)
{

    long unsigned int input[1][input_id.size()];
    for (int j=0;j<input_id.size();++j)
    {
        input[0][j] = input_id[j];
    }

    // 设置输入shape
    std::vector<std::vector<std::size_t>> inputShapes;
    inputShapes.push_back({1,input_id.size()});

    // 输入数据
    migraphx::parameter_map inputData;
    inputData[inputName]=migraphx::argument{migraphx::shape(inputShape.type(),inputShapes[0]),(long unsigned int*)input};

    // 推理
    std::vector<migraphx::argument> results = net.eval(inputData);

    // 获取输出节点的属性
    migraphx::argument result = results[0];
    migraphx::shape outputShape = result.get_shape();       // 输出节点的shape
    int numberOfOutput=outputShape.elements();              // 输出节点元素的个数
    float *data = (float *)result.data();                   // 输出节点数据指针

    // 保存推理结果
    long unsigned int n = 0;
    std::vector<Predictions> resultsOfPredictions(22557);
    for(int i=(input_id.size()-1)*22557; i<input_id.size()*22557; ++i)
    {
        resultsOfPredictions[n].index = n;
        resultsOfPredictions[n].predictionvalue = data[i];
        ++n;
    }

    // 对于[UNK]的概率设为无穷小，模型的预测结果不可能是[UNK]
    resultsOfPredictions[100].predictionvalue = -10000;

    // 排序
    std::sort(resultsOfPredictions.begin(), resultsOfPredictions.end(), CompareM);

    return resultsOfPredictions[0].index;
}

ErrorCode GPT2::Preprocessing(cuBERT::FullTokenizer tokenizer,
                             char *question,
                             std::vector<long unsigned int> &input_id)
{
    // 分词操作
    int max_seq_length =1024;
    std::vector<std::string> tokens_question;
    tokens_question.reserve(max_seq_length);
    tokenizer.tokenize(question, &tokens_question, max_seq_length);

    // 保存编码信息
    input_id.push_back(tokenizer.convert_token_to_id("[CLS]"));
    for (int i=0;i<tokens_question.size();++i) 
    {
        input_id.push_back(tokenizer.convert_token_to_id(tokens_question[i]));
    }

    return SUCCESS;
}

}