cmake_minimum_required(VERSION 3.18)
project(dots_tts_cpp LANGUAGES C CXX)

set(CMAKE_C_STANDARD 11)
set(CMAKE_CXX_STANDARD 17)

# ---- ggml library (from llama.cpp) ----
set(GGML_BUILD_COMMIT "dots_tts_cpp")
set(GGML_BUILD_NUMBER 0)
add_subdirectory(${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml ggml_build)

# ---- llama library (pre-built in llama.cpp/build/src) ----
set(LLAMA_BUILD_DIR ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/build)
find_library(LLAMA_LIB llama PATHS ${LLAMA_BUILD_DIR}/src NO_DEFAULT_PATH)
if(NOT LLAMA_LIB)
    message(FATAL_ERROR "llama library not found. Build llama.cpp first: cd ../llama/llama.cpp && cmake -B build && cmake --build build -j8 --target llama")
endif()
message(STATUS "Found llama: ${LLAMA_LIB}")

# ---- dots.tts.cpp (legacy target replaced by e2e_pipeline) ----
# dots_tts binary removed — use e2e_pipeline instead

# ---- llama library (pre-built in llama.cpp/build/src) ----

# ---- End-to-end pipeline ----
add_executable(e2e_pipeline
    src/e2e_pipeline.cpp
    src/gpt2_bpe_tokenizer.cpp
    src/gguf_extract.cpp
    src/modules/backbone/dit_forward.cpp
    src/modules/backbone/dit_attention.cpp
    src/modules/backbone/patchenc.cpp
    src/modules/vocoder/audiovae.cpp
    src/modules/vocoder/bigvgan_cpp.cpp
    src/modules/vocoder/lstm.cpp
    src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp
    src/llm_manual.cpp
    src/utils/dit_loader.cpp
    src/utils/patchenc_loader.cpp
    src/modules/speaker/campp.cpp
    src/modules/vocoder/audiovae_encoder.cpp
)
target_include_directories(e2e_pipeline PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/common
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(e2e_pipeline PRIVATE ${LLAMA_LIB} ggml pthread m openblas)

find_package(OpenMP)
if(OpenMP_CXX_FOUND)
    target_link_libraries(e2e_pipeline PRIVATE OpenMP::OpenMP_CXX)
    target_compile_options(e2e_pipeline PRIVATE ${OpenMP_CXX_FLAGS})
endif()

# ---- Standalone vocoder test ---- 
add_executable(test_vocoder
    test/test_vocoder.cpp
    src/modules/vocoder/bigvgan_cpp.cpp
    src/modules/vocoder/lstm.cpp
    src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp
)
target_include_directories(test_vocoder PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/common
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(test_vocoder PRIVATE ggml pthread m)

# ---- C++ Web demo server (replaces Flask) ----
add_executable(web_server tools/web_server.cpp)
target_link_libraries(web_server PRIVATE pthread)

# ---- Test dots LLM loading ----
add_executable(test_llm test/test_llm.cpp)
target_include_directories(test_llm PRIVATE
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/common
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(test_llm PRIVATE ${LLAMA_LIB} ggml pthread m)

# ---- Test C++ BPE tokenizer ----
add_executable(test_tokenizer test/test_tokenizer.cpp src/gpt2_bpe_tokenizer.cpp)
target_include_directories(test_tokenizer PRIVATE ${CMAKE_SOURCE_DIR}/include)
target_link_libraries(test_tokenizer PRIVATE pthread)

# ---- Test CAM++ speaker encoder ----
add_executable(test_campp test/test_campp.cpp src/modules/speaker/campp.cpp src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp)
target_include_directories(test_campp PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(test_campp PRIVATE ggml pthread m)

# ---- Test AudioVAE encoder ----
add_executable(test_encoder test/test_encoder.cpp src/modules/vocoder/audiovae_encoder.cpp src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp)
target_include_directories(test_encoder PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(test_encoder PRIVATE ggml pthread m)

# ---- Test time_embedder isolated ----
add_executable(test_time_embed test/test_time_embed.cpp
    src/modules/backbone/dit_forward.cpp
    src/modules/backbone/dit_attention.cpp
    src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp
    src/utils/dit_loader.cpp)
target_include_directories(test_time_embed PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(test_time_embed PRIVATE ggml pthread m)

# ---- DiT calibration test ----
add_executable(test_dit_calib test/test_dit_calibrate.cpp
    src/modules/backbone/dit_forward.cpp
    src/modules/backbone/dit_attention.cpp
    src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp
    src/utils/dit_loader.cpp)
target_include_directories(test_dit_calib PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(test_dit_calib PRIVATE ggml pthread m)

# ---- Manual block 0 test ----
add_executable(test_manual_b0 test/test_manual_b0.cpp
    src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp
    src/utils/dit_loader.cpp)
target_include_directories(test_manual_b0 PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(test_manual_b0 PRIVATE ggml pthread m)

# ---- 18-block stress test ----
add_executable(test_18blocks test/test_18blocks.cpp
    src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp src/utils/dit_loader.cpp)
target_include_directories(test_18blocks PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(test_18blocks PRIVATE ggml pthread m)

# DiT time value test
add_executable(test_dit_tvals test/test_dit_tvals.cpp
    src/modules/backbone/dit_forward.cpp
    src/modules/backbone/dit_attention.cpp
    src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp src/utils/dit_loader.cpp)
target_include_directories(test_dit_tvals PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_dit_tvals PRIVATE ggml pthread m)

# Pipeline simulation test
add_executable(test_pipeline_sim test/test_pipeline_sim.cpp
    src/modules/backbone/dit_forward.cpp
    src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp src/utils/dit_loader.cpp src/utils/patchenc_loader.cpp)
target_include_directories(test_pipeline_sim PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(test_pipeline_sim PRIVATE ggml pthread m)

# ---- Manual LLM test ----
add_executable(test_manual_llm test/test_manual_llm.cpp
    src/utils/safetensors.cpp src/llm_manual.cpp)
target_include_directories(test_manual_llm PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_manual_llm PRIVATE ggml pthread m)

# ---- KV cache correctness test ----
add_executable(test_kv_cache test/test_kv_cache_correctness.cpp
    src/utils/safetensors.cpp src/llm_manual.cpp)
target_include_directories(test_kv_cache PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_kv_cache PRIVATE ggml pthread m)

# ---- KV cache step test ----
add_executable(test_kv_step test/test_kv_step.cpp
    src/utils/safetensors.cpp src/llm_manual.cpp)
target_include_directories(test_kv_step PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_kv_step PRIVATE ggml pthread m)

# ---- PatchEncoder input test ----
add_executable(test_pe_input test/test_pe_input.cpp
    src/utils/safetensors.cpp src/utils/dit_loader.cpp src/utils/patchenc_loader.cpp
    src/modules/backbone/patchenc.cpp)
target_include_directories(test_pe_input PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_pe_input PRIVATE ggml pthread m)

# ---- PE dump layer 0 test ----
add_executable(test_pe_dump_l0 test/test_pe_dump_l0.cpp
    src/utils/safetensors.cpp src/utils/dit_loader.cpp src/utils/patchenc_loader.cpp
    src/modules/backbone/patchenc.cpp)
target_include_directories(test_pe_dump_l0 PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_pe_dump_l0 PRIVATE ggml pthread m)

# ---- PE comparison test ----
add_executable(test_pe_cmp test/test_pe_cmp.cpp
    src/utils/safetensors.cpp src/utils/dit_loader.cpp src/utils/patchenc_loader.cpp
    src/modules/backbone/patchenc.cpp)
target_include_directories(test_pe_cmp PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_pe_cmp PRIVATE ggml pthread m)

# ---- Conv debug test ----
add_executable(test_conv_debug test/test_conv_debug.cpp
    src/utils/safetensors.cpp src/utils/dit_loader.cpp src/utils/patchenc_loader.cpp
    src/modules/backbone/patchenc.cpp)
target_include_directories(test_conv_debug PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_conv_debug PRIVATE ggml pthread m)

# ---- PE layers dump test ----
add_executable(test_pe_layers test/test_pe_layers.cpp
    src/utils/safetensors.cpp src/utils/dit_loader.cpp src/utils/patchenc_loader.cpp
    src/modules/backbone/patchenc.cpp)
target_include_directories(test_pe_layers PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_pe_layers PRIVATE ggml pthread m)

# ---- Clean comparison test ----
add_executable(test_clean_cmp test/test_clean_cmp.cpp
    src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp src/utils/dit_loader.cpp
    src/modules/vocoder/bigvgan_cpp.cpp)
target_include_directories(test_clean_cmp PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(test_clean_cmp PRIVATE ggml pthread m)

# ---- Test BigVGAN vs Python reference ----
add_executable(test_bigvgan_cmp test/test_bigvgan_cmp.cpp src/modules/vocoder/bigvgan_cpp.cpp src/modules/vocoder/lstm.cpp src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp)
target_include_directories(test_bigvgan_cmp PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(test_bigvgan_cmp PRIVATE ggml pthread m)

# Leaf mul test
add_executable(test_leaf_mul ../test/test_leaf_mul.cpp)
target_include_directories(test_leaf_mul PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_leaf_mul PRIVATE ggml pthread m)
add_executable(test_cond_cmp test/test_cond_cmp.cpp src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp src/utils/dit_loader.cpp)
target_include_directories(test_cond_cmp PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_cond_cmp PRIVATE ggml pthread m)
add_executable(test_pp test/test_pp.cpp src/modules/vocoder/bigvgan_cpp.cpp src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp)
target_include_directories(test_pp PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_pp PRIVATE ggml pthread m)

add_executable(test_bn test/test_bn.cpp src/modules/vocoder/bigvgan_cpp.cpp src/modules/vocoder/lstm.cpp src/utils/safetensors.cpp)
target_include_directories(test_bn PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_bn PRIVATE ggml pthread m)
add_executable(test_cp2 test/test_cp2.cpp src/modules/vocoder/bigvgan_cpp.cpp src/modules/vocoder/lstm.cpp src/utils/safetensors.cpp)
target_include_directories(test_cp2 PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_cp2 PRIVATE ggml pthread m)
add_executable(test_bv_full test/test_bv_full.cpp src/modules/vocoder/bigvgan_cpp.cpp src/modules/vocoder/lstm.cpp src/utils/safetensors.cpp)
target_include_directories(test_bv_full PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_bv_full PRIVATE ggml pthread m)



# BigVGAN comparison test

add_executable(test_bv_cmp2 test/test_bv_cmp2.cpp src/modules/vocoder/bigvgan_cpp.cpp src/modules/vocoder/lstm.cpp src/utils/safetensors.cpp)
target_include_directories(test_bv_cmp2 PRIVATE include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_bv_cmp2 PRIVATE ggml m pthread)


add_executable(test_bv_quick test/test_bv_quick.cpp src/modules/vocoder/bigvgan_cpp.cpp src/modules/vocoder/lstm.cpp src/utils/safetensors.cpp)
target_include_directories(test_bv_quick PRIVATE include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_bv_quick PRIVATE ggml m pthread)

add_executable(test_bv_stage test/test_bv_stage.cpp src/modules/vocoder/bigvgan_cpp.cpp src/modules/vocoder/lstm.cpp src/utils/safetensors.cpp)
target_include_directories(test_bv_stage PRIVATE include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src)
target_link_libraries(test_bv_stage PRIVATE ggml m pthread)

# ---- DiT comparison test (C++ vs Python on same input) ----
add_executable(test_dit_compare test/test_dit_compare.cpp
    src/utils/safetensors.cpp src/utils/dit_loader.cpp src/modules/vocoder/lstm.cpp)
target_include_directories(test_dit_compare PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(test_dit_compare PRIVATE ggml pthread m openblas)

# ---- E2E pipeline v4 (matching Python _generate_latents_stream) ----
add_executable(e2e_pipeline_v4
    src/e2e_pipeline_v4.cpp
    src/gpt2_bpe_tokenizer.cpp
    src/gguf_extract.cpp
    src/modules/backbone/dit_forward.cpp
    src/modules/backbone/dit_attention.cpp
    src/modules/backbone/patchenc.cpp
    src/modules/vocoder/audiovae.cpp
    src/modules/vocoder/bigvgan_cpp.cpp
    src/modules/vocoder/lstm.cpp
    src/utils/safetensors.cpp src/modules/vocoder/lstm.cpp
    src/llm_manual.cpp
    src/utils/dit_loader.cpp
    src/utils/patchenc_loader.cpp
    src/modules/speaker/campp.cpp
    src/modules/vocoder/audiovae_encoder.cpp
)
target_include_directories(e2e_pipeline_v4 PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/common
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(e2e_pipeline_v4 PRIVATE ${LLAMA_LIB} ggml pthread m openblas)
if(OpenMP_CXX_FOUND)
    target_link_libraries(e2e_pipeline_v4 PRIVATE OpenMP::OpenMP_CXX)
    target_compile_options(e2e_pipeline_v4 PRIVATE ${OpenMP_CXX_FLAGS})
endif()

# ---- BLAS linear comparison test ----
add_executable(test_blas_linear test/test_blas_linear.cpp
    src/utils/safetensors.cpp src/utils/dit_loader.cpp src/modules/vocoder/lstm.cpp)
target_include_directories(test_blas_linear PRIVATE
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/include
    ${CMAKE_SOURCE_DIR}/../llama/llama.cpp/ggml/src
)
target_link_libraries(test_blas_linear PRIVATE ggml pthread m openblas)
