/****************************************************************************
 *
 *    Copyright (c) 2026  by amlogic Corp.  All rights reserved.
 *
 *    The material in this file is confidential and contains trade secrets
 *    of amlogic Corporation. No part of this work may be disclosed,
 *    reproduced, copied, transmitted, or used in any way for any purpose,
 *    without the express written permission of amlogic Corporation.
 *
 ***************************************************************************/

/*-------------------------------------------
                Includes
-------------------------------------------*/
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <float.h>
#include <sys/time.h>
#include <dlfcn.h>

#include <iostream>
#include <sstream>
#include <fstream>
#include <string>
#include <vector>

#include "nnsdk2.h"

/*-------------------------------------------
                  Functions
-------------------------------------------*/

typedef struct
{
    uint32_t index;
    float score;
} top5_item_t;

typedef struct
{
    uint32_t count;
    top5_item_t items[5];
} top5_result_t;

typedef struct
{
    void *handle;
    int (*init)(void **context, void *model, uint32_t size, amlnn_init_config *init_config);
    int (*query)(void *context, amlnn_query_cmd query_cmd, void *info, uint32_t size);
    int (*inputs_set)(void *context, uint32_t n_inputs, amlnn_input inputs[]);
    int (*run)(void *context, amlnn_run_config *run_config);
    int (*outputs_get)(void *context, uint32_t n_outputs, amlnn_output outputs[]);
    int (*destroy)(void *context);
    int (*alloc_dma_mem)(void *context, amlnn_dma_memory *dma_memory, amlnn_dma_config *dma_config);
    int (*bind_dma_mem)(void *context, amlnn_dma_memory *dma_memory, amlnn_dma_config *dma_config);
    int (*unbind_dma_mem)(void *context);
    int (*free_dma_mem)(void *context, amlnn_dma_memory *dma_memory, amlnn_dma_config *dma_config);
    int (*set_softop_opt)(void *context, amlnn_softop_opt_request requests[], uint32_t request_num);
    int (*set_input_shapes)(void *context, uint32_t n_inputs, amlnn_tensor_attr tensor_attr[]);
} nnsdk_api_t;

static nnsdk_api_t g_nnsdk_api = {};

template <typename T>
static bool load_symbol(void *handle, const char *symbol_name, T &symbol)
{
    dlerror();
    symbol = reinterpret_cast<T>(dlsym(handle, symbol_name));
    const char *err = dlerror();
    if (err != NULL)
    {
        printf("dlsym %s fail! %s\n", symbol_name, err);
        return false;
    }
    return true;
}

static int load_nnsdk_api(const char *lib_path)
{
    g_nnsdk_api.handle = dlopen(lib_path, RTLD_NOW | RTLD_LOCAL);
    if (g_nnsdk_api.handle == NULL)
    {
        printf("dlopen %s fail! %s\n", lib_path, dlerror());
        return -1;
    }

    if (!load_symbol(g_nnsdk_api.handle, "amlnn_init", g_nnsdk_api.init) ||
        !load_symbol(g_nnsdk_api.handle, "amlnn_query", g_nnsdk_api.query) ||
        !load_symbol(g_nnsdk_api.handle, "amlnn_inputs_set", g_nnsdk_api.inputs_set) ||
        !load_symbol(g_nnsdk_api.handle, "amlnn_run", g_nnsdk_api.run) ||
        !load_symbol(g_nnsdk_api.handle, "amlnn_outputs_get", g_nnsdk_api.outputs_get) ||
        !load_symbol(g_nnsdk_api.handle, "amlnn_destroy", g_nnsdk_api.destroy) ||
        !load_symbol(g_nnsdk_api.handle, "amlnn_alloc_dma_mem", g_nnsdk_api.alloc_dma_mem) ||
        !load_symbol(g_nnsdk_api.handle, "amlnn_bind_dma_mem", g_nnsdk_api.bind_dma_mem) ||
        !load_symbol(g_nnsdk_api.handle, "amlnn_unbind_dma_mem", g_nnsdk_api.unbind_dma_mem) ||
        !load_symbol(g_nnsdk_api.handle, "amlnn_free_dma_mem", g_nnsdk_api.free_dma_mem) ||
        !load_symbol(g_nnsdk_api.handle, "amlnn_set_softop_opt", g_nnsdk_api.set_softop_opt) ||
        !load_symbol(g_nnsdk_api.handle, "amlnn_set_input_shapes", g_nnsdk_api.set_input_shapes))
    {
        dlclose(g_nnsdk_api.handle);
        g_nnsdk_api = {};
        return -1;
    }

    return 0;
}

