test_tiledbvcf.py

import numpy as np
import subprocess
import os
import pandas as pd
import pyarrow as pa
import glob
import shutil
import platform
import pytest
import tiledbvcf
import tiledb

# Directory containing this file
CONTAINING_DIR = os.path.abspath(os.path.dirname(__file__))

# Test inputs directory
TESTS_INPUT_DIR = os.path.abspath(
    os.path.join(CONTAINING_DIR, "../../../libtiledbvcf/test/inputs")
)


def _check_dfs(expected, actual):
    def assert_series(s1, s2):
        if np.issubdtype(s2.dtype, np.floating):
            assert np.isclose(s1, s2, equal_nan=True).all()
        elif np.issubdtype(s2.dtype, np.integer):
            assert s1.astype("int64").equals(s2.astype("int64"))
        else:
            assert s1.equals(s2)

    for k in expected:
        assert_series(expected[k], actual[k])

    for k in actual:
        assert_series(expected[k], actual[k])


def check_if_compatible(uri):
    try:
        with tiledb.open(uri):
            return True
    except tiledb.libtiledb.TileDBError as e:
        if "incompatible format version" in str(e).lower():
            raise pytest.skip.Exception(
                "Test skipped due to incompatible format version"
            )
        raise pytest.skip.Exception(f"Test skipped due to TileDB error: {str(e)}")


@pytest.fixture
def test_ds_v4():
    return tiledbvcf.Dataset(
        os.path.join(TESTS_INPUT_DIR, "arrays/v4/ingested_2samples")
    )


@pytest.fixture
def test_ds():
    return tiledbvcf.Dataset(
        os.path.join(TESTS_INPUT_DIR, "arrays/v3/ingested_2samples")
    )


@pytest.fixture
def test_ds_attrs():
    return tiledbvcf.Dataset(
        os.path.join(TESTS_INPUT_DIR, "arrays/v3/ingested_2samples_GT_DP_PL")
    )


def test_basic_count(test_ds):
    assert test_ds.count() == 14


def test_retrieve_attributes(test_ds):
    builtin_attrs = [
        "sample_name",
        "contig",
        "pos_start",
        "pos_end",
        "alleles",
        "id",
        "fmt",
        "info",
        "filters",
        "qual",
        "query_bed_end",
        "query_bed_start",
        "query_bed_line",
    ]
    assert sorted(test_ds.attributes(attr_type="builtin")) == sorted(builtin_attrs)

    info_attrs = [
        "info_BaseQRankSum",
        "info_ClippingRankSum",
        "info_DP",
        "info_DS",
        "info_END",
        "info_HaplotypeScore",
        "info_InbreedingCoeff",
        "info_MLEAC",
        "info_MLEAF",
        "info_MQ",
        "info_MQ0",
        "info_MQRankSum",
        "info_ReadPosRankSum",
    ]
    assert test_ds.attributes(attr_type="info") == info_attrs

    fmt_attrs = [
        "fmt_AD",
        "fmt_DP",
        "fmt_GQ",
        "fmt_GT",
        "fmt_MIN_DP",
        "fmt_PL",
        "fmt_SB",
    ]
    assert test_ds.attributes(attr_type="fmt") == fmt_attrs


def test_retrieve_samples(test_ds):
    assert test_ds.samples() == ["HG00280", "HG01762"]


def test_read_unsupported_regions_type(test_ds):
    unsupported_region = 3.14
    unsupported_type_error = f'"regions" parameter cannot have type: {type(unsupported_region)}'
    wrong_dimension_region = np.array([["1:12700-13400"], ["1:12700-13400"]])
    ndarray_wrong_dimension_error = f'"regions" parameter of type {type(wrong_dimension_region)} must be 1-dimensional'
    with pytest.raises(Exception, match=unsupported_type_error):
        test_ds.read(regions=unsupported_region)
    with pytest.raises(Exception, match=ndarray_wrong_dimension_error):
        test_ds.read(regions=wrong_dimension_region)
    with pytest.raises(Exception, match=unsupported_type_error):
        test_ds.read_arrow(regions=unsupported_region)
    with pytest.raises(Exception, match=ndarray_wrong_dimension_error):
        test_ds.read_arrow(regions=wrong_dimension_region)
    with pytest.raises(Exception, match=unsupported_type_error):
        for variant in test_ds.read_iter(regions=unsupported_region):
            print(variant)
    with pytest.raises(Exception, match=ndarray_wrong_dimension_error):
        for variant in test_ds.read_iter(regions=wrong_dimension_region):
            print(variant)


def test_read_attrs(test_ds_attrs):
    attrs = ["sample_name"]
    df = test_ds_attrs.read(attrs=attrs)
    assert df.columns.values.tolist() == attrs

    attrs = ["sample_name", "fmt_GT"]
    df = test_ds_attrs.read(attrs=attrs)
    assert df.columns.values.tolist() == attrs

    attrs = ["sample_name"]
    df = test_ds_attrs.read(attrs=attrs)
    assert df.columns.values.tolist() == attrs


def test_basic_reads(test_ds):
    expected_df = pd.DataFrame(
        {
            "sample_name": pd.Series(
                [
                    "HG00280",
                    "HG01762",
                    "HG00280",
                    "HG01762",
                    "HG00280",
                    "HG01762",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                ]
            ),
            "pos_start": pd.Series(
                [
                    12141,
                    12141,
                    12546,
                    12546,
                    13354,
                    13354,
                    13375,
                    13396,
                    13414,
                    13452,
                    13520,
                    13545,
                    17319,
                    17480,
                ],
                dtype=np.int32,
            ),
            "pos_end": pd.Series(
                [
                    12277,
                    12277,
                    12771,
                    12771,
                    13374,
                    13389,
                    13395,
                    13413,
                    13451,
                    13519,
                    13544,
                    13689,
                    17479,
                    17486,
                ],
                dtype=np.int32,
            ),
        }
    ).sort_values(ignore_index=True, by=["sample_name", "pos_start"])

    for use_arrow in [False, True]:
        func = test_ds.read_arrow if use_arrow else test_ds.read

        df = func(attrs=["sample_name", "pos_start", "pos_end"])
        if use_arrow:
            df = df.to_pandas()

        _check_dfs(
            expected_df,
            df.sort_values(ignore_index=True, by=["sample_name", "pos_start"]),
        )

    # Region intersection
    df = test_ds.read(
        attrs=["sample_name", "pos_start", "pos_end"], regions=["1:12700-13400"]
    )
    expected_df = pd.DataFrame(
        {
            "sample_name": pd.Series(
                ["HG00280", "HG01762", "HG00280", "HG01762", "HG00280", "HG00280"]
            ),
            "pos_start": pd.Series(
                [12546, 12546, 13354, 13354, 13375, 13396], dtype=np.int32
            ),
            "pos_end": pd.Series(
                [12771, 12771, 13374, 13389, 13395, 13413], dtype=np.int32
            ),
        }
    ).sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    _check_dfs(
        expected_df, df.sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    )
    df = test_ds.read_arrow(
        attrs=["sample_name", "pos_start", "pos_end"], regions=["1:12700-13400"]
    ).to_pandas()
    _check_dfs(
        expected_df, df.sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    )

    # Regions as string
    df = test_ds.read(
        attrs=["sample_name", "pos_start", "pos_end"], regions="1:12700-13400"
    )
    _check_dfs(
        expected_df, df.sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    )
    df = test_ds.read_arrow(
        attrs=["sample_name", "pos_start", "pos_end"], regions="1:12700-13400"
    ).to_pandas()
    _check_dfs(
        expected_df, df.sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    )

    # Regions as numpy.ndarray
    df = test_ds.read(
        attrs=["sample_name", "pos_start", "pos_end"], regions=np.array(["1:12700-13400"])
    )
    _check_dfs(
        expected_df, df.sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    )
    df = test_ds.read_arrow(
        attrs=["sample_name", "pos_start", "pos_end"], regions=np.array(["1:12700-13400"])
    ).to_pandas()
    _check_dfs(
        expected_df, df.sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    )

    # Region and sample intersection
    df = test_ds.read(
        attrs=["sample_name", "pos_start", "pos_end"],
        regions=["1:12700-13400"],
        samples=["HG01762"],
    )
    expected_df = pd.DataFrame(
        {
            "sample_name": pd.Series(["HG01762", "HG01762"]),
            "pos_start": pd.Series([12546, 13354], dtype=np.int32),
            "pos_end": pd.Series([12771, 13389], dtype=np.int32),
        }
    ).sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    _check_dfs(
        expected_df, df.sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    )

    # Sample only
    df = test_ds.read(
        attrs=["sample_name", "pos_start", "pos_end"], samples=["HG01762"]
    )
    expected_df = pd.DataFrame(
        {
            "sample_name": pd.Series(["HG01762", "HG01762", "HG01762"]),
            "pos_start": pd.Series([12141, 12546, 13354], dtype=np.int32),
            "pos_end": pd.Series([12277, 12771, 13389], dtype=np.int32),
        }
    ).sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    _check_dfs(
        expected_df, df.sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    )


def test_multiple_counts(test_ds):
    assert test_ds.count() == 14
    assert test_ds.count() == 14
    assert test_ds.count(regions=["1:12700-13400"]) == 6
    assert test_ds.count(samples=["HG00280"], regions=["1:12700-13400"]) == 4
    assert test_ds.count() == 14
    assert test_ds.count(samples=["HG01762"]) == 3
    assert test_ds.count(samples=["HG00280"]) == 11


def test_empty_region(test_ds):
    assert test_ds.count(regions=["12:1-1000000"]) == 0


def test_missing_sample_raises_exception(test_ds):
    with pytest.raises(RuntimeError):
        test_ds.count(samples=["abcde"])


# TODO remove skip
@pytest.mark.skip
def test_bad_contig_raises_exception(test_ds):
    with pytest.raises(RuntimeError):
        test_ds.count(regions=["chr1:1-1000000"])
    with pytest.raises(RuntimeError):
        test_ds.count(regions=["1"])
    with pytest.raises(RuntimeError):
        test_ds.count(regions=["1:100-"])
    with pytest.raises(RuntimeError):
        test_ds.count(regions=["1:-100"])


def test_bad_attr_raises_exception(test_ds):
    with pytest.raises(RuntimeError):
        test_ds.read(attrs=["abcde"], regions=["1:12700-13400"])


def test_read_write_mode_exceptions():
    ds = tiledbvcf.Dataset(os.path.join(TESTS_INPUT_DIR, "arrays/v3/ingested_2samples"))
    samples = [os.path.join(TESTS_INPUT_DIR, s) for s in ["small.bcf", "small2.bcf"]]

    with pytest.raises(Exception):
        ds.create_dataset()

    with pytest.raises(Exception):
        ds.ingest_samples(samples)

    ds = tiledbvcf.Dataset(
        os.path.join(TESTS_INPUT_DIR, "arrays/v3/ingested_2samples"), mode="w"
    )
    with pytest.raises(Exception):
        ds.count()


