# Copyright (c) 2022 - 2024 Advanced Micro Devices, Inc. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import os, sys, struct
import datetime, pytz
import nnir as ir

def generateLicenseForScript(f):
    f.write( \
"""################################################################################
#
# MIT License
#
# Copyright (c) 2022 - 2024 Advanced Micro Devices, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
################################################################################

# This file is generated by nnir_to_zendnn.py on %s
""" % (datetime.datetime.now(tz=pytz.timezone('America/Los_Angeles')).isoformat()))

def generateLicenseForCPP(f):
        f.write( \
"""/*
#
# MIT License
#
# Copyright (c) 2022 - 2024 Advanced Micro Devices, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
*/

// This file is generated by nnir_to_zendnn.py on %s
""" % (datetime.datetime.now(tz=pytz.timezone('America/Los_Angeles')).isoformat()))

def generateCMakeFiles(graph,outputFolder):
    fileName = outputFolder + '/CMakeLists.txt'
    print('creating ' + fileName + ' ...')
    with open(fileName, 'w') as f:
        generateLicenseForScript(f)
        f.write( \
"""
cmake_minimum_required(VERSION 3.10)

# ROCM Path
if(DEFINED ENV{ROCM_PATH})
    set(ROCM_PATH $ENV{ROCM_PATH} CACHE PATH "Default ROCm installation path")
elseif(ROCM_PATH)
    message("-- INFO:ROCM_PATH Set -- ${ROCM_PATH}")
else()
    set(ROCM_PATH /opt/rocm CACHE PATH "Default ROCm installation path")
endif()
# Set AMD Clang as default compiler
if(NOT DEFINED CMAKE_CXX_COMPILER AND EXISTS "${ROCM_PATH}/lib/llvm/bin/amdclang++")
    set(CMAKE_CXX_COMPILER ${ROCM_PATH}/lib/llvm/bin/amdclang++)
endif()

# Set Default Compiler & Standard
# aocc-linux-compiler -- ENV:ZENDNN_AOCC_COMP_PATH
if(NOT DEFINED ZENDNN_AOCC_COMP_PATH)
    if(NOT DEFINED ENV{ZENDNN_AOCC_COMP_PATH})
        message(FATAL_ERROR "ZENDNN_AOCC_COMP_PATH NOT FOUND -- aocc-compiler-X.X.XX install path must be set to ZENDNN_AOCC_COMP_PATH${ColourReset}")
    else()
        set(ZENDNN_AOCC_COMP_PATH $ENV{ZENDNN_AOCC_COMP_PATH} CACHE PATH "Path to which ZENDNN_AOCC_COMP_PATH has been installed")
    endif()
endif()
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED On)
set(CMAKE_CXX_EXTENSIONS ON)
set(CMAKE_C_COMPILER ${ZENDNN_AOCC_COMP_PATH}/bin/clang)
set(CMAKE_CXX_COMPILER ${ZENDNN_AOCC_COMP_PATH}/bin/clang++)

project(zendnn_app)

# OpenMP
find_package(OpenMP REQUIRED)
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
# -fopenmp -- Enable OpenMP
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
set(LINK_LIBRARY_LIST ${LINK_LIBRARY_LIST} OpenMP::OpenMP_CXX)

# aocl-linux-aocc -- ENV:ZENDNN_BLIS_PATH
if(NOT DEFINED ZENDNN_BLIS_PATH)
    if(NOT DEFINED ENV{ZENDNN_BLIS_PATH})
        message(FATAL_ERROR "ZENDNN_BLIS_PATH NOT FOUND -- aocl-linux-aocc-X.X.XX install path must be set to ZENDNN_BLIS_PATH${ColourReset}")
    else()
        set(ZENDNN_BLIS_PATH $ENV{ZENDNN_BLIS_PATH} CACHE PATH "Path to which ZENDNN_BLIS_PATH has been installed")
    endif()
endif()

# amd Zen DNN Git Source
if(NOT DEFINED ZENDNN_GIT_ROOT)
    if(NOT DEFINED ENV{ZENDNN_GIT_ROOT})
        message(FATAL_ERROR "ZENDNN_GIT_ROOT NOT FOUND -- ZEN DNN Git Source path must be set to ZENDNN_GIT_ROOT${ColourReset}")
    else()
        set(ZENDNN_GIT_ROOT $ENV{ZENDNN_GIT_ROOT} CACHE PATH "Path to which ZENDNN_GIT_ROOT has been installed")
    endif()
endif()

# set global compiler flags for the project
# -O3 -- Optimize output file
# -fPIC -- Generate position-independent code if possible.
# -march -- Generate code for given CPU. [znver2 -- Zen Version 2]
# -Wreturn-type -- Warn whenever a function's return type defaults to "int" (C), or about inconsistent return types (C++).
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DLIBM_ENABLE=1 -DBIAS_ENABLED=1 -DZENDNN_ENABLE=1 -DZENDNN_X64=1")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -march=znver2 -Wreturn-type")

# INCLUDE FILES
# ZENDNN_GIT_ROOT includes
include_directories(${ZENDNN_GIT_ROOT}/inc ${ZENDNN_GIT_ROOT}/src ${ZENDNN_GIT_ROOT}/src/common ${ZENDNN_GIT_ROOT}/src/cpu)
# ZENDNN_BLIS_PATH includes
include_directories(${ZENDNN_BLIS_PATH}/include)
# GNU Includes
include_directories(/usr/lib/x86_64-linux-gnu/include)
# test utilities
include_directories(${ZENDNN_GIT_ROOT}/tests/api_tests)

# Link Library Directories
link_directories(${ZENDNN_GIT_ROOT}/_out/lib)
link_directories(${ZENDNN_BLIS_PATH}/lib)
link_directories(/usr/lib/x86_64-linux-gnu/lib)

file(GLOB My_Source_Files ./source/*.cpp)
add_executable(${PROJECT_NAME} ${My_Source_Files})

set(LINK_LIBRARY_LIST ${LINK_LIBRARY_LIST} amdZenDNN blis-mt alm)

# OpenCV -- Display Component
find_package(OpenCV REQUIRED)
# OpenCV 3/4 Support
if(${OpenCV_VERSION_MAJOR} EQUAL 3 OR ${OpenCV_VERSION_MAJOR} EQUAL 4)
	message("-- ${PROJECT_NAME} -- OpenCV Version Supported -- Version-${OpenCV_VERSION_MAJOR}.${OpenCV_VERSION_MINOR}.X")
	if(${OpenCV_VERSION_MAJOR} EQUAL 4)
	    target_compile_definitions(${PROJECT_NAME} PUBLIC USE_OPENCV_4=1)
    else()
	    target_compile_definitions(${PROJECT_NAME} PUBLIC USE_OPENCV_4=0)
    endif()
	include_directories(${OpenCV_INCLUDE_DIRS})
    #set(LINK_LIBRARY_LIST ${LINK_LIBRARY_LIST} ${OpenCV_LIBRARIES}) # TBD -- Investigate SegFault
else()
	message(FATAL_ERROR "${PROJECT_NAME} -- OpenCV -- Version-${OpenCV_VERSION_MAJOR}.${OpenCV_VERSION_MINOR}.X Not Supported")
endif()

target_link_libraries(${PROJECT_NAME} ${LINK_LIBRARY_LIST})

# Bring in weights & images to build folder
FILE(GLOB IMAGE_LIST ${CMAKE_CURRENT_SOURCE_DIR}/data/images/*.png)
FILE(GLOB IMAGE_BIN_LIST ${CMAKE_CURRENT_SOURCE_DIR}/data/images/*.bin)
file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/data/weights.bin DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/data/)
file(COPY ${IMAGE_LIST} DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/images/)
file(COPY ${IMAGE_BIN_LIST} DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/images/)

message("-- ${PROJECT_NAME} -- Using Compiler - Path:" ${CMAKE_CXX_COMPILER} "\t
        Version:" ${CMAKE_CXX_COMPILER_VERSION} "\tCompiler:" ${CMAKE_CXX_COMPILER_ID}${ColourReset})
message("-- ${PROJECT_NAME} -- CMAKE_CXX_FLAGS:${CMAKE_CXX_FLAGS}${ColourReset}")
message("-- ${PROJECT_NAME} -- Link Libraries: ${LINK_LIBRARY_LIST}${ColourReset}")
""")

