// FILE Segmentator_OpenVINO.h

#pragma once

#include <ie_api.h>
#include <ie_core.hpp>
#include <ie_iinfer_request.hpp>
#include <ie_input_info.hpp>
#include <ie_precision.hpp>
#include <inference_engine.hpp>

#include <map>

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

#include "library.h"
#include "openVINOUtils.h"
#include "DnnUtils.h"

using namespace InferenceEngine;

class DLL_EXPORT Segmentator_OpenVIN {
public:
    Segmentator_OpenVINO();

    std::multimap<AnatomicalStructure, std::vector<cv::Point>> run(const cv::Mat& image);

protected:

    void preprocess(const cv::Mat& inputImage);
    void forward();
	std::multimap<AnatomicalStructure, std::vector<cv::Point>> postprocess();

	// Network associated with one hardware device
	InferenceEngine::ExecutableNetwork executable_network_;

	InferenceEngine::InputsDataMap inputInfo_;
	InferenceEngine::InferRequest infer_request_;

	std::string imageInputName_;

	size_t netBatchSize_;
	size_t netInputHeight_;
	size_t netInputWidth_;

	float confThreshold_;
	float maskThreshold_;
	float overlapThreshold_;

	int image_width_;
	int image_height_;
};


///////////////////////////////////////////////////////////////////////////////////

// FILE Segmentator_OpenVINO.cpp


#include "Segmentator_OpenVINO.h"
#include "openVINOUtils.h"

using namespace InferenceEngine;
using namespace std;

Segmentator_OpenVINO::Segmentator_OpenVINO()
{

    InferenceEngine::Core core;
    InferenceEngine::CNNNetwork network = core.ReadNetwork("C:/Users/joseph/Documents/Local_DNN_models/mask_rcnn/106/frozen_inference_graph.xml");

    /** Take information about all topology inputs **/
    inputInfo_ = network.getInputsInfo();
    /** Iterate over all input info**/
    for (const auto& inputInfoItem : inputInfo_) {
        if (inputInfoItem.second->getTensorDesc().getDims().size() == 4) { // first input contains images
            imageInputName_ = inputInfoItem.first;
            inputInfoItem.second->setPrecision(InferenceEngine::Precision::U8);
        }
        else if (inputInfoItem.second->getTensorDesc().getDims().size() == 2) { // second input contains image info
            inputInfoItem.second->setPrecision(InferenceEngine::Precision::FP32);
        }
        else {
            throw std::logic_error(
                "Unsupported input shape with size = " + std::to_string(inputInfoItem.second->getTensorDesc().getDims().size()));
        }
    }

    /** network dimensions for image input **/
    const InferenceEngine::TensorDesc& inputDesc = inputInfo_[imageInputName_]->getTensorDesc();
    IE_ASSERT(inputDesc.getDims().size() == 4);

    netBatchSize_ = getTensorBatch(inputDesc);
    netInputHeight_ = getTensorHeight(inputDesc);
    netInputWidth_ = getTensorWidth(inputDesc);

    image_width_ = 480;
    image_height_ = 480;

    maskThreshold_ = 0.5;
    confThreshold_ = 0.5;
    overlapThreshold_ = 0.75;

    /** Take information about all topology outputs **/
    InferenceEngine::OutputsDataMap outputInfo(network.getOutputsInfo());

    /** Iterate over all output info**/
    for (auto& item : outputInfo) {
        item.second->setPrecision(InferenceEngine::Precision::FP32);
    }

    executable_network_ = core.LoadNetwork(network, "CPU");
}

std::multimap<AnatomicalStructure, std::vector<cv::Point>> Segmentator_OpenVINO::run(const cv::Mat& image)
{
    printNetCharactestic();
    preprocess(image);
    forward();
    return = postprocess();
}


