#include <CL/sycl.hpp>
#include <iostream>
#include <vector>
#include <string>
#include <dirent.h>

#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>

#define BATCHSIZE 2
using namespace cv;
using namespace std::chrono;
using namespace sycl;

// Provide the compiler with information such as input model, target hardware, etc
#pragma tensor_compute model(ResNet) input(resnet50_relay_bs2.mlir) subkind(tc_spirv) target_spec(VASTAI_OAK :.*, NVIDIA_CUDA :.*, HUAWEI_ASCEND :.*, CAMBRICON_MLU :.*, nullptr)

const size_t array_size = 1000;
size_t imageLen = 0;

constexpr cl::sycl::access::mode sycl_read = cl::sycl::access::mode::read;
constexpr cl::sycl::access::mode sycl_write = cl::sycl::access::mode::write;

typedef struct
{
  int index;
  float value;
} sort_st;

class nnet_resnet50
{
public:
    nnet_resnet50():sharp_width(224), sharp_height(224) {}
    ~nnet_resnet50() {}

    void get_file_names(std::string basePath, std::vector<std::string> &files);
    int resnet_tensor_compute(std::vector<std::string> fileNames, size_t count);
    cv::Mat nnet_preprocess(cv::Mat src, int sharp_width, int sharp_height, float *data);
private:
    size_t sharp_width;
    size_t sharp_height;
    uint32_t batch = BATCHSIZE;
};

void nnet_resnet50::get_file_names(std::string basePath, std::vector<std::string> &files)
{
    DIR *dir;
    struct dirent *ptr;

    if ((dir = opendir(basePath.c_str())) == NULL)
    {
        std::cerr << "Open dir error input Image error....\n";
        exit(1);
    }

    while ((ptr = readdir(dir)) != NULL)
    {
        if (strcmp(ptr->d_name, ".") == 0 || strcmp(ptr->d_name, "..") == 0)
        {
            continue;
        }
        else if (ptr->d_type == DT_REG) //DT_REG: This is a regular file.
        {
            std::string a = ptr->d_name;
            int pe = a.find_last_of(".");
	        std::string pic_name = a.substr(pe + 1);
            if (pic_name == "JPEG" || pic_name == "jpg")
            {
                std::string tmpname = basePath + "/" + ptr->d_name;
                files.push_back(tmpname);
            }
        }
        else if (ptr->d_type == DT_DIR) //DT_DIR: This is a directory.
        {
	        std::string base = basePath + "/" + ptr->d_name;
            get_file_names(base, files);
        }
    }
    closedir(dir);
    return;
}

static bool compares(sort_st a, sort_st b)
{
    return a.value > b.value;
}

cl::sycl::queue deviceQueue;

template <typename T, typename TT>
void testApp(T *input, TT *output)
{
    // Define arrays of model inputs and outputs
    cl::sycl::buffer<T, 1> bufferA(input, imageLen * BATCHSIZE, {sycl::property::buffer::use_host_ptr()});
    cl::sycl::buffer<TT, 1> bufferC(output, array_size * BATCHSIZE, {sycl::property::buffer::use_host_ptr()});

    deviceQueue.submit([&](cl::sycl::handler &cgh)
        {
            auto accessorA = bufferA.template get_access<sycl_read>(cgh);
            auto accessorC = bufferC.template get_access<sycl_write>(cgh);

            std::vector<void *> inputs;
            std::vector<void *> outputs;

            inputs.push_back(&accessorA);
            outputs.push_back(&accessorC);
            cgh.tensor_compute_task<class ResNet>(inputs, outputs);
        }).wait();

    host_accessor h_c(bufferC, read_only);
}

