# Set the minimum CMake version and policies for highest tested version
cmake_minimum_required(VERSION 3.15...3.30)

# Set up the project and ensure there is a working C++ compiler
project(rigel_ext LANGUAGES C CXX)

# Warn if the user invokes CMake directly
if (NOT SKBUILD)
  message(WARNING "\
  This CMake file is meant to be executed using 'scikit-build-core'. \
  Running it directly will almost certainly not produce the desired \
  result. If you are a user trying to install this package, use: \
  ===================================================================== \
   $ pip install . \
  ===================================================================== \
  For development with incremental rebuilds: \
  ===================================================================== \
   $ pip install nanobind scikit-build-core[pyproject] \
   $ pip install --no-build-isolation -ve . \
  =====================================================================")
endif()

# Don't use compiler extensions (e.g. -std=gnu++17 -> -std=c++17)
set(CMAKE_CXX_EXTENSIONS OFF)

# Release build by default
if (NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
  set(CMAKE_BUILD_TYPE Release CACHE STRING "Choose the type of build." FORCE)
  set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
endif()

# --- Find Python and nanobind ------------------------------------------------

# Try to import all Python components potentially needed by nanobind
if (CMAKE_VERSION VERSION_LESS 3.18)
  set(DEV_MODULE Development)
else()
  set(DEV_MODULE Development.Module)
endif()

find_package(Python 3.12
  REQUIRED COMPONENTS Interpreter ${DEV_MODULE}
  OPTIONAL_COMPONENTS Development.SABIModule)

# Import nanobind through CMake's find_package mechanism
execute_process(
  COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
  OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE nanobind_ROOT)
find_package(nanobind CONFIG REQUIRED)

# --- Build the _cgranges_impl extension ---------------------------------------

# Compile the nanobind extension with optimizations for speed (NOMINSIZE)
# since cgranges overlap queries are a hot path.
nanobind_add_module(
  _cgranges_impl
  NOMINSIZE
  STABLE_ABI
  src/rigel/cgranges/cgranges_bind.cpp
  src/rigel/cgranges/cgranges.c
)

# The C wrapper needs to find cgranges.h and khash.h
target_include_directories(_cgranges_impl PRIVATE
  ${CMAKE_CURRENT_SOURCE_DIR}/src/rigel/cgranges
)

# Install into the rigel package directory so it's importable as
# rigel._cgranges_impl
install(TARGETS _cgranges_impl LIBRARY DESTINATION rigel)

# --- Build the _scoring_impl extension ----------------------------------------

# C++ scoring hot-path kernels (compute_fragment_weight, etc.)
nanobind_add_module(
  _scoring_impl
  NOMINSIZE
  STABLE_ABI
  src/rigel/native/scoring.cpp
)

# Install into the rigel package directory so it's importable as
# rigel._scoring_impl
install(TARGETS _scoring_impl LIBRARY DESTINATION rigel)

# --- Build the _resolve_impl extension ----------------------------------------

# C++ fragment resolution kernel (resolve_fragment + cgranges queries)
nanobind_add_module(
  _resolve_impl
  NOMINSIZE
  STABLE_ABI
  src/rigel/native/resolve.cpp
  src/rigel/cgranges/cgranges.c
)

# Needs cgranges.h, khash.h, and shared native headers
target_include_directories(_resolve_impl PRIVATE
  ${CMAKE_CURRENT_SOURCE_DIR}/src/rigel/cgranges
  ${CMAKE_CURRENT_SOURCE_DIR}/src/rigel/native
)

# Install into the rigel package directory so it's importable as
# rigel._resolve_impl
install(TARGETS _resolve_impl LIBRARY DESTINATION rigel)

# --- Build the _bam_impl extension --------------------------------------------

# C++ BAM scanner using htslib — replaces the Python BAM parsing hot path.
# Reads BAM → resolves fragments → buffers results in one C++ call.

# Find htslib (pkg-config or cmake find)
# Note: IMPORTED_TARGET is intentionally omitted — it forces transitive dep
# resolution via pkg-config (e.g. bzip2.pc), which is not always available in
# conda cross-compilation builds.  We link against the shared libhts directly,
# so no -lbz2 etc. are needed at link time.
find_package(PkgConfig)
if(PkgConfig_FOUND)
  pkg_check_modules(HTSLIB htslib)
endif()

# Helper macro to detect htslib in a given directory prefix
macro(_find_htslib_in_prefix _prefix _label)
  if(NOT HTSLIB_FOUND AND EXISTS "${_prefix}/include/htslib/hts.h")
    foreach(_lib libhts.dylib libhts.so libhts.a)
      if(EXISTS "${_prefix}/lib/${_lib}")
        set(HTSLIB_FOUND TRUE)
        set(HTSLIB_INCLUDE_DIRS "${_prefix}/include")
        set(HTSLIB_LIBRARY_DIRS "${_prefix}/lib")
        set(HTSLIB_LIBRARIES "hts")
        message(STATUS "htslib found via ${_label}: ${_prefix}")
        break()
      endif()
    endforeach()
  endif()
endmacro()

# Fallback: conda-build host env ($PREFIX), then activated env ($CONDA_PREFIX)
if(DEFINED ENV{PREFIX})
  _find_htslib_in_prefix("$ENV{PREFIX}" "conda PREFIX")
endif()
if(DEFINED ENV{CONDA_PREFIX})
  _find_htslib_in_prefix("$ENV{CONDA_PREFIX}" "CONDA_PREFIX")
endif()

if(HTSLIB_FOUND)
  message(STATUS "htslib found — building _bam_impl C++ BAM scanner")
else()
  message(FATAL_ERROR
    "htslib NOT found.  _bam_impl is required.\n"
    "Install htslib (conda install -c conda-forge htslib) and re-run cmake.")
endif()

nanobind_add_module(
  _bam_impl
  NOMINSIZE
  STABLE_ABI
  src/rigel/native/bam_scanner.cpp
  src/rigel/native/calibration/accumulator.cpp
  src/rigel/cgranges/cgranges.c
)

target_include_directories(_bam_impl PRIVATE
  ${CMAKE_CURRENT_SOURCE_DIR}/src/rigel/cgranges
  ${CMAKE_CURRENT_SOURCE_DIR}/src/rigel/native
)

# Link htslib via variables set by pkg_check_modules or the manual fallback
target_include_directories(_bam_impl PRIVATE ${HTSLIB_INCLUDE_DIRS})
target_link_directories(_bam_impl PRIVATE ${HTSLIB_LIBRARY_DIRS})
target_link_libraries(_bam_impl PRIVATE ${HTSLIB_LIBRARIES})

install(TARGETS _bam_impl LIBRARY DESTINATION rigel)

# --- Build the _em_impl extension ---------------------------------------------

# C++ EM solver — replaces the Python EM hot path (_em_step, _vbem_step,
# _build_equiv_classes, SQUAREM loop) with a single native call.
# Pure C++17 + nanobind + numpy, no external dependencies.
nanobind_add_module(
  _em_impl
  NOMINSIZE
  STABLE_ABI
  src/rigel/native/em_solver.cpp
)
# Keep debug symbols for profiling (dsym/DWARF) even in Release builds
target_compile_options(_em_impl PRIVATE -g)

install(TARGETS _em_impl LIBRARY DESTINATION rigel)

# --- SIMD / Architecture Detection -------------------------------------------
#
# Detect CPU architecture and report available SIMD instruction sets.
# The C++ header simd_detect.h uses standard compiler-defined macros
# (__ARM_NEON, __AVX2__, etc.) for compile-time detection, plus CPUID
# for runtime dispatch on x86.  With -march=native the compiler sets
# these macros automatically; for portable builds the C++ code uses
# function-level __attribute__((target(...))) for multi-versioned SIMD.
#
# This block just reports what's available and sets RIGEL_ARCH for
# any downstream CMake logic.

if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64|ARM64")
    set(RIGEL_ARCH "ARM64")
    message(STATUS "SIMD: ARM64 — NEON is baseline (always available)")
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|AMD64|amd64")
    set(RIGEL_ARCH "X86_64")
    include(CheckCXXCompilerFlag)
    check_cxx_compiler_flag("-mavx2 -mfma" _rigel_has_avx2_fma)
    check_cxx_compiler_flag("-mavx512f" _rigel_has_avx512f)
    if(_rigel_has_avx512f)
        message(STATUS "SIMD: x86_64 — SSE2 + AVX2/FMA + AVX-512F available")
    elseif(_rigel_has_avx2_fma)
        message(STATUS "SIMD: x86_64 — SSE2 + AVX2/FMA available")
    else()
        message(STATUS "SIMD: x86_64 — SSE2 only")
    endif()
else()
    set(RIGEL_ARCH "UNKNOWN")
    message(STATUS "SIMD: ${CMAKE_SYSTEM_PROCESSOR} — no explicit SIMD detection")
endif()

# --- Compiler optimizations ---------------------------------------------------
#
# Phase 2 performance flags: -O3, -march=native, FMA contraction, and LTO.
# Applied per-target to control fast-math scope.
#
# _em_impl gets -ffp-contract=fast (enables FMA) but NOT full -ffast-math,
# because -ffast-math allows the compiler to reassociate floating-point
# additions, which would break the Kahan compensated summation used in
# the E-step kernel for numerical stability.
#
# _scoring_impl gets full -ffast-math (no compensated summation there).
# All other targets get -O3 -march=native without fast-math.

# RIGEL_PORTABLE: when ON, omit -march=native so wheels run on any machine
# of the same architecture. CI wheel builds set this; local dev leaves it OFF.
option(RIGEL_PORTABLE "Build portable wheels (no -march=native)" OFF)
option(
  RIGEL_PROFILE_NATIVE
  "Build native extensions with profiler-friendly symbols/frame pointers and disable LTO"
  OFF
)

if(CMAKE_CXX_COMPILER_ID MATCHES "Clang|GNU")
    if(RIGEL_PORTABLE)
        set(RIGEL_OPT_FLAGS -O3 -DNDEBUG)
    else()
        set(RIGEL_OPT_FLAGS -O3 -march=native -DNDEBUG)
    endif()

    target_compile_options(_em_impl      PRIVATE ${RIGEL_OPT_FLAGS} -ffp-contract=fast)
    target_compile_options(_scoring_impl PRIVATE ${RIGEL_OPT_FLAGS} -ffast-math)
    target_compile_options(_resolve_impl PRIVATE ${RIGEL_OPT_FLAGS})
    target_compile_options(_cgranges_impl PRIVATE ${RIGEL_OPT_FLAGS})
    if(TARGET _bam_impl)
        target_compile_options(_bam_impl PRIVATE ${RIGEL_OPT_FLAGS})
    endif()

    if(RIGEL_PROFILE_NATIVE)
      foreach(_rigel_native_target
          _em_impl _scoring_impl _resolve_impl _cgranges_impl _bam_impl)
        if(TARGET ${_rigel_native_target})
          target_compile_options(
            ${_rigel_native_target}
            PRIVATE -g -fno-omit-frame-pointer
          )
        endif()
      endforeach()
    endif()
endif()

# Link-Time Optimization (LTO): enables cross-TU inlining and dead code
# elimination across compilation units.
include(CheckIPOSupported)
check_ipo_supported(RESULT ipo_supported)
if(ipo_supported AND NOT RIGEL_PROFILE_NATIVE)
    set_property(TARGET _em_impl      PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
    set_property(TARGET _scoring_impl PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
    set_property(TARGET _resolve_impl PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
    set_property(TARGET _cgranges_impl PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
    if(TARGET _bam_impl)
        set_property(TARGET _bam_impl PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
    endif()
elseif(RIGEL_PROFILE_NATIVE)
  message(STATUS "Native profiling build: debug symbols/frame pointers enabled; LTO disabled")
endif()
