from __future__ import annotations

import numpy as np

from scripts.prepare_cam_retrieval_dataset import (
    dataset_config,
    pack_bits_to_words_le,
    stratified_indices,
    words_le_to_int,
)


def test_dataset_config_resolves_cifar10_and_cifar100() -> None:
    assert dataset_config("cifar10") == ("uoft-cs/cifar10", "label")
    assert dataset_config("cifar100") == ("uoft-cs/cifar100", "fine_label")


def test_dataset_config_rejects_unknown_dataset() -> None:
    try:
        dataset_config("mnist")
    except ValueError as exc:
        assert "dataset must be cifar10 or cifar100" in str(exc)
    else:
        raise AssertionError("expected ValueError")


def test_stratified_indices_balances_labels_and_is_deterministic() -> None:
    labels = [0, 0, 0, 1, 1, 1, 2, 2, 2]

    first = stratified_indices(labels, total=6, seed=123)
    second = stratified_indices(labels, total=6, seed=123)

    assert first == second
    selected_labels = [labels[i] for i in first]
    assert selected_labels.count(0) == 2
    assert selected_labels.count(1) == 2
    assert selected_labels.count(2) == 2


def test_stratified_indices_fills_remainder() -> None:
    labels = [0, 0, 0, 1, 1, 1]

    indices = stratified_indices(labels, total=5, seed=7)

    assert len(indices) == 5
    assert len(set(indices)) == 5


def test_pack_bits_to_words_le_roundtrip() -> None:
    bits = np.zeros((2, 128), dtype=np.uint8)
    bits[0, 0] = 1
    bits[0, 65] = 1
    bits[1, 63] = 1
    bits[1, 127] = 1

    words = pack_bits_to_words_le(bits, hash_bits=128)

    assert words.dtype == np.uint64
    assert words.shape == (2, 2)
    assert words_le_to_int(words[0]) == (1 << 0) | (1 << 65)
    assert words_le_to_int(words[1]) == (1 << 63) | (1 << 127)