static void unload_nnsdk_api()
{
    if (g_nnsdk_api.handle != NULL)
    {
        dlclose(g_nnsdk_api.handle);
    }
    g_nnsdk_api = {};
}

static void dump_tensor_attr(amlnn_tensor_attr *attr)
{
    printf("  index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, size_with_stride=%d, fmt=%s, type=%s, "
           "zp=%d, scale=%f\n",
           attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
           attr->n_elems, attr->size, attr->size_with_stride, get_format_string(attr->fmt), get_type_string(attr->type),
           attr->zp, attr->scale);
}

static void dump_input_dynamic_range(amlnn_input_range *dyn_range)
{
    std::string range_str = "";
    for (int n = 0; n < dyn_range->shape_number; ++n)
    {
        range_str += n == 0 ? "[" : ",[";
        range_str += dyn_range->n_dims < 1 ? "" : std::to_string(dyn_range->dyn_range[n][0]);
        for (int i = 1; i < dyn_range->n_dims; ++i)
        {
            range_str += ", " + std::to_string(dyn_range->dyn_range[n][i]);
        }
        range_str += "]";
    }

    printf("  index=%d, name=%s, shape_number=%d, range=[%s], fmt = %s\n", dyn_range->index, dyn_range->name,
           dyn_range->shape_number, range_str.c_str(), get_format_string(dyn_range->fmt));
}

static std::vector<std::string> split(const std::string &str, const std::string &pattern)
{
    std::vector<std::string> res;
    if (str == "")
        return res;
    std::string strs = str + pattern;
    size_t pos = strs.find(pattern);
    while (pos != strs.npos)
    {
        std::string temp = strs.substr(0, pos);
        res.push_back(temp);
        strs = strs.substr(pos + 1, strs.size());
        pos = strs.find(pattern);
    }
    return res;
}

static uint8_t *load_model(const char *filename, uint32_t *model_size)
{
    FILE *fp = fopen(filename, "rb");
    if (fp == nullptr)
    {
        printf("fopen %s fail!\n", filename);
        return NULL;
    }
    fseek(fp, 0, SEEK_END);
    uint32_t model_len = ftell(fp);
    uint8_t *model = (uint8_t *)malloc(model_len);
    fseek(fp, 0, SEEK_SET);
    if (model_len != fread(model, 1, model_len, fp))
    {
        printf("fread %s fail!\n", filename);
        free(model);
        return NULL;
    }
    *model_size = model_len;
    if (fp)
    {
        fclose(fp);
    }
    return model;
}

static void *read_file(const char *file_path, uint32_t *file_size)
{
    FILE *fp = NULL;
    uint32_t size = 0;
    void *buf = NULL;
    fp = fopen(file_path, "rb");
    if (NULL == fp)
    {
        printf("open file fail!\n");
        return NULL;
    }

    fseek(fp, 0, SEEK_END);
    size = ftell(fp);
    rewind(fp);

    buf = malloc(sizeof(uint8_t) * size);

    fread(buf, 1, size, fp);

    fclose(fp);

    *file_size = size;
    return buf;
}

template <typename T>
static top5_result_t collect_top5(const T *buf, uint32_t num)
{
    uint32_t i, j, k;
    top5_result_t result = {};
    float max_prob[5] = {-FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX, -FLT_MAX};
    result.count = num < 5 ? num : 5;

    for (i = 0; i < num; i++)
    {
        float score = static_cast<float>(buf[i]);
        for (j = 0; j < 5; j++)
        {
            if (score > max_prob[j])
            {
                for (k = 4; k > j; k--)
                {
                    max_prob[k] = max_prob[k - 1];
                    result.items[k] = result.items[k - 1];
                }
                max_prob[j] = score;
                result.items[j].index = i;
                result.items[j].score = score;
                break;
            }
        }
    }

    return result;
}

static top5_result_t process_top5_f32(float *buf, uint32_t num)
{
    return collect_top5(buf, num);
}

static top5_result_t process_top5_i8(int8_t *buf, uint32_t num)
{
    return collect_top5(buf, num);
}

static top5_result_t process_top5_u8(uint8_t *buf, uint32_t num)
{
    return collect_top5(buf, num);
}

