//
// Copyright © 2022-2023 Arm Ltd and Contributors. All rights reserved.
// SPDX-License-Identifier: MIT
//

#include "GpuFsaBackend.hpp"
#include "GpuFsaBackendContext.hpp"
#include "GpuFsaBackendDefaultAllocator.hpp"
#include "GpuFsaBackendId.hpp"
#include "GpuFsaLayerSupport.hpp"
#include "GpuFsaTensorHandleFactory.hpp"
#include "GpuFsaWorkloadFactory.hpp"

#include <armnn/backends/IBackendContext.hpp>
#include <armnn/backends/IMemoryManager.hpp>
#include <Optimizer.hpp>

#include <aclCommon/BaseMemoryManager.hpp>

#include <arm_compute/runtime/CL/CLBufferAllocator.h>

namespace armnn
{

const BackendId& GpuFsaBackend::GetIdStatic()
{
    static const BackendId s_Id{GpuFsaBackendId()};
    return s_Id;
}

IBackendInternal::IMemoryManagerUniquePtr GpuFsaBackend::CreateMemoryManager() const
{
    if (m_UsingCustomAllocator)
    {
        return std::make_unique<GpuFsaMemoryManager>(m_CustomAllocator);
    }
    return std::make_unique<GpuFsaMemoryManager>(std::make_unique<arm_compute::CLBufferAllocator>());
}

IBackendInternal::IWorkloadFactoryPtr GpuFsaBackend::CreateWorkloadFactory(
    const IBackendInternal::IMemoryManagerSharedPtr& memoryManager) const
{
    return std::make_unique<GpuFsaWorkloadFactory>(PolymorphicPointerDowncast<GpuFsaMemoryManager>(memoryManager));
}

IBackendInternal::IWorkloadFactoryPtr GpuFsaBackend::CreateWorkloadFactory(
    TensorHandleFactoryRegistry& registry) const
{
    std::shared_ptr<GpuFsaMemoryManager> memoryManager;
    if (m_UsingCustomAllocator)
    {
        memoryManager = std::make_shared<GpuFsaMemoryManager>(m_CustomAllocator);
    }
    else
    {
        memoryManager = std::make_shared<GpuFsaMemoryManager>(std::make_unique<arm_compute::CLBufferAllocator>());
    }

    std::unique_ptr<ITensorHandleFactory> factory = std::make_unique<GpuFsaTensorHandleFactory>(memoryManager);

    registry.RegisterMemoryManager(memoryManager);
    registry.RegisterFactory(std::move(factory));

    return std::make_unique<GpuFsaWorkloadFactory>(PolymorphicPointerDowncast<GpuFsaMemoryManager>(memoryManager));
}

IBackendInternal::IWorkloadFactoryPtr GpuFsaBackend::CreateWorkloadFactory(
    TensorHandleFactoryRegistry& registry,
    const ModelOptions& modelOptions,
    MemorySourceFlags inputFlags,
    MemorySourceFlags outputFlags) const
{
    IgnoreUnused(modelOptions);

    // To allow force import if inputFlags/outputFlags are Undefined, set it as Malloc
    if (inputFlags == static_cast<MemorySourceFlags>(MemorySource::Undefined))
    {
        inputFlags = static_cast<MemorySourceFlags>(MemorySource::Malloc);
    }
    if (outputFlags == static_cast<MemorySourceFlags>(MemorySource::Undefined))
    {
        outputFlags = static_cast<MemorySourceFlags>(MemorySource::Malloc);
    }

    std::shared_ptr<GpuFsaMemoryManager> memoryManager;
    if (m_UsingCustomAllocator)
    {
        memoryManager = std::make_shared<GpuFsaMemoryManager>(m_CustomAllocator);
    }
    else
    {
        memoryManager = std::make_shared<GpuFsaMemoryManager>(std::make_unique<arm_compute::CLBufferAllocator>());
    }

    std::unique_ptr<ITensorHandleFactory> factory = std::make_unique<GpuFsaTensorHandleFactory>(memoryManager);

    registry.RegisterMemoryManager(memoryManager);
    registry.RegisterFactory(std::move(factory));

    return std::make_unique<GpuFsaWorkloadFactory>(PolymorphicPointerDowncast<GpuFsaMemoryManager>(memoryManager));
}

std::vector<ITensorHandleFactory::FactoryId> GpuFsaBackend::GetHandleFactoryPreferences() const
{
    return std::vector<ITensorHandleFactory::FactoryId> { GpuFsaTensorHandleFactory::GetIdStatic() };
}

void GpuFsaBackend::RegisterTensorHandleFactories(TensorHandleFactoryRegistry& registry)
{
    std::shared_ptr<GpuFsaMemoryManager> memoryManager;
    if (m_UsingCustomAllocator)
    {
        memoryManager = std::make_shared<GpuFsaMemoryManager>(m_CustomAllocator);
    }
    else
    {
        memoryManager = std::make_shared<GpuFsaMemoryManager>(std::make_unique<arm_compute::CLBufferAllocator>());
    }

    std::unique_ptr<ITensorHandleFactory> factory = std::make_unique<GpuFsaTensorHandleFactory>(memoryManager);
    registry.RegisterMemoryManager(memoryManager);
    registry.RegisterFactory(std::move(factory));

}

void GpuFsaBackend::RegisterTensorHandleFactories(TensorHandleFactoryRegistry& registry,
                                                  MemorySourceFlags inputFlags,
                                                  MemorySourceFlags outputFlags)
{
    // To allow force import if inputFlags/outputFlags are Undefined, set it as Malloc
    if (inputFlags == static_cast<MemorySourceFlags>(MemorySource::Undefined))
    {
        inputFlags = static_cast<MemorySourceFlags>(MemorySource::Malloc);
    }
    if (outputFlags == static_cast<MemorySourceFlags>(MemorySource::Undefined))
    {
        outputFlags = static_cast<MemorySourceFlags>(MemorySource::Malloc);
    }

    std::shared_ptr<GpuFsaMemoryManager> memoryManager;
    if (m_UsingCustomAllocator)
    {
        memoryManager = std::make_shared<GpuFsaMemoryManager>(m_CustomAllocator);
    }
    else
    {
        memoryManager = std::make_shared<GpuFsaMemoryManager>(std::make_unique<arm_compute::CLBufferAllocator>());
    }

    std::unique_ptr<ITensorHandleFactory> factory = std::make_unique<GpuFsaTensorHandleFactory>(memoryManager);
    registry.RegisterMemoryManager(memoryManager);
    registry.RegisterFactory(std::move(factory));
}

IBackendInternal::IBackendContextPtr GpuFsaBackend::CreateBackendContext(const IRuntime::CreationOptions& options) const
{
    return IBackendContextPtr{new GpuFsaBackendContext{options}};
}

IBackendInternal::IBackendProfilingContextPtr GpuFsaBackend::CreateBackendProfilingContext(
    const IRuntime::CreationOptions&, IBackendProfilingPtr&)
{
    return IBackendProfilingContextPtr{};
}

IBackendInternal::ILayerSupportSharedPtr GpuFsaBackend::GetLayerSupport() const
{
    static ILayerSupportSharedPtr layerSupport{new GpuFsaLayerSupport};
    return layerSupport;
}

std::unique_ptr<ICustomAllocator> GpuFsaBackend::GetDefaultAllocator() const
{
    return std::make_unique<GpuFsaBackendDefaultAllocator>();
}

OptimizationViews GpuFsaBackend::OptimizeSubgraphView(const SubgraphView& subgraph,
                                                      const ModelOptions& modelOptions) const
{
    OptimizationViews optimizationViews(modelOptions);
    optimizationViews.AddUntouchedSubgraph(SubgraphView(subgraph));
    return optimizationViews;
}

} // namespace armnn