/********************************************************************************
 * 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 <memory>

#include "aidge/graph/Node.hpp"
#include "aidge/graph/GraphView.hpp"
#include "aidge/utils/Recipies.hpp"

// Graph Regex
#include "aidge/graphmatching/GRegex.hpp"
#include "aidge/graphmatching/NodeRegex.hpp"


namespace Aidge {
    void removeFlatten(std::set<std::shared_ptr<Node>> nodes) {
        assert(nodes.size() == 2 && "Wrong number of nodes to replace\n");
        std::shared_ptr<Node> flatten;
        for (const auto& element : nodes) {
            assert((element->type() == "FC" || element->type() == "Flatten") && "Wrong type for the nodes to replace");
            if (element->type() == "Flatten"){
                flatten = element;
            }
        }
        auto g = std::make_shared<GraphView>();
        // TODO : avoid using replace_with and use a remove method instead
        g->add(std::set<std::shared_ptr<Node>>({flatten}));
        g->replaceWith({});
    }

    void removeFlatten(std::shared_ptr<GraphView> graphView){
        std::map<std::string,NodeRegex*> nodesRegex ;
        nodesRegex["Flatten"] = new NodeRegex("Flatten");
        nodesRegex["FC"] = new NodeRegex("FC");
        std::vector<std::string> seqRegex;
        seqRegex.push_back("Flatten->FC;");
        GRegex GReg(nodesRegex, seqRegex);
        Match matches = GReg.match(graphView);
        std::vector<std::set<std::shared_ptr<Node>>> matchNodes = matches.getMatchNodes();
        for (size_t i = 0; i < matches.getNbMatch(); ++i) {
            removeFlatten(matchNodes[i]);
        }
    }
}