#include <cstring>
#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <getopt.h>
#include <sys/stat.h>
#include <opencv2/opencv.hpp>
#include <migraphx/onnx.hpp>
#include <migraphx/gpu/target.hpp>
#include <migraphx/gpu/hip.hpp>
#include <migraphx/generate.hpp>
#include <migraphx/quantization.hpp>

using namespace std;
using namespace cv;
using namespace cv::dnn;

static struct option long_options[] = {
    {"models", required_argument, NULL, 'm'},
    {"input", required_argument, NULL, 'i'},
    {NULL, 0, NULL, 0}
};

struct Classifier
{
    migraphx::program net;
    cv::Size inputSize;
    std::string inputName;
    migraphx::shape inputShape;
};

std::vector<float> ComputeSoftmax(const std::vector<float>& results)
{
    float maxValue=-3.40e+38F;
    for(int i=0;i<results.size();++i)
    {
        if(results[i]>maxValue)
        {
            maxValue=results[i];
        }
    }

    std::vector<float> softmaxResults(results.size());
    float sum=0.0;
    for(int i=0;i<results.size();++i)
    {
        softmaxResults[i]= exp((float)(results[i] - maxValue));
        sum+=softmaxResults[i];
    }
    for(int i=0;i<results.size();++i)
    {
       softmaxResults[i]= softmaxResults[i]/sum;
    }

    return softmaxResults;

}

void InitVit(std::string Model, struct Classifier *classifier)
{
    // parse onnx
    (*classifier).net = migraphx::parse_onnx(Model);

    std::pair<std::string, migraphx::shape> inputAttribute = *((*classifier).net.get_parameter_shapes().begin());
    (*classifier).inputName = inputAttribute.first;
    (*classifier).inputShape = inputAttribute.second;
    (*classifier).inputSize = cv::Size((*classifier).inputShape.lens()[3], (*classifier).inputShape.lens()[2]);

    // compile net
    migraphx::target gpuTarget = migraphx::gpu::target{};
    migraphx::compile_options options;
    options.device_id = 0;
    options.offload_copy = true;
    (*classifier).net.compile(gpuTarget,options);
    fprintf(stdout, "succeed to compile model: %s\n", Model.c_str());

    // run once for warmup
    migraphx::parameter_map inputData;
    inputData[(*classifier).inputName] = migraphx::generate_argument((*classifier).inputShape);
    (*classifier).net.eval(inputData);
}

migraphx::parameter_map preprocess(cv::Mat srcImage, struct Classifier classifier)
{
    cv::Mat inputBlob;
    migraphx::parameter_map inputData;
    cv::dnn::blobFromImage(srcImage, inputBlob, 1/127.5, classifier.inputSize, {127.5, 127.5, 127.5}, true, false);
    inputData[classifier.inputName] = migraphx::argument{classifier.inputShape, (float*)inputBlob.data};
    return inputData;
}

void postprocess(migraphx::argument result, int *n, string inputdir)
{
    const char* labels[] = {"daisy", "dandelion", "roses", "sunflowers", "tulips"};
    migraphx::shape outputShape = result.get_shape();
    float *logits = (float *)result.data();

    std::vector<float> logit;
    for(int j=0; j<outputShape.elements(); ++j)
    {
        logit.push_back(logits[j]);
    }

    std::vector<float> probs = ComputeSoftmax(logit);
    for (int j = 0; j < outputShape.elements(); ++j)
    {
        if (probs[j] >= 0.5)
        {
            fprintf(stdout, "labels: %s, confidence: %.3f\n", labels[j], probs[j]);
            if (inputdir.find(labels[j]))
                (*n) += 1;
        }
    }
}
void postprocess_single(migraphx::argument result, int *n, cv::Mat &srcImage)
{
    const char* labels[] = {"daisy", "dandelion", "roses", "sunflowers", "tulips"};
    migraphx::shape outputShape = result.get_shape();
    float *logits = (float *)result.data();

    std::vector<float> logit;
    for(int j=0; j<outputShape.elements(); ++j)
    {
        logit.push_back(logits[j]);
    }

    std::vector<float> probs = ComputeSoftmax(logit);
    for (int j = 0; j < outputShape.elements(); ++j)
    {
        if (probs[j] >= 0.5)
        {
            char text[20];
            char text1[20];
            fprintf(stdout, "labels: %s, confidence: %.3f\n", labels[j], probs[j]);
            snprintf(text, sizeof(text), "labels: %s", labels[j]);
            snprintf(text1, sizeof(text1), "confidence: %.3f", probs[j]);
            cv::putText(srcImage, text, cv::Point(8, 15), cv::FONT_HERSHEY_PLAIN, 1.0, Scalar(0, 0, 255), 1);
            cv::putText(srcImage, text1, cv::Point(8, 25), cv::FONT_HERSHEY_PLAIN, 1.0, Scalar(0, 0, 255), 1);
            cv::imwrite("result.jpg", srcImage);
        }
    }
}
int main(int argc, char *argv[])
{
    if (argc < 3 || argc > 3)
    {
        fprintf(stdout, "Two args are required: ./a --models=/path_to_model --input=/path_to_imgs\n");
        return -1;
    }
    int opt, index;
    std::string  Model, inputs;
    while ((opt = getopt_long(argc, argv, "m:i:", long_options, NULL)) != -1)
    {
        switch (opt)
        {
            case 'm':
                Model = optarg;
                fprintf(stdout, "Run Model: %s\n", Model.c_str());
                break;
            case 'i':
                inputs = optarg;
                fprintf(stdout, "Image Path: %s\n", inputs.c_str());
                break;
            case '?':
                fprintf(stdout, "argvs is wrong, use: ./a --models=/path_to_model --input=/path_to_imgs\n");
                return 0;
            default:
                return 0;
        }
    }

    struct Classifier classifier;
    InitVit(Model, &classifier);
    fprintf(stdout, "succeed to Init classifier net.\n");

    struct stat s;
    int result = stat(inputs.c_str(), &s);
    if (S_IFDIR & s.st_mode)
    {
        vector<String> srcImages;
        glob(inputs, srcImages, false);
        int n = 0.0; size_t i;
        for (i = 0; i < srcImages.size(); i++)
        {
            fprintf(stdout, "Inference for image[%d]:\n", i);
            cv::Mat srcImage = cv::imread(srcImages[i], 1);
            migraphx::parameter_map inputData = preprocess(srcImage, classifier);

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

            postprocess(results[0], &n, inputs);
        }
        printf("All images:%d, match images:%d, Accuracy: %.3f%\n", i, n, ((float)n/i)*100);
    }
 
    if (S_IFREG & s.st_mode)
    {
        int n = 0;
        cv::Mat srcImage = cv::imread(inputs.c_str(), 1);
        migraphx::parameter_map inputData = preprocess(srcImage, classifier);

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

        postprocess_single(results[0], &n, srcImage);
    }

    return 0;
}