load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
load("@rules_shell//shell:sh_test.bzl", "sh_test")
load("@tflm_pip_deps//:requirements.bzl", "requirement")
load(
    "//python:py_rules.bzl",
    "tflm_py_binary",
    "tflm_py_library",
    "tflm_py_test",
)
load(
    "//tensorflow/lite/micro:build_def.bzl",
    "INCOMPATIBLE_WITH_WINDOWS",
    "tflm_cc_library",
    "tflm_cc_test",
)
load(
    "//third_party/flatbuffers:build_defs.bzl",
    "flatbuffer_cc_library",
    "flatbuffer_py_library",
)
load("view_tests.bzl", "generate_view_tests")

package(
    default_visibility = [
        "//visibility:public",
    ],
)

py_library(
    name = "compression",
    srcs = ["__init__.py"],
    deps = [
        ":compress_lib",
        ":spec",
        ":spec_builder",
    ],
)

flatbuffer_cc_library(
    # Generates the header-only library "metadata_generated.h", used to read
    # the metadata flatbuffer.
    name = "metadata_cc",
    srcs = ["metadata.fbs"],
)

tflm_cc_library(
    # The header-only library generated by flatc in ":metadata_cc" is saved to
    # the source tree and committed to git as "metadata_saved.h", which is used
    # by code which builds via the Make build system, which has no means of
    # generating the header on the fly. Code which builds via both bazel and
    # Make should #include the saved header and use this target in its bazel
    # BUILD deps. Code built exclusively via bazel would typically depend
    # directly on ":metadata_cc", which would generate a header from the schema
    # on the fly, during the build.
    #
    # When the schema definition "metadata.fbs" is changed, this saved header
    # should be updated by running the script "./metadata_saved_update.sh",
    # outside of bazel (because bazel cannot modify the source tree). The
    # script regenerates the header from the schema and copies it to the source
    # tree as "metadata_saved.h".
    #
    # Committing the generated file risks inconsistency between the schema and
    # the saved header, so consistency ensured by the unit test
    # ":metadata_saved_test".
    #
    name = "metadata_saved",
    hdrs = ["metadata_saved.h"],
)

sh_test(
    # Ensures consistency between the schema and the saved generated header.
    # Fails if they mismatch, in which case, ./metadata_saved_update.sh should
    # be run. See :metadata_saved above.
    name = "metadata_saved_test",
    size = "small",
    srcs = ["metadata_saved_test.sh"],
    args = [
        "$(location metadata_saved.h)",
        "$(location :metadata_cc_srcs)",
    ],
    data = [
        "metadata_saved.h",
        ":metadata_cc_srcs",
    ],
    target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
)

tflm_cc_test(
    name = "metadata_test_cc",
    size = "small",
    srcs = ["metadata_test.cc"],
    target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
    deps = [
        ":metadata_saved",
        "//tensorflow/lite/micro:hexdump",
        "//tensorflow/lite/micro/testing:micro_test",
        "@flatbuffers//:runtime_cc",
    ],
)

flatbuffer_py_library(
    # Generates the Python module "metadata_py_generated", used to read the
    # metadata flatbuffer.
    name = "metadata_py",
    srcs = ["metadata.fbs"],
)

tflm_py_test(
    name = "metadata_test_py",
    size = "small",
    srcs = ["metadata_test.py"],
    main = "metadata_test.py",
    target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
    deps = [
        "metadata_py",
        requirement("flatbuffers"),
    ],
)

py_library(
    name = "compress_lib",
    srcs = [
        "compress.py",
    ],
    deps = [
        ":compressor",
        ":decode_insert",
        ":huffman",
        ":lut",
        ":model_editor",
        ":pruning",
        ":spec",
        "//tensorflow/lite/micro/tools:tflite_flatbuffer_align",
        requirement("absl_py"),
    ],
)

py_binary(
    name = "compress",
    srcs = [
        "compress.py",
    ],
    target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
    deps = [
        ":compress_lib",
    ],
)

py_test(
    name = "compress_test",
    size = "small",
    srcs = [
        "compress_test.py",
    ],
    tags = [
        "noasan",
        "nomsan",  # Sanitizer symbols don't work with Python extension libs
        "noubsan",
    ],
    target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
    deps = [
        ":compress",
        ":compressor",
        ":decode_insert",
        ":model_editor",
        ":spec",
        "//tensorflow/lite/python:schema_py",
        requirement("numpy"),
    ],
)

tflm_py_library(
    name = "verify",
    srcs = ["verify.py"],
    deps = [
        ":model_editor",
        ":tensor_type",
        "//python/tflite_micro:runtime",
        requirement("numpy"),
    ],
)

tflm_py_test(
    name = "verify_test",
    size = "small",
    srcs = ["verify_test.py"],
    tags = [
        "noasan",
        "nomsan",
        "noubsan",
    ],
    deps = [
        ":model_editor",
        ":verify",
        "//tensorflow/lite/python:schema_py",
        requirement("numpy"),
    ],
)

tflm_py_test(
    name = "compression_integration_test",
    size = "small",
    srcs = ["compression_integration_test.py"],
    tags = [
        "noasan",
        "nomsan",
        "noubsan",
    ],
    target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
    deps = [
        ":compress_lib",
        ":decode_insert",
        ":model_editor",
        ":spec",
        ":verify",
        "//python/tflite_micro:runtime",
        "//tensorflow/lite/python:schema_py",
        requirement("numpy"),
    ],
)

tflm_py_test(
    name = "proprietary_integration_test",
    size = "small",
    srcs = ["proprietary_integration_test.py"],
    tags = [
        "manual",
        "noasan",
        "nomsan",
        "noubsan",
    ],
    target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
    deps = [
        ":compress_lib",
        ":spec",
        ":verify",
        requirement("pyyaml"),
    ],
)

