#include <iostream>
#include <opencv2/opencv.hpp>
#include <vector>
#include "memx/accl/MxAcclMT.h"
#include <string>
#include <numeric>
#include <fstream>

namespace fs = std::filesystem;

// Imagenet folder
const fs::path imagenetPath = "imagenet100/";
std::string baseFilename = "ILSVRC2012_val_";

// model file 
const fs::path modelPath = "resnet.dfp";

//model info
MX::Types::MxModelInfo model_info;

std::vector<int> ground_truth  = {  65,970,230,809,516,57,334,415,674,332,109,286,370,757,595,147,473,23,478,517,334,173,
                                    948,727,23,846,270,167,55,858,324,573,150,981,586,887,32,398,777,74,516,756,129,198,
                                    256,725,565,167,717,394,92,29,844,591,358,468,259,994,872,588,474,183,107,46,842,390,
                                    101,887,870,841,467,149,21,476,80,424,159,275,175,461,970,160,788,58,479,498,369,28,487,
                                    50,270,383,366,780,373,705,330,142,949,349 };

int totalSamples = 0;
int counter = 100;
int correctTop1 = 0;
int correctTop5 = 0;

cv::Mat preprocessImage(cv::Mat& img) {
    
    cv::resize(img, img, cv::Size(224, 224));

    if (img.channels() == 1) {
        cv::cvtColor(img, img, cv::COLOR_GRAY2BGR);
    }

    cv::Mat img_float;
    img.convertTo(img_float, CV_32F);

    return img_float; 
}

void printScore(const std::vector<float*>& ofmaps){

    totalSamples+=1;
    std::vector<int> indices(1000);

    std::iota(indices.begin(), indices.end(), 0);

    std::sort(indices.begin(), indices.end(),
        [&ofmaps](int i1, int i2) {
            return ofmaps[0][i1] > ofmaps[0][i2];
        }
    );

    int trueIndex = ground_truth[totalSamples-1];

    if (indices[0] == trueIndex) {
            correctTop1++;
    }
    // Check top-5 accuracy
    if (std::find(indices.begin(), indices.begin() + 5, trueIndex) != indices.begin() + 5) {
        correctTop5++;
    }
}

int main(){

    // accl to connect to accelerator
    MX::Runtime::MxAcclMT* accl = new MX::Runtime::MxAcclMT(modelPath);

    int stream_label=0;

    model_info = accl->get_model_info(0);

    std::cout << "Connected stream \n\n\n";

    //  a vector to hold input data pointers and reserve space based on the number of input feature maps
    std::vector<float*> input_data;
    input_data.reserve(model_info.num_in_featuremaps);

    // a vector to hold output feature map data pointers
    std::vector<float*> ofmap;

    // Reserve space in the output vector based on the number of output feature maps
    ofmap.reserve(model_info.num_out_featuremaps);

    // Allocate memory for each output feature map and store pointers in the output vector
    for(int j=0; j<model_info.num_out_featuremaps; ++j){
        float * fmap = new float[model_info.out_featuremap_sizes[j]];
        ofmap.push_back(fmap);
    }

    for(int i=1; i <= counter; i++){
        
        cv::Mat inframe;
        std::stringstream ss;
        ss << baseFilename << std::setw(8) << std::setfill('0') << i << ".JPEG";
        fs::path fullPath = imagenetPath / ss.str();
        inframe = cv::imread(fullPath.string());

        int recvinf_label;

        if(!inframe.empty()){

            cv::Mat preProcframe = preprocessImage(inframe);

            // Add the preprocessed frame data as an input feature map
            input_data.push_back((float*)preProcframe.data);
            
            // Send the input data to the accelerator
            accl->send_input(input_data, model_info.model_index, stream_label, 0);

            // Receive the processed output from the accelerator
            accl->receive_output(ofmap, model_info.model_index, recvinf_label, 0);

            // Evaluate the output and print the scores
            printScore(ofmap);

            // Clear the input data vector for reuse
            input_data.clear();
        }
    }

    for (auto& fmap : ofmap) {
        delete[] fmap;
        fmap = NULL;
    }  
    
    // Print accuracy
    std::cout << "Top-1 Accuracy: " << static_cast<double>(correctTop1) / totalSamples * 100.0 << "%" << std::endl;
    std::cout << "Top-5 Accuracy: " << static_cast<double>(correctTop5) / totalSamples * 100.0 << "%" << std::endl;

    return 1;
}
