#!/usr/bin/env python

import resource_to_cpp
from platform_methods import run_in_subprocess

Import("env")
Import("env_modules")

env_oidn = env_modules.Clone()

# Thirdparty source files
thirdparty_dir = "#thirdparty/oidn/"
thirdparty_sources = [
    "core/api.cpp",
    "core/device.cpp",
    "core/filter.cpp",
    "core/network.cpp",
    "core/autoencoder.cpp",
    "core/transfer_function.cpp",
    "weights/rtlightmap_hdr.cpp",
    "mkl-dnn/src/common/batch_normalization.cpp",
    "mkl-dnn/src/common/concat.cpp",
    "mkl-dnn/src/common/convolution.cpp",
    "mkl-dnn/src/common/convolution_pd.cpp",
    "mkl-dnn/src/common/deconvolution.cpp",
    "mkl-dnn/src/common/eltwise.cpp",
    "mkl-dnn/src/common/engine.cpp",
    "mkl-dnn/src/common/inner_product.cpp",
    "mkl-dnn/src/common/inner_product_pd.cpp",
    "mkl-dnn/src/common/lrn.cpp",
    "mkl-dnn/src/common/memory.cpp",
    "mkl-dnn/src/common/memory_desc_wrapper.cpp",
    "mkl-dnn/src/common/mkldnn_debug.cpp",
    "mkl-dnn/src/common/mkldnn_debug_autogenerated.cpp",
    "mkl-dnn/src/common/pooling.cpp",
    "mkl-dnn/src/common/primitive.cpp",
    "mkl-dnn/src/common/primitive_attr.cpp",
    "mkl-dnn/src/common/primitive_desc.cpp",
    "mkl-dnn/src/common/primitive_exec_types.cpp",
    "mkl-dnn/src/common/primitive_iterator.cpp",
    "mkl-dnn/src/common/query.cpp",
    "mkl-dnn/src/common/reorder.cpp",
    "mkl-dnn/src/common/rnn.cpp",
    "mkl-dnn/src/common/scratchpad.cpp",
    "mkl-dnn/src/common/shuffle.cpp",
    "mkl-dnn/src/common/softmax.cpp",
    "mkl-dnn/src/common/stream.cpp",
    "mkl-dnn/src/common/sum.cpp",
    "mkl-dnn/src/common/utils.cpp",
    "mkl-dnn/src/common/verbose.cpp",
    "mkl-dnn/src/cpu/cpu_barrier.cpp",
    "mkl-dnn/src/cpu/cpu_concat.cpp",
    "mkl-dnn/src/cpu/cpu_engine.cpp",
    "mkl-dnn/src/cpu/cpu_memory.cpp",
    "mkl-dnn/src/cpu/cpu_reducer.cpp",
    "mkl-dnn/src/cpu/cpu_reorder.cpp",
    "mkl-dnn/src/cpu/cpu_sum.cpp",
    "mkl-dnn/src/cpu/jit_avx2_conv_kernel_f32.cpp",
    "mkl-dnn/src/cpu/jit_avx2_convolution.cpp",
    "mkl-dnn/src/cpu/jit_avx512_common_conv_kernel.cpp",
    "mkl-dnn/src/cpu/jit_avx512_common_conv_winograd_kernel_f32.cpp",
    "mkl-dnn/src/cpu/jit_avx512_common_convolution.cpp",
    "mkl-dnn/src/cpu/jit_avx512_common_convolution_winograd.cpp",
    "mkl-dnn/src/cpu/jit_avx512_core_fp32_wino_conv_2x3.cpp",
    "mkl-dnn/src/cpu/jit_avx512_core_fp32_wino_conv_4x3.cpp",
    "mkl-dnn/src/cpu/jit_avx512_core_fp32_wino_conv_4x3_kernel.cpp",
    "mkl-dnn/src/cpu/jit_sse42_conv_kernel_f32.cpp",
    "mkl-dnn/src/cpu/jit_sse42_convolution.cpp",
    "mkl-dnn/src/cpu/jit_transpose_src_utils.cpp",
    "mkl-dnn/src/cpu/jit_uni_eltwise.cpp",
    "mkl-dnn/src/cpu/jit_uni_pool_kernel_f32.cpp",
    "mkl-dnn/src/cpu/jit_uni_pooling.cpp",
    "mkl-dnn/src/cpu/jit_uni_reorder.cpp",
    "mkl-dnn/src/cpu/jit_uni_reorder_utils.cpp",
    "mkl-dnn/src/cpu/jit_utils/jit_utils.cpp",
    "mkl-dnn/src/cpu/jit_utils/jitprofiling/jitprofiling.c",
    "common/platform.cpp",
    "common/thread.cpp",
    "common/tensor.cpp",
]
thirdparty_sources = [thirdparty_dir + file for file in thirdparty_sources]

thirdparty_include_dirs = [
    "",
    "include",
    "mkl-dnn/include",
    "mkl-dnn/src",
    "mkl-dnn/src/common",
    "mkl-dnn/src/cpu/xbyak",
    "mkl-dnn/src/cpu",
]
thirdparty_include_dirs = [thirdparty_dir + file for file in thirdparty_include_dirs]


env_oidn.Prepend(CPPPATH=thirdparty_include_dirs)
env_oidn.Append(
    CPPDEFINES=[
        "MKLDNN_THR=MKLDNN_THR_SEQ",
        "OIDN_STATIC_LIB",
        "__STDC_CONSTANT_MACROS",
        "__STDC_LIMIT_MACROS",
        "DISABLE_VERBOSE",
        "MKLDNN_ENABLE_CONCURRENT_EXEC",
        "NDEBUG",
    ]
)

env_thirdparty = env_oidn.Clone()
env_thirdparty.disable_warnings()
env_thirdparty.add_source_files(env.modules_sources, thirdparty_sources)

weights_in_path = thirdparty_dir + "weights/rtlightmap_hdr.tza"
weights_out_path = thirdparty_dir + "weights/rtlightmap_hdr.cpp"

env_thirdparty.Depends(weights_out_path, weights_in_path)
env_thirdparty.CommandNoCache(weights_out_path, weights_in_path, resource_to_cpp.tza_to_cpp)

env_oidn.add_source_files(env.modules_sources, "denoise_wrapper.cpp")
env_modules.add_source_files(env.modules_sources, ["register_types.cpp", "lightmap_denoiser.cpp"])