def test_incomplete_reads():
    # Using undocumented "0 MB" budget to test incomplete reads.
    uri = os.path.join(TESTS_INPUT_DIR, "arrays/v3/ingested_2samples")
    cfg = tiledbvcf.ReadConfig(memory_budget_mb=0)
    test_ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)

    df = test_ds.read(attrs=["pos_end"], regions=["1:12700-13400"])
    assert not test_ds.read_completed()
    assert len(df) == 2
    _check_dfs(
        pd.DataFrame.from_dict({"pos_end": np.array([12771, 12771], dtype=np.int32)}),
        df,
    )

    df = test_ds.continue_read()
    assert not test_ds.read_completed()
    assert len(df) == 2
    _check_dfs(
        pd.DataFrame.from_dict({"pos_end": np.array([13374, 13389], dtype=np.int32)}),
        df,
    )

    df = test_ds.continue_read()
    assert test_ds.read_completed()
    assert len(df) == 2
    _check_dfs(
        pd.DataFrame.from_dict({"pos_end": np.array([13395, 13413], dtype=np.int32)}),
        df,
    )

    # test incomplete via read_arrow
    table = test_ds.read_arrow(attrs=["pos_end"], regions=["1:12700-13400"])
    assert not test_ds.read_completed()
    assert len(table) == 2
    _check_dfs(
        pd.DataFrame.from_dict({"pos_end": np.array([12771, 12771], dtype=np.int32)}),
        table.to_pandas(),
    )

    table = test_ds.continue_read_arrow()
    assert not test_ds.read_completed()
    assert len(table) == 2
    _check_dfs(
        pd.DataFrame.from_dict({"pos_end": np.array([13374, 13389], dtype=np.int32)}),
        table.to_pandas(),
    )

    table = test_ds.continue_read_arrow()
    assert test_ds.read_completed()
    assert len(table) == 2
    _check_dfs(
        pd.DataFrame.from_dict({"pos_end": np.array([13395, 13413], dtype=np.int32)}),
        table.to_pandas(),
    )


def test_incomplete_read_generator():
    # Using undocumented "0 MB" budget to test incomplete reads.
    uri = os.path.join(TESTS_INPUT_DIR, "arrays/v3/ingested_2samples")
    cfg = tiledbvcf.ReadConfig(memory_budget_mb=0)
    test_ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)
    expected_df = pd.DataFrame.from_dict(
            {
                "pos_end": np.array(
                    [12771, 12771, 13374, 13389, 13395, 13413], dtype=np.int32
                )
            }
        )

    # NOTE: Running multiple test shows that the iterator can be reused

    # Regions as string
    dfs = []
    for df in test_ds.read_iter(attrs=["pos_end"], regions="1:12700-13400"):
        dfs.append(df)
    overall_df = pd.concat(dfs, ignore_index=True)
    assert len(overall_df) == 6
    _check_dfs(expected_df, overall_df)

    # Regions as list
    dfs = []
    for df in test_ds.read_iter(attrs=["pos_end"], regions=["1:12700-13400"]):
        dfs.append(df)
    overall_df = pd.concat(dfs, ignore_index=True)
    assert len(overall_df) == 6
    _check_dfs(expected_df, overall_df)

    # Regions as numpy.ndarray
    dfs = []
    for df in test_ds.read_iter(attrs=["pos_end"], regions=np.array(["1:12700-13400"])):
        dfs.append(df)
    overall_df = pd.concat(dfs, ignore_index=True)
    assert len(overall_df) == 6
    _check_dfs(expected_df, overall_df)


def test_read_filters(test_ds):
    df = test_ds.read(
        attrs=["sample_name", "pos_start", "pos_end", "filters"],
        regions=["1:12700-13400"],
    )
    expected_df = pd.DataFrame(
        {
            "sample_name": pd.Series(
                ["HG00280", "HG01762", "HG00280", "HG01762", "HG00280", "HG00280"]
            ),
            "pos_start": pd.Series(
                [12546, 12546, 13354, 13354, 13375, 13396], dtype=np.int32
            ),
            "pos_end": pd.Series(
                [12771, 12771, 13374, 13389, 13395, 13413], dtype=np.int32
            ),
            "filters": pd.Series(
                map(
                    lambda lst: np.array(lst, dtype=object),
                    [None, None, ["LowQual"], None, None, None],
                )
            ),
        }
    ).sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    _check_dfs(
        expected_df, df.sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    )


def test_read_var_length_filters(tmp_path):
    uri = os.path.join(tmp_path, "dataset")
    ds = tiledbvcf.Dataset(uri, mode="w")
    samples = [os.path.join(TESTS_INPUT_DIR, s) for s in ["varLenFilter.vcf.gz"]]
    ds.create_dataset()
    ds.ingest_samples(samples)

    ds = tiledbvcf.Dataset(uri, mode="r")
    df = ds.read(["pos_start", "filters"])

    expected_df = pd.DataFrame(
        {
            "pos_start": pd.Series(
                [
                    12141,
                    12546,
                    13354,
                    13375,
                    13396,
                    13414,
                    13452,
                    13520,
                    13545,
                    17319,
                    17480,
                ],
                dtype=np.int32,
            ),
            "filters": pd.Series(
                map(
                    lambda lst: np.array(lst, dtype=object),
                    [
                        ["PASS"],
                        ["PASS"],
                        ["ANEUPLOID", "LowQual"],
                        ["PASS"],
                        ["PASS"],
                        ["ANEUPLOID", "LOWQ", "LowQual"],
                        ["PASS"],
                        ["PASS"],
                        ["PASS"],
                        ["LowQual"],
                        ["PASS"],
                    ],
                )
            ),
        }
    ).sort_values(ignore_index=True, by=["pos_start"])

    _check_dfs(expected_df, df.sort_values(ignore_index=True, by=["pos_start"]))


def test_read_alleles(test_ds):
    df = test_ds.read(
        attrs=["sample_name", "pos_start", "pos_end", "alleles"],
        regions=["1:12100-13360", "1:13500-17350"],
    )
    expected_df = pd.DataFrame(
        {
            "sample_name": pd.Series(
                [
                    "HG00280",
                    "HG01762",
                    "HG00280",
                    "HG01762",
                    "HG00280",
                    "HG01762",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                ]
            ),
            "pos_start": pd.Series(
                [12141, 12141, 12546, 12546, 13354, 13354, 13452, 13520, 13545, 17319],
                dtype=np.int32,
            ),
            "pos_end": pd.Series(
                [12277, 12277, 12771, 12771, 13374, 13389, 13519, 13544, 13689, 17479],
                dtype=np.int32,
            ),
            "alleles": pd.Series(
                map(
                    lambda lst: np.array(lst, dtype=object),
                    [
                        ["C", "<NON_REF>"],
                        ["C", "<NON_REF>"],
                        ["G", "<NON_REF>"],
                        ["G", "<NON_REF>"],
                        ["T", "<NON_REF>"],
                        ["T", "<NON_REF>"],
                        ["G", "<NON_REF>"],
                        ["G", "<NON_REF>"],
                        ["G", "<NON_REF>"],
                        ["T", "<NON_REF>"],
                    ],
                )
            ),
        }
    ).sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    _check_dfs(
        expected_df, df.sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    )


def test_read_multiple_alleles(tmp_path):
    uri = os.path.join(tmp_path, "dataset")
    ds = tiledbvcf.Dataset(uri, mode="w")
    samples = [os.path.join(TESTS_INPUT_DIR, s) for s in ["small3.bcf", "small.bcf"]]
    ds.create_dataset()
    ds.ingest_samples(samples)

    ds = tiledbvcf.Dataset(uri, mode="r")
    df = ds.read(
        attrs=["sample_name", "pos_start", "alleles", "id", "filters"],
        regions=["1:70100-1300000"],
    )
    expected_df = pd.DataFrame(
        {
            "sample_name": pd.Series(["HG00280", "HG00280"]),
            "pos_start": pd.Series([866511, 1289367], dtype=np.int32),
            "alleles": pd.Series(
                map(
                    lambda lst: np.array(lst, dtype=object),
                    [["T", "CCCCTCCCT", "C", "CCCCTCCCTCCCT", "CCCCT"], ["CTG", "C"]],
                )
            ),
            "id": pd.Series([".", "rs1497816"]),
            "filters": pd.Series(
                map(
                    lambda lst: np.array(lst, dtype=object),
                    [["LowQual"], ["LowQual"]],
                )
            ),
        }
    ).sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    _check_dfs(
        expected_df, df.sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    )


def test_read_var_len_attrs(test_ds):
    df = test_ds.read(
        attrs=["sample_name", "pos_start", "pos_end", "fmt_DP", "fmt_PL"],
        regions=["1:12100-13360", "1:13500-17350"],
    )
    expected_df = pd.DataFrame(
        {
            "sample_name": pd.Series(
                [
                    "HG00280",
                    "HG01762",
                    "HG00280",
                    "HG01762",
                    "HG00280",
                    "HG01762",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                ]
            ),
            "pos_start": pd.Series(
                [12141, 12141, 12546, 12546, 13354, 13354, 13452, 13520, 13545, 17319],
                dtype=np.int32,
            ),
            "pos_end": pd.Series(
                [12277, 12277, 12771, 12771, 13374, 13389, 13519, 13544, 13689, 17479],
                dtype=np.int32,
            ),
            "fmt_DP": pd.Series([0, 0, 0, 0, 15, 64, 10, 6, 0, 0], dtype=np.int32),
            "fmt_PL": pd.Series(
                map(
                    lambda lst: np.array(lst, dtype=np.int32),
                    [
                        [0, 0, 0],
                        [0, 0, 0],
                        [0, 0, 0],
                        [0, 0, 0],
                        [0, 24, 360],
                        [0, 66, 990],
                        [0, 21, 210],
                        [0, 6, 90],
                        [0, 0, 0],
                        [0, 0, 0],
                    ],
                )
            ),
        }
    ).sort_values(ignore_index=True, by=["sample_name", "pos_start"])

    _check_dfs(
        expected_df, df.sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    )


def test_sample_args(test_ds, tmp_path):
    sample_file = os.path.join(tmp_path, "1_sample.txt")
    with open(sample_file, "w") as file:
        file.write("HG00280")

    region = ["1:12141-12141"]
    df1 = test_ds.read(["sample_name"], regions=region, samples=["HG00280"])
    df2 = test_ds.read(["sample_name"], regions=region, samples_file=sample_file)
    _check_dfs(df1, df2)

    with pytest.raises(TypeError):
        test_ds.read(
            attrs=["sample_name"],
            regions=region,
            samples=["HG00280"],
            samples_file=sample_file,
        )


