#include <library/cpp/testing/gtest/gtest.h>
#include <library/cpp/testing/common/env.h>
#include <tensorflow/cc/saved_model/loader.h>
#include <tensorflow/cc/saved_model/tag_constants.h>

#include <utility>

namespace tf = tensorflow;

// EfficientNet fullstack test
TEST(basic_tests, check_fullstack_import_inference)
{
    const tf::string TESTCONV_TF_MODEL_RESOURCE = BinaryPath("maps/wikimap/mapspro/services/mrc/eye/experiments/signs_detection/learning/faster_rcnn_keras/tests/data/Desktop/model");

    tf::SavedModelBundleLite model_bundle;
    const tf::SessionOptions session_options = tf::SessionOptions();
    const tf::RunOptions run_options = tf::RunOptions();
    const std::unordered_set tag_consts({tf::string(tf::kSavedModelTagServe)});
    tf::Status status = tf::LoadSavedModel(session_options, run_options, TESTCONV_TF_MODEL_RESOURCE, tag_consts, &model_bundle);
    
    ASSERT_THAT(status, tf::Status::OK());
    
    
    tf::Tensor input_tensor(tf::DT_FLOAT,
      tf::TensorShape({1, 224, 224, 3}));

    std::vector<tf::Tensor> output_tensor(1, tf::Tensor(tf::DT_FLOAT,
      tf::TensorShape({1, 10})));

    tf::Status runStatus = model_bundle.GetSession()->Run({{"serving_default_efficientnetb0_input:0", input_tensor}}, {"StatefulPartitionedCall:0"}, {}, &output_tensor);

    EXPECT_THAT(runStatus, tf::Status::OK());
}
