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

#include "graph_transformer_tests.hpp"

using namespace vpu;

class VPU_MergePermuteTest : public GraphTransformerTest {
 protected:
    PassSet pipeline;
    Model model;
    using PermuteDims = PermutationDimsMap;
 public:
    void SetUp() override {
        GraphTransformerTest::SetUp();
        InitCompileEnv();
    }

    void CreateModelWithTwoPermutes(const DimsOrder& layout, const PermuteDims& firstPermute, const PermuteDims& secondPermute) {
        model = CreateModel();

        const DataDesc dataDesc(DataType::FP16, layout, MakeStubDims(layout));
        auto input = model->addInputData("Input", dataDesc);
        model->attrs().set<int>("numInputs", 1);

        auto output = model->addOutputData("Output", dataDesc);
        model->attrs().set<int>("numOutputs", 1);

        auto data1 = model->addNewData("data1", dataDesc);

        stageBuilder->addPermuteStage(model,
                                      "Permute1",
                                      nullptr,
                                      input,
                                      data1,
                                      firstPermute
                                    );

        stageBuilder->addPermuteStage(model,
                                      "Permute2",
                                      nullptr,
                                      data1,
                                      output,
                                      secondPermute
                                    );
        InitPipeline();
    }

    void CreateModelWithThreePermutes(const DimsOrder& layout, const PermuteDims& firstPermute, const PermuteDims& secondPermute, const PermuteDims& thirdPermute) {
        model = CreateModel();

        const DataDesc dataDesc(DataType::FP16, layout, MakeStubDims(layout));
        auto input = model->addInputData("Input", dataDesc);
        model->attrs().set<int>("numInputs", 1);

        auto output = model->addOutputData("Output", dataDesc);
        model->attrs().set<int>("numOutputs", 1);

        auto data1 = model->addNewData("data1", dataDesc);
        auto data2 = model->addNewData("data2", dataDesc);

        stageBuilder->addPermuteStage(model,
                                      "Permute1",
                                      nullptr,
                                      input,
                                      data1,
                                      firstPermute
                                    );

        stageBuilder->addPermuteStage(model,
                                      "Permute2",
                                      nullptr,
                                      data1,
                                      data2,
                                      secondPermute
                                    );

        stageBuilder->addPermuteStage(model,
                                      "Permute3",
                                      nullptr,
                                      data2,
                                      output,
                                      thirdPermute
                                    );
        InitPipeline();
    }

    void CreateModelWithPermuteAndReorder(const DimsOrder& layoutIn, const DimsOrder& layoutOut, const PermuteDims& firstPermute) {
        model = CreateModel();

        const DataDesc dataDescIn (DataType::FP16, layoutIn,  MakeStubDims(layoutIn));
        const DataDesc dataDescOut(DataType::FP16, layoutOut, MakeStubDims(layoutIn));  // That's not a typo: we use same Dim-to-size mapping for in and out.
        auto input = model->addInputData("Input", dataDescIn);
        model->attrs().set<int>("numInputs", 1);

        auto output = model->addOutputData("Output", dataDescOut);
        model->attrs().set<int>("numOutputs", 1);

        auto data1 = model->addNewData("data1", dataDescIn);

        stageBuilder->addPermuteStage(model,
                                      "Permute",
                                      nullptr,
                                      input,
                                      data1,
                                      firstPermute
                                    );

        stageBuilder->addReorderStage(model,
                                      "Reorder",
                                      nullptr,
                                      data1,
                                      output
                                    );
        InitPipeline();
    }

    void InitPipeline() {
        pipeline = PassSet();
        pipeline.addPass(passManager->dumpModel("initial"));
        pipeline.addPass(passManager->adjustDataLayout());
        pipeline.addPass(passManager->dumpModel("adjustDataLayout"));
        pipeline.addPass(passManager->mergePermuteStages());
        pipeline.addPass(passManager->dumpModel("mergePermuteStages"));
        pipeline.addPass(passManager->adjustDataLocation());
        pipeline.addPass(passManager->dumpModel("adjustDataLocation"));
        pipeline.addPass(passManager->finalCheck());
    }

    // Create fake dimension size array, using dimension count from layout
    DimValues MakeStubDims(const DimsOrder& layout)
    {
        const auto perm = layout.toPermutation();
        const int numDms = layout.numDims();
        DimValues dims;
        for (int ind = 0; ind < numDms; ++ind) {
            dims.set(perm[ind], 2 + ind * 2); // 2, 4, 6 ...
        }
        return dims;
    }