def test_read_null_attrs(tmp_path):
    uri = os.path.join(tmp_path, "dataset")
    ds = tiledbvcf.Dataset(uri, mode="w")
    samples = [os.path.join(TESTS_INPUT_DIR, s) for s in ["small3.bcf", "small.bcf"]]
    ds.create_dataset()
    ds.ingest_samples(samples)

    ds = tiledbvcf.Dataset(uri, mode="r")
    df = ds.read(
        attrs=[
            "sample_name",
            "pos_start",
            "pos_end",
            "info_BaseQRankSum",
            "info_DP",
            "fmt_DP",
            "fmt_MIN_DP",
        ],
        regions=["1:12700-13400", "1:69500-69800"],
    )
    expected_df = pd.DataFrame(
        {
            "sample_name": pd.Series(
                [
                    "HG00280",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                    "HG01762",
                    "HG01762",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                    "HG00280",
                ]
            ),
            "pos_start": pd.Series(
                [
                    12546,
                    13354,
                    13375,
                    13396,
                    12546,
                    13354,
                    69371,
                    69511,
                    69512,
                    69761,
                    69762,
                    69771,
                ],
                dtype=np.int32,
            ),
            "pos_end": pd.Series(
                [
                    12771,
                    13374,
                    13395,
                    13413,
                    12771,
                    13389,
                    69510,
                    69511,
                    69760,
                    69761,
                    69770,
                    69834,
                ],
                dtype=np.int32,
            ),
            "info_BaseQRankSum": pd.Series(
                [
                    None,
                    None,
                    None,
                    None,
                    None,
                    None,
                    None,
                    np.array([-0.787], dtype=np.float32),
                    None,
                    np.array([1.97], dtype=np.float32),
                    None,
                    None,
                ]
            ),
            "info_DP": pd.Series(
                [
                    None,
                    None,
                    None,
                    None,
                    None,
                    None,
                    None,
                    np.array([89], dtype=np.int32),
                    None,
                    np.array([24], dtype=np.int32),
                    None,
                    None,
                ]
            ),
            "fmt_DP": pd.Series(
                [0, 15, 6, 2, 0, 64, 180, 88, 97, 24, 23, 21], dtype=np.int32
            ),
            "fmt_MIN_DP": pd.Series([0, 14, 3, 1, 0, 30, 20, None, 24, None, 23, 19]),
        }
    ).sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    _check_dfs(
        expected_df, df.sort_values(ignore_index=True, by=["sample_name", "pos_start"])
    )


def test_read_config():
    uri = os.path.join(TESTS_INPUT_DIR, "arrays/v3/ingested_2samples")
    cfg = tiledbvcf.ReadConfig()
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)

    cfg = tiledbvcf.ReadConfig(
        memory_budget_mb=512,
        region_partition=(0, 3),
        tiledb_config=["sm.tile_cache_size=0", "sm.compute_concurrency_level=1"],
    )
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)

    with pytest.raises(TypeError):
        cfg = tiledbvcf.ReadConfig(abc=123)

    # Expect an exception when passing both cfg and tiledb_config
    with pytest.raises(Exception):
        cfg = tiledbvcf.ReadConfig()
        tiledb_config = {"foo": "bar"}
        ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg, tiledb_config=tiledb_config)


# This test is skipped because running it in the same process as all the normal
# tests will cause it to fail (the first context created in a process determines
# the number of TBB threads allowed).
@pytest.mark.skip
def test_tbb_threads_config():
    uri = os.path.join(TESTS_INPUT_DIR, "arrays/v3/ingested_2samples")
    cfg = tiledbvcf.ReadConfig(tiledb_config=["sm.num_tbb_threads=3"])
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)

    cfg = tiledbvcf.ReadConfig(tiledb_config=["sm.num_tbb_threads=4"])
    with pytest.raises(RuntimeError):
        ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)


def test_read_limit():
    uri = os.path.join(TESTS_INPUT_DIR, "arrays/v3/ingested_2samples")
    cfg = tiledbvcf.ReadConfig(limit=3)
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)
    df = ds.read(
        attrs=["sample_name", "pos_start", "pos_end", "fmt_DP", "fmt_PL"],
        regions=["1:12100-13360", "1:13500-17350"],
    )
    assert len(df) == 3


def test_region_partitioned_read():
    uri = os.path.join(TESTS_INPUT_DIR, "arrays/v3/ingested_2samples")

    cfg = tiledbvcf.ReadConfig(region_partition=(0, 2))
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)
    df = ds.read(
        attrs=["sample_name", "pos_start", "pos_end"],
        regions=["1:12000-13000", "1:17000-18000"],
    )
    assert len(df) == 4

    cfg = tiledbvcf.ReadConfig(region_partition=(1, 2))
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)
    df = ds.read(
        attrs=["sample_name", "pos_start", "pos_end"],
        regions=["1:12000-13000", "1:17000-18000"],
    )
    assert len(df) == 2

    # Too many partitions still produces results
    cfg = tiledbvcf.ReadConfig(region_partition=(1, 3))
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)
    df = ds.read(
        attrs=["sample_name", "pos_start", "pos_end"],
        regions=["1:12000-13000", "1:17000-18000"],
    )
    assert len(df) == 2

    # Error: index >= num partitions
    cfg = tiledbvcf.ReadConfig(region_partition=(2, 2))
    with pytest.raises(RuntimeError):
        ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)


def test_sample_partitioned_read():
    uri = os.path.join(TESTS_INPUT_DIR, "arrays/v3/ingested_2samples")

    cfg = tiledbvcf.ReadConfig(sample_partition=(0, 2))
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)
    df = ds.read(
        attrs=["sample_name", "pos_start", "pos_end"], regions=["1:12000-18000"]
    )
    assert len(df) == 11
    assert (df.sample_name == "HG00280").all()

    cfg = tiledbvcf.ReadConfig(sample_partition=(1, 2))
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)
    df = ds.read(
        attrs=["sample_name", "pos_start", "pos_end"], regions=["1:12000-18000"]
    )
    assert len(df) == 3
    assert (df.sample_name == "HG01762").all()

    # Error: too many partitions
    cfg = tiledbvcf.ReadConfig(sample_partition=(1, 3))
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)
    with pytest.raises(RuntimeError):
        df = ds.read(
            attrs=["sample_name", "pos_start", "pos_end"], regions=["1:12000-18000"]
        )

    # Error: index >= num partitions
    cfg = tiledbvcf.ReadConfig(sample_partition=(2, 2))
    with pytest.raises(RuntimeError):
        ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)


def test_sample_and_region_partitioned_read():
    uri = os.path.join(TESTS_INPUT_DIR, "arrays/v3/ingested_2samples")

    cfg = tiledbvcf.ReadConfig(region_partition=(0, 2), sample_partition=(0, 2))
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)
    df = ds.read(
        attrs=["sample_name", "pos_start", "pos_end"],
        regions=["1:12000-13000", "1:17000-18000"],
    )
    assert len(df) == 2
    assert (df.sample_name == "HG00280").all()

    cfg = tiledbvcf.ReadConfig(region_partition=(0, 2), sample_partition=(1, 2))
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)
    df = ds.read(
        attrs=["sample_name", "pos_start", "pos_end"],
        regions=["1:12000-13000", "1:17000-18000"],
    )
    assert len(df) == 2
    assert (df.sample_name == "HG01762").all()

    cfg = tiledbvcf.ReadConfig(region_partition=(1, 2), sample_partition=(0, 2))
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)
    df = ds.read(
        attrs=["sample_name", "pos_start", "pos_end"],
        regions=["1:12000-13000", "1:17000-18000"],
    )
    assert len(df) == 2
    assert (df.sample_name == "HG00280").all()

    cfg = tiledbvcf.ReadConfig(region_partition=(1, 2), sample_partition=(1, 2))
    ds = tiledbvcf.Dataset(uri, mode="r", cfg=cfg)
    df = ds.read(
        attrs=["sample_name", "pos_start", "pos_end"],
        regions=["1:12000-13000", "1:17000-18000"],
    )
    assert len(df) == 0


@pytest.mark.skipif(os.environ.get("CI") != "true", reason="CI only")
def test_large_export_correctness():
    uri = "s3://tiledb-inc-demo-data/tiledbvcf-arrays/v4/vcf-samples-20"

    ds = tiledbvcf.Dataset(uri)
    df = ds.read(
        attrs=[
            "sample_name",
            "contig",
            "pos_start",
            "pos_end",
            "query_bed_start",
            "query_bed_end",
        ],
        samples=["v2-DjrIAzkP", "v2-YMaDHIoW", "v2-usVwJUmo", "v2-ZVudhauk"],
        bed_file=os.path.join(
            TESTS_INPUT_DIR, "E001_15_coreMarks_dense_filtered.bed.gz"
        ),
    )

    # total number of exported records
    assert df.shape[0] == 1172081

    # number of unique exported records
    record_index = ["sample_name", "contig", "pos_start"]
    assert df[record_index].drop_duplicates().shape[0] == 1168430


def test_basic_ingest(tmp_path):
    # Create the dataset
    uri = os.path.join(tmp_path, "dataset")
    ds = tiledbvcf.Dataset(uri, mode="w")
    samples = [os.path.join(TESTS_INPUT_DIR, s) for s in ["small.bcf", "small2.bcf"]]
    ds.create_dataset()
    ds.ingest_samples(samples)

    # Open it back in read mode and check some queries
    ds = tiledbvcf.Dataset(uri, mode="r")
    assert ds.count() == 14
    assert ds.count(regions=["1:12700-13400"]) == 6
    assert ds.count(samples=["HG00280"], regions=["1:12700-13400"]) == 4


def test_disable_ingestion_tasks(tmp_path):
    # Create the dataset
    uri = os.path.join(tmp_path, "dataset")
    ds = tiledbvcf.Dataset(uri, mode="w")
    samples = [os.path.join(TESTS_INPUT_DIR, s) for s in ["small.bcf", "small3.bcf"]]
    ds.create_dataset(
        enable_allele_count=False, enable_variant_stats=False, enable_sample_stats=False
    )
    ds.ingest_samples(samples)

    # TODO: remove this workaround when sc-19721 is resolved
    if platform.system() != "Linux":
        return

    # Validate that stats arrays were not created
    ac_uri = os.path.join(tmp_path, "dataset", "allele_count")
    vs_uri = os.path.join(tmp_path, "dataset", "variant_stats")
    ss_uri = os.path.join(tmp_path, "dataset", "sample_stats")

    assert not os.path.exists(ac_uri)
    assert not os.path.exists(vs_uri)
    assert not os.path.exists(ss_uri)