tflm_py_test(
    name = "proprietary_integration_smoke_test",
    size = "small",
    srcs = [
        "compression_integration_test.py",
        "proprietary_integration_smoke_test.py",
        "proprietary_integration_test.py",
    ],
    tags = [
        "noasan",
        "nomsan",
        "noubsan",
    ],
    target_compatible_with = select({
        "//:with_compression_enabled": [],
        "//conditions:default": ["@platforms//:incompatible"],
    }),
    deps = [
        ":compress_lib",
        ":decode_insert",
        ":model_editor",
        ":spec",
        ":verify",
        "//python/tflite_micro:runtime",
        "//tensorflow/lite/python:schema_py",
        requirement("numpy"),
        requirement("pyyaml"),
    ],
)

tflm_py_library(
    name = "spec",
    srcs = ["spec.py"],
    deps = [
        requirement("pyyaml"),
    ],
)

py_test(
    name = "spec_test",
    size = "small",
    srcs = ["spec_test.py"],
    target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
    deps = [
        ":spec",
    ],
)

py_library(
    name = "spec_builder",
    srcs = ["spec_builder.py"],
    deps = [
        ":spec",
    ],
)

py_test(
    name = "spec_builder_test",
    size = "small",
    srcs = ["spec_builder_test.py"],
    target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
    deps = [
        ":spec",
        ":spec_builder",
    ],
)

tflm_py_library(
    name = "tensor_type",
    srcs = ["tensor_type.py"],
    deps = [
        "//tensorflow/lite/python:schema_py",
        requirement("numpy"),
    ],
)

tflm_py_test(
    name = "tensor_type_test",
    size = "small",
    srcs = ["tensor_type_test.py"],
    deps = [
        ":tensor_type",
        "//tensorflow/lite/python:schema_py",
        requirement("numpy"),
    ],
)

tflm_py_library(
    name = "model_editor",
    srcs = ["model_editor.py"],
    deps = [
        ":tensor_type",
        "//tensorflow/lite/python:schema_py",
        requirement("flatbuffers"),
        requirement("numpy"),
    ],
)

tflm_py_test(
    name = "model_editor_test",
    size = "small",
    srcs = ["model_editor_test.py"],
    deps = [
        ":model_editor",
        "//tensorflow/lite/python:schema_py",
        requirement("numpy"),
    ],
)

tflm_py_library(
    name = "decode",
    srcs = ["decode.py"],
)

tflm_py_test(
    name = "decode_test",
    size = "small",
    srcs = ["decode_test.py"],
    deps = [
        ":decode",
    ],
)

tflm_py_library(
    name = "compressor",
    srcs = ["compressor.py"],
    deps = [
        ":decode",
        ":model_editor",
        ":spec",
    ],
)

tflm_py_library(
    name = "lut",
    srcs = ["lut.py"],
    deps = [
        ":compressor",
        ":decode",
        ":model_editor",
        ":spec",
        requirement("bitarray"),
        requirement("numpy"),
    ],
)

tflm_py_test(
    name = "lut_test",
    size = "small",
    srcs = ["lut_test.py"],
    tags = [
        "noasan",
        "nomsan",
        "noubsan",
    ],
    deps = [
        ":compressor",
        ":decode",
        ":lut",
        ":model_editor",
        ":spec",
        "//tensorflow/lite/python:schema_py",
        requirement("numpy"),
    ],
)

tflm_py_library(
    name = "huffman",
    srcs = ["huffman.py"],
    deps = [
        ":compressor",
        ":decode",
        ":model_editor",
        ":spec",
    ],
)

tflm_py_library(
    name = "pruning",
    srcs = ["pruning.py"],
    deps = [
        ":compressor",
        ":decode",
        ":model_editor",
        ":spec",
    ],
)

tflm_py_library(
    name = "decode_insert",
    srcs = ["decode_insert.py"],
    deps = [
        ":compressor",
        ":model_editor",
        "//tensorflow/lite/python:schema_py",
    ],
)

tflm_py_test(
    name = "decode_insert_test",
    size = "small",
    srcs = ["decode_insert_test.py"],
    tags = [
        "noasan",
        "nomsan",
        "noubsan",
    ],
    deps = [
        ":compressor",
        ":decode",
        ":decode_insert",
        ":model_editor",
        "//tensorflow/lite/python:schema_py",
        requirement("numpy"),
    ],
)

tflm_py_test(
    name = "decode_insert_runtime_test",
    size = "small",
    srcs = ["decode_insert_runtime_test.py"],
    tags = [
        "noasan",
        "nomsan",
        "noubsan",
    ],
    deps = [
        ":decode_insert",
        ":lut",
        ":model_editor",
        ":spec",
        "//python/tflite_micro:runtime",
        "//tensorflow/lite/python:schema_py",
        requirement("numpy"),
    ],
)

tflm_py_binary(
    name = "view",
    srcs = [
        "view.py",
    ],
    target_compatible_with = INCOMPATIBLE_WITH_WINDOWS,
    deps = [
        ":metadata_py",
        ":tensor_type",
        "//tensorflow/lite/python:schema_py",
        requirement("absl_py"),
        requirement("bitarray"),
        requirement("numpy"),
        requirement("prettyprinter"),
    ],
)

generate_view_tests([
    "//tensorflow/lite/micro/models:keyword_scrambled.tflite",
    "//tensorflow/lite/micro/models:keyword_scrambled_8bit.tflite",
    "//tensorflow/lite/micro/models:person_detect.tflite",
    "//tensorflow/lite/micro/models:person_detect_vela.tflite",
])