def generateDeployCPP(graph,fileName):
    print('creating ' + fileName + ' ...')
    with open(fileName, 'w') as f:
        generateLicenseForCPP(f)
        f.write( \
"""
#include <assert.h>

#include <chrono>
#include <numeric>
#include <vector>
#include <unordered_map>
#include <iostream>
#include <fstream>
#include <algorithm>

#include "test_utils.hpp"
#include "zendnn_logging.hpp"

#include <opencv2/opencv.hpp>

using namespace zendnn;
using namespace std;

#if USE_OPENCV_4
#define CV_BGR2GRAY COLOR_BGR2GRAY
#endif

memory::dim product(const memory::dims &dims)
{
    return std::accumulate(dims.begin(), dims.end(), (memory::dim)1,
                           std::multiplies<memory::dim>());
}

// Check function status
#define ERROR_CHECK_STATUS(call)                                                               \
    {                                                                                          \
        int status = (call);                                                                   \
        if (status != 0)                                                                       \
        {                                                                                      \
            printf("ERROR: failed with status = (%d) at " __FILE__ "#%d\\n", status, __LINE__); \
            return status;                                                                     \
        }                                                                                      \
    }

// Intialize weights & bias tensors
static int initializeTensor(std::vector<float> *tensor, size_t tensorSize, FILE *fp, const char *binaryFilename)
{
    size_t itemsize = sizeof(float);
    unsigned int h[2] = {0};

    fread(h, 1, sizeof(h), fp);
    if (h[0] != 0xf00dd1e1 || (size_t)h[1] != (tensorSize * itemsize))
    {
        printf("ERROR: invalid data (magic,size)=(0x%x,%x) in %s at byte position %lu -- expected size is %ld\\n", 
                h[0], h[1], binaryFilename, ftell(fp) - sizeof(h), tensorSize * itemsize);
        return -1;
    }

    void *ptr = tensor->data();
    size_t n = fread(ptr, itemsize, tensorSize, fp);
    if (n != tensorSize)
    {
        printf("ERROR: expected char[%zu], but got char[%zu] in %s\\n", tensorSize * itemsize, n * itemsize, binaryFilename);
        return -1;
    }

    return 0;
}

// Model Zen DNN Implementation
int model_setup(engine::kind engine_kind, const char *binaryFilename, const char *imageFilename, int times)
{

    using tag = memory::format_tag;
    using dt = memory::data_type;

    // Check weights file
    FILE *fp__variables = fopen(binaryFilename, "rb");
    if (!fp__variables)
    {
        printf("ERROR: unable to open: %s\\n", binaryFilename);
        return -1;
    }
    {
        unsigned int magic = 0;
        fread(&magic, 1, sizeof(magic), fp__variables);
        if (magic != 0xf00dd1e0)
        {
            fclose(fp__variables);
            printf("ERROR: invalid file magic in %s\\n", binaryFilename);
            return -1;
        }
    }

    // Initialize engine and stream
    engine eng(engine_kind, 0);
    stream s(eng);

    // Create network
    std::vector<primitive> net;
    std::vector<std::unordered_map<int, memory>> net_args;
""")
        for tensor in graph.inputs:
            f.write( \
"""
    // Allocate input data %s
    std::vector<float> user_src(%s);
    memory::dims %s_src_tensor_dims = {%s};

    //Create memory that describes data layout
    auto user_src_memory = memory({{%s_src_tensor_dims}, dt::f32, tag::nchw}, eng);
    write_to_zendnn_memory(user_src.data(), user_src_memory);
"""% (tensor.name, ' * '.join([str(v) for v in tensor.shape]), 
        tensor.name, ', '.join([str(v) for v in tensor.shape]), tensor.name))
        for tensor in graph.outputs:
            f.write( \
"""
    // Allocate output data %s
    std::vector<float> user_dst(%s);
    memory::dims %s_dims = {%s};
"""% (tensor.name, ' * '.join([str(v) for v in tensor.shape]), tensor.name, ', '.join([str(v) for v in tensor.shape])))
        for tensor in graph.inputs:
            f.write( \
"""
    // Load Input Image - Binary Files
    std::ifstream input_binary_image(imageFilename, std::ios::binary);

    input_binary_image.seekg(0, input_binary_image.end);
    int elementCount = (input_binary_image.tellg() / sizeof(float)); // total number of elements
    input_binary_image.seekg(0, input_binary_image.beg);

    if (elementCount == (%s))
    {
        input_binary_image.read(reinterpret_cast<char *>(user_src.data()), user_src.size() * sizeof(float));
    }
    else
    {
        printf("ERROR: invalid Binary Image File -- %%s: Total Pixels:%s Received: %%d\\n", imageFilename, elementCount);
        return -1;
    }

    // Allocate initializers
"""% (' * '.join([str(v) for v in tensor.shape]), ' * '.join([str(v) for v in tensor.shape])))
        for tensor in graph.initializers:
            i = 0
            while i < len(tensor.shape):
                if tensor.shape[-1] == 1:
                    del tensor.shape[-1]
                i += 1
            f.write( \
"""    memory::dims %s_dims = {%s};
"""%(tensor.name, ', '.join([str(v) for v in tensor.shape])))
        f.write( 
"""
    // Allocate locals
""")
        for tensor in graph.locals:
            j = 0
            while j < len(tensor.shape):
                if tensor.shape[-1] == 1:
                    del tensor.shape[-1]
                j += 1
            f.write( \
"""    memory::dims %s_dims = {%s};
"""%(tensor.name, ', '.join([str(v) for v in tensor.shape])))
        layerNumber = 1
        for node in graph.nodes:
            if node.type == 'conv':
                # TODO: add support for group convolution
                pads = node.attr.get('pads')
                strides = node.attr.get('strides')
                f.write( \
"""
    // Start: Model Layer %d - conv
    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - conv Setup");
    // Conv set dimensions
    memory::dims %s_padding = {%s};
    memory::dims %s_strides = {%s};
    // Allocate buffers for weights and bias
    std::vector<float> %s(product(%s_dims));
    ERROR_CHECK_STATUS(initializeTensor(&%s, product(%s_dims), fp__variables, binaryFilename));
"""%(layerNumber, layerNumber, node.outputs[0], ', '.join([str(v) for v in pads]), 
    node.outputs[0], ', '.join([str(v) for v in strides]), node.inputs[1], node.inputs[1], 
    node.inputs[1], node.inputs[1]))
                if len(node.inputs) == 3:
                    f.write( \
"""
    std::vector<float> %s(product(%s_dims));
    ERROR_CHECK_STATUS(initializeTensor(&%s, product(%s_dims), fp__variables, binaryFilename));
"""%(node.inputs[2], node.inputs[2], node.inputs[2], node.inputs[2]))
                f.write( \
"""
    // Create memory that describes data layout in the buffers
    auto %s_user_weights_memory = memory({{%s_dims}, dt::f32, tag::oihw}, eng);
    write_to_zendnn_memory(%s.data(), %s_user_weights_memory);
"""%(node.outputs[0], node.inputs[1], node.inputs[1], node.outputs[0]))
                if len(node.inputs) == 3:
                    f.write( \
"""
    auto %s_user_bias_memory = memory({{%s_dims[1]}, dt::f32, tag::x}, eng);
    write_to_zendnn_memory(%s.data(), %s_user_bias_memory);

    // Create convolution memory descriptors with layout tag::any
"""%(node.outputs[0], node.inputs[2], node.inputs[2], node.outputs[0]))
                if layerNumber == 1:
                    f.write( \
"""
    auto %s_src_md = memory::desc({%s_src_tensor_dims}, dt::f32, tag::any);
"""%(node.outputs[0], node.inputs[0]))
                if len(node.inputs) == 3:
                    f.write( \
"""
    auto %s_bias_md = memory::desc({%s_dims[1]}, dt::f32, tag::any);
"""%(node.outputs[0], node.inputs[2]))
                f.write( \
"""
    auto %s_weights_md = memory::desc({%s_dims}, dt::f32, tag::any);
    auto %s_dst_md = memory::desc({%s_dims}, dt::f32, tag::any);

    // Create convolution descriptor
"""%(node.outputs[0], node.inputs[1], node.outputs[0], node.outputs[0]))
                if layerNumber == 1:
                    if len(node.inputs) == 3:
                        f.write( \
"""
    auto %s_desc = convolution_forward::desc(prop_kind::forward_inference,
                                            algorithm::convolution_gemm, %s_src_md, %s_weights_md,
                                            %s_bias_md, %s_dst_md, %s_strides, %s_padding,
                                            %s_padding);
"""%(node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0],
    node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0]))
                    else:
                        f.write( \
"""
    auto %s_desc = convolution_forward::desc(prop_kind::forward_inference,
                                            algorithm::convolution_gemm, %s_src_md, %s_weights_md,
                                            %s_dst_md, %s_strides, %s_padding, %s_padding);
"""%(node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0],
    node.outputs[0], node.outputs[0], node.outputs[0]))
                else:
                    if len(node.inputs) == 3:
                        f.write( \
"""
    auto %s_desc = convolution_forward::desc(prop_kind::forward_inference,
                                            algorithm::convolution_gemm, %s_dst_memory.get_desc(), %s_weights_md,
                                            %s_bias_md, %s_dst_md, %s_strides, %s_padding,
                                            %s_padding);
    // End: Create convolution descriptor
"""%(node.outputs[0], node.inputs[0], node.outputs[0], node.outputs[0],
    node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0]))
                    else:
                        f.write( \
"""
    auto %s_desc = convolution_forward::desc(prop_kind::forward_inference,
                                            algorithm::convolution_gemm, %s_dst_memory.get_desc(), %s_weights_md,
                                            %s_dst_md, %s_strides, %s_padding, %s_padding);
    // End: Create convolution descriptor
"""%(node.outputs[0], node.inputs[0], node.outputs[0], node.outputs[0],
    node.outputs[0], node.outputs[0], node.outputs[0]))
                f.write( \
"""
    // Create a convolution primitive descriptor
    auto %s_prim_desc = convolution_forward::primitive_desc(%s_desc, eng);
"""%(node.outputs[0], node.outputs[0]))
                if layerNumber == 1:
                    f.write( \
"""
    // Check data and weights formats - Reorder
    auto %s_src_memory = user_src_memory;
    if (%s_prim_desc.src_desc() != user_src_memory.get_desc())
    {
        %s_src_memory = memory(%s_prim_desc.src_desc(), eng);
        net.push_back(reorder(user_src_memory, %s_src_memory));
        net_args.push_back({{ZENDNN_ARG_FROM, user_src_memory},
                            {ZENDNN_ARG_TO, %s_src_memory}});
    }

    auto %s_weights_memory = %s_user_weights_memory;
    if (%s_prim_desc.weights_desc() != %s_user_weights_memory.get_desc())
    {
        %s_weights_memory = memory(%s_prim_desc.weights_desc(), eng);
        reorder(%s_user_weights_memory, %s_weights_memory)
            .execute(s, %s_user_weights_memory, %s_weights_memory);
    }
"""%(node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0],
    node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0],
    node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0],
    node.outputs[0]))
                else:
                    f.write( \
"""
    // Check weights formats - Reorder
    auto %s_weights_memory = %s_user_weights_memory;
    if (%s_prim_desc.weights_desc() != %s_user_weights_memory.get_desc())
    {
        %s_weights_memory = memory(%s_prim_desc.weights_desc(), eng);
        reorder(%s_user_weights_memory, %s_weights_memory)
            .execute(s, %s_user_weights_memory, %s_weights_memory);
    }
"""%(node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0],
    node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0]))
                f.write( \
"""
    // Create a memory primitive for output
    auto %s_dst_memory = memory(%s_prim_desc.dst_desc(), eng);
"""%(node.outputs[0], node.outputs[0]))
                if layerNumber == 1:
                    if len(node.inputs) == 3:
                        f.write( \
"""
    // Create a convolution primitive and add it to the net
    net.push_back(convolution_forward(%s_prim_desc));
    net_args.push_back({{ZENDNN_ARG_SRC, %s_src_memory},
                        {ZENDNN_ARG_WEIGHTS, %s_weights_memory},
                        {ZENDNN_ARG_BIAS, %s_user_bias_memory},
                        {ZENDNN_ARG_DST, %s_dst_memory}});
"""%(node.outputs[0], node.outputs[0], node.outputs[0],
    node.outputs[0], node.outputs[0]))
                    else:
                        f.write( \
"""
    // Create a convolution primitive and add it to the net
    net.push_back(convolution_forward(%s_prim_desc));
    net_args.push_back({{ZENDNN_ARG_SRC, %s_src_memory},
                        {ZENDNN_ARG_WEIGHTS, %s_weights_memory},
                        {ZENDNN_ARG_DST, %s_dst_memory}});
"""%(node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0]))
                else:
                    if len(node.inputs) == 3:
                        f.write( \
"""
    // Create a convolution primitive and add it to the net
    net.push_back(convolution_forward(%s_prim_desc));
    net_args.push_back({{ZENDNN_ARG_SRC, %s_dst_memory},
                        {ZENDNN_ARG_WEIGHTS, %s_weights_memory},
                        {ZENDNN_ARG_BIAS, %s_user_bias_memory},
                        {ZENDNN_ARG_DST, %s_dst_memory}});
"""%(node.outputs[0], node.inputs[0], node.outputs[0],
    node.outputs[0], node.outputs[0]))
                    else:
                        f.write( \
"""
    // Create a convolution primitive and add it to the net
    net.push_back(convolution_forward(%s_prim_desc));
    net_args.push_back({{ZENDNN_ARG_SRC, %s_dst_memory},
                        {ZENDNN_ARG_WEIGHTS, %s_weights_memory},
                        {ZENDNN_ARG_DST, %s_dst_memory}});
"""%(node.outputs[0], node.inputs[0], node.outputs[0], node.outputs[0]))
                f.write( \
"""
    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - conv Setup Complete");
    // End: Model Layer %d - conv
"""%(layerNumber, layerNumber))
                layerNumber += 1
            elif node.type == 'max_pool' or node.type == 'avg_pool':
                pads = node.attr.get('pads')
                strides = node.attr.get('strides')
                kernel = node.attr.get('kernel_shape')
                f.write( \
"""
    // Start: Model Layer %d - pool
    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - pool Setup");
    // Pool set dimensions
    memory::dims %s_padding = {%s};
    memory::dims %s_strides = {%s};
    memory::dims %s_kernel = {%s};

    // Pool - set dst memory desc
    auto %s_dst_md = memory::desc({%s_dims}, dt::f32, tag::any);

    // Create pooling primitive
    auto %s_desc = pooling_forward::desc(prop_kind::forward_inference,
                                            algorithm::%s, %s_dst_memory.get_desc(), %s_dst_md,
                                            %s_strides, %s_kernel, %s_padding, %s_padding);
    auto %s_pd = pooling_forward::primitive_desc(%s_desc, eng);

    // Pool - dst memory
    auto %s_dst_memory = memory(%s_pd.dst_desc(), eng);

    // Create a pooling primitive and add it to the net
    net.push_back(pooling_forward(%s_pd));
    net_args.push_back({{ZENDNN_ARG_SRC, %s_dst_memory},
                        {ZENDNN_ARG_DST, %s_dst_memory}});
    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - pool Setup Complete");
    // End: Model Layer %d - pool

"""%(layerNumber, layerNumber, node.outputs[0], ', '.join([str(v) for v in pads]), node.outputs[0], 
    ', '.join([str(v) for v in strides]), node.outputs[0], ', '.join([str(v) for v in kernel]), node.outputs[0], 
    node.outputs[0], node.outputs[0], "pooling_max" if node.type == 'max_pool' else "pooling_avg",
    node.inputs[0], node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], 
    node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], 
    node.outputs[0], node.inputs[0], node.outputs[0], layerNumber, layerNumber))
                layerNumber += 1
            elif node.type == 'gemm':
                weights_shape = graph.tensor_dict[node.inputs[1]].shape
                f.write( \
"""
    // Start: Model Layer %d - inner product
    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - inner product Setup");

    // Weights and bias
    std::vector<float> %s_weights(product(%s_dims));
    ERROR_CHECK_STATUS(initializeTensor(&%s_weights, product(%s_dims), fp__variables, binaryFilename));
    std::vector<float> %s_bias(product(%s_dims));
    ERROR_CHECK_STATUS(initializeTensor(&%s_bias, product(%s_dims), fp__variables, binaryFilename));

    // Create memory that describes data layout in the buffers
    auto %s_user_weights_memory = memory({{%s_dims}, dt::f32, tag::%s}, eng);
    write_to_zendnn_memory(%s_weights.data(), %s_user_weights_memory);
    auto %s_user_bias_memory = memory({{%s_dims[1]}, dt::f32, tag::x}, eng);
    write_to_zendnn_memory(%s_bias.data(), %s_user_bias_memory);

    // Create ip memory descriptors with layout tag::any
    auto %s_weights_md = memory::desc({%s_dims}, dt::f32, tag::any);
    auto %s_bias_md = memory::desc({%s_dims[1]}, dt::f32, tag::any);
    auto %s_dst_md = memory::desc({%s_dims}, dt::f32, tag::any);

    // Create inner product primitive descriptor
    auto %s_desc = inner_product_forward::desc(prop_kind::forward_inference,
                                                %s_dst_memory.get_desc(), %s_weights_md, %s_bias_md, 
                                                %s_dst_md);
    auto %s_prim_desc = inner_product_forward::primitive_desc(%s_desc, eng);
"""%(layerNumber, layerNumber, node.inputs[1], node.inputs[1], node.inputs[1], 
    node.inputs[1], node.inputs[2], node.inputs[2], node.inputs[2], node.inputs[2], 
    node.inputs[1], node.inputs[1], 'oihw' if len(weights_shape) == 4 else 'nc', 
    node.inputs[1], node.inputs[1], node.inputs[2],
    node.inputs[2], node.inputs[2], node.inputs[2], node.inputs[1], node.inputs[1], 
    node.inputs[2], node.inputs[2], node.outputs[0], node.outputs[0], node.outputs[0], 
    node.inputs[0], node.inputs[1], node.inputs[2], node.outputs[0], node.outputs[0], 
    node.outputs[0]))
                f.write( \
"""
    // Check weights formats - Reorder
    auto %s_weights_memory = %s_user_weights_memory;
    if (%s_prim_desc.weights_desc() != %s_user_weights_memory.get_desc())
    {
        %s_weights_memory = memory(%s_prim_desc.weights_desc(), eng);
        reorder(%s_user_weights_memory, %s_weights_memory)
            .execute(s, %s_user_weights_memory, %s_weights_memory);
    }

    // IP - dst memory
    auto %s_dst_memory = memory(%s_prim_desc.dst_desc(), eng);

    // Create a IP primitive and add it to the net
    net.push_back(inner_product_forward(%s_prim_desc));
    net_args.push_back({{ZENDNN_ARG_SRC, %s_dst_memory},
                        {ZENDNN_ARG_WEIGHTS, %s_weights_memory},
                        {ZENDNN_ARG_BIAS, %s_user_bias_memory},
                        {ZENDNN_ARG_DST, %s_dst_memory}});

    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - inner product Setup Complete");
    // End: Model Layer %d - inner product
"""%(node.inputs[1], node.inputs[1], node.outputs[0], node.inputs[1], node.inputs[1],
    node.outputs[0], node.inputs[1], node.inputs[1], node.inputs[1], node.inputs[1],
    node.outputs[0], node.outputs[0], node.outputs[0], node.inputs[0], node.inputs[1], 
    node.inputs[2], node.outputs[0], layerNumber, layerNumber))
                layerNumber += 1
            elif node.type == 'relu':
                alpha = node.attr.get('alpha')
                f.write( \
"""
    // Start: Model Layer %d - relu
    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - relu Setup");
    const float %s_alpha = %f;
    const float %s_beta = 0.0f;

    auto %s_dst_md = memory::desc({%s_dims}, dt::f32, tag::any);

    // Create relu primitive descriptor
    auto %s_desc = eltwise_forward::desc(prop_kind::forward_inference,
                                            algorithm::eltwise_relu, %s_dst_memory.get_desc(),
                                            %s_alpha, %s_beta);
    auto %s_prim_desc = eltwise_forward::primitive_desc(%s_desc, eng);

    // Relu - dst memory
    auto %s_dst_memory = memory(%s_prim_desc.dst_desc(), eng);

    // Create a relu primitive and add it to the net
    net.push_back(eltwise_forward(%s_prim_desc));
    net_args.push_back({{ZENDNN_ARG_SRC, %s_dst_memory},
                        {ZENDNN_ARG_DST, %s_dst_memory}});
    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - relu Setup Complete");
    // End: Model Layer %d - relu
"""%(layerNumber, layerNumber, node.outputs[0], alpha, node.outputs[0], node.outputs[0], 
    node.outputs[0],node.outputs[0], node.inputs[0], node.outputs[0], node.outputs[0], 
    node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], 
    node.inputs[0], node.outputs[0], layerNumber, layerNumber))
                layerNumber += 1
            elif node.type == 'lrn':
                size = node.attr.get('size')
                alpha = node.attr.get('alpha')
                beta = node.attr.get('beta')
                bias = node.attr.get('bias')
                f.write( \
"""
    // Start: Model Layer %d - LRN
    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - LRN Setup");
    const float %s_alpha = %f;
    const float %s_beta = %f;
    const float %s_bias = %f;
    const unsigned long %s_size = %lu;

    // Create LRN primitive descriptor
    auto %s_desc = lrn_forward::desc(prop_kind::forward_inference,
                                    algorithm::%s, 
                                    %s_dst_memory.get_desc(), %s_size, %s_alpha, %s_beta, %s_bias);

    auto %s_prim_desc = lrn_forward::primitive_desc(%s_desc, eng);

    // LRN - dst memory
    auto %s_dst_memory = memory(%s_prim_desc.dst_desc(), eng);

    // Create a LRN primitive and add it to the net
    net.push_back(lrn_forward(%s_prim_desc));
    net_args.push_back({{ZENDNN_ARG_SRC, %s_dst_memory},
                        {ZENDNN_ARG_DST, %s_dst_memory}});
    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - LRN Setup Complete");
    // End: Model Layer %d - LRN
"""%(layerNumber, layerNumber, node.outputs[0], alpha, node.outputs[0], beta, node.outputs[0], 
    bias, node.outputs[0], size, node.outputs[0] , 
    "lrn_across_channels" if node.attr.get('mode') == 0 else "lrn_within_channel",
    node.inputs[0], node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0],
    node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0],
    node.inputs[0], node.outputs[0], layerNumber, layerNumber))
                layerNumber += 1
            elif node.type == 'batch_norm':
                epsilon = node.attr.get('epsilon')
                f.write( \
"""
    // Start: Model Layer %d - Batch Normalization
    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - Batch Normalization Setup");
    const float %s_epsilon = %ef;

    // Allocate buffers for mean, variance, scale and bias
    std::vector<float> %s(product(%s_dims));
    ERROR_CHECK_STATUS(initializeTensor(&%s, product(%s_dims), fp__variables, binaryFilename));
    std::vector<float> %s(product(%s_dims));
    ERROR_CHECK_STATUS(initializeTensor(&%s, product(%s_dims), fp__variables, binaryFilename));
    std::vector<float> %s(product(%s_dims));
    ERROR_CHECK_STATUS(initializeTensor(&%s, product(%s_dims), fp__variables, binaryFilename));
    std::vector<float> %s(product(%s_dims));
    ERROR_CHECK_STATUS(initializeTensor(&%s, product(%s_dims), fp__variables, binaryFilename));

    // Allocate memory for scaleshift tensor
    memory::dims %s_scale_shift_dims = {2, %s};    
    std::vector<float> %s_scale_shift_data(product(%s_scale_shift_dims));

    auto %s_mid = %s_scale_shift_data.begin() + %s_scale_shift_dims[1];
    
    copy(%s.begin(), %s.end(), %s_scale_shift_data.begin());
    copy(%s.begin(), %s.end(), %s_mid);

    // Create memory that describes data layout in the buffers
    auto %s_user_mean_memory = memory({{%s_dims}, dt::f32, tag::nc}, eng);
    write_to_zendnn_memory(%s.data(), %s_user_mean_memory);
    auto %s_user_variance_memory = memory({{%s_dims}, dt::f32, tag::nc}, eng);
    write_to_zendnn_memory(%s.data(), %s_user_variance_memory);
    auto %s_user_scale_memory = memory({{%s_dims}, dt::f32, tag::nc}, eng);
    write_to_zendnn_memory(%s.data(), %s_user_scale_memory);
    auto %s_user_bias_memory = memory({{%s_dims}, dt::f32, tag::nc}, eng);
    write_to_zendnn_memory(%s.data(), %s_user_bias_memory);
    auto %s_scale_shift_memory = memory({{%s_scale_shift_dims}, dt::f32, tag::nc}, eng);
    write_to_zendnn_memory(%s_scale_shift_data.data(), %s_scale_shift_memory);

    // Create memory descriptors and memory objects. 
    auto %s_dst_md = memory::desc({%s_dims}, dt::f32, tag::any);

    auto %s_desc = batch_normalization_forward::desc(
            prop_kind::forward_inference, %s_dst_memory.get_desc(), %s_epsilon,
            normalization_flags::use_scale_shift);

    // Create primitive descriptor.
    auto %s_prim_desc = batch_normalization_forward::primitive_desc(%s_desc, eng);

    // Create a memory primitive for output
    auto %s_dst_memory = memory(%s_prim_desc.dst_desc(), eng);

    // Create a batch normalization primitive and add it to the net
    net.push_back(batch_normalization_forward(%s_prim_desc));
    net_args.push_back({{ZENDNN_ARG_SRC, %s_dst_memory},
                        {ZENDNN_ARG_MEAN, %s_user_mean_memory},
                        {ZENDNN_ARG_VARIANCE, %s_user_variance_memory},
                        {ZENDNN_ARG_SCALE_SHIFT, %s_scale_shift_memory},
                        {ZENDNN_ARG_DST, %s_dst_memory}});

    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - Batch Normalization Setup Complete");
    // End: Model Layer %d - Batch Normalization
"""%(layerNumber, layerNumber, node.outputs[0], epsilon, 
    node.inputs[1], node.inputs[1], node.inputs[1], node.inputs[1],
    node.inputs[2], node.inputs[2], node.inputs[2], node.inputs[2],
    node.inputs[3], node.inputs[3], node.inputs[3], node.inputs[3],
    node.inputs[4], node.inputs[4], node.inputs[4], node.inputs[4],
    node.outputs[0], str(graph.tensor_dict[node.inputs[0]].shape[1]), 
    node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], 
    node.outputs[0], node.inputs[1], node.inputs[1], node.outputs[0], 
    node.inputs[2], node.inputs[2], node.outputs[0], node.inputs[3], 
    node.inputs[3], node.inputs[3], node.inputs[3], node.inputs[4], 
    node.inputs[4], node.inputs[4], node.inputs[4], node.inputs[1], 
    node.inputs[1], node.inputs[1], node.inputs[1], node.inputs[2], 
    node.inputs[2], node.inputs[2], node.inputs[2], node.outputs[0], 
    node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], 
    node.outputs[0], node.outputs[0], node.inputs[0], node.outputs[0], 
    node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], 
    node.outputs[0], node.inputs[0], node.inputs[3],node.inputs[4], 
    node.outputs[0], node.outputs[0], layerNumber, layerNumber))
                layerNumber += 1
            elif node.type == 'sum' or node.type == 'add':
                num_src = len(node.inputs)
                f.write( \
"""
    // Start: Model Layer %d - Sum
    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - Sum Setup");

    // Sum number of inputs.
    const int %s_num_src = %d;

    // Scaling factors -- supports only 1
    std::vector<float> %s_scales(%s_num_src, 1);

    // Create an array of memory descriptors and memory objects for src tensors.
    std::vector<memory::desc> %s_src_md;
    std::vector<memory> %s_src_mem;

"""%(layerNumber, layerNumber, node.outputs[0], num_src, node.outputs[0], 
    node.outputs[0], node.outputs[0], node.outputs[0]))
                for i in range(num_src):
                    f.write( \
"""    
    %s_src_md.push_back(%s_dst_md);
    %s_src_mem.push_back(%s_dst_memory);

"""%(node.outputs[0], node.inputs[i], node.outputs[0], node.inputs[i]))
                f.write( \
"""
    // Create primitive descriptor.
    auto %s_pd = sum::primitive_desc(%s_scales, %s_src_md, eng);

    // Create memory object for dst.
    auto %s_dst_memory = memory(%s_pd.dst_desc(), eng);

    // Adding sum primitive to net 
    net.push_back(sum(%s_pd));
    net_args.push_back({{ZENDNN_ARG_DST, %s_dst_memory}});
"""%(node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], 
    node.outputs[0], node.outputs[0], node.outputs[0]))
                for i in range(num_src):
                    f.write( \
"""
    net_args.push_back({{ZENDNN_ARG_MULTIPLE_SRC + %d, %s_src_mem[%d]}});

    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - Sum Setup Complete");
    // End: Model Layer %d - Sum
"""%(i, node.outputs[0], i, layerNumber, layerNumber))
                layerNumber += 1
            elif node.type == 'softmax':
                axis = node.attr.get('axis')
                f.write( \
"""
    // Start: Model Layer %d - softMax
    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - softmax Setup");

    // Softmax axis.
    const int sm_axis = %d;

    // Create softmax primitive
    auto %s_desc = softmax_forward::desc(prop_kind::forward_inference, %s_dst_memory.get_desc(), sm_axis);
    auto %s_pd = softmax_forward::primitive_desc(%s_desc, eng);

    // softmax - dst memory
    auto %s_dst_memory = memory(%s_pd.dst_desc(), eng);

    // Create a softmax primitive and add it to the net
    net.push_back(softmax_forward(%s_pd));
    net_args.push_back({{ZENDNN_ARG_SRC, %s_dst_memory},
                        {ZENDNN_ARG_DST, %s_dst_memory}});
    zendnnInfo(ZENDNN_TESTLOG, "Model Layer %d - softmax Setup Complete");
    // End: Model Layer %d- softMax
"""%(layerNumber, layerNumber, axis, node.outputs[0], node.inputs[0], node.outputs[0], 
    node.outputs[0], node.outputs[0], node.outputs[0], node.outputs[0], node.inputs[0], node.outputs[0], layerNumber, layerNumber))
                layerNumber += 1
        f.write( \
"""
    {
        unsigned int magic = 0;
        fread(&magic, 1, sizeof(magic), fp__variables);
        fclose(fp__variables);
        if (magic != 0xf00dd1e2)
        {
            printf("ERROR: invalid eoff magic in %s\\n", binaryFilename);
            return -1;
        }
    }

    // Execute primitives
    // For this example, the net is executed multiple times and each execution is timed individually.
    float AverageTime = 0;
    for (int j = 0; j < times; ++j)
    {
        auto start_exec_time = std::chrono::high_resolution_clock::now();
        assert(net.size() == net_args.size() && "something is missing");
        for (size_t i = 0; i < net.size(); ++i)
        {
            net.at(i).execute(s, net_args.at(i));
        }
        // Wait for the computation to finalize.
        s.wait();
        auto end_exec_time = std::chrono::high_resolution_clock::now();
        auto exec_duration = end_exec_time - start_exec_time;
        auto exec_in_millis = std::chrono::duration_cast<std::chrono::milliseconds>(exec_duration);
        printf("Execution:%d -- \t%.8f ms\\n", j, (float)exec_in_millis.count());

        // Read output from engine
        auto start_mem_time = std::chrono::high_resolution_clock::now();
""")
        for tensor in graph.outputs:
            f.write ( \
"""
        read_from_zendnn_memory(user_dst.data(), %s_dst_memory);
"""%(tensor.name))
        f.write( \
"""
        auto end_mem_time = std::chrono::high_resolution_clock::now();
        auto exec_mem_duration = end_mem_time - start_mem_time;
        auto mem_in_millis = std::chrono::duration_cast<std::chrono::milliseconds>(exec_mem_duration);
        //printf("Mem Transfer:%d -- \t%.8f ms\\n", j, (float)mem_in_millis.count());

        AverageTime += (float)exec_in_millis.count() + (float)mem_in_millis.count();
    }
    printf("\\nAvg Inference Time -- \t%.8f ms\\n",(float)(AverageTime/times));
        return 0;
}

int main(int argc, const char **argv)
{
    zendnnInfo(ZENDNN_TESTLOG, "zendnn_app test starts");

    // check command-line usage
    if (argc != 3)
    {
        printf(
            "\\n"
            "Usage: ZENDNN_LOG_OPTS=ALL:5 ZENDNN_VERBOSE=1 ./zendnn_app [weights.bin] [imageName]\\n"
            "\\n"
            "   <weights.bin>: weights file to be used for the inference\\n"
            "   <imageName>: image file to be used for the inference\\n"
            "\\n");
        return -1;
    }

    engine::kind engine_kind = zendnn::engine::kind::cpu;
    const char *weights = argv[1];
    const char *image = argv[2];
    int NumExecution = 10;

    try
    {
        auto begin = chrono::duration_cast<chrono::milliseconds>(
                         chrono::steady_clock::now().time_since_epoch())
                         .count();

        int app_status = model_setup(engine_kind, weights, image, NumExecution);

        auto end = chrono::duration_cast<chrono::milliseconds>(
                       chrono::steady_clock::now().time_since_epoch())
                       .count();
        zendnnInfo(ZENDNN_TESTLOG, "Use time ", (end - begin) / (NumExecution + 0.0));
    }

    catch (error &e)
    {
        std::cerr << "status: " << e.status << std::endl;
        std::cerr << "message: " << e.message << std::endl;
    }
    zendnnInfo(ZENDNN_TESTLOG, "zendnn_app test ends\\n");
    return 0;
}
""")

