/*
 * Copyright (C) 2019 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#define LOG_TAG "Operations"

#include <algorithm>
#include <cfloat>
#include <limits>
#include <vector>

#include "OperationResolver.h"
#include "Tracing.h"
#include "nnapi/Validation.h"

#ifdef NN_INCLUDE_CPU_IMPLEMENTATION
#include <tensorflow/lite/kernels/internal/optimized/legacy_optimized_ops.h>
#include <tensorflow/lite/kernels/internal/optimized/optimized_ops.h>

#include "CpuOperationUtils.h"
#endif  // NN_INCLUDE_CPU_IMPLEMENTATION

namespace android {
namespace nn {

namespace softmax {

constexpr char kOperationName[] = "SOFTMAX";

constexpr uint32_t kNumInputs = 3;
constexpr uint32_t kInputTensor = 0;
constexpr uint32_t kBetaScalar = 1;
constexpr uint32_t kAxisScalar = 2;

constexpr uint32_t kNumOutputs = 1;
constexpr uint32_t kOutputTensor = 0;

#ifdef NN_INCLUDE_CPU_IMPLEMENTATION
namespace {

inline bool softmaxSlowFloat32(const float* inputData, const Shape& inputShape, const float beta,
                               int32_t axis, float* outputData, const Shape& outputShape) {
    NNTRACE_TRANS("softmaxFloatSlow32");
    const uint32_t outerSize = getNumberOfElements(inputShape, 0, axis);
    const uint32_t axisSize = getSizeOfDimension(inputShape, axis);
    const uint32_t innerSize =
            getNumberOfElements(inputShape, axis + 1, getNumberOfDimensions(inputShape));
    for (uint32_t outer = 0; outer < outerSize; ++outer) {
        const float* inputBeg = inputData + outer * axisSize * innerSize;
        const float* inputEnd = inputBeg + axisSize * innerSize;
        float* outputBeg = outputData + outer * axisSize * innerSize;
        for (uint32_t inner = 0; inner < innerSize; ++inner, ++inputBeg, ++inputEnd, ++outputBeg) {
            // Find max
            float maxValue = -FLT_MAX;
            for (const float* p = inputBeg; p < inputEnd; p += innerSize) {
                maxValue = std::max(maxValue, *p);
            }
            // Compute sum
            float sum = 0.0f;
            for (const float* p = inputBeg; p < inputEnd; p += innerSize) {
                sum += std::exp((*p - maxValue) * beta);
            }
            // Compute result
            float* pOut = outputBeg;
            for (const float* p = inputBeg; p < inputEnd; p += innerSize, pOut += innerSize) {
                *pOut = std::exp((*p - maxValue) * beta) / sum;
            }
        }
    }
    return true;
}

bool softmaxFloat32(const float* inputData, const Shape& inputShape, const float beta, int32_t axis,
                    float* outputData, const Shape& outputShape) {
    int32_t ndim = getNumberOfDimensions(inputShape);
    NN_CHECK(handleNegativeAxis(inputShape, &axis));
    // TFLite optimized implementation only supports computation along the last axis
    if (axis == ndim - 1) {
        NNTRACE_COMP("optimized_ops::Softmax::float");
        tflite::SoftmaxParams param = {.beta = beta};
        tflite::optimized_ops::Softmax(param, convertShapeToTflshape(inputShape), inputData,
                                       convertShapeToTflshape(outputShape), outputData);
        return true;
    } else {
        return softmaxSlowFloat32(inputData, inputShape, beta, axis, outputData, outputShape);
    }
}

bool softmaxFloat16(const _Float16* inputData, const Shape& inputShape, const float beta,
                    int32_t axis, _Float16* outputData, const Shape& outputShape) {
    NNTRACE_TRANS("softmaxFloat16");
    std::vector<float> inputData_float32(getNumberOfElements(inputShape));
    convertFloat16ToFloat32(inputData, &inputData_float32);
    std::vector<float> outputData_float32(getNumberOfElements(outputShape));

    softmaxFloat32(inputData_float32.data(), inputShape, beta, axis, outputData_float32.data(),
                   outputShape);
    convertFloat32ToFloat16(outputData_float32, outputData);

    return true;
}

template <typename T>
bool softmaxQuant8Impl(const T* inputData, const Shape& inputShape, const float beta, int32_t axis,
                       int32_t inputMultiplier, int32_t inputLeftShift, float diffMin,
                       T* outputData, const Shape& outputShape) {
    NNTRACE_TRANS("softmaxQuant8");
    // The representation chosen for the input to the exp() function is Q5.26.
    // We need to leave extra space since values that we skip might be as large as
    // -32 before multiplying by input_beta_multiplier, and therefore as large as
    // -16 afterwards.  Note that exp(-8) is definitely not insignificant to
    // accumulation, but exp(-16) definitely is.
    static const int32_t kScaledDiffIntegerBits = 5;
    static const int kAccumulationIntegerBits = 12;
    using FixedPointScaledDiff = gemmlowp::FixedPoint<int32_t, kScaledDiffIntegerBits>;
    using FixedPointAccum = gemmlowp::FixedPoint<int32_t, kAccumulationIntegerBits>;
    using FixedPoint0 = gemmlowp::FixedPoint<int32_t, 0>;

    const uint32_t outerSize = getNumberOfElements(inputShape, 0, axis);
    const uint32_t axisSize = getSizeOfDimension(inputShape, axis);
    const uint32_t innerSize =
            getNumberOfElements(inputShape, axis + 1, getNumberOfDimensions(inputShape));
    for (uint32_t outer = 0; outer < outerSize; ++outer) {
        const T* inputBeg = inputData + outer * axisSize * innerSize;
        const T* inputEnd = inputBeg + axisSize * innerSize;
        T* outputBeg = outputData + outer * axisSize * innerSize;
        for (uint32_t inner = 0; inner < innerSize; ++inner, ++inputBeg, ++inputEnd, ++outputBeg) {
            // Find max
            T maxValue = std::is_same_v<T, int8_t> ? -128 : 0;
            for (const T* p = inputBeg; p < inputEnd; p += innerSize) {
                maxValue = std::max(maxValue, *p);
            }

            // Compute sum
            FixedPointAccum sum_of_exps = FixedPointAccum::Zero();
            for (const T* p = inputBeg; p < inputEnd; p += innerSize) {
                int32_t input_diff = static_cast<int32_t>(*p) - maxValue;
                if (input_diff >= diffMin) {
                    const int32_t input_diff_rescaled =
                            tflite::MultiplyByQuantizedMultiplierGreaterThanOne(
                                    input_diff, inputMultiplier, inputLeftShift);
                    const auto scaled_diff_f8 = FixedPointScaledDiff::FromRaw(input_diff_rescaled);
                    sum_of_exps = sum_of_exps + gemmlowp::Rescale<kAccumulationIntegerBits>(
                                                        exp_on_negative_values(scaled_diff_f8));
                }
            }

            uint32_t fixed_sum_of_exps = static_cast<uint32_t>(sum_of_exps.raw());
            int32_t headroom_plus_one = tflite::CountLeadingZeros(fixed_sum_of_exps);
            // This is the number of bits to the left of the binary point above 1.0.
            // Consider fixed_sum_of_exps=1.25.  In that case shifted_scale=0.8 and
            // no later adjustment will be needed.
            int32_t num_bits_over_unit = kAccumulationIntegerBits - headroom_plus_one;
            int32_t shifted_sum_minus_one = static_cast<int32_t>(
                    (fixed_sum_of_exps << headroom_plus_one) - (static_cast<uint32_t>(1) << 31));

            FixedPoint0 shifted_scale = gemmlowp::one_over_one_plus_x_for_x_in_0_1(
                    FixedPoint0::FromRaw(shifted_sum_minus_one));

            // Compute result
            constexpr int32_t q_min = std::numeric_limits<T>::min();
            constexpr int32_t q_max = std::numeric_limits<T>::max();
            T* pOut = outputBeg;
            for (const T* p = inputBeg; p < inputEnd; p += innerSize, pOut += innerSize) {
                int32_t input_diff = static_cast<int32_t>(*p) - maxValue;
                if (input_diff >= diffMin) {
                    const int32_t input_diff_rescaled =
                            tflite::MultiplyByQuantizedMultiplierGreaterThanOne(
                                    input_diff, inputMultiplier, inputLeftShift);
                    const auto scaled_diff_f8 = FixedPointScaledDiff::FromRaw(input_diff_rescaled);

                    FixedPoint0 exp_in_0 = exp_on_negative_values(scaled_diff_f8);
                    int32_t unsat_output = gemmlowp::RoundingDivideByPOT(
                            (shifted_scale * exp_in_0).raw(), num_bits_over_unit + 31 - 8);
                    if (std::is_same_v<T, int8_t>) {
                        unsat_output -= 128;
                    }

                    *pOut = static_cast<T>(std::max(std::min(unsat_output, q_max), q_min));

                } else {
                    *pOut = std::is_same_v<T, int8_t> ? -128 : 0;
                }
            }
        }
    }
    return true;
}

template <typename T>
bool softmaxQuant8(const T* inputData, const Shape& inputShape, const float beta, int32_t axis,
                   T* outputData, const Shape& outputShape) {
    int32_t ndim = getNumberOfDimensions(inputShape);
    NN_CHECK(handleNegativeAxis(inputShape, &axis));

    if ((inputShape.type == OperandType::TENSOR_QUANT8_ASYMM && outputShape.offset != 0) ||
        (inputShape.type == OperandType::TENSOR_QUANT8_ASYMM_SIGNED &&
         outputShape.offset != -128) ||
        outputShape.scale != 1.f / 256) {
        LOG(ERROR) << "incorrect scale / offset for output";
        return false;
    }

    static const int32_t kScaledDiffIntegerBits = 5;
    const double input_beta_real_multiplier =
            std::min(1.0 * beta * inputShape.scale * (1 << (31 - kScaledDiffIntegerBits)),
                     (1LL << 31) - 1.0);

    int32_t inputMultiplier = 0, inputLeftShift = 0;
    if (!QuantizeMultiplierGreaterThanOne(input_beta_real_multiplier, &inputMultiplier,
                                          &inputLeftShift)) {
        return false;
    }
    int32_t diffMin = -CalculateInputRadius(kScaledDiffIntegerBits, inputLeftShift);

    return softmaxQuant8Impl(inputData, inputShape, beta, axis, inputMultiplier, inputLeftShift,
                             diffMin, outputData, outputShape);
}

}  // namespace
#endif  // NN_INCLUDE_CPU_IMPLEMENTATION

Result<Version> validate(const IOperationValidationContext* context) {
    NN_RET_CHECK(context->getNumInputs() == kNumInputs ||
                 context->getNumInputs() == kNumInputs - 1);
    NN_RET_CHECK_EQ(context->getNumOutputs(), kNumOutputs);
    auto inputType = context->getInputType(kInputTensor);
    std::vector<OperandType> inExpectedTypes;
    auto minSupportedVersion = Version::ANDROID_OC_MR1;
    if (inputType == OperandType::TENSOR_FLOAT32 || inputType == OperandType::TENSOR_QUANT8_ASYMM) {
        minSupportedVersion = Version::ANDROID_OC_MR1;
        inExpectedTypes = {inputType, OperandType::FLOAT32};
    } else if (inputType == OperandType::TENSOR_FLOAT16) {
        minSupportedVersion = Version::ANDROID_Q;
        inExpectedTypes = {inputType, OperandType::FLOAT16};
    } else if (inputType == OperandType::TENSOR_QUANT8_ASYMM_SIGNED) {
        minSupportedVersion = Version::ANDROID_R;
        inExpectedTypes = {inputType, OperandType::FLOAT32};
    } else {
        NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation " << kOperationName;
    }
    const auto inputRank = getNumberOfDimensions(context->getInputShape(kInputTensor));
    if (inputRank != 0) {
        NN_RET_CHECK_LE(inputRank, 4);
    }
    if (context->getNumInputs() == kNumInputs) {
        minSupportedVersion = combineVersions(minSupportedVersion, Version::ANDROID_Q);
        inExpectedTypes.push_back(OperandType::INT32);
    } else {
        if (inputRank != 2 && inputRank != 4 && inputRank != 0) {
            minSupportedVersion = combineVersions(minSupportedVersion, Version::ANDROID_Q);
        }
    }
    NN_RET_CHECK(validateInputTypes(context, inExpectedTypes));
    NN_RET_CHECK(validateOutputTypes(context, {inputType}));
    return minSupportedVersion;
}

#ifdef NN_INCLUDE_CPU_IMPLEMENTATION
bool prepare(IOperationExecutionContext* context) {
    Shape input = context->getInputShape(kInputTensor);
    float beta = (input.type == OperandType::TENSOR_FLOAT16)
                         ? context->getInputValue<_Float16>(kBetaScalar)
                         : context->getInputValue<float>(kBetaScalar);
    NN_RET_CHECK_LE(getNumberOfDimensions(input), 4);
    NN_RET_CHECK_GT(beta, 0.0f);
    Shape output = context->getOutputShape(kOutputTensor);
    output.dimensions = input.dimensions;
    return context->setOutputShape(kOutputTensor, output);
}

bool execute(IOperationExecutionContext* context) {
    // Bypass execution in the case of zero-sized input.
    if (getNumberOfElements(context->getOutputShape(kOutputTensor)) == 0) return true;
    int32_t axis = (context->getNumInputs() == kNumInputs)
                           ? context->getInputValue<int32_t>(kAxisScalar)
                           : -1;
    switch (context->getInputType(kInputTensor)) {
        case OperandType::TENSOR_FLOAT16:
            return softmaxFloat16(context->getInputBuffer<_Float16>(kInputTensor),
                                  context->getInputShape(kInputTensor),
                                  context->getInputValue<_Float16>(kBetaScalar), axis,
                                  context->getOutputBuffer<_Float16>(kOutputTensor),
                                  context->getOutputShape(kOutputTensor));
        case OperandType::TENSOR_FLOAT32:
            return softmaxFloat32(context->getInputBuffer<float>(kInputTensor),
                                  context->getInputShape(kInputTensor),
                                  context->getInputValue<float>(kBetaScalar), axis,
                                  context->getOutputBuffer<float>(kOutputTensor),
                                  context->getOutputShape(kOutputTensor));
        case OperandType::TENSOR_QUANT8_ASYMM:
            return softmaxQuant8(context->getInputBuffer<uint8_t>(kInputTensor),
                                 context->getInputShape(kInputTensor),
                                 context->getInputValue<float>(kBetaScalar), axis,
                                 context->getOutputBuffer<uint8_t>(kOutputTensor),
                                 context->getOutputShape(kOutputTensor));
        case OperandType::TENSOR_QUANT8_ASYMM_SIGNED:
            return softmaxQuant8(context->getInputBuffer<int8_t>(kInputTensor),
                                 context->getInputShape(kInputTensor),
                                 context->getInputValue<float>(kBetaScalar), axis,
                                 context->getOutputBuffer<int8_t>(kOutputTensor),
                                 context->getOutputShape(kOutputTensor));
        default:
            NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation " << kOperationName;
    }
}
#endif  // NN_INCLUDE_CPU_IMPLEMENTATION

}  // namespace softmax

NN_REGISTER_OPERATION(SOFTMAX, "SOFTMAX", softmax::validate, softmax::prepare, softmax::execute,
                      .allowZeroSizedInput = true);

}  // namespace nn
}  // namespace android