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

#include "SpaceToDepth.hpp"

#include <DataLayoutIndexed.hpp>

using namespace armnnUtils;

namespace {
    unsigned int GetOffset(const armnn::TensorShape& shape,
        unsigned int c,
        unsigned int h,
        unsigned int w,
        unsigned int b,
        const DataLayoutIndexed& dataLayout)
    {
        if (dataLayout.GetDataLayout() == armnn::DataLayout::NHWC)
        {
            return ((b * shape[dataLayout.GetHeightIndex()] + h) * shape[dataLayout.GetWidthIndex()] + w) *
                shape[dataLayout.GetChannelsIndex()] + c;
        }
        else
        {
            return ((b * shape[dataLayout.GetChannelsIndex()] + c) * shape[dataLayout.GetHeightIndex()] + h) *
                shape[dataLayout.GetWidthIndex()] + w;
        }
    }
}

namespace armnn
{

void SpaceToDepth(const TensorInfo& inputInfo,
                  const TensorInfo& outputInfo,
                  const SpaceToDepthDescriptor& params,
                  Decoder<float>& inputData,
                  Encoder<float>& outputData)
{
    DataLayoutIndexed dataLayout = params.m_DataLayout;

    const TensorShape& inputShape = inputInfo.GetShape();
    const TensorShape& outputShape = outputInfo.GetShape();

    const unsigned int inputBatchSize = inputShape[0];
    const unsigned int inputChannels = inputShape[dataLayout.GetChannelsIndex()];

    const unsigned int outputHeight = outputShape[dataLayout.GetHeightIndex()];
    const unsigned int outputWidth = outputShape[dataLayout.GetWidthIndex()];
    const unsigned int outputChannels = outputShape[dataLayout.GetChannelsIndex()];

    const unsigned int blockSize = params.m_BlockSize;

    if (blockSize == 0)
    {
        throw InvalidArgumentException(
            "Input shape must be divisible by block size in all spatial dimensions: Block size is"
            " equal to zero");
    }

    for (unsigned int outChannelIndex = 0; outChannelIndex < outputChannels; outChannelIndex++)
    {
        unsigned int inChannelIndex = outChannelIndex % inputChannels;

        unsigned int shiftW = (outChannelIndex / inputChannels) % blockSize;
        unsigned int shiftH = (outChannelIndex / inputChannels) / blockSize;

        for (unsigned int outH = 0; outH < outputHeight; outH++)
        {
            for (unsigned int outW = 0; outW < outputWidth; outW++)
            {
                for (unsigned int inBatchIndex = 0; inBatchIndex < inputBatchSize; inBatchIndex++)
                {
                    unsigned int inOffset = GetOffset(inputShape,
                        inChannelIndex,
                        (outH * blockSize + shiftH),
                        (outW * blockSize + shiftW),
                        inBatchIndex,
                        dataLayout);

                    unsigned int outOffset = GetOffset(outputShape,
                        outChannelIndex,
                        outH,
                        outW,
                        inBatchIndex,
                        dataLayout);

                    outputData += outOffset;
                    inputData += inOffset;
                    outputData.Set(inputData.Get());
                    inputData -= inOffset;
                    outputData -= outOffset;
                }
            }
        }
    }
}

void SpaceToDepth(const TensorInfo& inputInfo,
    const TensorInfo& outputInfo,
    const SpaceToDepthDescriptor& params,
    Decoder<float>& inputData,
    Encoder<float>& outData);

} //namespace armnn