def test_ingestion_tasks(tmp_path):
    # Create the dataset
    uri = os.path.join(tmp_path, "dataset")
    ds = tiledbvcf.Dataset(uri, mode="w")
    samples = [os.path.join(TESTS_INPUT_DIR, s) for s in ["small.bcf", "small3.bcf"]]
    ds.create_dataset(enable_allele_count=True, enable_variant_stats=True)
    ds.ingest_samples(samples)

    # TODO: remove this workaround when sc-19721 is resolved
    if platform.system() != "Linux":
        return

    # query allele_count array with TileDB
    ac_uri = tiledb.Group(uri)["allele_count"].uri

    check_if_compatible(ac_uri)

    contig = "1"
    region = slice(69896)
    with tiledb.open(ac_uri) as A:
        df = A.query(attrs=["alt", "count"], dims=["pos"]).df[contig, region]

    assert df["pos"].array == 69896
    assert df["alt"].array == "C"
    assert df["count"].array == 1

    # query variant_stats array with TileDB
    vs_uri = tiledb.Group(uri)["variant_stats"].uri

    contig = "1"
    region = slice(12140)
    with tiledb.open(vs_uri) as A:
        df = A.query(attrs=["allele", "ac"], dims=["pos"]).df[contig, region]

    assert df["pos"].array == 12140
    assert df["allele"].array == "C"
    assert df["ac"].array == 4

    # Test raw sample_stats

    expected_df = pd.DataFrame(
        {
            "sample": ["HG00280", "HG01762"],
            "dp_sum": [879, 64],
            "dp_sum2": [56375, 4096],
            "dp_count": [68, 2],
            "dp_min": [0, 0],
            "dp_max": [180, 64],
            "gq_sum": [1489, 99],
            "gq_sum2": [79129, 9801],
            "gq_count": [68, 2],
            "gq_min": [0, 0],
            "gq_max": [99, 99],
            "n_records": [70, 3],
            "n_called": [70, 3],
            "n_not_called": [0, 0],
            "n_hom_ref": [64, 3],
            "n_het": [3, 0],
            "n_singleton": [4, 0],
            "n_snp": [7, 0],
            "n_insertion": [2, 0],
            "n_deletion": [1, 0],
            "n_transition": [6, 0],
            "n_transversion": [1, 0],
            "n_star": [0, 0],
            "n_multiallelic": [5, 0],
        }
    ).astype("uint64", errors="ignore")

    ss_uri = tiledb.Group(uri)["sample_stats"].uri
    with tiledb.open(ss_uri) as A:
        df = A.df[:]

    # Convert to uint64 for comparison to expected_df
    df = df.astype("uint64", errors="ignore")

    assert df.equals(expected_df)

    # Test sample_qc
    expected_qc = pd.DataFrame(
        {
            "sample": ["HG00280", "HG01762"],
            "dp_mean": [12.92647, 32.0],
            "dp_stddev": [25.728399, 32.0],
            "dp_min": [0, 0],
            "dp_max": [180, 64],
            "gq_mean": [21.897058, 49.5],
            "gq_stddev": [26.156845, 49.5],
            "gq_min": [0, 0],
            "gq_max": [99, 99],
            "call_rate": [1.0, 1.0],
            "n_called": [70, 3],
            "n_not_called": [0, 0],
            "n_hom_ref": [64, 3],
            "n_het": [3, 0],
            "n_hom_var": [3, 0],
            "n_non_ref": [6, 0],
            "n_singleton": [4, 0],
            "n_snp": [7, 0],
            "n_insertion": [2, 0],
            "n_deletion": [1, 0],
            "n_transition": [6, 0],
            "n_transversion": [1, 0],
            "n_star": [0, 0],
            "r_ti_tv": [6.0, np.nan],
            "r_het_hom_var": [1.0, np.nan],
            "r_insertion_deletion": [2.0, np.nan],
            "n_records": [70, 3],
            "n_multiallelic": [5, 0],
        }
    )

    qc = tiledbvcf.sample_qc(uri)
    _check_dfs(expected_qc, qc)


def test_incremental_ingest(tmp_path):
    uri = os.path.join(tmp_path, "dataset")
    ds = tiledbvcf.Dataset(uri, mode="w")
    ds.create_dataset()
    ds.ingest_samples([os.path.join(TESTS_INPUT_DIR, "small.bcf")])
    ds.ingest_samples([os.path.join(TESTS_INPUT_DIR, "small2.bcf")])

    # Open it back in read mode and check some queries
    ds = tiledbvcf.Dataset(uri, mode="r")
    assert ds.count() == 14
    assert ds.count(regions=["1:12700-13400"]) == 6
    assert ds.count(samples=["HG00280"], regions=["1:12700-13400"]) == 4


def test_ingest_disable_merging(tmp_path):
    # Create the dataset
    uri = os.path.join(tmp_path, "dataset_disable_merging")

    cfg = tiledbvcf.ReadConfig(memory_budget_mb=1024)
    attrs = ["sample_name", "contig", "pos_start", "pos_end"]

    ds = tiledbvcf.Dataset(uri, mode="w")
    samples = [
        os.path.join(TESTS_INPUT_DIR, s) for s in ["v2-DjrIAzkP-downsampled.vcf.gz"]
    ]
    ds.create_dataset()
    ds.ingest_samples(samples, contig_fragment_merging=False)

    # Open it back in read mode and check some queries
    ds = tiledbvcf.Dataset(uri, cfg=cfg, mode="r", verbose=False)
    df = ds.read(attrs=attrs)
    assert ds.count() == 246
    assert ds.count(regions=["chrX:9032893-9032893"]) == 1

    # Create the dataset
    uri = os.path.join(tmp_path, "dataset_merging_separate")
    ds2 = tiledbvcf.Dataset(uri, mode="w", verbose=False)
    samples = [
        os.path.join(TESTS_INPUT_DIR, s) for s in ["v2-DjrIAzkP-downsampled.vcf.gz"]
    ]
    ds2.create_dataset()
    ds2.ingest_samples(samples, contigs_to_keep_separate=["chr1"])

    # Open it back in read mode and check some queries
    ds2 = tiledbvcf.Dataset(uri, cfg=cfg, mode="r", verbose=False)
    df2 = ds2.read(attrs=attrs)
    assert df.equals(df2)

    assert ds.count() == 246
    assert ds.count(regions=["chrX:9032893-9032893"]) == 1


def test_ingest_merging_separate(tmp_path):
    # Create the dataset
    uri = os.path.join(tmp_path, "dataset_merging_separate")
    ds = tiledbvcf.Dataset(uri, mode="w")
    samples = [
        os.path.join(TESTS_INPUT_DIR, s) for s in ["v2-DjrIAzkP-downsampled.vcf.gz"]
    ]
    ds.create_dataset()
    ds.ingest_samples(samples, contigs_to_keep_separate=["chr1"])

    # Open it back in read mode and check some queries
    ds = tiledbvcf.Dataset(uri, mode="r")
    assert ds.count() == 246
    assert ds.count(regions=["chrX:9032893-9032893"]) == 1


def test_ingest_merging(tmp_path):
    # Create the dataset
    uri = os.path.join(tmp_path, "dataset_merging")
    ds = tiledbvcf.Dataset(uri, mode="w")
    samples = [
        os.path.join(TESTS_INPUT_DIR, s) for s in ["v2-DjrIAzkP-downsampled.vcf.gz"]
    ]
    ds.create_dataset()
    ds.ingest_samples(samples, contigs_to_allow_merging=["chr1", "chr2"])

    # Open it back in read mode and check some queries
    ds = tiledbvcf.Dataset(uri, mode="r")
    assert ds.count() == 246
    assert ds.count(regions=["chrX:9032893-9032893"]) == 1


def test_ingest_mode_merged(tmp_path):
    # tiledbvcf.config_logging("debug")
    # Create the dataset
    uri = os.path.join(tmp_path, "dataset_merging")
    ds = tiledbvcf.Dataset(uri, mode="w")
    samples = [
        os.path.join(TESTS_INPUT_DIR, s) for s in ["v2-DjrIAzkP-downsampled.vcf.gz"]
    ]
    ds.create_dataset()
    # ingest only merged contigs (pseudo-contigs)
    ds.ingest_samples(samples, contig_mode="merged")

    # Open it back in read mode and check some queries
    ds = tiledbvcf.Dataset(uri, mode="r")
    assert ds.count() == 19
    assert ds.count(regions=["chrX:9032893-9032893"]) == 0


@pytest.fixture
def test_stats_bgzipped_inputs(tmp_path):
    tmp_path_contents = os.listdir(tmp_path)
    if "stats" in tmp_path_contents:
        shutil.rmtree(os.path.join(tmp_path, "stats"))
    shutil.copytree(
        os.path.join(TESTS_INPUT_DIR, "stats"), os.path.join(tmp_path, "stats")
    )
    raw_inputs = glob.glob(os.path.join(tmp_path, "stats", "*.vcf"))
    # print(f"raw inputs: {raw_inputs}")
    for vcf_file in raw_inputs:
        subprocess.run(
            "bcftools view --no-version -Oz -o " + vcf_file + ".gz " + vcf_file,
            shell=True,
            check=True,
        )
    bgzipped_inputs = glob.glob(os.path.join(tmp_path, "stats", "*.gz"))
    for vcf_file in bgzipped_inputs:
        assert subprocess.run("bcftools index " + vcf_file, shell=True).returncode == 0
    if "outputs" in tmp_path_contents:
        shutil.rmtree(os.path.join(tmp_path, "outputs"))
    if "stats_test" in tmp_path_contents:
        shutil.rmtree(os.path.join(tmp_path, "stats_test"))
    return bgzipped_inputs


@pytest.fixture
def test_stats_sample_names(test_stats_bgzipped_inputs):
    assert len(test_stats_bgzipped_inputs) == 8
    return [os.path.basename(file).split(".")[0] for file in test_stats_bgzipped_inputs]


@pytest.fixture
def test_stats_v3_ingestion(tmp_path, test_stats_bgzipped_inputs):
    assert len(test_stats_bgzipped_inputs) == 8
    # print(f"bgzipped inputs: {test_stats_bgzipped_inputs}")
    ds = tiledbvcf.Dataset(uri=os.path.join(tmp_path, "stats_test"), mode="w")
    ds.create_dataset(
        enable_variant_stats=True, enable_allele_count=True, variant_stats_version=3
    )
    ds.ingest_samples(test_stats_bgzipped_inputs)
    ds = tiledbvcf.Dataset(uri=os.path.join(tmp_path, "stats_test"), mode="r")
    return ds