static void print_top5_result(const char *run_mode, uint32_t output_index, const top5_result_t &result)
{
    printf("[%s] output %u top5:\n", run_mode, output_index);
    for (uint32_t i = 0; i < result.count; i++)
    {
        printf("%3u: %8.6f\n", result.items[i].index, result.items[i].score);
    }
}

static void print_top5_result_with_tag(const char *run_mode, const char *tag, uint32_t output_index, const top5_result_t &result)
{
    printf("[%s][%s] output %u top5:\n", run_mode, tag, output_index);
    for (uint32_t i = 0; i < result.count; i++)
    {
        printf("%3u: %8.6f\n", result.items[i].index, result.items[i].score);
    }
}


/*-------------------------------------------
                  Main Function
-------------------------------------------*/
int base_api(int argc, char **argv)
{
    int ret = 0;

    const char *model_path = argv[1];
    const char *input_path = argv[2];

    if (argc != 4)
    {
        printf("Usage: %s <adla model_path> <input_path> \n", argv[0]);
        return -1;
    }

    // Get sdk version
    amlnn_sdk_version sdk_ver;
    ret = g_nnsdk_api.query(NULL, AMLNN_QUERY_SDK_VERSION, &sdk_ver, sizeof(sdk_ver));
    if (ret != AMLNN_SUCCESS)
    {
        printf("amlnn_query fail! ret=%d\n", ret);
        return -1;
    }
    printf("version info: sdk2-[%s], adla-[%s], delegate-[%s]\n", sdk_ver.sdk2_api_version, sdk_ver.adla_kmd_version, sdk_ver.delegate_version);

    // Load ADLA Model
    void *ctx = NULL;
    amlnn_init_config init_config;
    memset(&init_config, 0, sizeof(amlnn_init_config));
    init_config.backend_type = AMLNN_BACKEND_ADLA_NPU;
    ret = g_nnsdk_api.init(&ctx, (void *)model_path, 0, &init_config);
    if (ret < 0)
    {
        printf("amlnn_init fail! ret=%d\n", ret);
        return -1;
    }

    // Get Model Input Output Info
    amlnn_input_output_num io_num;
    ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret != AMLNN_SUCCESS)
    {
        printf("amlnn_query fail! ret=%d\n", ret);
        return -1;
    }
    printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);

    printf("input tensors:\n");
    amlnn_tensor_attr input_attrs[io_num.n_input];
    memset(input_attrs, 0, sizeof(input_attrs));
    for (uint32_t i = 0; i < io_num.n_input; i++)
    {
        input_attrs[i].index = i;
        ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(amlnn_tensor_attr));
        if (ret != AMLNN_SUCCESS)
        {
            printf("amlnn_query fail! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&(input_attrs[i]));
    }

    printf("output tensors:\n");
    amlnn_tensor_attr output_attrs[io_num.n_output];
    memset(output_attrs, 0, sizeof(output_attrs));
    for (uint32_t i = 0; i < io_num.n_output; i++)
    {
        output_attrs[i].index = i;
        ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(amlnn_tensor_attr));
        if (ret != AMLNN_SUCCESS)
        {
            printf("amlnn_query fail! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&(output_attrs[i]));
    }

    // Load input data
    void *input_data;
    uint32_t input_size;
    uint32_t file_size;
    input_data = read_file(input_path, &file_size);
    input_size = input_attrs[0].n_elems * sizeof(uint8_t);

    // Set input
    printf("amlnn_inputs_set\n");
    amlnn_input inputs[1];
    memset(inputs, 0, sizeof(inputs));
    inputs[0].index = 0;
    inputs[0].buf = input_data;
    inputs[0].size = input_size;
    ret = g_nnsdk_api.inputs_set(ctx, io_num.n_input, inputs);
    if (ret < 0)
    {
        printf("amlnn_inputs_set fail! ret=%d\n", ret);
        return -1;
    }

    // Run
    printf("amlnn_run\n");
    ret = g_nnsdk_api.run(ctx, NULL);
    if (ret < 0)
    {
        printf("amlnn_run fail! ret=%d\n", ret);
        return -1;
    }

    // Get Output
    printf("amlnn_outputs_get\n");
    amlnn_output outputs[io_num.n_output];
    memset(outputs, 0, sizeof(outputs));
    for (uint32_t i = 0; i < io_num.n_output; i++)
    {
        outputs[i].is_float = 1;
        outputs[i].index = i;
    }
    ret = g_nnsdk_api.outputs_get(ctx, io_num.n_output, outputs);
    if (ret < 0)
    {
        printf("amlnn_outputs_get fail! ret=%d\n", ret);
        return -1;
    }

    // Post Process
    for (uint32_t i = 0; i < io_num.n_output; i++)
    {
        float *buffer = (float *)outputs[i].buf;
        uint32_t sz = outputs[i].size / sizeof(float);

        top5_result_t result = process_top5_f32(buffer, sz);
        print_top5_result("base", i, result);
    }

    // Destroy
    ret = g_nnsdk_api.destroy(ctx);
    if (ret < 0)
    {
        printf("amlnn_destroy fail! ret=%d\n", ret);
        return -1;
    }

    if (input_data)
    {
        free(input_data);
    }

    return 0;
}

int dma_api(int argc, char **argv)
{
    int ret = 0;

    char *model_path = argv[1];
    char *input_path = argv[2];

    if (argc != 4)
    {
        printf("Usage: %s <adla model_path> <input_path> \n", argv[0]);
        return -1;
    }

    // Load ADLA Model
    void *ctx = NULL;
    uint8_t *model = NULL;
    uint32_t model_len = 0;
    model = load_model(model_path, &model_len);
    amlnn_init_config init_config;
    memset(&init_config, 0, sizeof(amlnn_init_config));
    ret = g_nnsdk_api.init(&ctx, (void *)model, model_len, &init_config);
    if (ret < 0)
    {
        printf("amlnn_init fail! ret=%d\n", ret);
        return -1;
    }

    // Get Model Input Output Info
    amlnn_input_output_num io_num;
    ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret != AMLNN_SUCCESS)
    {
        printf("amlnn_query fail! ret=%d\n", ret);
        return -1;
    }
    printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);

    printf("input tensors:\n");
    amlnn_tensor_attr input_attrs[io_num.n_input];
    memset(input_attrs, 0, io_num.n_input * sizeof(amlnn_tensor_attr));
    for (uint32_t i = 0; i < io_num.n_input; i++)
    {
        input_attrs[i].index = i;
        // query info
        ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(amlnn_tensor_attr));
        if (ret < 0)
        {
            printf("amlnn_query fail! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&input_attrs[i]);
    }

    printf("output tensors:\n");
    amlnn_tensor_attr output_attrs[io_num.n_output];
    memset(output_attrs, 0, io_num.n_output * sizeof(amlnn_tensor_attr));
    for (uint32_t i = 0; i < io_num.n_output; i++)
    {
        output_attrs[i].index = i;
        // query info
        ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(amlnn_tensor_attr));
        if (ret != AMLNN_SUCCESS)
        {
            printf("amlnn_query fail! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&output_attrs[i]);
    }

    // Alloc input dma memory
    printf("amlnn_alloc_dma_mem input\n");
    amlnn_dma_config dma_config_input;
    memset(&dma_config_input, 0, sizeof(amlnn_dma_config));
    dma_config_input.cacheable = 1;
    dma_config_input.access_mode = AMLNN_EXT_MEM_ACCESS_READ_WRITE;
    dma_config_input.buffer_type = AMLNN_EXT_BUFFER_VIRTUAL_ADDR;
    dma_config_input.io_type = AMLNN_EXT_IO_INPUT;

    amlnn_dma_memory dma_memory_input[1];
    memset(dma_memory_input, 0, sizeof(dma_memory_input));
    dma_memory_input[0].index = 0;
    dma_memory_input[0].size = input_attrs[0].size_with_stride;
    ret = g_nnsdk_api.alloc_dma_mem(ctx, &dma_memory_input[0], &dma_config_input);
    if (ret != AMLNN_SUCCESS)
    {
        printf("amlnn_alloc_dma_mem fail! ret=%d\n", ret);
        return -1;
    }

    // Bind input dma memory
    printf("amlnn_bind_dma_mem input\n");
    ret = g_nnsdk_api.bind_dma_mem(ctx, &dma_memory_input[0], &dma_config_input);
    if (ret != AMLNN_SUCCESS)
    {
        printf("amlnn_bind_dma_mem fail! ret=%d\n", ret);
        return -1;
    }

    // Alloc output dma memory
    printf("amlnn_alloc_dma_mem output\n");
    amlnn_dma_config dma_config_output;
    memset(&dma_config_output, 0, sizeof(amlnn_dma_config));
    dma_config_output.cacheable = 1;
    dma_config_output.access_mode = AMLNN_EXT_MEM_ACCESS_READ_WRITE;
    dma_config_output.buffer_type = AMLNN_EXT_BUFFER_VIRTUAL_ADDR;
    dma_config_output.io_type = AMLNN_EXT_IO_OUTPUT;

    amlnn_dma_memory dma_memory_output[1];
    memset(dma_memory_output, 0, sizeof(dma_memory_output));
    dma_memory_output[0].index = 0;
    dma_memory_output[0].size = output_attrs[0].size_with_stride;
    ret = g_nnsdk_api.alloc_dma_mem(ctx, &dma_memory_output[0], &dma_config_output);
    if (ret != AMLNN_SUCCESS)
    {
        printf("amlnn_alloc_dma_mem fail! ret=%d\n", ret);
        return -1;
    }

    // Bind output dma memory
    printf("amlnn_bind_dma_mem output\n");
    ret = g_nnsdk_api.bind_dma_mem(ctx, &dma_memory_output[0], &dma_config_output);
    if (ret != AMLNN_SUCCESS)
    {
        printf("amlnn_bind_dma_mem fail! ret=%d\n", ret);
        return -1;
    }

    // Load input data
    void *input_data;
    uint32_t input_size;
    uint32_t file_size;
    input_data = read_file(input_path, &file_size);
    input_size = input_attrs[0].n_elems * sizeof(int8_t);

    // Copy input data to dma mem
    memcpy(dma_memory_input[0].viraddr, input_data, input_size);

    // Run
    printf("amlnn_run\n");
    ret = g_nnsdk_api.run(ctx, NULL);
    if (ret < 0)
    {
        printf("amlnn_run fail! ret=%d\n", ret);
        return -1;
    }

    // Post Process
    for (uint32_t i = 0; i < io_num.n_output; i++)
    {
        uint8_t *buffer = (uint8_t *)dma_memory_output[i].viraddr;
        uint32_t sz = output_attrs[i].n_elems;

        top5_result_t result = process_top5_u8(buffer, sz);
        print_top5_result("dma", i, result);
    }

    // Unbind all dma memory
    ret = g_nnsdk_api.unbind_dma_mem(ctx);
    if (ret < 0)
    {
        printf("amlnn_unbind_dma_mem fail! ret=%d\n", ret);
        return -1;
    }

    // Free input dma memory
    ret = g_nnsdk_api.free_dma_mem(ctx, &dma_memory_input[0], &dma_config_input);
    if (ret < 0)
    {
        printf("amlnn_free_dma_mem fail! ret=%d\n", ret);
        return -1;
    }

    // Free output dma memory
    ret = g_nnsdk_api.free_dma_mem(ctx, &dma_memory_output[0], &dma_config_output);
    if (ret < 0)
    {
        printf("amlnn_free_dma_mem fail! ret=%d\n", ret);
        return -1;
    }

    // Destroy
    ret = g_nnsdk_api.destroy(ctx);
    if (ret < 0)
    {
        printf("amlnn_destroy fail! ret=%d\n", ret);
        return -1;
    }

    if (model)
    {
        free(model);
    }

    if (input_data)
    {
        free(input_data);
    }

    return 0;
}

int softop_acc_api(int argc, char **argv)
{
    int ret = 0;

    const char *model_path = argv[1];
    const char *input_path = argv[2];

    if (argc != 4)
    {
        printf("Usage: %s <adla model_path> <input_path> \n", argv[0]);
        return -1;
    }

    // Load ADLA Model
    void *ctx = NULL;
    uint8_t *model = NULL;
    uint32_t model_len = 0;
    model = load_model(model_path, &model_len);
    amlnn_init_config init_config;
    memset(&init_config, 0, sizeof(amlnn_init_config));
    ret = g_nnsdk_api.init(&ctx, (void *)model, model_len, &init_config);
    if (ret < 0)
    {
        printf("amlnn_init fail! ret=%d\n", ret);
        return -1;
    }

    // Get Model Input Output Info
    amlnn_input_output_num io_num;
    ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret != AMLNN_SUCCESS)
    {
        printf("amlnn_query fail! ret=%d\n", ret);
        return -1;
    }
    printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);

    printf("input tensors:\n");
    amlnn_tensor_attr input_attrs[io_num.n_input];
    memset(input_attrs, 0, sizeof(input_attrs));
    for (uint32_t i = 0; i < io_num.n_input; i++)
    {
        input_attrs[i].index = i;
        ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(amlnn_tensor_attr));
        if (ret != AMLNN_SUCCESS)
        {
            printf("amlnn_query fail! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&(input_attrs[i]));
    }

    printf("output tensors:\n");
    amlnn_tensor_attr output_attrs[io_num.n_output];
    memset(output_attrs, 0, sizeof(output_attrs));
    for (uint32_t i = 0; i < io_num.n_output; i++)
    {
        output_attrs[i].index = i;
        ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(amlnn_tensor_attr));
        if (ret != AMLNN_SUCCESS)
        {
            printf("amlnn_query fail! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&(output_attrs[i]));
    }

    // Set software operator acceleration
    amlnn_softop_opt_request softop_opt_requests[] =
    {
        {
           .softop_type = AMLNN_Unknown,
           .acc_type = AMLNN_SOFTOP_ACC_OPENMP,
        },
    };
    ret = g_nnsdk_api.set_softop_opt(ctx, softop_opt_requests, sizeof(softop_opt_requests) / sizeof(amlnn_softop_opt_request));
    if (ret != AMLNN_SUCCESS)
    {
        printf("amlnn_set_softop_opt fail! ret=%d\n", ret);
        return -1;
    }

    // Load input data
    void *input_data;
    uint32_t input_size;
    uint32_t file_size;
    input_data = read_file(input_path, &file_size);
    input_size = input_attrs[0].n_elems * sizeof(uint8_t);

    // Set input
    printf("amlnn_inputs_set\n");
    amlnn_input inputs[1];
    memset(inputs, 0, sizeof(inputs));
    inputs[0].index = 0;
    inputs[0].buf = input_data;
    inputs[0].size = input_size;
    ret = g_nnsdk_api.inputs_set(ctx, io_num.n_input, inputs);
    if (ret < 0)
    {
        printf("amlnn_inputs_set fail! ret=%d\n", ret);
        return -1;
    }

    // Run
    printf("amlnn_run\n");
    ret = g_nnsdk_api.run(ctx, NULL);
    if (ret < 0)
    {
        printf("amlnn_run fail! ret=%d\n", ret);
        return -1;
    }

    // Get Output
    printf("amlnn_outputs_get\n");
    amlnn_output outputs[io_num.n_output];
    memset(outputs, 0, sizeof(outputs));
    for (uint32_t i = 0; i < io_num.n_output; i++)
    {
        outputs[i].is_float = 1;
        outputs[i].index = i;
    }
    ret = g_nnsdk_api.outputs_get(ctx, io_num.n_output, outputs);
    if (ret < 0)
    {
        printf("amlnn_outputs_get fail! ret=%d\n", ret);
        return -1;
    }

    // Post Process
    for (uint32_t i = 0; i < io_num.n_output; i++)
    {
        float *buffer = (float *)outputs[i].buf;
        uint32_t sz = outputs[i].size / sizeof(float);

        top5_result_t result = process_top5_f32(buffer, sz);
        print_top5_result("softop_acc", i, result);
    }

    // Destroy
    ret = g_nnsdk_api.destroy(ctx);
    if (ret < 0)
    {
        printf("amlnn_destroy fail! ret=%d\n", ret);
        return -1;
    }

    if (model)
    {
        free(model);
    }

    if (input_data)
    {
        free(input_data);
    }

    return 0;
}

int dynamic_shape_api(int argc, char **argv)
{
    int ret = 0;

    char *model_path = argv[1];
    char *input_paths = argv[2];

    if (argc != 4)
    {
        printf("Usage: %s <adla model_path> <input_paths> \n", argv[0]);
        return -1;
    }

    std::vector<std::string> input_paths_split = split(input_paths, "#");

    // Load ADLA Model
    void *ctx = NULL;
    amlnn_init_config init_config;
    memset(&init_config, 0, sizeof(amlnn_init_config));
    init_config.backend_type = AMLNN_BACKEND_ADLA_NPU;
    ret = g_nnsdk_api.init(&ctx, (void *)model_path, 0, &init_config);
    if (ret < 0)
    {
        printf("amlnn_init fail! ret=%d\n", ret);
        return -1;
    }

    // Get Model Input Output Info
    amlnn_input_output_num io_num;
    ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret != AMLNN_SUCCESS)
    {
        printf("amlnn_query fail! ret=%d\n", ret);
        return -1;
    }
    printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);

    printf("input tensors:\n");
    amlnn_tensor_attr input_attrs[io_num.n_input];
    memset(input_attrs, 0, io_num.n_input * sizeof(amlnn_tensor_attr));
    for (uint32_t i = 0; i < io_num.n_input; i++)
    {
        input_attrs[i].index = i;
        // query info
        ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(amlnn_tensor_attr));
        if (ret < 0)
        {
            printf("amlnn_query fail! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&input_attrs[i]);
    }

    printf("output tensors:\n");
    amlnn_tensor_attr output_attrs[io_num.n_output];
    memset(output_attrs, 0, io_num.n_output * sizeof(amlnn_tensor_attr));
    for (uint32_t i = 0; i < io_num.n_output; i++)
    {
        output_attrs[i].index = i;
        // query info
        ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(amlnn_tensor_attr));
        if (ret != AMLNN_SUCCESS)
        {
            printf("amlnn_query fail! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&output_attrs[i]);
    }

    printf("dynamic inputs shape range:\n");
    amlnn_input_range shape_range[io_num.n_input];
    memset(shape_range, 0, io_num.n_input * sizeof(amlnn_input_range));
    for (uint32_t i = 0; i < io_num.n_input; i++)
    {
        shape_range[i].index = i;
        ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_INPUT_DYNAMIC_RANGE, &shape_range[i], sizeof(amlnn_input_range));
        if (ret != AMLNN_SUCCESS)
        {
            printf("amlnn_query fail! ret=%d\n", ret);
            return -1;
        }
        dump_input_dynamic_range(&shape_range[i]);
    }

    void *input_data[io_num.n_input];
    uint32_t input_size[io_num.n_input];
    uint32_t file_size[io_num.n_input];
    for (uint32_t i = 0; i < io_num.n_input; i++)
    {
        input_data[i] = NULL;
        input_size[i] = input_attrs[i].n_elems * sizeof(uint8_t);
    }
    uint32_t shape_num = shape_range[0].shape_number;

    // Load input data
    if (io_num.n_input != input_paths_split.size())
    {
        return -1;
    }
    for (uint32_t i = 0; i < io_num.n_input; i++)
    {
        printf("%s\n", input_paths_split[i].c_str());

        input_data[i] = read_file(input_paths_split[i].c_str(), &file_size[i]);
    }

    for (uint32_t s = 0; s < shape_num; ++s)
    {
        for (uint32_t i = 0; i < io_num.n_input; i++)
        {
            for (uint32_t j = 0; j < input_attrs[i].n_dims; ++j)
            {
                input_attrs[i].dims[j] = shape_range[i].dyn_range[s][j];
            }
        }

        ret = g_nnsdk_api.set_input_shapes(ctx, io_num.n_input, input_attrs);
        if (ret < 0)
        {
            fprintf(stderr, "amlnn_set_input_shapes fail! ret=%d\n", ret);
            return -1;
        }

        printf("current input tensors:\n");
        amlnn_tensor_attr cur_input_attrs[io_num.n_input];
        memset(cur_input_attrs, 0, io_num.n_input * sizeof(amlnn_tensor_attr));
        for (uint32_t i = 0; i < io_num.n_input; i++)
        {
            cur_input_attrs[i].index = i;
            // query info
            ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_CURRENT_INPUT_ATTR, &(cur_input_attrs[i]), sizeof(amlnn_tensor_attr));
            if (ret < 0)
            {
                printf("amlnn_query fail! ret=%d\n", ret);
                return -1;
            }
            dump_tensor_attr(&cur_input_attrs[i]);
        }

        printf("current output tensors:\n");
        amlnn_tensor_attr cur_output_attrs[io_num.n_output];
        memset(cur_output_attrs, 0, io_num.n_output * sizeof(amlnn_tensor_attr));
        for (uint32_t i = 0; i < io_num.n_output; i++)
        {
            cur_output_attrs[i].index = i;
            // query info
            ret = g_nnsdk_api.query(ctx, AMLNN_QUERY_CURRENT_OUTPUT_ATTR, &(cur_output_attrs[i]), sizeof(amlnn_tensor_attr));
            if (ret != AMLNN_SUCCESS)
            {
                printf("amlnn_query fail! ret=%d\n", ret);
                return -1;
            }
            dump_tensor_attr(&cur_output_attrs[i]);
        }

        uint32_t input_size_pad[io_num.n_input];
        void *input_data_pad[io_num.n_input];
        std::vector<std::vector<uint8_t>> input_buf_tmp(io_num.n_input);
        for (uint32_t i = 0; i < io_num.n_input; i++)
        {
            input_size_pad[i] = cur_input_attrs[i].n_elems * sizeof(uint8_t);
            input_buf_tmp[i].resize(input_size_pad[i], 0);

            uint32_t multiplier = input_size_pad[i] / input_size[i];
            for (uint32_t j = 0; j < multiplier; j++)
            {
                memcpy(input_buf_tmp[i].data() + (j * input_size[i]), input_data[i], input_size[i]);
            }
            input_data_pad[i] = input_buf_tmp[i].data();
        }

        printf("amlnn_inputs_set\n");
        amlnn_input inputs[io_num.n_input];
        memset(inputs, 0, io_num.n_input * sizeof(amlnn_input));
        for (uint32_t i = 0; i < io_num.n_input; i++)
        {
            inputs[i].index = i;
            inputs[i].buf = input_data_pad[i];
            inputs[i].size = input_size_pad[i];
        }
        ret = g_nnsdk_api.inputs_set(ctx, io_num.n_input, inputs);
        if (ret < 0)
        {
            printf("amlnn_inputs_set fail! ret=%d\n", ret);
            return -1;
        }

        printf("amlnn_run\n");
        ret = g_nnsdk_api.run(ctx, NULL);
        if (ret < 0)
        {
            printf("amlnn_run fail! ret=%d\n", ret);
            return -1;
        }

        printf("amlnn_outputs_get\n");
        amlnn_output outputs[io_num.n_output];
        memset(outputs, 0, io_num.n_output * sizeof(amlnn_output));
        for (uint32_t i = 0; i < io_num.n_output; ++i)
        {
            outputs[i].is_float = 1;
            outputs[i].index = i;
        }
        ret = g_nnsdk_api.outputs_get(ctx, io_num.n_output, outputs);
        if (ret < 0)
        {
            printf("amlnn_outputs_get fail! ret=%d\n", ret);
            return ret;
        }

        for (uint32_t i = 0; i < io_num.n_output; i++)
        {
            float *buffer = (float *)outputs[i].buf;
            uint32_t sz = outputs[i].size / sizeof(float);

            top5_result_t result = process_top5_f32(buffer, sz);
            std::string tag = "shape_" + std::to_string(s);
            print_top5_result_with_tag("dynamic_shape", tag.c_str(), i, result);
        }
    }

    // Destroy
    ret = g_nnsdk_api.destroy(ctx);
    if (ret < 0)
    {
        printf("amlnn_destroy fail! ret=%d\n", ret);
        return -1;
    }

    for (uint32_t i = 0; i < io_num.n_input; i++)
    {
        free(input_data[i]);
    }

    return 0;
}

