// Copyright (C) 2018-2019 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#include <iomanip>
#include <vector>
#include <memory>
#include <string>
#include <cstdlib>

#include <opencv2/opencv.hpp>
#include <inference_engine.hpp>
#include <samples/classification_results.h>

using namespace InferenceEngine;

int main(int argc, char *argv[]) {
    try {
        // ------------------------------ Parsing and validation of input args ---------------------------------
        if (argc != 4) {
            std::cout << "Usage : ./hello_request_classification <path_to_model> <path_to_image> <device_name>"
                      << std::endl;
            return EXIT_FAILURE;
        }

        const std::string input_model{argv[1]};
        const std::string input_image_path{argv[2]};
        const std::string device_name{argv[3]};
        // -----------------------------------------------------------------------------------------------------

        // --------------------------- 1. Load Plugin for inference engine -------------------------------------
        InferencePlugin plugin = PluginDispatcher().getPluginByDevice(device_name);
        // -----------------------------------------------------------------------------------------------------

        // --------------------------- 2. Read IR Generated by ModelOptimizer (.xml and .bin files) ------------
        CNNNetReader network_reader;
        network_reader.ReadNetwork(input_model);
        network_reader.ReadWeights(input_model.substr(0, input_model.size() - 4) + ".bin");
        network_reader.getNetwork().setBatchSize(1);
        CNNNetwork network = network_reader.getNetwork();
        // -----------------------------------------------------------------------------------------------------

        // --------------------------- 3. Configure input & output ---------------------------------------------

        // --------------------------- Prepare input blobs -----------------------------------------------------
        /** Taking information about all topology inputs **/
        InputsDataMap input_info(network.getInputsInfo());
        /** Iterating over all input info**/
        for (auto &item : input_info) {
            InputInfo::Ptr input_data = item.second;
            input_data->setPrecision(Precision::U8);
            input_data->setLayout(Layout::NCHW);
        }

        // ------------------------------ Prepare output blobs -------------------------------------------------
        /** Taking information about all topology outputs **/
        OutputsDataMap output_info(network.getOutputsInfo());
        /** Iterating over all output info**/
        for (auto &item : output_info) {
            DataPtr output_data = item.second;
            if (!output_data) {
                throw std::runtime_error("Output data pointer is invalid");
            }
            output_data->setPrecision(Precision::FP32);
        }
        // -----------------------------------------------------------------------------------------------------

        // --------------------------- 4. Loading model to the plugin ------------------------------------------
        ExecutableNetwork executable_network = plugin.LoadNetwork(network, {});
        // -----------------------------------------------------------------------------------------------------

        // --------------------------- 5. Create infer request -------------------------------------------------
        InferRequest async_infer_request = executable_network.CreateInferRequest();
        // -----------------------------------------------------------------------------------------------------

        // --------------------------- 6. Prepare input --------------------------------------------------------
        for (auto &item : input_info) {
            cv::Mat image = cv::imread(input_image_path);

            auto input_name = item.first;
            InputInfo::Ptr input_data = item.second;

            /** Getting input blob **/
            Blob::Ptr input = async_infer_request.GetBlob(input_name);
            auto input_buffer = input->buffer().as<PrecisionTrait<Precision::U8>::value_type *>();

            /** Fill input tensor with planes. First b channel, then g and r channels **/
            if (image.empty()) throw std::logic_error("Invalid image at path: " + input_image_path);

            /* Resize and copy data from the image to the input blob */
            cv::resize(image, image, cv::Size(input_data->getTensorDesc().getDims()[3], input_data->getTensorDesc().getDims()[2]));
            auto dims = input->getTensorDesc().getDims();
            size_t channels_number = dims[1];
            size_t image_size = dims[3] * dims[2];
            for (size_t pid = 0; pid < image_size; ++pid) {
                for (size_t ch = 0; ch < channels_number; ++ch) {
                    input_buffer[ch * image_size + pid] = image.at<cv::Vec3b>(pid)[ch];
                }
            }
        }
        // -----------------------------------------------------------------------------------------------------

        // --------------------------- 7. Do inference ---------------------------------------------------------
        const int max_number_of_iterations = 10;
        int iterations = max_number_of_iterations;
        /** Set callback function for calling on completion of async request **/
        async_infer_request.SetCompletionCallback(
                [&] {
                    std::cout << "Completed " << max_number_of_iterations - iterations + 1 << " async request"
                              << std::endl;
                    if (--iterations) {
                        /** Start async request (max_number_of_iterations - 1) more times **/
                        async_infer_request.StartAsync();
                    }
                });
        /** Start async request for the first time **/
        async_infer_request.StartAsync();
        /** Wait all repetition of async requests **/
        for (int i = 0; i < max_number_of_iterations; i++) {
            async_infer_request.Wait(IInferRequest::WaitMode::RESULT_READY);
        }
        // -----------------------------------------------------------------------------------------------------

        // --------------------------- 8. Process output -------------------------------------------------------
        for (auto &item : output_info) {
            auto output_name = item.first;
            Blob::Ptr output = async_infer_request.GetBlob(output_name);;
            // Print classification results
            ClassificationResult classificationResult(output, {input_image_path});
            classificationResult.print();
        }
        // -----------------------------------------------------------------------------------------------------
    } catch (const std::exception & ex) {
        std::cerr << ex.what() << std::endl;
        return EXIT_FAILURE;
    }
    return EXIT_SUCCESS;
}