# Ok to skip is missing bcftools in Windows CI job
@pytest.mark.skipif(
    os.environ.get("CI") == "true"
    and platform.system() == "Windows"
    and shutil.which("bcftools") is None,
    reason="no bcftools",
)
def test_ingest_with_stats_v3(
    tmp_path, test_stats_v3_ingestion, test_stats_sample_names
):
    data_frame = test_stats_v3_ingestion.read(
        samples=test_stats_sample_names,
        attrs=["contig", "pos_start", "id", "qual", "info_TILEDB_IAF", "sample_name"],
        set_af_filter="<0.2",
    )
    assert data_frame.shape == (1, 8)
    assert data_frame.query("sample_name == 'second'")["qual"].iloc[0] == pytest.approx(
        343.73
    )
    assert (
        data_frame[data_frame["sample_name"] == "second"]["info_TILEDB_IAF"].iloc[0][0]
        == 0.9375
    )
    data_frame = test_stats_v3_ingestion.read(
        samples=test_stats_sample_names,
        attrs=["contig", "pos_start", "id", "qual", "info_TILEDB_IAF", "sample_name"],
        scan_all_samples=True,
    )
    assert (
        data_frame[
            (data_frame["sample_name"] == "second") & (data_frame["pos_start"] == 4)
        ]["info_TILEDB_IAF"].iloc[0][0]
        == 0.9375
    )

    ######################
    # read_variant_stats #
    ######################

    # test errors
    no_parameter_error = '"region" or "regions" parameter is required'
    exclusive_parameter_error = (
        '"region" and "regions" parameters are mutually exclusive'
    )
    format_error = '"region" parameter must have format "<contig>:<start>-<end>"'
    empty_contig_error = "Region contig cannot be empty"
    base_1_error = "Regions must be 1-based"
    interval_error = '"100-1" is not a valid region interval'
    with pytest.raises(Exception, match=no_parameter_error):
        test_stats_v3_ingestion.read_variant_stats()
    with pytest.raises(Exception, match=no_parameter_error):
        test_stats_v3_ingestion.read_variant_stats_arrow()
    with pytest.raises(Exception, match=exclusive_parameter_error):
        test_stats_v3_ingestion.read_variant_stats("chr1:1-100", regions=["chr1:1-100"])
    with pytest.raises(Exception, match=exclusive_parameter_error):
        test_stats_v3_ingestion.read_variant_stats_arrow(
            "chr1:1-100", regions=["chr1:1-100"]
        )
    with pytest.raises(Exception, match=format_error):
        test_stats_v3_ingestion.read_variant_stats(regions=[""])
    with pytest.raises(Exception, match=format_error):
        test_stats_v3_ingestion.read_variant_stats_arrow(regions=[""])
    with pytest.raises(Exception, match=format_error):
        test_stats_v3_ingestion.read_variant_stats(regions=["chr1"])
    with pytest.raises(Exception, match=format_error):
        test_stats_v3_ingestion.read_variant_stats_arrow(regions=["chr1"])
    with pytest.raises(Exception, match=format_error):
        test_stats_v3_ingestion.read_variant_stats(regions=["chr1:-"])
    with pytest.raises(Exception, match=format_error):
        test_stats_v3_ingestion.read_variant_stats_arrow(regions=["chr1:-"])
    with pytest.raises(Exception, match=empty_contig_error):
        test_stats_v3_ingestion.read_variant_stats(regions=[":1-100"])
    with pytest.raises(Exception, match=empty_contig_error):
        test_stats_v3_ingestion.read_variant_stats_arrow(regions=[":1-100"])
    with pytest.raises(Exception, match=base_1_error):
        test_stats_v3_ingestion.read_variant_stats(regions=["chr1:0-100"])
    with pytest.raises(Exception, match=base_1_error):
        test_stats_v3_ingestion.read_variant_stats_arrow(regions=["chr1:0-100"])
    with pytest.raises(Exception, match=interval_error):
        test_stats_v3_ingestion.read_variant_stats(regions=["chr1:100-1"])
    with pytest.raises(Exception, match=interval_error):
        test_stats_v3_ingestion.read_variant_stats_arrow(regions=["chr1:100-1"])

    # test empty region
    assert test_stats_v3_ingestion.read_variant_stats(regions=["chr3:1-10000"]).empty

    # test types and deprecated region parameter
    region1 = "chr1:1-10000"
    df = test_stats_v3_ingestion.read_variant_stats(region1)
    tbl = test_stats_v3_ingestion.read_variant_stats_arrow(region1)
    assert isinstance(df, pd.DataFrame)
    assert isinstance(tbl, pa.Table)
    assert df.shape == (13, 6)
    assert df.equals(tbl.to_pandas())
    df = test_stats_v3_ingestion.read_variant_stats(regions=[region1])
    tbl = test_stats_v3_ingestion.read_variant_stats_arrow(regions=[region1])
    assert isinstance(df, pd.DataFrame)
    assert isinstance(tbl, pa.Table)
    assert df.shape == (13, 6)
    assert df.equals(tbl.to_pandas())

    # test a region on a different contig
    region2 = "chr2:1-10000"
    df = test_stats_v3_ingestion.read_variant_stats(regions=[region2])
    tbl = test_stats_v3_ingestion.read_variant_stats_arrow(regions=[region2])
    assert df.shape == (2, 6)
    assert df.equals(tbl.to_pandas())

    # test multiple regions from different contigs and their ordering
    regions = [region1, region2]
    contigs = ["chr1"] * 13 + ["chr2"] * 2
    df = test_stats_v3_ingestion.read_variant_stats(regions=regions)
    assert df.shape == (15, 6)
    assert contigs == list(df["contig"].values)
    df2 = test_stats_v3_ingestion.read_variant_stats(regions=reversed(regions))
    assert df.equals(df2)
    tbl = test_stats_v3_ingestion.read_variant_stats_arrow(regions=regions)
    tbl2 = test_stats_v3_ingestion.read_variant_stats_arrow(regions=reversed(regions))
    assert tbl.equals(tbl2)
    assert df.equals(tbl.to_pandas())
    assert df2.equals(tbl2.to_pandas())

    # test overlapping regions on different contigs and their order
    region1 = "chr1:1-1"
    df = test_stats_v3_ingestion.read_variant_stats(regions=[region1])
    assert df.shape == (2, 6)
    region2 = "chr1:1-2"
    df = test_stats_v3_ingestion.read_variant_stats(regions=[region2])
    assert df.shape == (5, 6)
    region3 = "chr1:3-4"
    df = test_stats_v3_ingestion.read_variant_stats(regions=[region3])
    assert df.shape == (6, 6)
    region4 = "chr1:2-5"
    df = test_stats_v3_ingestion.read_variant_stats(regions=[region4])
    assert df.shape == (11, 6)
    regions_chr1 = [region1, region2, region3, region4]
    df = test_stats_v3_ingestion.read_variant_stats(regions=regions_chr1)
    df2 = test_stats_v3_ingestion.read_variant_stats(regions=reversed(regions_chr1))
    assert df.shape == (13, 6)
    assert df.equals(df2)
    region5 = "chr2:1-1"
    df = test_stats_v3_ingestion.read_variant_stats(regions=[region5])
    assert df.shape == (1, 6)
    region6 = "chr2:3-3"
    df = test_stats_v3_ingestion.read_variant_stats(regions=[region6])
    assert df.shape == (1, 6)
    regions_chr2 = [region5, region6]
    df = test_stats_v3_ingestion.read_variant_stats(regions=regions_chr2)
    df2 = test_stats_v3_ingestion.read_variant_stats(regions=reversed(regions_chr2))
    assert df.shape == (2, 6)
    assert df.equals(df2)
    regions = regions_chr1 + regions_chr2
    df = test_stats_v3_ingestion.read_variant_stats(regions=regions)
    df2 = test_stats_v3_ingestion.read_variant_stats(regions=reversed(regions))
    assert df.shape == (15, 6)
    assert contigs == list(df["contig"].values)
    assert df.equals(df2)
    regions = regions_chr2 + regions_chr1
    df = test_stats_v3_ingestion.read_variant_stats(regions=regions)
    df2 = test_stats_v3_ingestion.read_variant_stats(regions=reversed(regions))
    assert df.shape == (15, 6)
    assert contigs == list(df["contig"].values)
    assert df.equals(df2)

    # test scan_all_samples
    ac = [8, 8, 5, 6, 5, 4, 4, 4, 4, 1, 15, 1, 2, 2, 2]
    an = [16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 3, 3, 2, 2]
    af = [
        0.5,
        0.5,
        0.3125,
        0.375,
        0.3125,
        0.25,
        0.25,
        0.25,
        0.25,
        0.0625,
        0.9375,
        0.33333334,
        0.6666667,
        1.0,
        1.0,
    ]
    df = test_stats_v3_ingestion.read_variant_stats(regions=regions)
    assert ac == list(df["ac"].values)
    assert an == list(df["an"].values)
    assert af == list(df["af"].values)
    ac = [8, 8, 5, 6, 5, 4, 4, 4, 4, 1, 15, 1, 2, 2, 2]
    an = [16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16]
    af = [
        0.5,
        0.5,
        0.3125,
        0.375,
        0.3125,
        0.25,
        0.25,
        0.25,
        0.25,
        0.0625,
        0.9375,
        0.0625,
        0.125,
        0.125,
        0.125,
    ]
    df = test_stats_v3_ingestion.read_variant_stats(
        regions=regions,
        scan_all_samples=True,
    )
    assert ac == list(df["ac"].values)
    assert an == list(df["an"].values)
    assert af == list(df["af"].values)

    # test drop_ref
    alleles = [
        "T,C",
        "ref",
        "G,GTTTA",
        "G,T",
        "ref",
        "C,A",
        "C,G",
        "C,T",
        "ref",
        "G,GTTTA",
        "ref",
        "C,T",
        "ref",
        "G,GTTTA",
        "G,GTTTA",
    ]
    df = test_stats_v3_ingestion.read_variant_stats(regions=regions)
    assert alleles == list(df["alleles"].values)
    alleles = [
        "T,C",
        "G,GTTTA",
        "G,T",
        "C,A",
        "C,G",
        "C,T",
        "G,GTTTA",
        "C,T",
        "G,GTTTA",
        "G,GTTTA",
    ]
    df = test_stats_v3_ingestion.read_variant_stats(
        regions=regions,
        drop_ref=True,
    )
    assert alleles == list(df["alleles"].values)

    ######################
    # read_allele_count #
    ######################

    # test errors
    with pytest.raises(Exception, match=no_parameter_error):
        test_stats_v3_ingestion.read_allele_count()
    with pytest.raises(Exception, match=no_parameter_error):
        test_stats_v3_ingestion.read_allele_count_arrow()
    with pytest.raises(Exception, match=exclusive_parameter_error):
        test_stats_v3_ingestion.read_allele_count("chr1:1-100", regions=["chr1:1-100"])
    with pytest.raises(Exception, match=exclusive_parameter_error):
        test_stats_v3_ingestion.read_allele_count_arrow(
            "chr1:1-100", regions=["chr1:1-100"]
        )
    with pytest.raises(Exception, match=format_error):
        test_stats_v3_ingestion.read_allele_count(regions=[""])
    with pytest.raises(Exception, match=format_error):
        test_stats_v3_ingestion.read_allele_count_arrow(regions=[""])
    with pytest.raises(Exception, match=format_error):
        test_stats_v3_ingestion.read_allele_count(regions=["chr1"])
    with pytest.raises(Exception, match=format_error):
        test_stats_v3_ingestion.read_allele_count_arrow(regions=["chr1"])
    with pytest.raises(Exception, match=format_error):
        test_stats_v3_ingestion.read_allele_count(regions=["chr1:-"])
    with pytest.raises(Exception, match=format_error):
        test_stats_v3_ingestion.read_allele_count_arrow(regions=["chr1:-"])
    with pytest.raises(Exception, match=empty_contig_error):
        test_stats_v3_ingestion.read_allele_count(regions=[":1-100"])
    with pytest.raises(Exception, match=empty_contig_error):
        test_stats_v3_ingestion.read_allele_count_arrow(regions=[":1-100"])
    with pytest.raises(Exception, match=base_1_error):
        test_stats_v3_ingestion.read_allele_count(regions=["chr1:0-100"])
    with pytest.raises(Exception, match=base_1_error):
        test_stats_v3_ingestion.read_allele_count_arrow(regions=["chr1:0-100"])
    with pytest.raises(Exception, match=interval_error):
        test_stats_v3_ingestion.read_allele_count(regions=["chr1:100-1"])
    with pytest.raises(Exception, match=interval_error):
        test_stats_v3_ingestion.read_allele_count_arrow(regions=["chr1:100-1"])

    # test empty region
    assert test_stats_v3_ingestion.read_allele_count(regions=["chr3:1-10000"]).empty

    # test types and deprecated region parameter
    region1 = "chr1:1-10000"
    pos = (0, 1, 1, 2, 2, 2, 3)
    count = (8, 5, 3, 4, 2, 2, 1)
    df = test_stats_v3_ingestion.read_allele_count(region1)
    tbl = test_stats_v3_ingestion.read_allele_count_arrow(region1)
    assert isinstance(df, pd.DataFrame)
    assert isinstance(tbl, pa.Table)
    assert df.shape == (7, 7)
    assert df.equals(tbl.to_pandas())
    assert sum(df["pos"] == pos) == 7
    assert sum(df["count"] == count) == 7
    df = test_stats_v3_ingestion.read_allele_count(regions=[region1])
    tbl = test_stats_v3_ingestion.read_allele_count_arrow(regions=[region1])
    assert isinstance(df, pd.DataFrame)
    assert isinstance(tbl, pa.Table)
    assert df.shape == (7, 7)
    assert df.equals(tbl.to_pandas())
    assert sum(df["pos"] == pos) == 7
    assert sum(df["count"] == count) == 7

    # test a region on a different contig
    region2 = "chr2:1-10000"
    df = test_stats_v3_ingestion.read_allele_count(regions=[region2])
    tbl = test_stats_v3_ingestion.read_allele_count_arrow(regions=[region2])
    assert df.shape == (2, 7)
    assert df.equals(tbl.to_pandas())

    # test multiple regions from different contigs and their ordering
    regions = [region1, region2]
    contigs = ["chr1"] * 7 + ["chr2"] * 2
    df = test_stats_v3_ingestion.read_allele_count(regions=regions)
    assert df.shape == (9, 7)
    assert contigs == list(df["contig"].values)
    df2 = test_stats_v3_ingestion.read_allele_count(regions=reversed(regions))
    assert df.equals(df2)
    tbl = test_stats_v3_ingestion.read_allele_count_arrow(regions=regions)
    tbl2 = test_stats_v3_ingestion.read_allele_count_arrow(regions=reversed(regions))
    assert tbl.equals(tbl2)
    assert df.equals(tbl.to_pandas())
    assert df2.equals(tbl2.to_pandas())

    # test overlapping regions on different contigs and their order
    region1 = "chr1:1-1"
    df = test_stats_v3_ingestion.read_allele_count(regions=[region1])
    assert df.shape == (1, 7)
    region2 = "chr1:1-2"
    df = test_stats_v3_ingestion.read_allele_count(regions=[region2])
    assert df.shape == (3, 7)
    region3 = "chr1:3-4"
    df = test_stats_v3_ingestion.read_allele_count(regions=[region3])
    assert df.shape == (4, 7)
    region4 = "chr1:2-5"
    df = test_stats_v3_ingestion.read_allele_count(regions=[region4])
    assert df.shape == (6, 7)
    regions_chr1 = [region1, region2, region3, region4]
    df = test_stats_v3_ingestion.read_allele_count(regions=regions_chr1)
    df2 = test_stats_v3_ingestion.read_allele_count(regions=reversed(regions_chr1))
    assert df.shape == (7, 7)
    assert df.equals(df2)
    region5 = "chr2:1-1"
    df = test_stats_v3_ingestion.read_allele_count(regions=[region5])
    assert df.shape == (1, 7)
    region6 = "chr2:3-3"
    df = test_stats_v3_ingestion.read_allele_count(regions=[region6])
    assert df.shape == (1, 7)
    regions_chr2 = [region5, region6]
    df = test_stats_v3_ingestion.read_allele_count(regions=regions_chr2)
    df2 = test_stats_v3_ingestion.read_allele_count(regions=reversed(regions_chr2))
    assert df.shape == (2, 7)
    assert df.equals(df2)
    regions = regions_chr1 + regions_chr2
    df = test_stats_v3_ingestion.read_allele_count(regions=regions)
    df2 = test_stats_v3_ingestion.read_allele_count(regions=reversed(regions))
    assert df.shape == (9, 7)
    assert contigs == list(df["contig"].values)
    assert df.equals(df2)
    regions = regions_chr2 + regions_chr1
    df = test_stats_v3_ingestion.read_allele_count(regions=regions)
    df2 = test_stats_v3_ingestion.read_allele_count(regions=reversed(regions))
    assert df.shape == (9, 7)
    assert contigs == list(df["contig"].values)
    assert df.equals(df2)

    #########################
    # read_allele_frequency #
    #########################

    region = "chr1:1-10000"
    df = tiledbvcf.allele_frequency.read_allele_frequency(
        os.path.join(tmp_path, "stats_test"), region
    )
    assert df.pos.is_monotonic_increasing
    df["an_check"] = (df.ac / df.af).round(0).astype("int32")
    assert df.an_check.equals(df.an)
    df = test_stats_v3_ingestion.read_variant_stats(region)
    assert df.shape == (13, 6)


