#pragma once

#include <maps/libs/common/include/exception.h>

#include <opencv2/opencv.hpp>
#include <tensorflow/core/public/session.h>

#include <string>
#include <vector>

namespace maps {
namespace wiki {
namespace tf_inferencer {

typedef std::vector<cv::Mat> ImagesBatch;

class TensorFlowInferencer {
public:
    /**
    * @param path to tensorflow model in protobuf format
    */
    TensorFlowInferencer(const std::string& path);
    static TensorFlowInferencer fromResource(const std::string& resourceName);
    TensorFlowInferencer(const tensorflow::GraphDef& graphDef);

    /**
        * @brief Launch inference of network loaded in constructor
        *
        * @param inputLayerName  name of the input layer
        * @param inputImage      data for the input layer
        * @param outputLayerName name of the output layer
        *
        * @return output from the network
        */
    tensorflow::Tensor inference(const std::string &inputLayerName,
                         const cv::Mat &inputImage,
                         const std::string &outputLayerName) const;

    std::vector<tensorflow::Tensor>
    inference(const std::string &inputLayerName,
              const cv::Mat &inputImage,
              const std::vector<std::string> &outputLayerNames) const;

    std::vector<tensorflow::Tensor>
    inference(const std::vector<std::pair<std::string, cv::Mat>> &inputLayerImages,
              const std::vector<std::string> &outputLayerNames) const;

    /*
        All images in batch must have same sizes and amount of channels
    */
    std::vector<tensorflow::Tensor>
    inference(const std::string &inputLayerName,
              const ImagesBatch &inputImagesBatch,
              const std::vector<std::string> &outputLayerNames) const;

    tensorflow::Tensor inference(const std::string &outputLayerName) const;

    std::vector<tensorflow::Tensor>
    inference(const std::vector<std::pair<std::string, tensorflow::Tensor>> &inputLayerTensors,
              const std::vector<std::string> &outputLayerNames) const;
private:
    std::unique_ptr<tensorflow::Session> session_;

    std::vector<tensorflow::Tensor>
    inference(const std::string &inputLayerName,
              const tensorflow::Tensor &input,
              const std::vector<std::string> &outputLayerNames) const;
};


cv::Mat tensorToImage(const tensorflow::Tensor& tensor);

tensorflow::Tensor cvMatToTensor(const cv::Mat& mat, bool addBatchDimension, bool singleColumnAsDim, bool singleChannelAsDim);

template<class T>
std::vector<T> tensorToVector(const tensorflow::Tensor& tensor)
{
    REQUIRE(1 == tensor.dims(),
            "Tensor has unacceptable dimensions count" << tensor.dims());

    auto eigenTensor = tensor.vec<T>();

    std::vector<T> result;
    size_t classesNum = eigenTensor.dimension(0);
    result.reserve(classesNum);

    for (size_t i = 0; i < classesNum; ++i) {
        result.emplace_back(eigenTensor(i));
    }

    return result;
}

} //namespace tf_inferencer
} //namespace wiki
} //namespace maps
