/*
 * Copyright (C) 2017 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.
 */

#include "CpuOperationUtils.h"
#include "OperationResolver.h"

#include "tensorflow/lite/kernels/internal/optimized/legacy_optimized_ops.h"
#include "tensorflow/lite/kernels/internal/reference/legacy_reference_ops.h"

#include "Tracing.h"

namespace android {
namespace nn {
namespace concatenation {

constexpr char kOperationName[] = "CONCATENATION";

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

namespace {

template <typename T>
bool concatenation(const std::vector<const T*>& inputDataPtrs,
                   const std::vector<Shape>& inputShapes, int32_t axis, T* outputData,
                   const Shape& outputShape) {
    NNTRACE_TRANS("concatenation");
    int num_inputs = inputShapes.size();
    std::vector<tflite::Dims<4>*> inputDimsPtr(num_inputs);
    std::vector<tflite::Dims<4> > inputDims(num_inputs);
    for (int i = 0; i < num_inputs; i++) {
        inputDims[i] = convertShapeToDims(inputShapes[i]);
        inputDimsPtr[i] = &inputDims[i];
    }
    NNTRACE_COMP_SWITCH("optimized_ops::Concatenation");
    tflite::optimized_ops::Concatenation<tflite::FusedActivationFunctionType::kNone, T>(
            getNumberOfDimensions(outputShape) - axis - 1, inputDataPtrs.data(),
            inputDimsPtr.data(), num_inputs, outputData, convertShapeToDims(outputShape));

    return true;
}

template <>
bool concatenation<uint8_t>(const std::vector<const uint8_t*>& inputDataPtrs,
                            const std::vector<Shape>& inputShapes, int32_t axis,
                            uint8_t* outputData, const Shape& outputShape) {
    NNTRACE_TRANS("concatenationQuant8");
    int num_inputs = inputShapes.size();
    std::vector<float> inputScales(num_inputs);
    std::vector<int32> inputOffsets(num_inputs);
    std::vector<tflite::Dims<4>*> inputDimsPtr(num_inputs);
    std::vector<tflite::Dims<4> > inputDims(num_inputs);
    for (int i = 0; i < num_inputs; i++) {
        inputScales[i] = inputShapes[i].scale;
        inputOffsets[i] = inputShapes[i].offset;
        inputDims[i] = convertShapeToDims(inputShapes[i]);
        inputDimsPtr[i] = &inputDims[i];
    }

    NNTRACE_COMP_SWITCH("reference_ops::Concatenation");
    tflite::reference_ops::Concatenation(
            getNumberOfDimensions(outputShape) - axis - 1, inputDataPtrs.data(),
            inputDimsPtr.data(), inputOffsets.data(), inputScales.data(), num_inputs, outputData,
            convertShapeToDims(outputShape), outputShape.offset, outputShape.scale);

    return true;
}

template <typename T>
inline bool concatenation(IOperationExecutionContext* context) {
    uint32_t inputCount = context->getNumInputs() - 1;
    std::vector<const T*> inputDatas;
    std::vector<Shape> inputShapes;
    for (uint32_t i = 0; i < inputCount; ++i) {
        const T* buffer = context->getInputBuffer<T>(i);
        if (buffer == nullptr) continue;
        inputDatas.push_back(buffer);
        inputShapes.push_back(context->getInputShape(i));
    }
    return concatenation(inputDatas, inputShapes, context->getInputValue<int32_t>(inputCount),
                         context->getOutputBuffer<T>(kOutputTensor),
                         context->getOutputShape(kOutputTensor));
}

}  // namespace

bool validate(const IOperationValidationContext* context) {
    uint32_t inputCount = context->getNumInputs();
    NN_RET_CHECK_GE(inputCount, 2);
    NN_RET_CHECK_EQ(context->getNumOutputs(), kNumOutputs);
    const OperandType inputType = context->getInputType(0);
    if (inputType == OperandType::TENSOR_FLOAT32 || inputType == OperandType::TENSOR_QUANT8_ASYMM) {
        NN_RET_CHECK(validateHalVersion(context, HalVersion::V1_0));
    } else if (inputType == OperandType::TENSOR_FLOAT16) {
        NN_RET_CHECK(validateHalVersion(context, HalVersion::V1_2));
    } else {
        NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation " << kOperationName;
    }
    std::vector<OperandType> inExpectedTypes(inputCount - 1, inputType);
    inExpectedTypes.push_back(OperandType::INT32);
    if (context->getHalVersion() < HalVersion::V1_2 &&
        inputType == OperandType::TENSOR_QUANT8_ASYMM) {
        const Shape& output = context->getOutputShape(kOutputTensor);
        for (uint32_t i = 0; i < inputCount - 1; ++i) {
            const Shape& input = context->getInputShape(i);
            NN_RET_CHECK_EQ(input.scale, output.scale);
            NN_RET_CHECK_EQ(input.offset, output.offset);
        }
    }
    return validateInputTypes(context, inExpectedTypes) &&
           validateOutputTypes(context, {inputType});
}

bool prepare(IOperationExecutionContext* context) {
    uint32_t numInputs = context->getNumInputs();
    NN_RET_CHECK_GE(numInputs, 2);
    const Shape& input0 = context->getInputShape(0);
    uint32_t numDimensions = getNumberOfDimensions(input0);
    int32_t axis = context->getInputValue<int32_t>(numInputs - 1);
    NN_RET_CHECK_GE(axis, 0);
    NN_RET_CHECK_LT(axis, numDimensions);

    uint32_t sumAxis = getSizeOfDimension(input0, axis);
    for (uint32_t i = 1; i < numInputs - 1; ++i) {
        const Shape& input = context->getInputShape(i);
        NN_RET_CHECK_EQ(getNumberOfDimensions(input), numDimensions);
        NN_RET_CHECK(input.type == input0.type);
        for (uint32_t d = 0; d < numDimensions; ++d) {
            if (d == axis) {
                sumAxis += getSizeOfDimension(input, axis);
            } else {
                NN_RET_CHECK_EQ(getSizeOfDimension(input0, d), getSizeOfDimension(input, d));
            }
        }
    }

    Shape output = context->getOutputShape(kOutputTensor);
    output.type = input0.type;
    output.dimensions = input0.dimensions;
    output.dimensions[axis] = sumAxis;
    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;
    switch (context->getInputType(0)) {
        case OperandType::TENSOR_FLOAT16:
            return concatenation<_Float16>(context);
        case OperandType::TENSOR_FLOAT32:
            return concatenation<float>(context);
        case OperandType::TENSOR_QUANT8_ASYMM:
            return concatenation<uint8_t>(context);
        default:
            NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation " << kOperationName;
    }
}

}  // namespace concatenation

NN_REGISTER_OPERATION(CONCATENATION, concatenation::kOperationName, concatenation::validate,
                      concatenation::prepare, concatenation::execute, .allowZeroSizedInput = true);

}  // namespace nn
}  // namespace android