#include <Classifier.h>
#include <Filesystem.h>
#include <SimpleLog.h>
#include <algorithm>
#include <CommonUtility.h>

namespace ortSamples
{

Classifier::Classifier()
{
}

Classifier::~Classifier()
{
    delete dcu_session;
    inputNamesPtr.clear();
    outputNamesPtr.clear();
    configurationFile.release();
}

ErrorCode Classifier::Initialize(InitializationParameterOfClassifier initializationParameterOfClassifier)
{
    // 读取配置文件
    std::string configFilePath=initializationParameterOfClassifier.configFilePath;
    if(Exists(configFilePath)==false)
    {
        LOG_ERROR(stdout, "no configuration file!\n");
        return CONFIG_FILE_NOT_EXIST;
    }
    if(!configurationFile.open(configFilePath, cv::FileStorage::READ))
    {
       LOG_ERROR(stdout, "fail to open configuration file\n");
       return FAIL_TO_OPEN_CONFIG_FILE;
    }
    LOG_INFO(stdout, "succeed to open configuration file\n");

    // 获取配置文件参数
    cv::FileNode netNode = configurationFile["Classifier"];
    std::string modelPath=(std::string)netNode["ModelPath"];

    // 初始化session
    OrtMIGraphXProviderOptions migraphx_options;
    migraphx_options.device_id = 0;
    migraphx_options.migraphx_fp16_enable = 1;
    migraphx_options.migraphx_int8_enable = 0;
    migraphx_options.dynamic_model = 0;
    migraphx_options.migraphx_profile_max_shapes = "";
    migraphx_options.migraphx_load_compiled_model=0;
    migraphx_options.migraphx_save_compiled_model=0;
    sessionOptions.AppendExecutionProvider_MIGraphX(migraphx_options);

    dcu_session = new Ort::Session(env, modelPath.c_str(), sessionOptions);
    return SUCCESS;
}

ErrorCode Classifier::Classify(const std::vector<cv::Mat> &srcImages,std::vector<std::vector<ResultOfPrediction>> &predictions)
{
    if(srcImages.size()==0||srcImages[0].empty()||srcImages[0].depth()!=CV_8U)
    {
        LOG_ERROR(stdout, "image error!\n");
        return IMAGE_ERROR;
    }
    
    // 数据预处理
    std::vector<cv::Mat> image;
    for(int i =0;i<srcImages.size();++i)
    {
        //BGR转换为RGB
        cv::Mat imgRGB;
        cv::cvtColor(srcImages[i], imgRGB, cv::COLOR_BGR2RGB);

        // 调整大小，使短边为256，保持长宽比
        cv::Mat shrink;
        float ratio = (float)256 / min(imgRGB.cols, imgRGB.rows);
        if(imgRGB.rows > imgRGB.cols)
        {
            cv::resize(imgRGB, shrink, cv::Size(256, int(ratio * imgRGB.rows)), 0, 0);
        }
        else
        {
            cv::resize(imgRGB, shrink, cv::Size(int(ratio * imgRGB.cols), 256), 0, 0);
        }

        // 裁剪中心窗口为224*224
        int start_x = shrink.cols/2 - 224/2;
        int start_y = shrink.rows/2 - 224/2;
        cv::Rect rect(start_x, start_y, 224, 224);
        cv::Mat images = shrink(rect);
        image.push_back(images);
    }

    // normalize并转换为NCHW
    cv::Mat inputBlob;
    Image2BlobParams image2BlobParams;
    image2BlobParams.scalefactor=cv::Scalar(1/58.395, 1/57.12, 1/57.375);
    image2BlobParams.mean=cv::Scalar(123.675, 116.28, 103.53);
    image2BlobParams.swapRB=false;
    blobFromImagesWithParams(image,inputBlob,image2BlobParams);

    // 获取模型输入输出信息
    Ort::AllocatorWithDefaultOptions allocator;
    for ( size_t i=0; i<dcu_session->GetInputCount(); i++)
    {
        auto input_name = dcu_session->GetInputNameAllocated(i , allocator);
        inputNamesPtr.push_back(std::move(input_name));
    }
    for ( size_t i=0; i<dcu_session->GetOutputCount(); i++)
    {
        auto out_name = dcu_session->GetOutputNameAllocated(i , allocator);
        outputNamesPtr.push_back(std::move(out_name));
    }
    std::vector<const char *> inputNames = {inputNamesPtr.data()->get()};
    std::vector<const char *> outputNames = {outputNamesPtr.data()->get()};

    float* input_data = (float*)inputBlob.data;
    std::array<float, 3 * 224 * 224> input_data_len{};

    Ort::MemoryInfo memoryInfo =Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);

    std::vector<Ort::Value> inputTensors;
    for(size_t i=0; i<inputNames.size(); i++)
    {
        Ort::TypeInfo inputTypeInfo = dcu_session->GetInputTypeInfo(i);
        auto inputTensorInfo = inputTypeInfo.GetTensorTypeAndShapeInfo();
        std::vector<int64_t> inputDims = inputTensorInfo.GetShape();
        inputTensors.push_back(Ort::Value::CreateTensor<float>(memoryInfo,input_data,input_data_len.size(), inputDims.data(), inputDims.size()));
    }
    
    // 进行推理
    auto output = dcu_session->Run(Ort::RunOptions{nullptr}, inputNames.data(), inputTensors.data(), inputNames.size(), outputNames.data(), outputNames.size());
    
    // 解析输出结果
    const float* pdata = output[0].GetTensorMutableData<float>();
    int numberOfClasses = 1000 ;
    for(int i=0;i<srcImages.size();++i)
    {
        int startIndex=numberOfClasses*i;
        std::vector<float> logit;
        for(int j=0;j<numberOfClasses;++j)
        {
            logit.push_back(pdata[startIndex+j]);
        }     
        std::vector<ResultOfPrediction> resultOfPredictions;
        for(int j=0;j<numberOfClasses;++j)
        {
            ResultOfPrediction prediction;
            prediction.label=j;
            prediction.confidence=logit[j];
            resultOfPredictions.push_back(prediction);
        }
        predictions.push_back(resultOfPredictions);
    }
    return SUCCESS;
}
}