//
// Copyright © 2017 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//

#include "SubGraphSelector.hpp"
#include "Graph.hpp"
#include <boost/assert.hpp>
#include <algorithm>
#include <unordered_map>

namespace armnn
{

namespace
{

struct LayerSelectionInfo
{
    using LayerInfoContainer = std::unordered_map<Layer*, LayerSelectionInfo>;
    static constexpr uint32_t InitialSplitId() { return 1; }

    LayerSelectionInfo(Layer* layer, const SubGraphSelector::LayerSelectorFunction& selector)
    : m_Layer{layer}
    , m_SplitId{0}
    , m_IsSelected{selector(*layer)}
    {
        // fill topology information by storing direct children
        for (auto&& slot = m_Layer->BeginOutputSlots(); slot != m_Layer->EndOutputSlots(); ++slot)
        {
            for (InputSlot* childLayerInputSlot : slot->GetConnections())
            {
                Layer& childLayer = childLayerInputSlot->GetOwningLayer();
                m_DirectChildren.push_back(&childLayer);
            }
        }
    }

    void MarkChildrenSplits(LayerInfoContainer& network,
                            uint32_t splitId,
                            bool prevSelected)
    {
        if (m_SplitId < splitId)
        {
            m_SplitId = splitId;
        }

        // introduce a new split point at all non-selected points, but only if the
        // previous point was selected. this prevents creating a new subgraph at
        // every non-selected layer
        if (!m_IsSelected && prevSelected)
        {
            ++m_SplitId;
        }

        for (auto& layer : m_DirectChildren)
        {
            auto it = network.find(layer);
            BOOST_ASSERT_MSG(it != network.end(), "All layers must be part of the topology.");
            if (it != network.end())
            {
                it->second.MarkChildrenSplits(network, m_SplitId, m_IsSelected);
            }
        }
    }

    bool IsInputLayer() const
    {
        return m_Layer->GetType() == armnn::LayerType::Input;
    }

    void CollectNonSelectedInputs(SubGraph::InputSlots& slots,
                                  const SubGraphSelector::LayerSelectorFunction& selector)
    {
        for (auto&& slot = m_Layer->BeginInputSlots(); slot != m_Layer->EndInputSlots(); ++slot)
        {
            OutputSlot* parentLayerOutputSlot = slot->GetConnectedOutputSlot();
            BOOST_ASSERT_MSG(parentLayerOutputSlot != nullptr, "The slots must be connected here.");
            if (parentLayerOutputSlot)
            {
                Layer& parentLayer = parentLayerOutputSlot->GetOwningLayer();
                if (selector(parentLayer) == false)
                {
                    slots.push_back(&(*slot));
                }
            }
        }
    }

    void CollectNonSelectedOutputSlots(SubGraph::OutputSlots& slots,
                                       const SubGraphSelector::LayerSelectorFunction& selector)
    {
        for (auto&& slot = m_Layer->BeginOutputSlots(); slot != m_Layer->EndOutputSlots(); ++slot)
        {
            for (InputSlot* childLayerInputSlot : slot->GetConnections())
            {
                Layer& childLayer = childLayerInputSlot->GetOwningLayer();
                if (selector(childLayer) == false)
                {
                    slots.push_back(&(*slot));
                }
            }
        }
    }

    std::vector<Layer*> m_DirectChildren;
    Layer* m_Layer;
    uint32_t m_SplitId;
    bool m_IsSelected;
};

} // namespace <anonymous>

SubGraphSelector::SubGraphs
SubGraphSelector::SelectSubGraphs(Graph& graph,
                                  const LayerSelectorFunction& selector)
{
    LayerSelectionInfo::LayerInfoContainer layerInfo;

    for (auto& layer : graph)
    {
        layerInfo.emplace(layer, LayerSelectionInfo{layer, selector});
    }

    uint32_t splitNo = LayerSelectionInfo::InitialSplitId();
    for (auto& info : layerInfo)
    {
        if (info.second.IsInputLayer())
        {
            // for each input layer we mark the graph where subgraph
            // splits need to happen because of the dependency between
            // the selected and non-selected nodes
            info.second.MarkChildrenSplits(layerInfo, splitNo, false);
        }
    }

    // Collect all selected layers keyed by split id into a map
    using SelectionInfoPtrs = std::vector<LayerSelectionInfo*>;
    std::unordered_map<uint32_t, SelectionInfoPtrs> splitMap;
    for (auto& info : layerInfo)
    {
        if (info.second.m_IsSelected)
        {
            auto it = splitMap.find(info.second.m_SplitId);
            if (it == splitMap.end())
            {
                splitMap.insert(std::make_pair(info.second.m_SplitId, SelectionInfoPtrs{&info.second}));
            }
            else
            {
                it->second.push_back(&info.second);
            }
        }
    }

    // Now each non-empty split id represents a subgraph
    SubGraphs result;
    for (auto& splitGraph : splitMap)
    {
        if (splitGraph.second.empty() == false)
        {
            SubGraph::OutputSlots outputs;
            SubGraph::InputSlots inputs;
            SubGraph::Layers layers;
            for (auto&& infoPtr : splitGraph.second)
            {
                infoPtr->CollectNonSelectedOutputSlots(outputs, selector);
                infoPtr->CollectNonSelectedInputs(inputs, selector);
                layers.insert(infoPtr->m_Layer);
            }
            result.emplace_back(
                std::make_unique<SubGraph>(
                    std::move(inputs),
                    std::move(outputs),
                    std::move(layers)));
        }
    }

    return result;
}

} // namespace armnn