//
// Copyright © 2017, 2019, 2021-2023 Arm Ltd and Contributors. All rights reserved.
// SPDX-License-Identifier: MIT
//

#pragma once

#include <armnn/BackendOptions.hpp>
#include <armnn/Deprecated.hpp>
#include <armnn/INetwork.hpp>
#include <armnn/backends/SubgraphView.hpp>
#include <vector>

namespace armnn
{

class OptimizationViews
{
public:
    OptimizationViews(const NetworkOptions& networkOptions = {}) : m_INetwork(INetwork::Create(networkOptions)) {}
    OptimizationViews(const OptimizationViews&) = delete;
    OptimizationViews& operator=(const OptimizationViews&) = delete;
    OptimizationViews(OptimizationViews&&) = default;
    OptimizationViews& operator=(OptimizationViews&&) = default;

    struct SubstitutionPair
    {
        /// Subgraph of Layers from the original graph which should be replaced
        SubgraphView m_SubstitutableSubgraph;

        /// A subgraph of new layers which will replace layers in m_SubstitutableSubgraph
        SubgraphView m_ReplacementSubgraph;
    };

    using Subgraphs = std::vector<SubgraphView>;
    using Substitutions = std::vector<SubstitutionPair>;

    void AddSubstitution(SubstitutionPair&& substitution)
    {
        m_SuccesfulOptimizations.emplace_back(substitution);
    }

    void AddFailedSubgraph(SubgraphView&& subgraph)
    {
        m_FailedOptimizations.emplace_back(subgraph);
    }

    void AddUntouchedSubgraph(SubgraphView&& subgraph)
    {
        m_UntouchedSubgraphs.emplace_back(subgraph);
    }

    void AddDeletedSubgraph(SubgraphView&& subgraph)
    {
        m_DeletedSubgraphs.emplace_back(subgraph);
    }

    const Substitutions& GetSubstitutions() const { return m_SuccesfulOptimizations; }
    const Subgraphs& GetFailedSubgraphs() const { return m_FailedOptimizations; }
    const Subgraphs& GetUntouchedSubgraphs() const { return m_UntouchedSubgraphs; }
    const Subgraphs& GetDeletedSubgraphs() const { return m_DeletedSubgraphs; }

    Substitutions& GetSubstitutions() { return m_SuccesfulOptimizations; }
    Subgraphs& GetFailedSubgraphs() { return m_FailedOptimizations; }
    Subgraphs& GetUntouchedSubgraphs() { return m_UntouchedSubgraphs; }

    bool Validate(const SubgraphView& originalSubgraph) const;

    INetwork* GetINetwork() { return m_INetwork.get(); }
    INetwork& GetINetworkRef() { return *m_INetwork; }

private:
    Substitutions m_SuccesfulOptimizations;     ///< Proposed substitutions from successful optimizations
    Subgraphs m_FailedOptimizations;            ///< Subgraphs from the original subgraph which cannot be supported
    Subgraphs m_UntouchedSubgraphs;             ///< Subgraphs from the original subgraph which remain unmodified
    Subgraphs m_DeletedSubgraphs;               ///< Subgraphs from the original subgraph which have been deleted

    /// INetworkPtr object used only as a container for any layer generated by the optimization process
    /// Also, can use to AddPrecompiledLayer to the SubstitutionPair
    INetworkPtr m_INetwork = INetwork::Create();
};

} //namespace armnn