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

#pragma once

#include "BaseIterator.hpp"
#include <backendsCommon/CpuTensorHandle.hpp>

namespace
{

// Helper functions ported from the Android code base
// Refer to: android/external/tensorflow/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.cc

void MatrixBatchVectorMultiplyAccumulate(armnn::Decoder<float>& matrix,
                                         uint32_t mRows,
                                         uint32_t mCols,
                                         armnn::Decoder<float>& vector,
                                         uint32_t nBatch,
                                         armnn::Encoder<float>& outResult)
{
    for (uint32_t b = 0; b < nBatch; b++)
    {
        for (uint32_t r = 0; r < mRows; r++)
        {
            vector += b * mCols;
            for (uint32_t c = 0; c < mCols; c++)
            {
                outResult.Set(outResult.Get() + matrix.Get() * vector.Get());
                ++matrix;
                ++vector;
            }
            outResult += 1;
            vector -= (b+1) * mCols;
        }
        matrix -= (mRows * mCols);
    }
    outResult -= (mRows * nBatch);
}

void VectorBatchVectorAssign(armnn::Decoder<float>& vector,
                             uint32_t vSize,
                             uint32_t nBatch,
                             armnn::Encoder<float>& outBatchVector)
{
    for (uint32_t b = 0; b < nBatch; b++)
    {
        for (uint32_t v = 0; v < vSize; v++)
        {
            outBatchVector.Set(vector.Get());
            ++outBatchVector;
            ++vector;
        }
        vector -= vSize;
    }
    outBatchVector -= (nBatch * vSize);
}

void VectorBatchVectorCwiseProductAccumulate(armnn::Decoder<float>& vector,
                                             uint32_t vSize,
                                             armnn::Decoder<float>& batchVector,
                                             uint32_t nBatch,
                                             armnn::Encoder<float>& outResult)
{
    for (uint32_t b = 0; b < nBatch; b++)
    {
        for (uint32_t v = 0; v < vSize; v++)
        {
            outResult.Set(outResult.Get() + vector.Get() * batchVector.Get());
            ++outResult;
            ++vector;
            ++batchVector;
        }
        vector -= vSize;
    }
    batchVector -= vSize * nBatch;
    outResult -= vSize * nBatch;
}

void Sub1Vector(armnn::Decoder<float>& vector,
                uint32_t vSize,
                armnn::Encoder<float>& result)
{
    for (uint32_t v = 0; v < vSize; v++)
    {
        result.Set(1.0f - vector.Get());
        ++vector;
        ++result;
    }
    vector -= vSize;
    result -= vSize;
}

void VectorVectorCwiseProduct(armnn::Decoder<float>& vector1,
                              armnn::Decoder<float>& vector2,
                              uint32_t vSize,
                              armnn::Encoder<float>& outResult)
{
    for (uint32_t v = 0; v < vSize; v++)
    {
        outResult.Set(vector1.Get() * vector2.Get());
        ++outResult;
        ++vector1;
        ++vector2;
    }
    outResult -= vSize;
    vector1 -= vSize;
    vector2 -= vSize;
}

void VectorVectorCwiseProductAccumulate(armnn::Decoder<float>& vector1,
                                        armnn::Decoder<float>& vector2,
                                        uint32_t vSize,
                                        armnn::Encoder<float>& outResult)
{
    for (uint32_t v = 0; v < vSize; v++)
    {
        outResult.Set(outResult.Get() + vector1.Get() * vector2.Get());
        ++outResult;
        ++vector1;
        ++vector2;
    }
    outResult -= vSize;
    vector1 -= vSize;
    vector2 -= vSize;
}

float Clip(float f,
           float absLimit)
{
    float result = (absLimit < f) ? absLimit : f;
    result = (-absLimit > result) ? -absLimit : result;
    return result;
}

void ClipVector(armnn::Decoder<float>& vector,
                uint32_t vSize,
                float absLimit,
                armnn::Encoder<float>& outResult)
{
    for (uint32_t v = 0; v < vSize; v++)
    {
        outResult.Set(Clip(vector.Get(), absLimit));
        ++vector;
        ++outResult;
    }
    vector -= vSize;
    outResult -= vSize;
}

void CopyVector(armnn::Decoder<float>& vector,
                uint32_t vSize,
                armnn::Encoder<float>& outResult)
{
    for (uint32_t v = 0; v < vSize; v++)
    {
        outResult.Set(vector.Get());
        ++outResult;
        ++vector;
    }
    outResult -= vSize;
    vector -= vSize;
}

void SetActivationParameters(uint32_t activation,
                             armnn::ActivationFunction& outArmnnActivation,
                             float& outA,
                             float& outB)
{
    switch (activation)
    {
    case 0: // None
        outA = 0;
        outB = 0;
        return;

    case 1: // Relu
        outArmnnActivation = armnn::ActivationFunction::ReLu;
        outA = 0;
        outB = 0;
        return;

    case 3: // Relu6
        outArmnnActivation = armnn::ActivationFunction::BoundedReLu;
        outA = 6;
        outB = 0;
        return;

    case 4: // Tanh
        outArmnnActivation = armnn::ActivationFunction::TanH;
        outA = 1;
        outB = 1;
        return;

    case 6: // Sigmoid
        outArmnnActivation = armnn::ActivationFunction::Sigmoid;
        outA = 0;
        outB = 0;
        return;

    default:
        throw armnn::Exception("Unsupported activation function: " + std::to_string(activation));
    }
}

std::unique_ptr<armnn::ScopedCpuTensorHandle> AssignScopedCpuTensorHandle(const armnn::ConstCpuTensorHandle* ptr)
{
    if (!ptr)
    {
        return nullptr;
    }

    return std::make_unique<armnn::ScopedCpuTensorHandle>(*ptr);
}

} // anonymous namespace