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

#include "OutputShapeUtils.hpp"

#include <algorithm>

namespace armnn_driver
{

using namespace armnn;

TensorShape InferPreluOutputShape(const TensorShape& inputShape, const TensorShape& alphaShape)
{
    // NOTE: The inferred PReLU output size will be the maximum size along each dimension
    // of input and alpha, starting with the trailing dimensions, and working its way forward.
    //
    // Example: inputShape={4, 1, 2}, alphaShape={5, 4, 3, 1} => outputShape={5, 4, 3, 2}

    const unsigned int numInputDims = inputShape.GetNumDimensions();
    const unsigned int numAlphaDims = alphaShape.GetNumDimensions();

    const unsigned int maxNumDims = std::max(numInputDims, numAlphaDims);

    TensorShape outputShape = TensorShape(maxNumDims);
    for (unsigned int reverseIdx = 1u; reverseIdx <= maxNumDims; ++reverseIdx)
    {
        const int inputIdx = numInputDims - reverseIdx;
        const int alphaIdx = numAlphaDims - reverseIdx;

        const unsigned int inputDimSize = inputIdx >= 0 ? inputShape[inputIdx] : 0u;
        const unsigned int alphaDimSize = alphaIdx >= 0 ? alphaShape[alphaIdx] : 0u;

        const unsigned int outputIdx = maxNumDims - reverseIdx;
        outputShape[outputIdx] = std::max(inputDimSize, alphaDimSize);
    }

    return outputShape;
}

} // namespace armnn_driver