int nnet_resnet50::resnet_tensor_compute(std::vector<std::string> fileNames, size_t count)
{
    auto alltime_start = system_clock::now();
    unsigned char *tempD;

    imageLen = this->sharp_width * this->sharp_height * 3;
    cl::sycl::cl_float *imageData = (cl::sycl::cl_float *)malloc_pinned(imageLen * batch * sizeof(cl::sycl::cl_float), deviceQueue);
    cl::sycl::cl_float *result = (cl::sycl::cl_float *)malloc_pinned(array_size * batch * sizeof(cl::sycl::cl_float), deviceQueue);

    for (int i = 0; i < count; i++)
    {
        std::string file_name = fileNames.back();
        fileNames.pop_back();
        std::cout << "\npic-name:" << file_name << std::endl;
        cv::Mat img = cv::imread(file_name.c_str(), cv::IMREAD_COLOR);
        cv::Mat dst;
        {
            float data[sharp_height * sharp_width * 3];
            std::cout << "sharp_height " << sharp_height << " , sharp_width " << sharp_width << std::endl;
            dst = nnet_preprocess(img, this->sharp_width, this->sharp_height, data);
            float *p = data;
            int pos = 0;
            for (int i = 0; i < imageLen * batch; ++i)
            {
                imageData[i] = *(p + pos);
                pos = (pos + 1) % imageLen;
            }

            auto hardprocess_starttime = system_clock::now();
            testApp(imageData, result);
        }

        std::cout << std::endl;
        std::vector<sort_st> result_vector(array_size * batch);
        for (int m = 0; m < batch; m++)
        {
            int index = 0;
            for (int n = m * (array_size); n < (m + 1) * (array_size); n++)
            {
                result_vector[index].value = result[n];
                result_vector[index].index = index;
                index++;
            }

            sort(result_vector.begin(), result_vector.end(), compares);
            std::cout << "Top 5 is: ";
            for (int i = 0; i < 5; ++i)
            {
                std::cout << result_vector[i].index << ": " << result_vector[i].value << " ";
            }
            std::cout << std::endl;
        }
    }

    free(imageData, deviceQueue);
    free(result, deviceQueue);
    duration<double> alltime_end = system_clock::now() - alltime_start;

    std::cout << std::endl << "Finished!\n";
    return 0;
}

// Set the image ROI
cv::Mat nnet_resnet50::nnet_preprocess(cv::Mat src, int sharp_width, int sharp_height, float *data)
{
    cv::Mat src_img;
    cv::cvtColor(src, src, cv::COLOR_BGR2RGB);

    size_t resize = 256;
    size_t resize_w = 0;
    size_t resize_h = 0;

    if (src.cols > src.rows)
    {
        resize_h = resize;
        resize_w = int(resize * float(src.cols) / float(src.rows));
    }
    else
    {
        resize_w = resize;
        resize_h = int(resize * float(src.rows) / float(src.cols));
    }

    size_t crop_x = int((resize_w - sharp_width) / 2);
    size_t crop_y = int((resize_h - sharp_height) / 2);

    cv::resize(src, src_img, cv::Size(resize_w, resize_h));
    Rect roi(crop_x, crop_y, sharp_width, sharp_height);

    std::vector<float> mean_value{0, 0, 0};
    std::vector<float> std_value{1, 1, 1};

    cv::Mat dst;
    uint32_t count = 0;
    src_img(roi).copyTo(dst);
    for (int i = 0; i < sharp_height; i++)
    {
        uchar *uc_pixel = dst.data + i * dst.step;
        for (int j = 0; j < sharp_width; j++)
        {
            data[count] = (uc_pixel[0] / 255. - mean_value[0]) / std_value[0];
            data[count + sharp_height * sharp_width] = (uc_pixel[1] / 255. - mean_value[1]) / std_value[1];
            data[count + 2 * sharp_height * sharp_width] = (uc_pixel[2] / 255. - mean_value[2]) / std_value[2];
            uc_pixel += 3;
            count++;
        }
    }
    return dst;
}

int main(int argc, char **argv)
{
    nnet_resnet50 *resnet50 = new nnet_resnet50();

    std::vector<std::string> fileNames;
    resnet50->get_file_names("/opt/share/pub_share/c_images/dataset/2012img", fileNames);
    size_t count = fileNames.size();
    resnet50->resnet_tensor_compute(fileNames, count);
    return 0;
}