typedef enum
{
    MODE_NONE,
    MODE_BASE,
    MODE_DMA,
    MODE_SOFTOP_ACC,
    MODE_DYNAMIC_SHAPE
} run_mode_t;

int main(int argc, char **argv)
{
    run_mode_t mode = MODE_NONE;
    int ret = 0;

    for (int i = 1; i < argc; i++)
    {
        if (strcmp(argv[i], "--base") == 0)
        {
            if (mode != MODE_NONE)
                goto multi_error;
            mode = MODE_BASE;
        }
        else if (strcmp(argv[i], "--dma") == 0)
        {
            if (mode != MODE_NONE)
                goto multi_error;
            mode = MODE_DMA;
        }
        else if (strcmp(argv[i], "--softop_acc") == 0)
        {
            if (mode != MODE_NONE)
                goto multi_error;
            mode = MODE_SOFTOP_ACC;
        }
        else if (strcmp(argv[i], "--dynamic_shape") == 0)
        {
            if (mode != MODE_NONE)
                goto multi_error;
            mode = MODE_DYNAMIC_SHAPE;
        }
    }

    if (mode == MODE_NONE)
    {
        printf("Error: Please specify one mode (--base, --dma, --softop_acc, --dynamic_shape).\n");
        return 1;
    }

    if (load_nnsdk_api("libnnsdk.so") != 0)
    {
        return 1;
    }

    switch (mode)
    {
    case MODE_BASE:
        ret = base_api(argc, argv);
        break;
    case MODE_DMA:
        ret = dma_api(argc, argv);
        break;
    case MODE_SOFTOP_ACC:
        ret = softop_acc_api(argc, argv);
        break;
    case MODE_DYNAMIC_SHAPE:
        ret = dynamic_shape_api(argc, argv);
        break;
    default:
        ret = 1;
        break;
    }

    unload_nnsdk_api();
    return ret == 0 ? 0 : 1;

multi_error:
    printf("Error: Multiple modes specified. You can only run ONE at a time.\n");
    return 1;
}