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

#include "shared_test_classes/single_layer/space_to_depth.hpp"
#include "test_utils/cpu_test_utils.hpp"
#include "shared_test_classes/base/ov_subgraph.hpp"


using namespace InferenceEngine;
using namespace CPUTestUtils;
using namespace ngraph::opset3;
using namespace ov::test;

namespace CPULayerTestsDefinitions {

using SpaceToDepthLayerCPUTestParamSet = std::tuple<
        InputShape,                                     // Input shape
        ElementType,                                    // Input element type
        SpaceToDepth::SpaceToDepthMode,                 // Mode
        std::size_t,                                    // Block size
        CPUSpecificParams
>;

class SpaceToDepthLayerCPUTest : public testing::WithParamInterface<SpaceToDepthLayerCPUTestParamSet>,
                                 virtual public ov::test::SubgraphBaseTest, public CPUTestsBase {
public:
    static std::string getTestCaseName(testing::TestParamInfo<SpaceToDepthLayerCPUTestParamSet> obj) {
        InputShape shapes;
        ElementType inType;
        SpaceToDepth::SpaceToDepthMode mode;
        std::size_t blockSize;
        CPUSpecificParams cpuParams;
        std::tie(shapes, inType, mode, blockSize, cpuParams) = obj.param;

        std::ostringstream results;
        results << "IS=" << CommonTestUtils::partialShape2str({shapes.first}) << "_";
        results << "TS=";
        for (const auto& item : shapes.second) {
            results << CommonTestUtils::vec2str(item) << "_";
        }
        results << "Prc=" << inType << "_";
        switch (mode) {
            case SpaceToDepth::SpaceToDepthMode::BLOCKS_FIRST:
                results << "BLOCKS_FIRST_";
                break;
            case SpaceToDepth::SpaceToDepthMode::DEPTH_FIRST:
                results << "DEPTH_FIRST_";
                break;
            default:
                throw std::runtime_error("Unsupported SpaceToDepthMode");
        }
        results << "BS=" << blockSize << "_";
        results << CPUTestsBase::getTestCaseName(cpuParams);

        return results.str();
    }
protected:
    void SetUp() override {
        InputShape shapes;
        ElementType inType;
        SpaceToDepth::SpaceToDepthMode mode;
        std::size_t blockSize;
        CPUSpecificParams cpuParams;
        std::tie(shapes, inType, mode, blockSize, cpuParams) = this->GetParam();

        std::tie(inFmts, outFmts, priority, selectedType) = cpuParams;
        if (selectedType.empty()) {
            selectedType = getPrimitiveType();
        }
        selectedType = selectedType + "_" + InferenceEngine::details::convertPrecision(inType).name();
        targetDevice = CommonTestUtils::DEVICE_CPU;
        init_input_shapes({shapes});

        auto params = ngraph::builder::makeDynamicParams(inType, inputDynamicShapes);
        auto d2s = ngraph::builder::makeSpaceToDepth(params[0], mode, blockSize);
        function = makeNgraphFunction(inType, params, d2s, "SpaceToDepthCPU");
    }
};

TEST_P(SpaceToDepthLayerCPUTest, CompareWithRefs) {
    SKIP_IF_CURRENT_TEST_IS_DISABLED()

    run();
    CPUTestsBase::CheckPluginRelatedResults(executableNetwork, "SpaceToDepth");
}

namespace {

const auto cpuParams_nChw16c = CPUSpecificParams {{nChw16c}, {nChw16c}, {}, {}};
const auto cpuParams_nCdhw16c = CPUSpecificParams {{nCdhw16c}, {nCdhw16c}, {}, {}};

const auto cpuParams_nChw8c = CPUSpecificParams {{nChw8c}, {nChw8c}, {}, {}};
const auto cpuParams_nCdhw8c = CPUSpecificParams {{nCdhw8c}, {nCdhw8c}, {}, {}};

const auto cpuParams_nhwc = CPUSpecificParams {{nhwc}, {nhwc}, {}, {}};
const auto cpuParams_ndhwc = CPUSpecificParams {{ndhwc}, {ndhwc}, {}, {}};

const auto cpuParams_nchw = CPUSpecificParams {{nchw}, {nchw}, {}, {}};
const auto cpuParams_ncdhw = CPUSpecificParams {{ncdhw}, {ncdhw}, {}, {}};

const std::vector<CPUSpecificParams> CPUParams4D = {
        cpuParams_nhwc,
        cpuParams_nchw
};

const std::vector<CPUSpecificParams> CPUParamsBlocked4D = {
        cpuParams_nChw16c,
        cpuParams_nChw8c,
        cpuParams_nhwc
};

const std::vector<CPUSpecificParams> CPUParams5D = {
        cpuParams_ndhwc,
        cpuParams_ncdhw
};

const std::vector<CPUSpecificParams> CPUParamsBlocked5D = {
        cpuParams_nCdhw16c,
        cpuParams_nCdhw8c,
        cpuParams_ndhwc
};

const std::vector<ElementType> inputElementType = {
        ElementType::f32,
        ElementType::bf16,
        ElementType::i8
};

const std::vector<SpaceToDepth::SpaceToDepthMode> SpaceToDepthModes = {
        SpaceToDepth::SpaceToDepthMode::BLOCKS_FIRST,
        SpaceToDepth::SpaceToDepthMode::DEPTH_FIRST
};

/* *========================* Static Shapes Tests *========================* */

namespace static_shapes {

const std::vector<ov::Shape> inputShapesBS2_4D = {{1, 16, 2, 2}, {1, 16, 4, 2},
                                                  {1, 32, 6, 8}, {2, 32, 4, 6},
                                                  {2, 48, 4, 4}, {2, 64, 8, 2}};

const std::vector<ov::Shape> inputShapesBS3_4D = {
    {1, 2, 3, 3}, {1, 3, 3, 6}, {1, 5, 6, 3}, {2, 5, 9, 3}, {3, 5, 6, 6}};

INSTANTIATE_TEST_SUITE_P(smoke_CPUSpaceToDepthBS2_4D, SpaceToDepthLayerCPUTest,
                         testing::Combine(
                                 testing::ValuesIn(static_shapes_to_test_representation(inputShapesBS2_4D)),
                                 testing::ValuesIn(inputElementType),
                                 testing::ValuesIn(SpaceToDepthModes),
                                 testing::Values(1, 2),
                                 testing::ValuesIn(filterCPUInfoForDevice(CPUParamsBlocked4D))),
                         SpaceToDepthLayerCPUTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_CPUSpaceToDepthStaticBS3_4D, SpaceToDepthLayerCPUTest,
                         testing::Combine(
                                 testing::ValuesIn(static_shapes_to_test_representation(inputShapesBS3_4D)),
                                 testing::ValuesIn(inputElementType),
                                 testing::ValuesIn(SpaceToDepthModes),
                                 testing::Values(1, 3),
                                 testing::ValuesIn(filterCPUInfoForDevice(CPUParams4D))),
                         SpaceToDepthLayerCPUTest::getTestCaseName);

const std::vector<ov::Shape> inputShapesBS2_5D = {
        {1, 16, 2, 2, 2}, {1, 16, 4, 4, 2}, {1, 32, 2, 6, 2}, {2, 32, 4, 2, 2}, {1, 48, 6, 2, 2}, {2, 64, 2, 2, 6}
};

const std::vector<ov::Shape> inputShapesBS3_5D = {
        {1, 2, 3, 3, 3}, {1, 2, 3, 6, 9}, {1, 5, 6, 3, 3}, {2, 5, 3, 9, 3}, {3, 5, 3, 3, 6}
};

INSTANTIATE_TEST_SUITE_P(smoke_CPUSpaceToDepthStaticBS2_5D, SpaceToDepthLayerCPUTest,
                         testing::Combine(
                                 testing::ValuesIn(static_shapes_to_test_representation(inputShapesBS2_5D)),
                                 testing::ValuesIn(inputElementType),
                                 testing::ValuesIn(SpaceToDepthModes),
                                 testing::Values(1, 2),
                                 testing::ValuesIn(filterCPUInfoForDevice(CPUParamsBlocked5D))),
                         SpaceToDepthLayerCPUTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_CPUSpaceToDepthStaticBS3_5D, SpaceToDepthLayerCPUTest,
                         testing::Combine(
                                 testing::ValuesIn(static_shapes_to_test_representation(inputShapesBS3_5D)),
                                 testing::ValuesIn(inputElementType),
                                 testing::ValuesIn(SpaceToDepthModes),
                                 testing::Values(1, 3),
                                 testing::ValuesIn(filterCPUInfoForDevice(CPUParams5D))),
                         SpaceToDepthLayerCPUTest::getTestCaseName);

} // namespace static_shapes
/* *========================* *==================* *========================* */


/* *========================* Dynamic Shapes Tests *========================* */
namespace dynamic_shapes {

const std::vector<InputShape> inputShapes4D = {
    {{-1, -1, -1, -1},                                                   // dynamic
     {{2, 12, 6, 12}, {1, 2, 18, 6}, {1, 10, 12, 24}, {2, 12, 6, 12}}},  // target

    {{-1, 32, -1, -1},                                                     // dynamic
     {{1, 32, 6, 12}, {2, 32, 24, 24}, {3, 32, 30, 6}, {2, 32, 24, 24}}},  // target

    {{{2, 5}, {1, 50}, {1, 100}, {1, 100}},                            // dynamic
     {{3, 5, 12, 36}, {2, 3, 6, 12}, {5, 2, 12, 18}, {2, 3, 6, 12}}},  // target
};

const std::vector<InputShape> inputShapes5D = {
    {{-1, -1, -1, -1, -1},                                                            // dynamic
     {{2, 2, 6, 12, 24}, {1, 4, 24, 24, 36}, {1, 7, 6, 30, 18}, {2, 2, 6, 12, 24}}},  // target

    {{{1, 3}, {5, 16}, {1, 60}, {1, 60}, {1, 60}},                // dynamic
     {{3, 5, 12, 6, 24}, {1, 6, 24, 18, 6}, {3, 5, 12, 6, 24}}},  // target
};

const std::vector<InputShape> inputShapesBlocked5D = {
    {{-1, 16, -1, -1, -1},                                                            // dynamic
     {{1, 16, 4, 6, 10}, {1, 16, 12, 8, 2}, {3, 16, 2, 14, 24}, {1, 16, 12, 8, 2}}},  // target

    {{{1, 3}, 32, {1, 32}, {1, 32}, {1, 32}},                                            // dynamic
     {{1, 32, 4, 16, 10}, {1, 32, 18, 6, 14}, {3, 32, 2, 14, 12}, {1, 32, 18, 6, 14}}},  // target
};

INSTANTIATE_TEST_SUITE_P(smoke_CPUSpaceToDepthDynamic4D, SpaceToDepthLayerCPUTest,
                         testing::Combine(
                                 testing::ValuesIn(inputShapes4D),
                                 testing::ValuesIn(inputElementType),
                                 testing::ValuesIn(SpaceToDepthModes),
                                 testing::Values(1, 2, 3),
                                 testing::ValuesIn(filterCPUInfoForDevice(CPUParams4D))),
                         SpaceToDepthLayerCPUTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_CPUSpaceToDepthDynamicBlocksFirstBlocked4D, SpaceToDepthLayerCPUTest,
                         testing::Combine(
                                 testing::Values(inputShapes4D[1]),
                                 testing::ValuesIn(inputElementType),
                                 testing::Values(SpaceToDepth::SpaceToDepthMode::BLOCKS_FIRST),
                                 testing::Values(1, 2, 3),
                                 testing::ValuesIn(filterCPUInfoForDevice(CPUParamsBlocked4D))),
                        SpaceToDepthLayerCPUTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_CPUSpaceToDepthDynamicDepthFirstBlocked4D, SpaceToDepthLayerCPUTest,
                         testing::Combine(
                                 testing::Values(inputShapes4D[1]),
                                 testing::ValuesIn(inputElementType),
                                 testing::Values(SpaceToDepth::SpaceToDepthMode::DEPTH_FIRST),
                                 testing::Values(1, 2),
                                 testing::ValuesIn(filterCPUInfoForDevice(CPUParamsBlocked4D))),
                         SpaceToDepthLayerCPUTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_CPUSpaceToDepthDynamic5D, SpaceToDepthLayerCPUTest,
                         testing::Combine(
                                 testing::ValuesIn(inputShapes5D),
                                 testing::ValuesIn(inputElementType),
                                 testing::ValuesIn(SpaceToDepthModes),
                                 testing::Values(1, 2, 3),
                                 testing::ValuesIn(filterCPUInfoForDevice(CPUParams5D))),
                         SpaceToDepthLayerCPUTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_CPUSpaceToDepthDynamicCPUSpecific5D, SpaceToDepthLayerCPUTest,
                         testing::Combine(
                                 testing::ValuesIn(inputShapesBlocked5D),
                                 testing::ValuesIn(inputElementType),
                                 testing::ValuesIn(SpaceToDepthModes),
                                 testing::Values(1, 2),
                                 testing::ValuesIn(filterCPUInfoForDevice(CPUParamsBlocked5D))),
                         SpaceToDepthLayerCPUTest::getTestCaseName);
} // namespace dynamic_shapes
/* *========================* *==================* *========================* */

} // namespace
} // namespace CPULayerTestsDefinitions