/********************************************************************************
 * Copyright (c) 2023 CEA-List
 *
 * This program and the accompanying materials are made available under the
 * terms of the Eclipse Public License 2.0 which is available at
 * http://www.eclipse.org/legal/epl-2.0.
 *
 * SPDX-License-Identifier: EPL-2.0
 *
 ********************************************************************************/

#include <catch2/catch_test_macros.hpp>
#include <set>

#include "aidge/data/Tensor.hpp"
#include "aidge/graph/GraphView.hpp"
#include "aidge/operator/GenericOperator.hpp"
#include "aidge/operator/FC.hpp"
#include "aidge/recipies/Recipies.hpp"

namespace Aidge {


TEST_CASE("[cpu/recipies] RemoveFlatten", "[RemoveFlatten][recipies]") {
    // generate the original GraphView
    auto flatten = GenericOperator("Flatten", 1, 0, 1, "myFlatten");
    auto fc = FC(10, 50, "myFC");

    flatten -> addChild(fc);

    auto g = std::make_shared<GraphView>();
    g->add({fc, flatten});

    // Check original graph
    // g -> save("before_remove_flatten");

    // use recipie
    removeFlatten(g);

    // Check transformed graph
    // g -> save("after_remove_flatten");

    REQUIRE(g->getOrderedInputs().size() == 1);
    REQUIRE(g->getOrderedOutputs().size() == 1);
    REQUIRE(g->getOrderedInputs()[0].first == fc);
    REQUIRE(g->getOrderedOutputs()[0].first == fc);
}

}  // namespace Aidge