void Segmentator_OpenVINO::preprocess(const cv::Mat& inputImage)
{
    cv::Mat image(inputImage);
    cv::resize(image, image, cv::Size(netInputHeight_, netInputWidth_));
    cv::cvtColor(image, image, cv::COLOR_RGBA2BGR);

    infer_request_ = executable_network_.CreateInferRequest();
    /** Iterate over all the input blobs **/
    for (const auto& inputInfoItem : inputInfo_) {
        InferenceEngine::Blob::Ptr input = infer_request_.GetBlob(inputInfoItem.first);
        std::cout << input->getTensorDesc().getLayout() << std::endl;
        /** Fill first input tensor with images. First b channel, then g and r
         * channels **/
        if (inputInfoItem.second->getTensorDesc().getDims().size() == 4) {
            /** Iterate over all input images **/
            matU8ToBlob<unsigned char>(image, input);
        }
        /** Fill second input tensor with image info **/
        if (inputInfoItem.second->getTensorDesc().getDims().size() == 2) {
            auto data = input->buffer()
                .as<InferenceEngine::PrecisionTrait<
                InferenceEngine::Precision::FP32>::value_type*>();
            data[0] = static_cast<float>(netInputHeight_); // height
            data[1] = static_cast<float>(netInputWidth_); // width
            data[2] = 1;
        }
    }

    InferenceEngine::Blob::Ptr imgBlob = wrapMat2Blob(image);

    // infer_request accepts input blob of any size
    infer_request_.SetBlob(imageInputName_, imgBlob);
}

void Segmentator_OpenVINO::forward()
{
    infer_request_.Infer();
}

std::multimap<AnatomicalStructure, std::vector<cv::Point>> Segmentator_OpenVINO::postprocess()
{
    if (!infer_request_) {
        std::cerr << "Infer request is uninitialized " << std::endl;
        return {};
    }
    else
        std::cout << "Infer request in fine" << std::endl;

    std::multimap<int, cv::Mat> previous_detection;
    std::multimap<AnatomicalStructure, std::vector<cv::Point>> detectedPolygons;
    cv::Mat res(image_height_, image_width_, CV_8UC1, cv::Scalar(0, 0, 0));

    const auto do_blob = infer_request_.GetBlob("reshape_do_2d");
    const auto do_data = do_blob->buffer().as<float*>();

    const TensorDesc& do_desc = do_blob->getTensorDesc();

    const auto masks_blob = infer_request_.GetBlob("masks");
    const auto masks_data = masks_blob->buffer().as<float*>();

    IE_ASSERT(do_blob->getTensorDesc().getDims().size() == 2);
    size_t BOX_DESCRIPTION_SIZE = do_blob->getTensorDesc().getDims().back();

    const TensorDesc& masksDesc = masks_blob->getTensorDesc();
    IE_ASSERT(masksDesc.getDims().size() == 4);
    size_t BOXES = getTensorBatch(masksDesc);
    size_t C = getTensorChannels(masksDesc);
    size_t H = getTensorHeight(masksDesc);
    size_t W = getTensorWidth(masksDesc);

    size_t box_stride = W * H * C;

    std::map<size_t, size_t> class_color;

    bool overlapping = false;

    /** Iterating over all boxes **/
    for (size_t box = 0; box < BOXES; ++box) {
        float* box_info = do_data + box * BOX_DESCRIPTION_SIZE;
        auto batch = static_cast<int>(box_info[0]);
        if (batch < 0)
            break;
        if (batch >= static_cast<int>(netBatchSize_))
            throw std::logic_error("Invalid batch ID within detection output box");
        float prob = box_info[2];
        std::cout << "\tProb: " << prob << std::endl;
        if (prob > confThreshold_) {
            float x1 = std::min(std::max(0.0f, box_info[3] * image_width_), static_cast<float>(image_width_));
            float y1 = std::min(std::max(0.0f, box_info[4] * image_height_), static_cast<float>(image_height_));
            float x2 = std::min(std::max(0.0f, box_info[5] * image_width_), static_cast<float>(image_width_));
            float y2 = std::min(std::max(0.0f, box_info[6] * image_height_), static_cast<float>(image_height_));
            int box_width = std::min(static_cast<int>(std::max(0.0f, x2 - x1)), image_width_);
            int box_height = std::min(static_cast<int>(std::max(0.0f, y2 - y1)), image_height_);
            auto class_id = static_cast<size_t>(box_info[1] + 1e-6f);

            float* mask_arr = masks_data + box_stride * box + H * W * (class_id - 1);
            cv::Mat objectMask(H, W, CV_32FC1, mask_arr);

            .......

        }
    }
    return detectedPolygons;
}