//
//  QNNBinary.cpp
//  MNN
//
//  Created by MNN on b'2025/04/10'.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#include "QNNBinary.hpp"
#include <algorithm>

namespace MNN {
namespace QNN {
#ifdef ENABLE_QNN_ONLINE_FINALIZE

static bool needChangeInputOrder(const std::string& binaryTypeName) {
    std::set<std::string> NoneedChangeSet = {"ElementWiseAdd",     "ElementWiseMultiply", "ElementWiseMinimum",
                                             "ElementWiseMaximum", "ElementWiseOr",       "ElementWiseEqual",
                                             "ElementWiseNotEqual"};
    return NoneedChangeSet.find(binaryTypeName) == NoneedChangeSet.end();
}

ErrorCode QNNBinary::onEncode(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
    int dim0 = inputs[0]->dimensions();
    int dim1 = inputs[1]->dimensions();
    int minDim = dim0 > dim1 ? dim1 : dim0;
    int fullIndex = dim0 > dim1 ? 0 : 1;

    // Broadcast binary with scalar.
    // By our experiments, this branch is faster than using Qnn binary operations directly, although Qnn binary operations supports scalar broadcasting.
    if(dim0 != dim1 && minDim == 0) {
        return this->onEncodeScalarOptimize(inputs, outputs, fullIndex);
    }
    
    if (dim0 != dim1 && TensorUtils::getDimType(inputs[0]) != TensorUtils::getDimType(inputs[1])) {
        fullIndex = TensorUtils::getDimType(inputs[0]) != TensorUtils::getDimType(outputs[0]) ? 1 : 0;
        return this->onEncodeBroadcast(inputs, outputs, fullIndex);
    }

    mNodeType = mBinaryTypeName;
    this->addNodeCommon(inputs, outputs);

    return NO_ERROR;
}

ErrorCode QNNBinary::onEncodeScalarOptimize(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs, int fullIndex) {
    std::vector<uint32_t> shape = getNHWCShape(inputs[fullIndex]);
    int idleIndex = 1 - fullIndex;
    Qnn_DataType_t dataType = mBackend->getNativeTensor(inputs[fullIndex])->v1.dataType;

    std::vector<uint32_t> dim(inputs[fullIndex]->dimensions(), 1);

    this->createStageTensor("stage_0", dataType, dim, inputs[fullIndex]); // mTempTensorWrappers[0]
    this->createStageTensor("stage_1", dataType, shape, inputs[fullIndex]); // mTempTensorWrappers[1]

    this->createParamTensor("multiples", QNN_DATATYPE_UINT_32, {(uint32_t)shape.size()},
                            shape.data()); // mParamTensorWrappers[0]

    // Reshape
    {
        CLEAR_BEFORE_ADDING_NODE;

        std::string name = mNodeName + "_Reshape";
        mNodeType = "Reshape";

        mInputs.push_back(*(mBackend->getNativeTensor(inputs[idleIndex]))); // idle input
        mOutputs.push_back(*(mTempTensorWrappers[0]->getNativeTensor())); // stage 0

        mBackend->addNodeToGraph(mOpConfigVersion, name.c_str(), mPackageName.c_str(), mNodeType.c_str(), mParams, mInputs, mOutputs);
    }

    // Tile
    {
        CLEAR_BEFORE_ADDING_NODE;

        std::string name = mNodeName + "_Tile";
        mNodeType = "Tile";

        mInputs.push_back(*(mTempTensorWrappers[0]->getNativeTensor())); // stage 0
        mParams.push_back(*(mParamTensorWrappers[0]->getNativeParam())); // multiples
        mOutputs.push_back(*(mTempTensorWrappers[1]->getNativeTensor())); // stage 1

        mBackend->addNodeToGraph(mOpConfigVersion, name.c_str(), mPackageName.c_str(), mNodeType.c_str(), mParams, mInputs, mOutputs);
    }

    // Binary
    // Ensure correct input order for non-commutative operations (Sub, Div, Pow, etc.)
    // input0Tensor corresponds to original inputs[0], input1Tensor corresponds to original inputs[1]
    const auto& input0Tensor = (fullIndex == 0) ? *(mBackend->getNativeTensor(inputs[fullIndex]))
                                                : *(mTempTensorWrappers[1]->getNativeTensor());
    const auto& input1Tensor = (fullIndex == 0) ? *(mTempTensorWrappers[1]->getNativeTensor())
                                                : *(mBackend->getNativeTensor(inputs[fullIndex]));
    {
        CLEAR_BEFORE_ADDING_NODE;

        mNodeType = mBinaryTypeName;

        if (needChangeInputOrder(mBinaryTypeName)) {
            mInputs.push_back(input0Tensor);
            mInputs.push_back(input1Tensor);
        } else {
            mInputs.push_back(*(mBackend->getNativeTensor(inputs[fullIndex]))); // full input
            mInputs.push_back(*(mTempTensorWrappers[1]->getNativeTensor()));    // stage 1
        }
        mOutputs.push_back(*(mBackend->getNativeTensor(outputs[0])));

        mBackend->addNodeToGraph(mOpConfigVersion, mNodeName.c_str(), mPackageName.c_str(), mNodeType.c_str(), mParams, mInputs, mOutputs);
    }
    return NO_ERROR;
}

ErrorCode QNNBinary::onEncodeBroadcast(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs, int fullIndex) {
    // Create resources.
    int idleIndex = 1 - fullIndex;
    int fullDim = inputs[fullIndex]->dimensions();
    int idleDim = inputs[idleIndex]->dimensions();
    int offset = fullDim - idleDim;
    std::vector<uint32_t> idleShape = getNHWCShape(inputs[idleIndex]);
    std::vector<uint32_t> permData(fullDim);
    Qnn_DataType_t dataType = mBackend->getNativeTensor(inputs[fullIndex])->v1.dataType;
    if(TensorUtils::getDescribe(inputs[idleIndex])->dimensionFormat == MNN_DATA_FORMAT_NCHW){
        permData[0] = 0;
        permData[fullDim - 1] = 1;
        for(int i = 1; i < fullDim - 1; ++i){
            permData[i] = i + 1;
        }
    } else{
        permData[0] = 0;
        permData[1] = fullDim - 1;
        for(int i = 2; i < fullDim; ++i){
            permData[i] = i - 1;
        }
    }
    this->createParamTensor("perm", QNN_DATATYPE_UINT_32, {(uint32_t) fullDim}, (void *) permData.data()); // mParamTensorWrappers[0]
    std::vector<uint32_t> shapeStageReshape(fullDim, 1);
    for (int i = 0; i < idleDim; i++) {shapeStageReshape[i + offset] = idleShape[i];}
    this->createStageTensor("stage_reshape", dataType, shapeStageReshape, inputs[idleIndex]); // mTempTensorWrappers[0]
    
    std::vector<uint32_t> shapeStagePerm(fullDim);
    for (int i = 0; i < fullDim; i++) {shapeStagePerm[i] = shapeStageReshape[permData[i]];}
    this->createStageTensor("stage_perm", dataType, shapeStagePerm, inputs[idleIndex]); // mTempTensorWrappers[1]
    // Reshape.
    this->addNodeCommonReshape("Reshape",
                               *(mBackend->getNativeTensor(inputs[idleIndex])),
                               *(mTempTensorWrappers[0]->getNativeTensor()));
    
    // Permute.
    this->addNodeCommonPermute("Permute",
                               *(mTempTensorWrappers[0]->getNativeTensor()),
                               *(mParamTensorWrappers[0]->getNativeParam()),
                               *(mTempTensorWrappers[1]->getNativeTensor()));
    // Binary broadcast.
    // Ensure correct input order for non-commutative operations (Sub, Div, Pow, etc.)
    {
        CLEAR_BEFORE_ADDING_NODE;
        mNodeType = mBinaryTypeName;
        if (needChangeInputOrder(mBinaryTypeName) && fullIndex == 1) {
            mInputs.push_back(*(mTempTensorWrappers[1]->getNativeTensor()));
            mInputs.push_back(*(mBackend->getNativeTensor(inputs[fullIndex])));
        } else {
            mInputs.push_back(*(mBackend->getNativeTensor(inputs[fullIndex])));
            mInputs.push_back(*(mTempTensorWrappers[1]->getNativeTensor()));
        }
        mOutputs.push_back(*(mBackend->getNativeTensor(outputs[0])));
        mBackend->addNodeToGraph(mOpConfigVersion, mNodeName.c_str(), mPackageName.c_str(), mNodeType.c_str(), mParams, mInputs, mOutputs);
    }
    return NO_ERROR;
}

class QNNBinaryCreator : public QnnBackend::Creator {
public:
    virtual QNNCommonExecution * onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, const MNN::Op* op,
                                Backend* backend) const override {
        MNN_ASSERT(inputs.size() == 2 && outputs.size() == 1);

        std::map<BinaryOpOperation, std::string> binaryMap{{BinaryOpOperation_ADD, "ElementWiseAdd"},
                                                           {BinaryOpOperation_SUB, "ElementWiseSubtract"},
                                                           {BinaryOpOperation_MUL, "ElementWiseMultiply"},
                                                           {BinaryOpOperation_DIV, "ElementWiseDivide"},
                                                           {BinaryOpOperation_POW, "ElementWisePower"},
                                                           {BinaryOpOperation_REALDIV, "ElementWiseDivide"},
                                                           {BinaryOpOperation_MINIMUM, "ElementWiseMinimum"},
                                                           {BinaryOpOperation_MAXIMUM, "ElementWiseMaximum"},
                                                           {BinaryOpOperation_GREATER, "ElementWiseGreater"},
                                                           {BinaryOpOperation_GREATER_EQUAL, "ElementWiseGreaterEqual"},
                                                           {BinaryOpOperation_LESS, "ElementWiseLess"},
                                                           {BinaryOpOperation_FLOORDIV, "ElementWiseFloorDiv"},
                                                           {BinaryOpOperation_LESS_EQUAL, "ElementWiseLessEqual"},
                                                           {BinaryOpOperation_FLOORMOD, "ElementWiseFmod"},
                                                           {BinaryOpOperation_EQUAL, "ElementWiseEqual"},
                                                           {BinaryOpOperation_MOD, "ElementWiseMod"},
                                                           {BinaryOpOperation_LOGICALOR, "ElementWiseOr"},
                                                           {BinaryOpOperation_NOTEQUAL, "ElementWiseNotEqual"}};

        std::map<EltwiseType, std::string> eltwiseMap {
            {EltwiseType_PROD, "ElementWiseMultiply"},
            {EltwiseType_SUM, "ElementWiseAdd"},
            {EltwiseType_SUB, "ElementWiseSubtract"},
            {EltwiseType_MAXIMUM, "ElementWiseMaximum"}
        };

        std::string binaryTypeName;
        if (op->type() == OpType_BinaryOp) {
            auto iter = binaryMap.find(static_cast<BinaryOpOperation>(op->main_as_BinaryOp()->opType()));
            if (iter == binaryMap.end()) {
                MNN_ERROR("MNN_QNN: Not supported Binary type.\n");
                return nullptr;
            }
            binaryTypeName = iter->second;
        } else {
            auto iter = eltwiseMap.find(op->main_as_Eltwise()->type());
            if (iter == eltwiseMap.end()) {
                MNN_ERROR("MNN_QNN: Not supported Eltwise type.\n");
                return nullptr;
            }
            binaryTypeName = iter->second;
        }

        return new QNNBinary(backend, op, binaryTypeName);
    }
};

REGISTER_QNN_OP_CREATOR(QNNBinaryCreator, OpType_BinaryOp)
REGISTER_QNN_OP_CREATOR(QNNBinaryCreator, OpType_Eltwise)
#endif
} // end namespace QNN
} // end namespace MNN