from itertools import islice

from pathlib import Path

import polars

import pytest

import edsnlp

def test_read_write(blank_nlp, text, df_notes_pandas):

    import polars

    df_notes_polars = polars.from_pandas(df_notes_pandas)

    reader = edsnlp.data.from_polars(

        df_notes_polars,

        converter="omop",

        nlp=blank_nlp,

    ).set_processing(backend="simple")

    doc = list(reader)[0]

    assert doc.text == text

    blank_nlp.add_pipe("eds.matcher", config={"terms": {"douleur": ["douleurs"]}})

    blank_nlp.add_pipe("eds.negation")

    docs = reader.map_pipeline(blank_nlp)

    writer: polars.DataFrame = docs.to_polars(

        converter="omop",

        span_attributes=["negation"],

        span_getter=["ents"],

    )

    res = writer.to_dicts()

    assert len(res) == 20

    assert sum(len(r["entities"]) for r in res) == 20

@pytest.mark.parametrize("num_cpu_workers", [0, 2])

def test_read_shuffle_loop(num_cpu_workers: int):

    data = polars.read_parquet(

        Path(__file__).parent.parent.resolve() / "resources" / "docs.parquet"

    )

    notes = (

        edsnlp.data.from_polars(

            data,

            shuffle="dataset",

            seed=42,

            loop=True,

        )

        .map(lambda x: x["note_id"])

        .set_processing(num_cpu_workers=num_cpu_workers)

    )

    notes = list(islice(notes, 6))

    assert notes == [

        "subfolder/doc-1",

        "subfolder/doc-3",

        "subfolder/doc-2",

        "subfolder/doc-1",

        "subfolder/doc-2",

        "subfolder/doc-3",

    ]