# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The Lance Authors

from pathlib import Path

import lance
import numpy as np
import pyarrow as pa
import pytest

torch = pytest.importorskip("torch")


@pytest.mark.parametrize("accelerator", [None, "cuda"])
def test_f16_embeddings(tmp_path: Path, accelerator: str):
    if not torch.cuda.is_available() and accelerator == "cuda":
        pytest.skip("CUDA not available")

    DIM = 16
    TOTAL = 256
    values = np.random.random(TOTAL * DIM).astype(np.float16)
    fsl = pa.FixedSizeListArray.from_arrays(values, DIM)
    data = pa.Table.from_arrays([fsl, np.arange(TOTAL)], names=["vec", "id"])

    ds = lance.write_dataset(data, tmp_path)
    assert ds.schema.field("vec").type.value_type == pa.float16()

    ds = ds.create_index(
        "vec",
        "IVF_PQ",
        replace=True,
        num_partitions=2,
        num_sub_vectors=2,
        accelerator=accelerator,
    )

    # Can use float32 to search
    query = np.random.random(DIM).astype(np.float32)
    rst32 = ds.to_table(nearest={"column": "vec", "q": query})

    # Can use float16 to search
    rst16 = ds.to_table(nearest={"column": "vec", "q": query.astype(np.float16)})

    rst_py = ds.to_table(
        nearest={
            "column": "vec",
            "q": query.tolist(),
        }
    )

    assert rst16 == rst32
    assert rst16 == rst_py
