//
//  QNNCommonExecution.cpp
//  MNN
//
//  Created by MNN on 2025/02/11.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#include "QNNCommonExecution.hpp"
namespace MNN {
namespace QNN {
#ifdef ENABLE_QNN_ONLINE_FINALIZE

// #define QNN_VERBOSE
QNNCommonExecution::QNNCommonExecution(Backend *backend, const Op *op) : Execution(backend), mOp(op) {
    mBackend = (QnnBackend *)backend;
}

ErrorCode QNNCommonExecution::onResize(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
    this->setNodeName(mOp, inputs, outputs);

    std::string nodeNameBase = MNN::EnumNameOpType(mOp->type());
    #ifdef QNN_VERBOSE
    MNN_PRINT("%s encoding start\n", nodeNameBase.c_str());
    #endif
    ErrorCode result = this->onEncode(inputs, outputs);
    if (result != NO_ERROR) {
        MNN_ERROR("Error %s encoding\n", nodeNameBase.c_str());
        return result;
    }

    #ifdef QNN_VERBOSE
    MNN_PRINT("%s encoding end\n", nodeNameBase.c_str());
    #endif
    this->clean();

    return NO_ERROR;
}

ErrorCode QNNCommonExecution::onEncode(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
    return NO_ERROR;
}

ErrorCode QNNCommonExecution::onExecute(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
    return NO_ERROR;
}

void QNNCommonExecution::setNodeName(const Op * op, const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
    if (nullptr != op->name()) {
        mNodeName = op->name()->str();
        for (int i=0; i<mNodeName.size(); ++i) {
            if (mNodeName[i] == '.' || mNodeName[i] == '/' || mNodeName[i] == ':') {
                mNodeName[i] = '_';
            }
        }
        return;
    }
    std::string nodeNameBase = MNN::EnumNameOpType(mOp->type());
    nodeNameBase += "_";
    std::string inputTag = "I_";
    std::string outputTag = "O_";

    for (int i = 0; i < inputs.size(); i++) {
        inputTag += std::to_string(mBackend->getTensorIdx(inputs[i]));
        inputTag += "_";
    }

    for (int j = 0; j < outputs.size() - 1; j++) {
        outputTag += std::to_string(mBackend->getTensorIdx(outputs[j]));
        outputTag += "_";
    }
    outputTag += std::to_string(mBackend->getTensorIdx(outputs[outputs.size() - 1]));

    mNodeName = nodeNameBase + inputTag + outputTag;
}

std::shared_ptr<QNNTensorWrapper> QNNCommonExecution::createStaticTensor(const std::string & name, Qnn_DataType_t dataType, const std::vector<uint32_t> & dimensions, const void * buffer, Qnn_QuantizeParams_t quantizeParam) {
    std::string tensorName = mNodeName + "_" + name;
    std::shared_ptr<QNNTensorWrapper> tensorWrapper = QNNTensorWrapper::createStaticTensor(tensorName, dataType, dimensions, buffer, quantizeParam);
    mBackend->addTensor(tensorWrapper->getNativeTensor());
    mTempTensorWrappers.push_back(tensorWrapper);
    return tensorWrapper;
}

std::shared_ptr<QNNTensorWrapper> QNNCommonExecution::createStaticFloatTensor(const std::string & name, Qnn_DataType_t dataType, const std::vector<uint32_t> & dimensions, const float * buffer, Qnn_QuantizeParams_t quantize) {
    std::string tensorName = mNodeName + "_" + name;
    std::shared_ptr<QNNTensorWrapper> tensorWrapper = QNNTensorWrapper::createStaticFloatTensor(tensorName, dataType, dimensions, buffer, quantize);
    mBackend->addTensor(tensorWrapper->getNativeTensor());
    mTempTensorWrappers.push_back(tensorWrapper);
    return tensorWrapper;
}

std::shared_ptr<QNNTensorWrapper> QNNCommonExecution::createStageTensor(const std::string & name, Qnn_DataType_t dataType, const std::vector<int> & dimensions, const Tensor* tensor) {
    std::vector<uint32_t> vec(dimensions.size());
    for (int i = 0; i < dimensions.size(); i++) {
        vec[i] = (uint32_t)dimensions[i];
    }
    return this->createStageTensor(name, dataType, vec, tensor);
}
std::shared_ptr<QNNTensorWrapper> QNNCommonExecution::createStageTensor(const std::string & name, Qnn_DataType_t dataType, const std::vector<uint32_t> & dimensions, const Tensor* tensor) {
    std::string tensorName = mNodeName + "_" + name;
    Qnn_QuantizeParams_t quantize = DEFAULT_QUANTIZE_PARAMS;
    Qnn_ScaleOffset_t tScaleOffsetEncoding;
    if(tensor != nullptr && TensorUtils::getDescribe(tensor)->applyQuant){
        quantize.encodingDefinition = QNN_DEFINITION_DEFINED;
        quantize.quantizationEncoding = QNN_QUANTIZATION_ENCODING_SCALE_OFFSET;
        tScaleOffsetEncoding.scale = mBackend->getNativeTensor(tensor)->v1.quantizeParams.scaleOffsetEncoding.scale;
        tScaleOffsetEncoding.offset = mBackend->getNativeTensor(tensor)->v1.quantizeParams.scaleOffsetEncoding.offset;
        quantize.scaleOffsetEncoding = tScaleOffsetEncoding;
    }
    std::shared_ptr<QNNTensorWrapper> tensorWrapper = QNNTensorWrapper::create(tensorName, QNN_TENSOR_TYPE_NATIVE, dataType, dimensions, quantize);
    mBackend->addTensor(tensorWrapper->getNativeTensor());
    mTempTensorWrappers.push_back(tensorWrapper);
    return tensorWrapper;
}

std::shared_ptr<QNNParamTensorWrapper> QNNCommonExecution::createParamTensor(const std::string & paramName, Qnn_DataType_t dataType, const std::vector<uint32_t> & dims, void * data, std::string postName) {
    MNN_ASSERT(!mNodeName.empty());
    std::string tensorName;
    if (postName.empty()) {
        tensorName = mNodeName + "_" + paramName + "_PARAM";
    } else {
        tensorName = mNodeName + "_" + paramName + "_" + postName + "_PARAM";
    }
    std::shared_ptr<QNNParamTensorWrapper> result = QNNParamTensorWrapper::create(paramName, tensorName, dataType, dims);

    void * dst = result->alloc();
    ::memcpy(dst, data, result->getNativeTensor()->v1.clientBuf.dataSize);

    mBackend->addTensor(result->getNativeTensor());

    mParamTensorWrappers.push_back(result);

    return mParamTensorWrappers.back();
}

std::shared_ptr<QNNParamScalarWrapper> QNNCommonExecution::createParamScalar(const std::string & name, bool data) {
    mParamScalarWrappers.push_back(QNNParamScalarWrapper::create(name, data));
    return mParamScalarWrappers.back();
}

std::shared_ptr<QNNParamScalarWrapper> QNNCommonExecution::createParamScalar(const std::string & name, uint32_t data) {
    mParamScalarWrappers.push_back(QNNParamScalarWrapper::create(name, data));
    return mParamScalarWrappers.back();
}

std::shared_ptr<QNNParamScalarWrapper> QNNCommonExecution::createParamScalar(const std::string & name, int data) {
    mParamScalarWrappers.push_back(QNNParamScalarWrapper::create(name, data));
    return mParamScalarWrappers.back();
}


std::shared_ptr<QNNParamScalarWrapper> QNNCommonExecution::createParamScalar(const std::string & name, float data) {
    mParamScalarWrappers.push_back(QNNParamScalarWrapper::create(name, data));
    return mParamScalarWrappers.back();
}

void QNNCommonExecution::addNodeCommon(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs, int inputSize, int outputSize) {
    for (int i = 0; i < mParamTensorWrappers.size(); i++) {
        mParams.push_back(*(mParamTensorWrappers[i]->getNativeParam()));
    }

    for (int j = 0; j < mParamScalarWrappers.size(); j++) {
        mParams.push_back(*(mParamScalarWrappers[j]->getNativeParam()));
    }
    if (0 == inputSize) {
        inputSize = inputs.size();
    }
    if (0 == outputSize) {
        outputSize = outputs.size();
    }
    for (int k = 0; k < inputSize; k++) {
        mInputs.push_back(*(mBackend->getNativeTensor(inputs[k])));
    }

    for (int l = 0; l < outputSize; l++) {
        mOutputs.push_back(*(mBackend->getNativeTensor(outputs[l])));
    }

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

void QNNCommonExecution::addNodeCommonPermute(const std::string & nodeNamePostfix, const Qnn_Tensor_t & input, const Qnn_Param_t & paramPerm, const Qnn_Tensor_t & output) {
    CLEAR_BEFORE_ADDING_NODE;

    std::string name = mNodeName + "_" + nodeNamePostfix;
    mNodeType = "Transpose";
    mInputs.push_back(input);
    mParams.push_back(paramPerm);
    mOutputs.push_back(output);

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

void QNNCommonExecution::addNodeCommonReshape(const std::string & nodeNamePostfix, const Qnn_Tensor_t & input, const Qnn_Tensor_t & output) {
    CLEAR_BEFORE_ADDING_NODE;

    std::string name = mNodeName + "_" + nodeNamePostfix;
    mNodeType = "Reshape";
    mInputs.push_back(input);
    mOutputs.push_back(output);

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

void QNNCommonExecution::clean() {
    mTempTensorWrappers.clear();
    mParamTensorWrappers.clear();
    mParamScalarWrappers.clear();
    mNodeName.clear();
    mNodeType.clear();
    mParams.clear();
    mInputs.clear();
    mOutputs.clear();
}
#endif
} // end namespace QNN
} // end namespace MNN