    const PermuteDims& GetPermuteDimsForStage(size_t index) {
        auto stages = model->getStages();
        auto stageIter = stages.begin();
        std::advance(stageIter, index);
        auto permuteStage = *stageIter;

        IE_ASSERT(permuteStage->type() == StageType::Permute);
        return permuteStage->attrs().get<PermutationDimsMap>("permutation");
    }

};
TEST_F(VPU_MergePermuteTest, InternalFunctions) {
    // Meaning of integer permutation vectors:
    // suppose you have vector {1, 2, 0}, expand it in mapping, where key would be vector index:
    // [0] = 1
    // [1] = 2
    // [2] = 0
    // The key is destination dimension, and the value is the source dimension.
    // For CHW, C==2, H==1, W==0 index, so permutation above would be equal to:
    // W <- H
    // H <- C
    // C <- W

    const auto order = DimsOrder::NCDHW;
    auto origMapping = PermuteDims{{Dim::N, Dim::C}, {Dim::C, Dim::D}, {Dim::D, Dim::H}, {Dim::H, Dim::W}, {Dim::W, Dim::N}};
    auto permVector = permuteMapToVector(origMapping, order, order);
    {
        ASSERT_EQ(permVector, (PermutationIndexVector{4, 0, 1, 2, 3}));
        auto identityCheckMapping = permuteVectorToMap(permVector, order, order);
        ASSERT_EQ(origMapping, identityCheckMapping);
    }

    {
        auto permVectorCombined = combinePermutationVectors(permVector, permVector);
        ASSERT_EQ(permVectorCombined, (PermutationIndexVector{3, 4, 0, 1, 2}));
        auto combinedMapping = permuteVectorToMap(permVectorCombined, order, order);
        auto identityCombinedMapping = PermuteDims{{Dim::N, Dim::D}, {Dim::C, Dim::H}, {Dim::D, Dim::W}, {Dim::H, Dim::N}, {Dim::W, Dim::C}};
        ASSERT_EQ(combinedMapping, identityCombinedMapping);

    }
    {
        auto permVectorCombined = combinePermutationVectors(PermutationIndexVector{0, 1, 2, 4, 3}, PermutationIndexVector{0, 1, 3, 2, 4});
        ASSERT_EQ(permVectorCombined, (PermutationIndexVector{0, 1, 4, 2, 3}));
    }

    {
        auto permVector1 = permuteMapToVector(PermuteDims{{Dim::H, Dim::C}, {Dim::W, Dim::H}, {Dim::C, Dim::W}},
                                              DimsOrder::CHW, DimsOrder::CHW);
        ASSERT_EQ(permVector1, (PermutationIndexVector{1, 2, 0}));
        auto permVector2 = permuteMapToVector(PermuteDims{{Dim::H, Dim::C}, {Dim::C, Dim::H}, {Dim::W, Dim::W}},
                                              DimsOrder::CHW, DimsOrder::CHW);
        ASSERT_EQ(permVector2, (PermutationIndexVector{0, 2, 1}));

        auto permVectorCombined = combinePermutationVectors(permVector1, permVector2);
        ASSERT_EQ(permVectorCombined, (PermutationIndexVector{1, 0, 2}));
    }
}

