# Copyright 2024 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

load(
    "@python_version_repo//:py_version.bzl",
    "HERMETIC_PYTHON_VERSION",
    "REQUIREMENTS",
)
load("@rules_python//python:pip.bzl", "compile_pip_requirements")
load("@rules_python//python:py_library.bzl", "py_library")
load("//jaxlib:jax.bzl", "all_py_deps")

licenses(["notice"])

REQUIREMENTS_TXT = (
    "requirements_lock_{}.txt".format(
        HERMETIC_PYTHON_VERSION.replace(".", "_").replace("-", "_")
    )
    if REQUIREMENTS == "//:requirements.txt"
    else REQUIREMENTS
)

REQUIREMENTS_SRCS = [
    "requirements.in",
    "test-requirements.txt",
    "nvidia-requirements.txt",
]

exports_files(["numpy.json"])

filegroup(
    name = "numpy_json",
    srcs = ["numpy.json"],
    visibility = ["//visibility:public"],
)

compile_pip_requirements(
    name = "requirements",
    srcs = REQUIREMENTS_SRCS,
    extra_args = [
        "--allow-unsafe",
        "--build-isolation",
        "--rebuild",
        "--index-url=https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple",
    ],
    generate_hashes = True,
    requirements_txt = REQUIREMENTS_TXT,
)

compile_pip_requirements(
    name = "requirements_nightly",
    srcs = REQUIREMENTS_SRCS,
    extra_args = [
        "--allow-unsafe",
        "--build-isolation",
        "--index-url=https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple",
        "--extra-index-url=https://pypi.anaconda.org/scientific-python-nightly-wheels/simple",
        "--pre",
        "--upgrade",
    ],
    generate_hashes = False,
    requirements_txt = REQUIREMENTS_TXT,
)

compile_pip_requirements(
    name = "requirements_dev",
    srcs = REQUIREMENTS_SRCS,
    extra_args = [
        "--allow-unsafe",
        "--build-isolation",
        "--upgrade",
        "--rebuild",
        "--index-url=https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple",
    ],
    generate_hashes = False,
    requirements_txt = REQUIREMENTS_TXT,
)

py_library(
    name = "all_py_deps",
    deps = all_py_deps([
        "zstandard",
        "tensorstore",
    ]),
)