# Copyright (c) Meta Platforms, Inc. and affiliates.

# @noautodeps

load("@fbcode_macros//build_defs:cpp_python_extension.bzl", "cpp_python_extension")
load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
load("//data_compression/experimental/zstrong:defs.bzl", "zs_cxxlibrary")

oncall("data_compression")

zs_cxxlibrary(
    name = "pybind_helpers",
    srcs = ["pybind_helpers.cpp"],
    headers = ["pybind_helpers.h"],
    deps = [
        "fbsource//third-party/fmt:fmt",
        "fbsource//third-party/pybind11:pybind11",
        "//data_compression/experimental/zstrong/tools:zstrong_cpp",
    ],
    exported_deps = [
        "//data_compression/experimental/zstrong:zstronglib",
    ],
)

zs_cxxlibrary(
    name = "zstrong_ml_native",
    srcs = ["zstrong_ml_pybind.cpp"],
    headers = ["zstrong_ml_pybind.h"],
    compiler_flags = ["-Wno-float-equal"],
    strict_conversions = False,
    deps = [
        "fbsource//third-party/pybind11:pybind11",
        "fbsource//third-party/pypi/numpy:numpy",
        ":pybind_helpers",
        "//data_compression/experimental/zstrong/tools:zstrong_ml",
    ],
)

cpp_python_extension(
    name = "zstrong_json",
    srcs = ["zstrong_json_pybind.cpp"],
    base_module = "",
    types = [
        "zstrong_json.pyi",
        "zstrong_json.ml.pyi",
    ],
    deps = [
        "fbsource//third-party/pybind11:pybind11",
        "fbsource//third-party/pypi/numpy:numpy",
        ":pybind_helpers",
        ":zstrong_ml_native",
        "//data_compression/experimental/zstrong/tools:zstrong_json",
        "//data_compression/experimental/zstrong/tools:zstrong_ml",
        "//folly:json",
    ],
)

python_library(
    name = "zstrong_ml",
    srcs = ["zstrong_ml.py"],
    base_module = "",
    deps = [
        "fbsource//third-party/pypi/pandas:pandas",
        ":zstrong_json",
    ],
    external_deps = ["xgboost"],
)
