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

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

#include <armnn/BackendRegistry.hpp>

#include <armnn/backends/IBackendContext.hpp>
#include <armnn/backends/IMemoryManager.hpp>
#include <backendsCommon/DefaultAllocator.hpp>

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

#include <algorithm>

namespace
{

bool IsLayerSupported(const armnn::Layer* layer)
{
    ARMNN_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)
{
    ARMNN_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
{

MockBackendInitialiser::MockBackendInitialiser()
{
    BackendRegistryInstance().Register(MockBackend::GetIdStatic(),
                                       []()
                                       {
                                           return IBackendInternalUniquePtr(new MockBackend);
                                       });
}

MockBackendInitialiser::~MockBackendInitialiser()
{
    try
    {
        BackendRegistryInstance().Deregister(MockBackend::GetIdStatic());
    }
    catch (...)
    {
        std::cerr << "could not deregister mock backend" << std::endl;
    }
}

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::IBackendProfilingContextPtr MockBackend::CreateBackendProfilingContext(
    const IRuntime::CreationOptions& options, IBackendProfilingPtr& backendProfiling)
{
    IgnoreUnused(options);
    std::shared_ptr<armnn::MockBackendProfilingContext> context =
        std::make_shared<MockBackendProfilingContext>(backendProfiling);
    MockBackendProfilingService::Instance().SetProfilingContextPtr(context);
    return context;
}

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

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::IConnectableLayers& subgraphLayers = subgraph.GetIConnectableLayers();

    // Parse the layers
    SubgraphView::IConnectableLayers supportedLayers;
    SubgraphView::IConnectableLayers unsupportedLayers;
    SubgraphView::IConnectableLayers untouchedLayers;
    std::for_each(subgraphLayers.begin(),
                  subgraphLayers.end(),
                  [&](IConnectableLayer* layer)
    {
        bool supported = IsLayerSupported(PolymorphicDowncast<Layer*>(layer));
        if (supported)
        {
            // Layer supported, check if it's optimizable
            bool optimizable = IsLayerOptimizable(PolymorphicDowncast<Layer*>(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)
        {
            ARMNN_ASSERT(supportedSubgraph != nullptr);

            CompiledBlobPtr blobPtr;
            BackendId backend = MockBackendId();

            IConnectableLayer* preCompiledLayer =
                optimizationViews.GetINetwork()->AddPrecompiledLayer(
                        PreCompiledDescriptor(supportedSubgraph->GetNumInputSlots(),
                                              supportedSubgraph->GetNumOutputSlots()),
                                              std::move(blobPtr),
                                              backend,
                                              nullptr);

            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)
        {
            ARMNN_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)
        {
            ARMNN_ASSERT(untouchedSubgraph != nullptr);

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

    return optimizationViews;
}

std::unique_ptr<ICustomAllocator> MockBackend::GetDefaultAllocator() const
{
    return std::make_unique<DefaultAllocator>();
}

} // namespace armnn