@pytest.mark.skipif(
    os.environ.get("CI") == "true"
    and platform.system() == "Windows"
    and shutil.which("bcftools") is None,
    reason="no bcftools",
)
def test_delete_samples(tmp_path, test_stats_v3_ingestion, test_stats_sample_names):
    #    assert test_stats_v3_ingestion.samples() == test_stats_sample_names
    assert "second" in test_stats_sample_names
    assert "fifth" in test_stats_sample_names
    assert "third" in test_stats_sample_names
    ds = tiledbvcf.Dataset(uri=os.path.join(tmp_path, "stats_test"), mode="w")
    # tiledbvcf.config_logging("trace")
    ds.delete_samples(["second", "fifth"])
    ds = tiledbvcf.Dataset(uri=os.path.join(tmp_path, "stats_test"), mode="r")
    sample_names = ds.samples()
    assert "second" not in sample_names
    assert "fifth" not in sample_names
    assert "third" in sample_names


# Ok to skip is missing bcftools in Windows CI job
@pytest.mark.skipif(
    os.environ.get("CI") == "true"
    and platform.system() == "Windows"
    and shutil.which("bcftools") is None,
    reason="no bcftools",
)
def test_ingest_with_stats_v2(tmp_path):
    # tiledbvcf.config_logging("debug")
    tmp_path_contents = os.listdir(tmp_path)
    if "stats" in tmp_path_contents:
        shutil.rmtree(os.path.join(tmp_path, "stats"))
    shutil.copytree(
        os.path.join(TESTS_INPUT_DIR, "stats"), os.path.join(tmp_path, "stats")
    )
    raw_inputs = glob.glob(os.path.join(tmp_path, "stats", "*.vcf"))
    # print(f"raw inputs: {raw_inputs}")
    for vcf_file in raw_inputs:
        subprocess.run(
            "bcftools view --no-version -Oz -o " + vcf_file + ".gz " + vcf_file,
            shell=True,
            check=True,
        )
    bgzipped_inputs = glob.glob(os.path.join(tmp_path, "stats", "*.gz"))
    # print(f"bgzipped inputs: {bgzipped_inputs}")
    for vcf_file in bgzipped_inputs:
        assert subprocess.run("bcftools index " + vcf_file, shell=True).returncode == 0
    if "outputs" in tmp_path_contents:
        shutil.rmtree(os.path.join(tmp_path, "outputs"))
    if "stats_test" in tmp_path_contents:
        shutil.rmtree(os.path.join(tmp_path, "stats_test"))
    # tiledbvcf.config_logging("trace")
    ds = tiledbvcf.Dataset(uri=os.path.join(tmp_path, "stats_test"), mode="w")
    ds.create_dataset(enable_variant_stats=True, enable_allele_count=True)
    ds.ingest_samples(bgzipped_inputs)
    ds = tiledbvcf.Dataset(uri=os.path.join(tmp_path, "stats_test"), mode="r")
    sample_names = [os.path.basename(file).split(".")[0] for file in bgzipped_inputs]
    data_frame = ds.read(
        samples=sample_names,
        attrs=["contig", "pos_start", "id", "qual", "info_TILEDB_IAF", "sample_name"],
        set_af_filter="<0.2",
    )
    assert data_frame.shape == (1, 8)
    assert data_frame.query("sample_name == 'second'")["qual"].iloc[0] == pytest.approx(
        343.73
    )
    assert (
        data_frame[data_frame["sample_name"] == "second"]["info_TILEDB_IAF"].iloc[0][0]
        == 0.9375
    )
    data_frame = ds.read(
        samples=sample_names,
        attrs=["contig", "pos_start", "id", "qual", "info_TILEDB_IAF", "sample_name"],
        scan_all_samples=True,
    )
    assert (
        data_frame[
            (data_frame["sample_name"] == "second") & (data_frame["pos_start"] == 4)
        ]["info_TILEDB_IAF"].iloc[0][0]
        == 0.9375
    )
    ds = tiledbvcf.Dataset(uri=os.path.join(tmp_path, "stats_test"), mode="r")
    df = ds.read_variant_stats("chr1:1-10000")
    assert df.shape == (13, 6)
    df = tiledbvcf.allele_frequency.read_allele_frequency(
        os.path.join(tmp_path, "stats_test"), "chr1:1-10000"
    )
    assert df.pos.is_monotonic_increasing
    df["an_check"] = (df.ac / df.af).round(0).astype("int32")
    assert df.an_check.equals(df.an)
    df = ds.read_variant_stats("chr1:1-10000")
    assert df.shape == (13, 6)
    df = ds.read_allele_count("chr1:1-10000")
    assert df.shape == (7, 7)
    assert sum(df["pos"] == (0, 1, 1, 2, 2, 2, 3)) == 7
    assert sum(df["count"] == (8, 5, 3, 4, 2, 2, 1)) == 7


