/**
 * \file InferenceHandler.cpp
 * \brief Fichier source, definit le corps des fonctions declarees dans InferenceHandler.h
 * \author Ayaan Toorab
 * \date 17 mai 2021
 */

#include "LibTensorRT.h"

//Util function to reorder the vectors

template <typename T>
std::vector<size_t> TensorRT_Handler::sort_indexes(const std::vector<T> &v) {

  // initialize original index locations
  std::vector<size_t> idx(v.size());
  std::iota(idx.begin(), idx.end(), 0);

  // sort indexes based on comparing values in v
  // using std::stable_sort instead of std::sort
  // to avoid unnecessary index re-orderings
  // when v contains elements of equal values 
  std::stable_sort(idx.begin(), idx.end(),
       [&v](size_t i1, size_t i2) {return v[i1] > v[i2];});

  return idx;
}

// TODO :: necessaire si je preprocess des le redimentionnement ?
bool TensorRT_Handler::processInput(const samplesCommon::BufferManager& buffers)
{

    // mean_.clear();
    // stddev_.clear();
    // mean_ = {0.485F,0.456F,0.406F};
    // stddev_ = {0.229F,0.224F,0.225F};

    float* hostDataBuff = static_cast<float*>(buffers.getHostBuffer("input_1"));

    for (int32_t i = 0, volImg = channel_ * inputH_ * inputW_; i < batchSize_; ++i)
    {
        for (uint32_t j = 0, volChl = inputH_ * inputW_; j < volChl; ++j)
        {
            // The color image to input should be in BGR order --> Les valeurs semblent convenir par rapport au benchmark maintenant trouver l'ordre
            for (int32_t c = 0; c < channel_; ++c){ 
                hostDataBuff[j * channel_ + c] = (static_cast<float>(image_.data[j * channel_ + 2 - c])/255 - mean_[c])/stddev_[c];//2-c before
            }

        }
    } 

    return true;
}

bool TensorRT_Handler::readPreprocessing(const std::string path_preprocessing_file){

    std::ifstream infile(path_preprocessing_file);

    int cnt_read = 0;
    std::string line, num_values, token, delimiter;
    size_t pos = 0;
    delimiter = ",";
    mean_.clear();
    stddev_.clear();

    while (std::getline(infile, line) && cnt_read<3)
    {
        
        if((line.starts_with("MEAN=") && cnt_read == 0) || (line.starts_with("STDDEV=") && cnt_read == 1)){
            num_values = line.substr(line.find("{"),line.find("}"));
            num_values.erase(remove(num_values.begin(), num_values.end(), '{'), num_values.end());
            num_values.erase(remove(num_values.begin(), num_values.end(), '}'), num_values.end());

            while ((pos = num_values.find(delimiter)) != std::string::npos) {
                token = num_values.substr(0, pos);
                if(cnt_read == 0){
                    mean_.push_back(std::stof(token+"f"));
                }else if(cnt_read == 1){
                    stddev_.push_back(std::stof(token+"f"));
                }
                num_values.erase(0, pos + delimiter.length());
            }
            token = num_values.substr(0, pos);
            if(cnt_read == 0){
                mean_.push_back(std::stof(token+"f"));
            }else if(cnt_read == 1){
                stddev_.push_back(std::stof(token+"f"));
            }

            cnt_read++;
            pos = 0;
        }else if(line.starts_with("IMAGEFORMAT=") && cnt_read == 2){

            num_values = line.substr(line.find("{"),line.find("}"));
            num_values.erase(remove(num_values.begin(), num_values.end(), '{'), num_values.end());
            num_values.erase(remove(num_values.begin(), num_values.end(), '}'), num_values.end());
            int cnt_line = 0;

            while ((pos = num_values.find(delimiter)) != std::string::npos) {

                token = num_values.substr(0, pos);

                switch(cnt_line){
                    case 0 :
                        inputW_ = std::stoi(token);
                        break;
                    case 1 :
                        inputH_ = std::stoi(token);
                        break;
                    case 2 :
                        channel_ = std::stoi(token);
                        break;
                    default :
                        break;
                }

                num_values.erase(0, pos + delimiter.length());
            }
            
            token = num_values.substr(0, pos);
            batchSize_ = std::stoi(token);

            cnt_read++;
            pos = 0;

        
        }else{
            mean_.clear();
            stddev_.clear();
            return false;
        }
    }

    return true;
}



// bool TensorRT_Handler::constructNetwork(SampleUniquePtr<nvinfer1::IBuilder>& builder,
//     SampleUniquePtr<nvinfer1::INetworkDefinition>& network, SampleUniquePtr<nvinfer1::IBuilderConfig>& config,
//     SampleUniquePtr<nvonnxparser::IParser>& parser)
// {
//     auto parsed = parser->parseFromFile(model_path_.c_str(),
//         static_cast<int>(sample::gLogger.getReportableSeverity()));
//     if (!parsed)
//     {
//         return false;
//     }

//     config->setFlag(BuilderFlag::kFP16);
//     // TODO --> Handle the case of pruning
//     // config->setFlag(BuilderFlag::kINT8);
//     // samplesCommon::setAllDynamicRanges(network.get(), 127.0f, 127.0f);

//     // Default value : -1
//     samplesCommon::enableDLA(builder.get(), config.get(), -1);

//     return true;
// }

