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

#include "MockBackend.hpp"
#include "MockBackendId.hpp"

#include <armnn/BackendRegistry.hpp>

#include <backendsCommon/IBackendContext.hpp>
#include <backendsCommon/IMemoryManager.hpp>

#include <Optimizer.hpp>
#include <SubgraphViewSelector.hpp>

#include <boost/cast.hpp>

#include <algorithm>

namespace
{

bool IsLayerSupported(const armnn::Layer* layer)
{
    BOOST_ASSERT(layer != nullptr);

    armnn::LayerType layerType = layer->GetType();
    switch (layerType)
    {
    case armnn::LayerType::Input:
    case armnn::LayerType::Output:
    case armnn::LayerType::Addition:
    case armnn::LayerType::Convolution2d:
        // Layer supported
        return true;
    default:
        // Layer unsupported
        return false;
    }
}

bool IsLayerSupported(const armnn::Layer& layer)
{
    return IsLayerSupported(&layer);
}

bool IsLayerOptimizable(const armnn::Layer* layer)
{
    BOOST_ASSERT(layer != nullptr);

    // A Layer is not optimizable if its name contains "unoptimizable"
    const std::string layerName(layer->GetName());
    bool optimizable = layerName.find("unoptimizable") == std::string::npos;

    return optimizable;
}

bool IsLayerOptimizable(const armnn::Layer& layer)
{
    return IsLayerOptimizable(&layer);
}

} // Anonymous namespace

