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

#pragma once

#include <armnn/Tensor.hpp>

#include <armnnUtils/DataLayoutIndexed.hpp>

namespace armnn
{

// Utility class providing access to raw tensor memory based on indices along each dimension.
template <typename DataType>
class TensorBufferArrayView
{
public:
    TensorBufferArrayView(const TensorShape& shape, DataType* data,
                          armnnUtils::DataLayoutIndexed dataLayout = DataLayout::NCHW)
        : m_Shape(shape)
        , m_Data(data)
        , m_DataLayout(dataLayout)
    {
        ARMNN_THROW_INVALIDARG_MSG_IF_FALSE(m_Shape.GetNumDimensions() == 4,
                                            "Only $d tensors are supported by TensorBufferArrayView.");
    }

    DataType& Get(unsigned int b, unsigned int c, unsigned int h, unsigned int w) const
    {
        return m_Data[m_DataLayout.GetIndex(m_Shape, b, c, h, w)];
    }

private:
    const TensorShape             m_Shape;
    DataType*                     m_Data;
    armnnUtils::DataLayoutIndexed m_DataLayout;
};

} //namespace armnn