# Ok to skip is missing bcftools in Windows CI job
@pytest.mark.skipif(
    os.environ.get("CI") == "true"
    and platform.system() == "Windows"
    and shutil.which("bcftools") is None,
    reason="no bcftools",
)
def test_ingest_polyploid(tmp_path):
    tmp_path_contents = os.listdir(tmp_path)
    if "polyploid" in tmp_path_contents:
        shutil.rmtree(os.path.join(tmp_path, "polyploid"))
    shutil.copytree(
        os.path.join(TESTS_INPUT_DIR, "polyploid"), os.path.join(tmp_path, "polyploid")
    )
    raw_inputs = glob.glob(os.path.join(tmp_path, "polyploid", "*.vcf"))
    # print(f"raw inputs: {raw_inputs}")
    for vcf_file in raw_inputs:
        subprocess.run(
            "bcftools view --no-version -Oz -o " + vcf_file + ".gz " + vcf_file,
            shell=True,
            check=True,
        )
    bgzipped_inputs = glob.glob(os.path.join(tmp_path, "polyploid", "*.gz"))
    # print(f"bgzipped inputs: {bgzipped_inputs}")
    for vcf_file in bgzipped_inputs:
        assert subprocess.run("bcftools index " + vcf_file, shell=True).returncode == 0
    if "polyploid" in tmp_path_contents:
        shutil.rmtree(os.path.join(tmp_path, "polyploid"))
    if "outputs" in tmp_path_contents:
        shutil.rmtree(os.path.join(tmp_path, "outputs"))
    if "polyploid_test" in tmp_path_contents:
        shutil.rmtree(os.path.join(tmp_path, "polyploid_test"))
    # tiledbvcf.config_logging("trace")
    ds = tiledbvcf.Dataset(uri=os.path.join(tmp_path, "polyploid_test"), mode="w")
    ds.create_dataset(enable_variant_stats=True)
    ds.ingest_samples(bgzipped_inputs)
    ds = tiledbvcf.Dataset(uri=os.path.join(tmp_path, "polyploid_test"), mode="r")
    sample_names = [os.path.basename(file).split(".")[0] for file in bgzipped_inputs]
    data_frame = ds.read(
        samples=sample_names,
        attrs=["contig", "pos_start", "id", "qual", "info_TILEDB_IAF", "sample_name"],
        set_af_filter="<0.8",
    )
    # print(data_frame)


def test_ingest_mode_separate(tmp_path):
    # tiledbvcf.config_logging("debug")
    # Create the dataset
    uri = os.path.join(tmp_path, "dataset_merging")
    ds = tiledbvcf.Dataset(uri, mode="w")
    samples = [
        os.path.join(TESTS_INPUT_DIR, s) for s in ["v2-DjrIAzkP-downsampled.vcf.gz"]
    ]
    ds.create_dataset()
    # ingest only merged contigs (pseudo-contigs)
    ds.ingest_samples(
        samples, contigs_to_keep_separate=["chr1"], contig_mode="separate"
    )

    # Open it back in read mode and check some queries
    ds = tiledbvcf.Dataset(uri, mode="r")
    assert ds.count() == 17
    assert ds.count(regions=["chrX:9032893-9032893"]) == 0


def test_vcf_attrs(tmp_path):
    # Create the dataset with vcf info and fmt attributes
    uri = os.path.join(tmp_path, "vcf_attrs_dataset")
    ds = tiledbvcf.Dataset(uri, mode="w")
    vcf_uri = os.path.join(TESTS_INPUT_DIR, "v2-DjrIAzkP-downsampled.vcf.gz")
    ds.create_dataset(vcf_attrs=vcf_uri)

    # Open it back in read mode and check attributes
    ds = tiledbvcf.Dataset(uri, mode="r")

    queryable_attrs = [
        "alleles",
        "contig",
        "filters",
        "fmt",
        "fmt_DP",
        "fmt_GQ",
        "fmt_GT",
        "fmt_MIN_DP",
        "fmt_PS",
        "fmt_SB",
        "fmt_STR_MAX_LEN",
        "fmt_STR_PERIOD",
        "fmt_STR_TIMES",
        "fmt_VAR_CONTEXT",
        "fmt_VAR_TYPE",
        "id",
        "info",
        "info_AC",
        "info_AC_AFR",
        "info_AC_AMR",
        "info_AC_Adj",
        "info_AC_CONSANGUINEOUS",
        "info_AC_EAS",
        "info_AC_FEMALE",
        "info_AC_FIN",
        "info_AC_Hemi",
        "info_AC_Het",
        "info_AC_Hom",
        "info_AC_MALE",
        "info_AC_NFE",
        "info_AC_OTH",
        "info_AC_POPMAX",
        "info_AC_SAS",
        "info_AF",
        "info_AF_AFR",
        "info_AF_AMR",
        "info_AF_Adj",
        "info_AF_EAS",
        "info_AF_FIN",
        "info_AF_NFE",
        "info_AF_OTH",
        "info_AF_SAS",
        "info_AGE_HISTOGRAM_HET",
        "info_AGE_HISTOGRAM_HOM",
        "info_AN",
        "info_AN_AFR",
        "info_AN_AMR",
        "info_AN_Adj",
        "info_AN_CONSANGUINEOUS",
        "info_AN_EAS",
        "info_AN_FEMALE",
        "info_AN_FIN",
        "info_AN_MALE",
        "info_AN_NFE",
        "info_AN_OTH",
        "info_AN_POPMAX",
        "info_AN_SAS",
        "info_BaseQRankSum",
        "info_CCC",
        "info_CSQ",
        "info_ClippingRankSum",
        "info_DB",
        "info_DOUBLETON_DIST",
        "info_DP",
        "info_DP_HIST",
        "info_DS",
        "info_END",
        "info_ESP_AC",
        "info_ESP_AF_GLOBAL",
        "info_ESP_AF_POPMAX",
        "info_FS",
        "info_GQ_HIST",
        "info_GQ_MEAN",
        "info_GQ_STDDEV",
        "info_HWP",
        "info_HaplotypeScore",
        "info_Hemi_AFR",
        "info_Hemi_AMR",
        "info_Hemi_EAS",
        "info_Hemi_FIN",
        "info_Hemi_NFE",
        "info_Hemi_OTH",
        "info_Hemi_SAS",
        "info_Het_AFR",
        "info_Het_AMR",
        "info_Het_EAS",
        "info_Het_FIN",
        "info_Het_NFE",
        "info_Het_OTH",
        "info_Het_SAS",
        "info_Hom_AFR",
        "info_Hom_AMR",
        "info_Hom_CONSANGUINEOUS",
        "info_Hom_EAS",
        "info_Hom_FIN",
        "info_Hom_NFE",
        "info_Hom_OTH",
        "info_Hom_SAS",
        "info_InbreedingCoeff",
        "info_K1_RUN",
        "info_K2_RUN",
        "info_K3_RUN",
        "info_KG_AC",
        "info_KG_AF_GLOBAL",
        "info_KG_AF_POPMAX",
        "info_MLEAC",
        "info_MLEAF",
        "info_MQ",
        "info_MQ0",
        "info_MQRankSum",
        "info_NCC",
        "info_NEGATIVE_TRAIN_SITE",
        "info_OLD_VARIANT",
        "info_POPMAX",
        "info_POSITIVE_TRAIN_SITE",
        "info_QD",
        "info_ReadPosRankSum",
        "info_VQSLOD",
        "info_clinvar_conflicted",
        "info_clinvar_measureset_id",
        "info_clinvar_mut",
        "info_clinvar_pathogenic",
        "info_culprit",
        "pos_end",
        "pos_start",
        "qual",
        "query_bed_end",
        "query_bed_line",
        "query_bed_start",
        "sample_name",
    ]

    assert ds.attributes(attr_type="info") == []
    assert ds.attributes(attr_type="fmt") == []
    assert sorted(ds.attributes()) == sorted(queryable_attrs)


@pytest.mark.parametrize("compress", [True, False])
def test_sample_compression(tmp_path, compress):
    # Create the dataset
    dataset_uri = os.path.join(tmp_path, "sample_compression")
    array_uri = os.path.join(dataset_uri, "data")
    ds = tiledbvcf.Dataset(dataset_uri, mode="w")
    ds.create_dataset(compress_sample_dim=compress)

    check_if_compatible(array_uri)

    # Check for the presence of the Zstd filter
    found_zstd = False
    with tiledb.open(array_uri) as A:
        for filter in A.domain.dim("sample").filters:
            found_zstd = found_zstd or "Zstd" in str(filter)

    assert found_zstd == compress


@pytest.mark.parametrize("level", [1, 4, 16, 22])
def test_compression_level(tmp_path, level):
    # Create the dataset
    dataset_uri = os.path.join(tmp_path, "compression_level")
    array_uri = os.path.join(dataset_uri, "data")
    ds = tiledbvcf.Dataset(dataset_uri, mode="w")
    ds.create_dataset(compression_level=level)

    check_if_compatible(array_uri)

    # Check for the expected compression level
    with tiledb.open(array_uri) as A:
        for i in range(A.schema.nattr):
            attr = A.schema.attr(i)
            for filter in attr.filters:
                if "Zstd" in str(filter):
                    assert filter.level == level


# Ok to skip is missing bcftools in Windows CI job
@pytest.mark.skipif(
    os.environ.get("CI") == "true"
    and platform.system() == "Windows"
    and shutil.which("bcftools") is None,
    reason="no bcftools",
)
def test_gvcf_export(tmp_path):
    # Compress the input VCFs
    vcf_inputs = glob.glob(os.path.join(TESTS_INPUT_DIR, "gvcf-export", "*.vcf"))
    for vcf_input in vcf_inputs:
        vcf_output = os.path.join(tmp_path, os.path.basename(vcf_input)) + ".gz"
        cmd = f"bcftools view --no-version -Oz -o {vcf_output} {vcf_input}"
        subprocess.run(cmd, shell=True, check=True)

    # Index the compressed VCFs
    vcf_files = glob.glob(os.path.join(tmp_path, "*.gz"))
    for vcf_file in vcf_files:
        cmd = f"bcftools index {vcf_file}"
        subprocess.run(cmd, shell=True, check=True)

    # Ingest the VCFs
    uri = os.path.join(tmp_path, "vcf.tdb")
    ds = tiledbvcf.Dataset(uri=uri, mode="w")
    ds.create_dataset()
    ds.ingest_samples(vcf_files)
    ds = tiledbvcf.Dataset(uri=uri, mode="r")

    # List of tests.
    tests = [
        {"region": "chr1:100-120", "samples": ["s0", "s1", "s2"]},
        {"region": "chr1:110-120", "samples": ["s0", "s1"]},
        {"region": "chr1:149-149", "samples": ["s0", "s1", "s3"]},
        {"region": "chr1:150-150", "samples": ["s0", "s1", "s3", "s4"]},
    ]

    # No IAF filtering or reporting
    for test in tests:
        df = ds.read(regions=test["region"])
        assert set(df["sample_name"].unique()) == set(test["samples"])

    attrs = [
        "sample_name",
        "contig",
        "pos_start",
        "alleles",
        "fmt_GT",
        "info_TILEDB_IAF",
    ]

    # IAF reporting
    for test in tests:
        df = ds.read(attrs=attrs, regions=test["region"])
        assert set(df["sample_name"].unique()) == set(test["samples"])

    # IAF filtering and reporting
    for test in tests:
        df = ds.read(attrs=attrs, regions=test["region"], set_af_filter="<=1.0")
        assert set(df["sample_name"].unique()) == set(test["samples"])