def generateBinary(graph,fileName):
    VARIABLES_FILE_MAGIC = 0xF00DD1E0
    VARIABLES_DATA_MAGIC = 0xF00DD1E1
    VARIABLES_EOFF_MAGIC = 0xF00DD1E2
    print('creating ' + fileName + ' ...')
    with open(fileName, 'wb') as f:
        f.write(struct.pack('I', VARIABLES_FILE_MAGIC))
        for tensor in graph.initializers:
            binary = graph.binaries[tensor.name]
            f.write(struct.pack('II', VARIABLES_DATA_MAGIC, len(binary)))
            f.write(binary)
        f.write(struct.pack('I', VARIABLES_EOFF_MAGIC))

def generateCode(graph,argmaxOutput,outputFolder):
    if not os.path.isdir(outputFolder):
        os.mkdir(outputFolder)
        os.mkdir(outputFolder + '/data')
        os.mkdir(outputFolder + '/source')
    generateCMakeFiles(graph,outputFolder)
    generateBinary(graph,outputFolder + '/data/weights.bin')
    generateDeployCPP(graph,outputFolder + '/source/zentest.cpp')

def main():
    usage = """
Usage: python nnir_to_zendnn.py <nnirInputFolder> <outputFolder>

"""
    pos = 1
    argmaxOutput = None
    while len(sys.argv[pos:]) >= 2 and sys.argv[pos][:2] == '--':
        if sys.argv[pos] != '--help':
            print('ERROR: invalid option: %s' % (sys.argv[pos]))
        print(usage)
        sys.exit(1)
        pos = pos + 2
    if len(sys.argv[pos:]) < 2:
        print(usage)
        sys.exit(1)
    inputFolder = sys.argv[pos]
    outputFolder = sys.argv[pos+1]
    print('reading IR model from ' + inputFolder + ' ...')
    graph = ir.IrGraph(True)
    graph.fromFile(inputFolder)
    print('creating C code in ' + outputFolder + ' ...')
    generateCode(graph,argmaxOutput,outputFolder)


if __name__ == '__main__':
    main()
