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

#include "test_utils/cpu_test_utils.hpp"
#include "test_utils/fusing_test_utils.hpp"
#include "ngraph_functions/builders.hpp"
#include "common_test_utils/common_utils.hpp"

#include <algorithm>
#include <cassert>

using namespace ngraph;
using namespace InferenceEngine;
using namespace CPUTestUtils;

namespace SubgraphTestsDefinitions {

using AlignMatMulInputRanksTestParams = std::tuple<std::pair<SizeVector, SizeVector>, // IS fully connected
                                       fusingSpecificParams>;

class AlignMatMulInputRanksTest : public testing::WithParamInterface<AlignMatMulInputRanksTestParams>, public CpuTestWithFusing,
                      virtual public LayerTestsUtils::LayerTestsCommon {
public:
    static std::string getTestCaseName(testing::TestParamInfo<AlignMatMulInputRanksTestParams> obj) {
        std::pair<SizeVector, SizeVector> supportedInputShapes;
        fusingSpecificParams fusingParams;
        std::tie(supportedInputShapes, fusingParams) = obj.param;
        SizeVector inputShapeA = supportedInputShapes.first; SizeVector inputShapeB = supportedInputShapes.second;

        std::ostringstream result;
        result << "IS_A=" << CommonTestUtils::vec2str(inputShapeA) << "_";
        result << "IS_B=" << CommonTestUtils::vec2str(inputShapeB) << "_";
        result << CpuTestWithFusing::getTestCaseName(fusingParams);

        return result.str();
    }

protected:
    void SetUp() override {
        targetDevice = CommonTestUtils::DEVICE_CPU;
        std::pair<SizeVector, SizeVector> inShapes;
        fusingSpecificParams fusingParams;
        std::tie(inShapes, fusingParams) = this->GetParam();

        if (inShapes.first.size() != inShapes.second.size())
            expectedNumOfReshapes++;  // one input will be unsqueezed
        if (inShapes.first.size() == 1 || inShapes.second.size() == 1)
            expectedNumOfReshapes++;  // output will be squeezed
        if (inShapes.first.size() == 1 && inShapes.second.size() == 1)
            expectedNumOfReshapes+=2; // both inputs unsqueezed and output squeezed

        if (inShapes.first.size() != 1 && inShapes.second.size() != 1) // no fusing through Reshape after output
            std::tie(postOpMgrPtr, fusedOps) = fusingParams;

        const auto ngPrec = element::f32;
        auto inputParams = builder::makeParams(ngPrec, {inShapes.first, inShapes.second});
        const auto outputNodes = helpers::convert2OutputVector(helpers::castOps2Nodes<op::Parameter>(inputParams));
        const auto matMul = builder::makeMatMul(outputNodes[0], outputNodes[1], false, false);

        function = makeNgraphFunction(ngPrec, inputParams, matMul, "AlignMatMulInputRanks");
    }

    int expectedNumOfReshapes = 0;
};

TEST_P(AlignMatMulInputRanksTest, CompareWithRefs) {
    SKIP_IF_CURRENT_TEST_IS_DISABLED()

    Run();
    CheckNodeOfTypeCount(executableNetwork, "Reshape", expectedNumOfReshapes); // Squeeze / Unsqueeze turns into Reshape
    CheckFusingResults(executableNetwork, "MatMul");
}

namespace {

const std::vector<std::pair<SizeVector, SizeVector>> supportedInputShapes = {
    {{4, 10, 5}, {1, 5, 10}},      // nothing to be done
    {{3}, {3}},                    // 3x1 * 1x3 -> 1
    {{18}, {1, 5, 18, 20}},        // 1x1x1x18 * 1x5x18x20 -> 1x5x20
    {{2, 3, 4, 4, 4, 10, 5}, {5}}, // 2x3x4x4x4x10x5 * 1x1x1x1x1x5x1 -> 1x1x1x1x1x5
    {{1, 18}, {1, 5, 18, 20}},
    {{1, 70, 18}, {1, 5, 18, 20}},
    {{7, 1, 10, 3, 2, 7}, {1, 7, 5}},
    {{2, 3, 4, 4, 4, 10, 5}, {5, 20}},
};

// verify fusing just in case
std::vector<fusingSpecificParams> fusingParamsSet {
        emptyFusingSpec,
        fusingElu,
};

INSTANTIATE_TEST_SUITE_P(smoke_Check, AlignMatMulInputRanksTest,
                         ::testing::Combine(::testing::ValuesIn(supportedInputShapes),
                                            ::testing::ValuesIn(fusingParamsSet)),
                         AlignMatMulInputRanksTest::getTestCaseName);

} // namespace

} // namespace SubgraphTestsDefinitions
