#!/usr/bin/env python

import resource_to_cpp

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

env_oidn = env_modules.Clone()

# Thirdparty source files

thirdparty_obj = []

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.gen.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",
    ]
)
env_oidn.AppendUnique(CPPDEFINES=["NDEBUG"])  # No assert() even in debug builds.

env_thirdparty = env_oidn.Clone()
env_thirdparty.disable_warnings()

if env["disable_exceptions"]:
    # OIDN hard-requires exceptions, so we re-enable them here.
    if env.msvc and ("_HAS_EXCEPTIONS", 0) in env_thirdparty["CPPDEFINES"]:
        env_thirdparty["CPPDEFINES"].remove(("_HAS_EXCEPTIONS", 0))
        env_thirdparty.AppendUnique(CCFLAGS=["/EHsc"])
    elif not env.msvc and "-fno-exceptions" in env_thirdparty["CCFLAGS"]:
        env_thirdparty["CCFLAGS"].remove("-fno-exceptions")

env_thirdparty.add_source_files(thirdparty_obj, thirdparty_sources)
env.modules_sources += thirdparty_obj

if env["platform"] == "windows" and not env.msvc:
    env_thirdparty.Append(CPPFLAGS=["-mstackrealign"])

weights_in_path = thirdparty_dir + "weights/rtlightmap_hdr.tza"
weights_out_path = thirdparty_dir + "weights/rtlightmap_hdr.gen.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)

# Godot source files

module_obj = []

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

# Needed to force rebuilding the module files when the thirdparty library is updated.
env.Depends(module_obj, thirdparty_obj)