def test_flag_export(tmp_path):
    # Create the dataset
    uri = os.path.join(tmp_path, "dataset")
    ds = tiledbvcf.Dataset(uri, mode="w")
    samples = [os.path.join(TESTS_INPUT_DIR, s) for s in ["small.vcf.gz"]]
    ds.create_dataset()
    ds.ingest_samples(samples)

    # Read info flags
    ds = tiledbvcf.Dataset(uri, mode="r")
    df = ds.read(attrs=["pos_start", "info_DB", "info_DS"])
    df = df.sort_values(by=["pos_start"])

    # Check if flags match the expected values
    expected_db = [1, 1, 1, 0, 0, 1]
    assert df["info_DB"].tolist() == expected_db

    expected_ds = [1, 1, 0, 0, 1, 1]
    assert df["info_DS"].tolist() == expected_ds


def test_bed_filestore(tmp_path, test_ds_v4):
    # tiledbvcf.config_logging("debug")

    expected_df = pd.DataFrame(
        {
            "sample_name": pd.Series(
                [
                    "HG00280",
                    "HG01762",
                    "HG00280",
                    "HG01762",
                    "HG00280",
                ]
            ),
            "pos_start": pd.Series(
                [
                    12141,
                    12141,
                    12546,
                    12546,
                    17319,
                ],
                dtype=np.int32,
            ),
            "pos_end": pd.Series(
                [
                    12277,
                    12277,
                    12771,
                    12771,
                    17479,
                ],
                dtype=np.int32,
            ),
        }
    ).sort_values(ignore_index=True, by=["sample_name", "pos_start"])

    # Create BED file
    bed_file = os.path.join(tmp_path, "test.bed")

    regions = [
        (1, 12000, 13000),
        (1, 17000, 17479),
    ]

    with open(bed_file, "w") as f:
        for region in regions:
            f.write(f"{region[0]}\t{region[1]}\t{region[2]}\n")

    # Create BED filestore from BED file
    bed_filestore = os.path.join(tmp_path, "test.bed.filestore")
    tiledb.Array.create(bed_filestore, tiledb.ArraySchema.from_file(bed_file))
    tiledb.Filestore.copy_from(bed_filestore, bed_file)

    # Create the dataset
    for use_arrow in [False, True]:
        func = test_ds_v4.read_arrow if use_arrow else test_ds_v4.read

        df = func(attrs=["sample_name", "pos_start", "pos_end"], bed_file=bed_filestore)
        if use_arrow:
            df = df.to_pandas()

        # print(df)

        _check_dfs(
            expected_df,
            df.sort_values(ignore_index=True, by=["sample_name", "pos_start"]),
        )


def test_bed_array(tmp_path, test_ds_v4):
    expected_df = pd.DataFrame(
        {
            "sample_name": pd.Series(
                [
                    "HG00280",
                    "HG01762",
                    "HG00280",
                    "HG01762",
                    "HG00280",
                ]
            ),
            "pos_start": pd.Series(
                [
                    12141,
                    12141,
                    12546,
                    12546,
                    17319,
                ],
                dtype=np.int32,
            ),
            "pos_end": pd.Series(
                [
                    12277,
                    12277,
                    12771,
                    12771,
                    17479,
                ],
                dtype=np.int32,
            ),
        }
    ).sort_values(ignore_index=True, by=["sample_name", "pos_start"])

    # Create bed array
    bed_array = os.path.join(tmp_path, "bed_array")
    tiledb.from_pandas(
        bed_array,
        pd.DataFrame(
            {
                "chrom": ["1", "1"],
                "chromStart": [12000, 17000],
                "chromEnd": [13000, 17479],
            }
        ),
        sparse=True,
        index_col=["chrom", "chromStart"],
    )

    # Add aliases to the array metadata
    with tiledb.Array(bed_array, "w") as A:
        A.meta["alias contig"] = "chrom"
        A.meta["alias start"] = "chromStart"
        A.meta["alias end"] = "chromEnd"

    # Create the dataset
    for use_arrow in [False, True]:
        func = test_ds_v4.read_arrow if use_arrow else test_ds_v4.read

        df = func(attrs=["sample_name", "pos_start", "pos_end"], bed_file=bed_array)
        if use_arrow:
            df = df.to_pandas()

        _check_dfs(
            expected_df,
            df.sort_values(ignore_index=True, by=["sample_name", "pos_start"]),
        )


def test_info_end(tmp_path):
    """
    This test checks that the info_END attribute is handled correctly, even when the
    VCF header incorrectly defines the END attribute as a string.

    The test also checks that info_END contains the original values from the VCF,
    including the missing values.
    """

    expected_end = pd.DataFrame(
        {
            "pos_end": pd.Series(
                [
                    12277,
                    12771,
                    13374,
                    13395,
                    13413,
                    13451,
                    13519,
                    13544,
                    13689,
                    17479,
                    17486,
                    30553,
                    35224,
                    35531,
                    35786,
                    69096,
                    69103,
                    69104,
                    69109,
                    69110,
                    69111,
                    69112,
                    69114,
                    69115,
                    69122,
                    69123,
                    69128,
                    69129,
                    69130,
                    69192,
                    69195,
                    69196,
                    69215,
                    69222,
                    69227,
                    69228,
                    69261,
                    69262,
                    69269,
                    69270,
                    69346,
                    69349,
                    69352,
                    69353,
                    69370,
                    69510,
                    69511,
                    69760,
                    69761,
                    69770,
                    69834,
                    69835,
                    69838,
                    69861,
                    69863,
                    69866,
                    69896,
                    69897,
                    69912,
                    69938,
                    69939,
                    69941,
                    69946,
                    69947,
                    69948,
                    69949,
                    69953,
                    70012,
                    866511,
                    1289369,
                ],
                dtype=np.int32,
            ),
            # Expected values are strings because the small3.vcf.gz defines END as a string
            "info_END": pd.Series(
                [
                    "12277",
                    "12771",
                    "13374",
                    "13395",
                    "13413",
                    "13451",
                    "13519",
                    "13544",
                    "13689",
                    "17479",
                    "17486",
                    "30553",
                    "35224",
                    "35531",
                    "35786",
                    "69096",
                    "69103",
                    "69104",
                    "69109",
                    "69110",
                    "69111",
                    "69112",
                    "69114",
                    "69115",
                    "69122",
                    "69123",
                    "69128",
                    "69129",
                    "69130",
                    "69192",
                    "69195",
                    "69196",
                    "69215",
                    "69222",
                    "69227",
                    "69228",
                    "69261",
                    "69262",
                    "69269",
                    None,
                    "69346",
                    "69349",
                    "69352",
                    "69353",
                    "69370",
                    "69510",
                    None,
                    "69760",
                    None,
                    "69770",
                    "69834",
                    "69835",
                    "69838",
                    "69861",
                    "69863",
                    "69866",
                    "69896",
                    None,
                    "69912",
                    "69938",
                    "69939",
                    "69941",
                    "69946",
                    "69947",
                    "69948",
                    "69949",
                    "69953",
                    "70012",
                    None,
                    None,
                ],
                dtype=object,
            ),
        }
    )

    # Ingest the data
    uri = os.path.join(tmp_path, "dataset")
    ds = tiledbvcf.Dataset(uri, mode="w")
    samples = [os.path.join(TESTS_INPUT_DIR, s) for s in ["small3.vcf.gz"]]
    ds.create_dataset()
    ds.ingest_samples(samples)

    # Read the data
    ds = tiledbvcf.Dataset(uri)
    df = ds.read(attrs=["sample_name", "pos_start", "pos_end", "info_END"])

    # Sort the results because VCF uses an unordered reader
    df.sort_values(ignore_index=True, by=["sample_name", "pos_start"], inplace=True)

    # Drop the columns that are not used for comparison
    df.drop(columns=["sample_name", "pos_start"], inplace=True)

    # Check the results
    _check_dfs(df, expected_end)


def test_context_manager():
    ds1_uri = os.path.join(TESTS_INPUT_DIR, "arrays/v4/ingested_2samples")
    expected_count1 = 14
    ds2_uri = os.path.join(TESTS_INPUT_DIR, "arrays/v3/synth-array")
    expected_count2 = 19565

    # Test the context manager
    with tiledbvcf.Dataset(ds1_uri) as ds:
        assert ds.count() == expected_count1

    with tiledbvcf.Dataset(ds2_uri) as ds:
        assert ds.count() == expected_count2

    # Open the datasets outside the context manager
    ds1 = tiledbvcf.Dataset(ds1_uri)
    assert ds1.count() == expected_count1

    ds2 = tiledbvcf.Dataset(ds2_uri)
    assert ds2.count() == expected_count2

    # Check that an exception is raised when trying to access a closed dataset
    ds1.close()
    with pytest.raises(Exception):
        assert ds1.count() == expected_count1

    assert ds2.count() == expected_count2

    ds2.close()
    with pytest.raises(Exception):
        assert ds2.count() == expected_count2


def test_delete_dataset(tmp_path):
    uri = os.path.join(tmp_path, "delete_dataset")

    with tiledbvcf.Dataset(uri, mode="w") as ds:
        ds.create_dataset()

    # Check that the dataset exists
    assert os.path.exists(uri)

    # Delete the dataset
    tiledbvcf.Dataset.delete(uri)

    # Check that the dataset does not exist
    assert not os.path.exists(uri)


def test_equality_old_new_format():
    old_ds = tiledbvcf.Dataset(os.path.join(TESTS_INPUT_DIR, "arrays/old_format"))
    new_ds = tiledbvcf.Dataset(os.path.join(TESTS_INPUT_DIR, "arrays/new_format"))

    assert old_ds.count() == new_ds.count()
    assert old_ds.samples() == new_ds.samples()
    assert old_ds.read().equals(new_ds.read())