namespace armnn
{

namespace
{

static BackendRegistry::StaticRegistryInitializer g_RegisterHelper
{
    BackendRegistryInstance(),
    MockBackend::GetIdStatic(),
    []()
    {
        return IBackendInternalUniquePtr(new MockBackend);
    }
};

}

const BackendId& MockBackend::GetIdStatic()
{
    static const BackendId s_Id{MockBackendId()};
    return s_Id;
}

IBackendInternal::IWorkloadFactoryPtr MockBackend::CreateWorkloadFactory(
    const IBackendInternal::IMemoryManagerSharedPtr& memoryManager) const
{
    return IWorkloadFactoryPtr{};
}

IBackendInternal::IBackendContextPtr MockBackend::CreateBackendContext(const IRuntime::CreationOptions&) const
{
    return IBackendContextPtr{};
}

IBackendInternal::IMemoryManagerUniquePtr MockBackend::CreateMemoryManager() const
{
    return IMemoryManagerUniquePtr{};
}

IBackendInternal::Optimizations MockBackend::GetOptimizations() const
{
    return Optimizations{};
}

IBackendInternal::ILayerSupportSharedPtr MockBackend::GetLayerSupport() const
{
    static ILayerSupportSharedPtr layerSupport{new MockLayerSupport};
    return layerSupport;
}

OptimizationViews MockBackend::OptimizeSubgraphView(const SubgraphView& subgraph) const
{
    // Prepare the optimization views
    OptimizationViews optimizationViews;

    // Get the layers of the input sub-graph
    const SubgraphView::Layers& subgraphLayers = subgraph.GetLayers();

    // Parse the layers
    SubgraphView::Layers supportedLayers;
    SubgraphView::Layers unsupportedLayers;
    SubgraphView::Layers untouchedLayers;
    std::for_each(subgraphLayers.begin(),
                  subgraphLayers.end(),
                  [&](Layer* layer)
    {
        bool supported = IsLayerSupported(layer);
        if (supported)
        {
            // Layer supported, check if it's optimizable
            bool optimizable = IsLayerOptimizable(layer);
            if (optimizable)
            {
                // Layer fully supported
                supportedLayers.push_back(layer);
            }
            else
            {
                // Layer supported but not optimizable
                untouchedLayers.push_back(layer);
            }
        }
        else
        {
            // Layer unsupported
            unsupportedLayers.push_back(layer);
        }
    });

    // Check if there are supported layers
    if (!supportedLayers.empty())
    {
        // Select the layers that are neither inputs or outputs, but that are optimizable
        auto supportedSubgraphSelector = [](const Layer& layer)
        {
            return layer.GetType() != LayerType::Input &&
                   layer.GetType() != LayerType::Output &&
                   IsLayerSupported(layer) &&
                   IsLayerOptimizable(layer);
        };

        // Apply the subgraph selector to the supported layers to group them into sub-graphs were appropriate
        SubgraphView mutableSubgraph(subgraph);
        SubgraphViewSelector::Subgraphs supportedSubgraphs =
                SubgraphViewSelector::SelectSubgraphs(mutableSubgraph, supportedSubgraphSelector);

        // Create a substitution pair for each supported sub-graph
        std::for_each(supportedSubgraphs.begin(),
                      supportedSubgraphs.end(),
                      [&optimizationViews](const SubgraphView::SubgraphViewPtr& supportedSubgraph)
        {
            BOOST_ASSERT(supportedSubgraph != nullptr);

            PreCompiledLayer* preCompiledLayer =
                optimizationViews.GetGraph().AddLayer<PreCompiledLayer>(
                        PreCompiledDescriptor(supportedSubgraph->GetNumInputSlots(),
                                              supportedSubgraph->GetNumOutputSlots()),
                        "pre-compiled");
            preCompiledLayer->SetBackendId(MockBackendId());

            SubgraphView substitutionSubgraph(*supportedSubgraph);
            SubgraphView replacementSubgraph(preCompiledLayer);

            optimizationViews.AddSubstitution({ substitutionSubgraph, replacementSubgraph });
        });
    }

    // Check if there are unsupported layers
    if (!unsupportedLayers.empty())
    {
        // Select the layers that are neither inputs or outputs, and are not optimizable
        auto unsupportedSubgraphSelector = [](const Layer& layer)
        {
            return layer.GetType() != LayerType::Input &&
                   layer.GetType() != LayerType::Output &&
                   !IsLayerSupported(layer);
        };

        // Apply the subgraph selector to the unsupported layers to group them into sub-graphs were appropriate
        SubgraphView mutableSubgraph(subgraph);
        SubgraphViewSelector::Subgraphs unsupportedSubgraphs =
                SubgraphViewSelector::SelectSubgraphs(mutableSubgraph, unsupportedSubgraphSelector);

        // Add each unsupported sub-graph to the list of failed sub-graphs in the optimizization views
        std::for_each(unsupportedSubgraphs.begin(),
                      unsupportedSubgraphs.end(),
                      [&optimizationViews](const SubgraphView::SubgraphViewPtr& unsupportedSubgraph)
        {
            BOOST_ASSERT(unsupportedSubgraph != nullptr);

            optimizationViews.AddFailedSubgraph(SubgraphView(*unsupportedSubgraph));
        });
    }

    // Check if there are untouched layers
    if (!untouchedLayers.empty())
    {
        // Select the layers that are neither inputs or outputs, that are supported but that and are not optimizable
        auto untouchedSubgraphSelector = [](const Layer& layer)
        {
            return layer.GetType() != LayerType::Input &&
                   layer.GetType() != LayerType::Output &&
                   IsLayerSupported(layer) &&
                   !IsLayerOptimizable(layer);
        };

        // Apply the subgraph selector to the untouched layers to group them into sub-graphs were appropriate
        SubgraphView mutableSubgraph(subgraph);
        SubgraphViewSelector::Subgraphs untouchedSubgraphs =
                SubgraphViewSelector::SelectSubgraphs(mutableSubgraph, untouchedSubgraphSelector);

        // Add each untouched sub-graph to the list of untouched sub-graphs in the optimizization views
        std::for_each(untouchedSubgraphs.begin(),
                      untouchedSubgraphs.end(),
                      [&optimizationViews](const SubgraphView::SubgraphViewPtr& untouchedSubgraph)
        {
            BOOST_ASSERT(untouchedSubgraph != nullptr);

            optimizationViews.AddUntouchedSubgraph(SubgraphView(*untouchedSubgraph));
        });
    }

    return optimizationViews;
}

} // namespace armnn