from itertools import islice
from pathlib import Path

import pytest
from typing_extensions import Literal

import edsnlp


@pytest.mark.parametrize("num_cpu_workers", [0, 2])
@pytest.mark.parametrize("shuffle", ["dataset"])
def test_read_shuffle_loop(
    num_cpu_workers: int,
    shuffle: Literal["dataset", "fragment"],
):
    input_file = (
        Path(__file__).parent.parent.resolve() / "training" / "rhapsodie_sample.conllu"
    )
    notes = edsnlp.data.read_conll(
        input_file,
        shuffle=shuffle,
        seed=42,
        loop=True,
    ).set_processing(num_cpu_workers=num_cpu_workers)
    notes = list(islice(notes, 6))
    assert len(notes) == 6
    # 32	ce	ce	PRON	_	Gender=Masc|Number=Sing|Person=3|PronType=Dem	30	obl:arg	_	_  # noqa: E501
    word_attrs = {
        "text": "ce",
        "lemma_": "ce",
        "pos_": "PRON",
        "dep_": "obl:arg",
        "morph": "Gender=Masc|Number=Sing|Person=3|PronType=Dem",
        "head": "profité",
    }
    word = notes[0][31]
    for attr, val in word_attrs.items():
        assert str(getattr(word, attr)) == val