from __future__ import annotations

from model.ref_model import (
    generate_grouped_flip_mask,
    match_top1_with_read_noise,
    xnor_popcount_score,
)


def test_grouped_flip_mask_full_rate_one_bit_per_64_bit_group():
    random_value = 0
    for group in range(8):
        bit_idx = group + 1
        sample = 0
        random_value |= bit_idx << (group * 14)
        random_value |= sample << (group * 14 + 6)

    mask = generate_grouped_flip_mask(
        random_value=random_value,
        hash_bits=512,
        noise_bits=8,
        rate_num=1,
        rate_den=1,
    )

    expected = 0
    for group in range(8):
        expected |= 1 << (group * 64 + group + 1)

    assert mask == expected
    assert mask.bit_count() == 8


def test_grouped_flip_mask_zero_rate_no_flips():
    mask = generate_grouped_flip_mask(
        random_value=(1 << 128) - 1,
        hash_bits=512,
        noise_bits=8,
        rate_num=0,
        rate_den=100,
    )
    assert mask == 0


def test_score_is_bit_match_popcount_not_hamming_distance():
    query = 0b1010
    stored = 0b1000
    assert xnor_popcount_score(query, stored, width=4) == 3


def test_read_noise_model_is_reproducible_after_reset_seed():
    rows = [0, (1 << 512) - 1, 0x1234, 0x5678, 0x9ABC, 0xDEF0, 0x1357, 0x2468]
    query = rows[2]
    kwargs = dict(
        query=query,
        rows=rows,
        width=512,
        lanes=8,
        noise_bits=8,
        rate_num=1,
        rate_den=100,
        seed=0x6A09_E667_F3BC_C909,
    )

    first = match_top1_with_read_noise(**kwargs)
    second = match_top1_with_read_noise(**kwargs)

    assert first.top1_index == second.top1_index
    assert first.top1_score == second.top1_score
    assert first.scores.tolist() == second.scores.tolist()