/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#include "sample_dynamic_batch.h"
#include <map>
#include <vector>

int main()
{
    INFO_LOG("SAMPLE start to execute.");

    const std::map<ge::AscendString, ge::AscendString> options{
            {"ge.graphRunMode", "0"}
    };
    std::unique_ptr<SampleDynamicBatch> sampleDynamicBatchPtr = nullptr;
    try {
        sampleDynamicBatchPtr = std::make_unique<SampleDynamicBatch>(options);
    } catch (std::runtime_error &e) {
        ERROR_LOG("SampleDynamicBatch creation failed");
        return FAILED;
    }
    if (!sampleDynamicBatchPtr) {
        ERROR_LOG("SampleDynamicBatch creation failed");
        return FAILED;
    }

    // resnet50_Opset16.onnx is static model, input node is x:[1, 3, 224, 224]
    const std::string modelPath = "../model/resnet50_Opset16.onnx";
    // parse onnx model and build ge graph
    auto ret = sampleDynamicBatchPtr->ParseModelAndBuildGraph(modelPath);
    if (ret != SUCCESS) {
        ERROR_LOG("Parse onnx model failed, model:%s", modelPath.c_str());
        return FAILED;
    }
    INFO_LOG("Parse model %s success", modelPath.c_str());

    // enable dynamic batch
    const std::map<ge::AscendString, ge::AscendString> graph_options{
        {"ge.inputShape", "x:-1,3,224,224"},
        {"ge.dynamicDims", "1;2;4;8"},
        {"ge.dynamicNodeType", "1"}
    };
    // set specific input tensor
    std::vector<ge::Tensor> input_tensors;
    const std::initializer_list<int64_t> dims{2, 3, 224, 224};
    const ge::Shape input_shape(dims); // batchSize:2
    const ge::Tensor input_tensor(ge::TensorDesc(input_shape, ge::FORMAT_NCHW, ge::DT_FLOAT));
    input_tensors.push_back(input_tensor);

    ret = sampleDynamicBatchPtr->CompileGraph(graph_options, input_tensors);
    if (ret != SUCCESS) {
        ERROR_LOG("SampleDynamicBatch compile graph failed");
        return FAILED;
    }

    const std::vector<std::string> testFiles = {"../data/dog1_1024_683.bin","../data/dog2_1024_683.bin"};
    ret = sampleDynamicBatchPtr->Process(testFiles, dims);
    if (ret != SUCCESS) {
        ERROR_LOG("SampleDynamicBatch process graph failed");
        return FAILED;
    }

    // print output result
    sampleDynamicBatchPtr->OutputModelResult();

    INFO_LOG("SAMPLE PASSED.");
    return SUCCESS;
}