int TensorRT_Handler::LoadCNN(std::string model_path) {

    //temp modif
    int DLACore = -1; 

    model_path_ = model_path;
    std::filesystem::path temp_model_path = model_path_;
    dir_model_path_ = temp_model_path.parent_path().generic_string();

    std::ifstream engineFile(model_path, std::ios::binary);
    if (!engineFile)
    {
        std::cout << "Error opening engine file: " << model_path << std::endl;
        mEngine_ = nullptr;
        return -1;
    }

    engineFile.seekg(0, engineFile.end);
    long int fsize = engineFile.tellg();
    engineFile.seekg(0, engineFile.beg);

    std::vector<char> engineData(fsize);
    engineFile.read(engineData.data(), fsize);
    if (!engineFile)
    {
        std::cout << "Error loading engine file: " << model_path << std::endl;
        mEngine_ = nullptr;
        return -1;
    }

    SampleUniquePtr<IRuntime> runtime{createInferRuntime(sample::gLogger.getTRTLogger())};
    if (DLACore != -1)
    {
        runtime->setDLACore(DLACore);
    }

    mEngine_ = std::shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(engineData.data(), fsize, nullptr));

    //Get preprocessing values for inference time
    std::filesystem::path path_preprocessing(model_path_);
    path_preprocessing = std::filesystem::absolute(path_preprocessing.parent_path().append(PREPROCESSING_FILE_NAME));
    readPreprocessing(path_preprocessing.generic_string());

    return 0;
}

void TensorRT_Handler::InferCNN() {

    // Create RAII buffer manager object
    samplesCommon::BufferManager buffers(mEngine_);

    auto context = SampleUniquePtr<nvinfer1::IExecutionContext>(mEngine_->createExecutionContext());
    if (!context)
    {
        return;
    }

    if (!processInput(buffers))
    {
        return;
    }

    cudaStream_t stream;
    CHECK(cudaStreamCreate(&stream));

    //TODO --> Preprocess with opencv to have the same behavious as in training mode

    // Memcpy from host input buffers to device input buffers
    //buffers.copyInputToDevice();
    buffers.copyInputToDeviceAsync(stream);

    //bool status = context->executeV2(buffers.getDeviceBindings().data());
    bool status = context->enqueueV2(buffers.getDeviceBindings().data(),stream,nullptr);
    if (!status)
    {
        return;
    }

    // Memcpy from device output buffers to host output buffers
    //buffers.copyOutputToHost();
    buffers.copyOutputToHostAsync(stream);

    float* output = static_cast<float*>(buffers.getHostBuffer("dense_2"));

    std::cout << std::endl;
    std::vector<float> probas_temp;

    for (int i = 0; i < nb_label_; i++)
    {
        probas_temp.push_back(output[i]);
    }

    std::cout << "TOP RESULTS GPU :" << std::endl;

    for (auto i: sort_indexes(probas_temp)) {
        std::cout << i << " : ";
        labels_index_.push_back(i);
        std::cout << probas_temp[i] << std::endl;
        probas_.push_back(probas_temp[i]);
        labels_.push_back(w_labels_[i]);
    }

    // Wait for the work in the stream to complete
    CHECK(cudaStreamSynchronize(stream));
    // Release stream
    CHECK(cudaStreamDestroy(stream));


}

void TensorRT_Handler::setImageToInfer(cv::Mat image)
{
    image_ = image;
}

void TensorRT_Handler::setLabels(const char ** w_labels, int nb_label){

    nb_label_ = nb_label;
    w_labels_.clear();
    if(flag_alloc_){
        std::cout << "[INFO] Delete dynamic arrays in setLabels" << std::endl; 
        delete[] res_int_;
        delete[] res_float_;
        flag_alloc_ = false;
    }

    for(int i=0; i<nb_label_;i++){
        w_labels_.push_back(std::string(w_labels[i]));
    }

    //One allocation for the return tabs
    res_int_ = new int [nb_label_];
    res_float_ = new float [nb_label_];
    flag_alloc_ = true;
}

void TensorRT_Handler::emptyReports() {
    probas_.clear();
    labels_.clear();
    labels_index_.clear();
}

void TensorRT_Handler::setImageName(std::string name) {
    image_path_ = name;
}

// std::tuple<std::vector<std::string>, std::vector<float>> InferenceHandler::PredictionCNN(cv::Mat image) {
//     setImageToInfer(image);
//     InferCNN();
//     std::tuple<std::vector<std::string>, std::vector<float>> report = std::make_tuple(labels_, probas_);
//     emptyReports();
//     return report;
// }

void TensorRT_Handler::PredictionCNN_C(cv::Mat image){
    setImageToInfer(image);
    InferCNN();
    std::memcpy(res_int_, &labels_index_.front(), nb_label_ * sizeof(int));
    std::memcpy(res_float_, &probas_.front(), nb_label_ * sizeof(float));
    emptyReports();
}

int * TensorRT_Handler::getLabels(){
    return res_int_;
}

float * TensorRT_Handler::getProbas(){
    return res_float_;
}

float TensorRT_Handler::getProbasIndex(int index){

    if(index < nb_label_){
        return probas_[index];
    }else{
        return probas_[0];
    }

}

void TensorRT_Handler::DestroyDynamicTabs(){
    if(flag_alloc_){
        delete[] res_int_;
        delete[] res_float_;
        flag_alloc_ = false;
    }
}

// void InferenceHandler::DisplayImage(){

//     cv::imshow("defaut",image_);
//     cv::waitKey(0);

// }