TEST_F(VPU_MergePermuteTest, TwoPermutes) {
    // TLDR: {H, C} + {W, H} => {W, C}.
    // How does permutation should combine?
    // according to stages/permute.cpp, parsePermute() stage, in  PermutationDimsMap first is destination dimension,
    // and the second value in pair - is source dimension.
    // for {Dim::W, Dim::H}, data in W output will be relevant to H dimension.
    // If we have {Dim::H, Dim::C} in first permute, and {Dim::W, Dim::H} in the second, the resulting permutation is
    // {Dim::W, Dim::C}  - basically, we taking the dest (first in pair) from the second permute,
    // and the source (second in pair) value  from the first permute.

    // clockwise permutes on 4 dims.
    CreateModelWithTwoPermutes(DimsOrder::NCHW,
                               PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::N}, {Dim::N, Dim::W}},
                               PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::N}, {Dim::N, Dim::W}});
    ASSERT_NO_THROW(pipeline.run(model));
    ASSERT_EQ(model->getStages().size(), 1U);
    ASSERT_EQ(GetPermuteDimsForStage(0), (PermuteDims{{Dim::W, Dim::C}, {Dim::H, Dim::N}, {Dim::C, Dim::W}, {Dim::N, Dim::H}}));

    // counter-clockwise permutes on 5 dims
    CreateModelWithTwoPermutes(DimsOrder::NCDHW,
                               PermuteDims{{Dim::N, Dim::C}, {Dim::C, Dim::D}, {Dim::D, Dim::H}, {Dim::H, Dim::W}, {Dim::W, Dim::N}},
                               PermuteDims{{Dim::N, Dim::C}, {Dim::C, Dim::D}, {Dim::D, Dim::H}, {Dim::H, Dim::W}, {Dim::W, Dim::N}});
    ASSERT_NO_THROW(pipeline.run(model));
    ASSERT_EQ(model->getStages().size(), 1U);
    ASSERT_EQ(GetPermuteDimsForStage(0), (PermuteDims{{Dim::N, Dim::D}, {Dim::C, Dim::H}, {Dim::D, Dim::W}, {Dim::H, Dim::N}, {Dim::W, Dim::C}}));

    // partial compensation (3 dim)
    CreateModelWithTwoPermutes(DimsOrder::CHW,
                               PermuteDims{{Dim::H, Dim::C}, {Dim::W, Dim::H}, {Dim::C, Dim::W}},
                               PermuteDims{{Dim::C, Dim::H}, {Dim::W, Dim::W}, {Dim::H, Dim::C}});
    ASSERT_NO_THROW(pipeline.run(model));
    ASSERT_EQ(model->getStages().size(), 1U);
    ASSERT_EQ(GetPermuteDimsForStage(0), (PermuteDims{{Dim::H, Dim::W}, {Dim::W, Dim::H}, {Dim::C, Dim::C}}));

    // two compensating permutes (3 dim)
    CreateModelWithTwoPermutes(DimsOrder::CHW,
                               PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::W}},
                               PermuteDims{{Dim::H, Dim::W}, {Dim::C, Dim::H}, {Dim::W, Dim::C}});
    ASSERT_NO_THROW(pipeline.run(model));
    ASSERT_EQ(model->getStages().size(), 1U);
    ASSERT_EQ(model->getStages().front()->type(), StageType::Copy);
}

TEST_F(VPU_MergePermuteTest, ThreePermutes) {
    // three compensating permutes (3 dim)
    CreateModelWithThreePermutes(DimsOrder::CHW,
                                 PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::W}},
                                 PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::W}},
                                 PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::W}});
    ASSERT_NO_THROW(pipeline.run(model));
    ASSERT_EQ(model->getStages().size(), 1U);
    ASSERT_EQ(model->getStages().front()->type(), StageType::Copy);

    // three merging permutes (5 dim)
    CreateModelWithThreePermutes(DimsOrder::NCDHW,
                                 PermuteDims{{Dim::N, Dim::C}, {Dim::C, Dim::N}, {Dim::D, Dim::D}, {Dim::H, Dim::H}, {Dim::W, Dim::W}},
                                 PermuteDims{{Dim::N, Dim::N}, {Dim::C, Dim::D}, {Dim::D, Dim::C}, {Dim::H, Dim::H}, {Dim::W, Dim::W}},
                                 PermuteDims{{Dim::N, Dim::N}, {Dim::C, Dim::C}, {Dim::D, Dim::H}, {Dim::H, Dim::D}, {Dim::W, Dim::W}});
    ASSERT_NO_THROW(pipeline.run(model));
    ASSERT_EQ(model->getStages().size(), 1U);
    ASSERT_EQ(GetPermuteDimsForStage(0), (PermuteDims{{Dim::H, Dim::N}, {Dim::C, Dim::D}, {Dim::N, Dim::C}, {Dim::D, Dim::H}, {Dim::W, Dim::W}}));
}

