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

#include <gtest/gtest.h>
#include "myriad_layers_reference_functions.hpp"
#include "myriad_layers_tests.hpp"

using namespace InferenceEngine;

#define NUM_ELEMS_PRIORS (4)
#define ERROR_BOUND (2.5e-3f)

struct PriorGridGeneratorParam {
    int               flatten;
    int               grid_w;
    int               grid_h;
    float             stride_w;
    float             stride_h;

    friend std::ostream& operator<<(std::ostream& os, PriorGridGeneratorParam const& tst)
    {
        return os << "grid width = " << tst.grid_w
                  << ", grid height = " << tst.grid_h
                  << ", step width = " << tst.stride_w
                  << ", step height = " << tst.stride_h;
    };
};

struct InputDims {
    tensor_test_params priors;
    tensor_test_params featureMap;
    tensor_test_params imData;

    InputDims(Dims priorDims, Dims featureMapDims, Dims imDataDims) :
                priors(priorDims),
                featureMap(featureMapDims),
                imData(imDataDims) {}
    
    InputDims() = default;
};

using ExpPriorGridGeneratorTestParams = std::tuple<InputDims, PriorGridGeneratorParam>;

typedef myriadLayerTestBaseWithParam<ExpPriorGridGeneratorTestParams> myriadLayersTestsExpPriorGridGenerator_smoke;

static void genPriors(InferenceEngine::Blob::Ptr rois,
                    const tensor_test_params& params,
                    const uint32_t numPriors) {
    ie_fp16 *roisBlobData = rois->buffer().as<ie_fp16*>();
    const int maxRangeWidth  = params.w * 4 / 5;
    const int maxRangeHeight = params.h * 4 / 5;

    for (int i = 0; i < numPriors; i++) {
        int x0 = std::rand() % maxRangeWidth;
        int x1 = x0 + (std::rand() % (params.w - x0 - 1)) + 1;
        int y0 = std::rand() % maxRangeHeight;
        int y1 = y0 + (std::rand() % (params.h - y0 - 1)) + 1;

        roisBlobData[i * NUM_ELEMS_PRIORS + 0] = PrecisionUtils::f32tof16((float)x0);
        roisBlobData[i * NUM_ELEMS_PRIORS + 1] = PrecisionUtils::f32tof16((float)y0);
        roisBlobData[i * NUM_ELEMS_PRIORS + 2] = PrecisionUtils::f32tof16((float)x1);
        roisBlobData[i * NUM_ELEMS_PRIORS + 3] = PrecisionUtils::f32tof16((float)y1);
    }
}

TEST_P(myriadLayersTestsExpPriorGridGenerator_smoke, ExpPriorGridGenerator) {
    InputDims inputTensorsDims = std::get<0>(GetParam());
    PriorGridGeneratorParam opParams = std::get<1>(GetParam());

    const auto numPriors = inputTensorsDims.priors.n;

    _config[VPU_CONFIG_KEY(DETECT_NETWORK_BATCH)] = CONFIG_VALUE(NO);

    IN_OUT_desc inputTensors, outputTensors;
    inputTensors.push_back({inputTensorsDims.priors.n, inputTensorsDims.priors.c});
    inputTensors.push_back({inputTensorsDims.featureMap.n,
                             inputTensorsDims.featureMap.c,
                             inputTensorsDims.featureMap.h,
                             inputTensorsDims.featureMap.w});
    inputTensors.push_back({inputTensorsDims.imData.n,
                             inputTensorsDims.imData.c,
                             inputTensorsDims.imData.h,
                             inputTensorsDims.imData.w});

    const int gridWidth  = opParams.grid_w ? opParams.grid_w : inputTensorsDims.featureMap.w;
    const int gridHeight = opParams.grid_h ? opParams.grid_h : inputTensorsDims.featureMap.h;

    outputTensors.push_back({numPriors * gridHeight * gridWidth, inputTensorsDims.priors.c});

    SetInputTensors(inputTensors);
    SetOutputTensors(outputTensors);

    std::map<std::string, std::string> layerParams = {
        {"flatten",  std::to_string(opParams.flatten)},
        {"h",        std::to_string(opParams.grid_h)},
        {"w",        std::to_string(opParams.grid_w)},
        {"stride_y", std::to_string(opParams.stride_h)},
        {"stride_x", std::to_string(opParams.stride_w)}
    };

    ASSERT_NO_FATAL_FAILURE(makeSingleLayerNetwork(LayerInitParams("ExperimentalDetectronPriorGridGenerator").params(layerParams)));

    /* Input data generating */
    for (auto blob : _inputMap) {
        if (blob.second == _inputMap.begin()->second) {
            genPriors(blob.second, inputTensorsDims.featureMap, numPriors);
        } else {
            GenRandomData(blob.second);
        }
    }

    std::vector<InferenceEngine::Blob::Ptr> refInputBlobs;
    std::vector<InferenceEngine::Blob::Ptr> refOutputBlobs;

    for (auto blob : _inputMap) {
        auto _refInputBlob = make_shared_blob<ie_fp16>({Precision::FP16,
                                                        blob.second->getTensorDesc().getDims(),
                                                        blob.second->getTensorDesc().getLayout()},
                                                        blob.second->buffer());
        refInputBlobs.push_back(_refInputBlob);
    }

    for (auto blob : _outputMap) {
        auto refOutputBlob = make_shared_blob<ie_fp16>({Precision::FP16,
                                                      blob.second->getTensorDesc().getDims(),
                                                      blob.second->getTensorDesc().getLayout()});
        refOutputBlob->allocate();
        refOutputBlobs.push_back(refOutputBlob);
    }

    ref_ExpPriorGridGenerator(refInputBlobs,
                              refOutputBlobs,
                              opParams.grid_w,
                              opParams.grid_h,
                              opParams.stride_w,
                              opParams.stride_h);

    ASSERT_TRUE(Infer());
    CompareCommonAbsolute(_outputMap.begin()->second, refOutputBlobs[0], ERROR_BOUND);
}

static std::vector<InputDims> s_ExpPriorGridGeneratorLayerInputs = {
    {
        InputDims(
            Dims({3, 4}),          // priors
            Dims({1, 128, 8, 8}),  // feature map
            Dims({1, 3, 480, 480}) // im_data
        )
    },
    {
        InputDims(
            Dims({3, 4}),           // priors
            Dims({1, 128, 60, 60}), // feature map
            Dims({1, 3, 480, 480})  // im_data
        )
    },
    {
        InputDims(
            Dims({3, 4}),             // priors
            Dims({1, 128, 120, 120}), // feature map
            Dims({1, 3, 480, 480})    // im_data
        )
    },
    {
        InputDims(
            Dims({64, 4}),          // priors
            Dims({1, 128, 16, 16}), // feature map
            Dims({1, 3, 480, 480})  // im_data
        )
    },
};

static std::vector<PriorGridGeneratorParam> s_ExpPriorGridGeneratorLayerParam = {
    {1, 0, 0, 16.0f, 16.0f}, {1, 0, 0, 8.0f, 8.0f}, {1, 0, 0, 4.0f, 4.0f},
    {1, 8, 8, 64.0f, 64.0f}, {1, 10, 16, 0.0f, 0.0f}
};