TEST_F(VPU_MergePermuteTest, TwoSequences) {
    model = CreateModel();
    const auto layout = DimsOrder::CHW;
    const DataDesc dataDesc(DataType::FP16, layout, MakeStubDims(layout));
    auto input = model->addInputData("Input", dataDesc);
    model->attrs().set<int>("numInputs", 1);

    auto output = model->addOutputData("Output", dataDesc);
    model->attrs().set<int>("numOutputs", 1);

    // sequence: input - Permute11 - data1 - Permute12 - data2 - Copy - data3 - Permute21 - data4 - Permute22 - output

    auto data1 = model->addNewData("data1", dataDesc);
    auto data2 = model->addNewData("data2", dataDesc);
    auto data3 = model->addNewData("data3", dataDesc);
    auto data4 = model->addNewData("data4", dataDesc);

    stageBuilder->addPermuteStage(model,
                                  "Permute11",
                                  nullptr,
                                  input,
                                  data1,
                                  PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::W}}
                                );
    stageBuilder->addPermuteStage(model,
                                  "Permute12",
                                  nullptr,
                                  data1,
                                  data2,
                                  PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::W}}
                                );

    stageBuilder->addCopyStage(model,
                               "Copy",
                               nullptr,
                               data2,
                               data3,
                               ""
                               );

    stageBuilder->addPermuteStage(model,
                                  "Permute21",
                                  nullptr,
                                  data3,
                                  data4,
                                  PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::W}}
                                );
    stageBuilder->addPermuteStage(model,
                                  "Permute22",
                                  nullptr,
                                  data4,
                                  output,
                                  PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::W}}
                                );
    InitPipeline();

    ASSERT_NO_THROW(pipeline.run(model));
    ASSERT_EQ(model->getStages().size(), 3U); // 2 Permute and 1 copy stage.
}

TEST_F(VPU_MergePermuteTest, OutputBeteweenPermutes) {
    model = CreateModel();
    const auto layout = DimsOrder::CHW;
    const DataDesc dataDesc(DataType::FP16, layout, MakeStubDims(layout));
    auto input = model->addInputData("Input", dataDesc);
    model->attrs().set<int>("numInputs", 1);

    auto output1 = model->addOutputData("Output1", dataDesc);
    auto output2 = model->addOutputData("Output2", dataDesc);
    model->attrs().set<int>("numOutputs", 2);

    // sequence: input - Permute1 - Output1 - Permute2 - Output2

    stageBuilder->addPermuteStage(model,
                                  "Permute1",
                                  nullptr,
                                  input,
                                  output1,
                                  PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::W}}
                                );
    stageBuilder->addPermuteStage(model,
                                  "Permute2",
                                  nullptr,
                                  output1,
                                  output2,
                                  PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::W}}
                                );
    InitPipeline();

    ASSERT_NO_THROW(pipeline.run(model));
    ASSERT_EQ(model->getStages().size(), 2U);
}

TEST_F(VPU_MergePermuteTest, DataSourceSplitBeteweenPermutes) {
    model = CreateModel();
    const auto layout = DimsOrder::CHW;
    const DataDesc dataDesc(DataType::FP16, layout, MakeStubDims(layout));
    auto input = model->addInputData("Input", dataDesc);
    model->attrs().set<int>("numInputs", 1);

    auto output1 = model->addOutputData("Output1", dataDesc);
    auto output2 = model->addOutputData("Output2", dataDesc);
    model->attrs().set<int>("numOutputs", 2);

    auto data1 = model->addNewData("data1", dataDesc);

    // sequence: input - Permute1 - data1 - Permute2 - Output1
    //                                    \ Copy     - Output2

    stageBuilder->addPermuteStage(model,
                                  "Permute1",
                                  nullptr,
                                  input,
                                  data1,
                                  PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::W}}
                                );
    stageBuilder->addPermuteStage(model,
                                  "Permute2",
                                  nullptr,
                                  data1,
                                  output1,
                                  PermuteDims{{Dim::W, Dim::H}, {Dim::H, Dim::C}, {Dim::C, Dim::W}}
                                );
    stageBuilder->addCopyStage(model, "Copy", nullptr, data1, output2, "");
    InitPipeline();

    ASSERT_NO_THROW(pipeline.run(model));
    ASSERT_EQ(model->getStages().size(), 3U);
}

TEST_F(VPU_MergePermuteTest, PermuteAndReorder) {
    CreateModelWithPermuteAndReorder(DimsOrder::CHW,
                                     DimsOrder::HWC,
                                     PermuteDims{{Dim::W, Dim::W}, {Dim::H, Dim::C}, {Dim::C, Dim::H}});
    ASSERT_NO_THROW(pipeline.run(model));
    ASSERT_EQ(model->getStages().size(), 1U);
    ASSERT_EQ(GetPermuteDimsForStage(0), (PermuteDims{{Dim::W, Dim::W}, {Dim::H, Dim::C}, {Dim::C, Dim::H}}));
}
