import pytest

import edsnlp

from edsnlp.utils.collections import ld_to_dl

try:

    import torch.nn

except ImportError:

    torch = None

def test_map_batches():

    items = [1, 2, 3, 4, 5]

    stream = edsnlp.data.from_iterable(items)

    stream = stream.map(lambda x: x + 1)  # 2, 3, 4, 5, 6

    stream = stream.map_batches(lambda x: [sum(x)])

    stream = stream.set_processing(

        num_cpu_workers=2,

        sort_chunks=False,

        batch_size=2,

    )

    res = list(stream)

    assert res == [6, 8, 6]  # 2+4, 3+5, 6

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

def test_flat_iterable(num_cpu_workers):

    items = [1, 2, 3, 4]

    stream = edsnlp.data.from_iterable(items)

    stream = stream.set_processing(num_cpu_workers=num_cpu_workers)

    stream = stream.map(lambda x: [x] * x)

    stream = stream.flatten()

    res = list(stream.to_iterable(converter=lambda x: x))

    assert sorted(res) == [1, 2, 2, 3, 3, 3, 4, 4, 4, 4]

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

@pytest.mark.skipif(torch is None, reason="torch not installed")

def test_map_gpu(num_gpu_workers):

    import torch

    def prepare_batch(batch, device):

        return {"tensor": torch.tensor(batch).to(device)}

    def forward(batch):

        return {"outputs": batch["tensor"] * 2}

    items = range(15)

    stream = edsnlp.data.from_iterable(items)

    if num_gpu_workers == 0:

        # this is just to fuse tests, and test map_gpu

        # following a map_batches without specifying a batch size

        stream = stream.map_batches(lambda x: x)

    stream = stream.map_gpu(prepare_batch, forward)

    stream = stream.set_processing(

        num_gpu_workers=num_gpu_workers,

        gpu_worker_devices=["cpu"] * num_gpu_workers,

        sort_chunks=False,

        batch_size=2,

    )

    res = ld_to_dl(stream)

    res = torch.cat(res["outputs"])

    assert set(res.tolist()) == {i * 2 for i in range(15)}

# fmt: off

@pytest.mark.parametrize(

    "sort,num_cpu_workers,batch_kwargs,expected",

    [

        (False, 1, {"batch_size": 10, "batch_by": "words"}, [3, 1, 3, 1, 3, 1]),  # noqa: E501

        (False, 1, {"batch_size": 10, "batch_by": "padded_words"}, [2, 1, 1, 2, 1, 1, 2, 1, 1]),  # noqa: E501

        (False, 1, {"batch_size": 10, "batch_by": "docs"}, [10, 2]),  # noqa: E501

        (False, 2, {"batch_size": 10, "batch_by": "words"}, [2, 1, 2, 1, 2, 1, 1, 1, 1]),  # noqa: E501

        (False, 2, {"batch_size": 10, "batch_by": "padded_words"}, [2, 1, 2, 1, 2, 1, 1, 1, 1]),  # noqa: E501

        (False, 2, {"batch_size": 10, "batch_by": "docs"}, [6, 6]),  # noqa: E501

        (True, 2, {"batch_size": 10, "batch_by": "padded_words"}, [3, 3, 2, 1, 1, 1, 1]),  # noqa: E501

        (False, 2, {"batch_size": "10 words"}, [2, 1, 2, 1, 2, 1, 1, 1, 1]),  # noqa: E501

    ],

)

# fmt: on

def test_map_with_batching(sort, num_cpu_workers, batch_kwargs, expected):

    nlp = edsnlp.blank("eds")

    nlp.add_pipe(

        "eds.matcher",

        config={

            "terms": {

                "foo": ["This", "is", "a", "sentence", ".", "Short", "snippet", "too"],

            }

        },

        name="matcher",

    )

    samples = [

        "This is a sentence.",

        "Short snippet",

        "Short snippet too",

        "This is a very very long sentence that will make more than 10 words",

    ] * 3

    stream = edsnlp.data.from_iterable(samples)

    if sort:

        stream = stream.map_batches(lambda x: sorted(x, key=len), batch_size=1000)

    stream = stream.map_pipeline(nlp)

    stream = stream.map_batches(len)

    stream = stream.set_processing(

        num_cpu_workers=num_cpu_workers,

        **batch_kwargs,

        chunk_size=1000,  # deprecated

        split_into_batches_after="matcher",

        show_progress=True,

    )

    assert list(stream) == expected

def test_repr(frozen_ml_nlp, tmp_path):

    items = ["ceci est un test", "ceci est un autre test"]

    stream = (

        edsnlp.data.from_iterable(items, converter=frozen_ml_nlp.make_doc)

        .map(lambda x: x)

        .map_pipeline(frozen_ml_nlp, batch_size=2)

        .map_batches(lambda b: sorted(b, key=len))

        .set_processing(num_cpu_workers=2)

        .write_json(tmp_path / "out_test.jsonl", lines=True, execute=False)

    )

    assert "Stream" in repr(stream)

@pytest.mark.parametrize("shuffle_reader", [True, False])

def test_shuffle_before_generator(shuffle_reader):

    def gen_fn(x):

        yield x

        yield x

    items = [1, 2, 3, 4, 5]

    stream = edsnlp.data.from_iterable(items)

    stream = stream.map(lambda x: x)

    stream = stream.shuffle(seed=42, shuffle_reader=shuffle_reader)

    stream = stream.map(gen_fn)

    assert stream.reader.shuffle == ("dataset" if shuffle_reader else False)

    assert len(stream.ops) == (2 if shuffle_reader else 5)

    res = list(stream)

    assert res == [4, 4, 2, 2, 3, 3, 5, 5, 1, 1]

def test_shuffle_after_generator():

    def gen_fn(x):

        yield x

        yield x

    items = [1, 2, 3, 4, 5]

    stream = edsnlp.data.from_iterable(items)

    stream = stream.map(lambda x: x)

    stream = stream.map(gen_fn)

    stream = stream.shuffle(seed=43)

    assert stream.reader.shuffle == "dataset"

    assert len(stream.ops) == 5

    res = list(stream)

    assert res == [1, 2, 4, 3, 1, 3, 5, 5, 4, 2]

def test_shuffle_frozen_ml_pipeline(run_in_test_dir, frozen_ml_nlp):

    stream = edsnlp.data.read_parquet("../resources/docs.parquet", converter="omop")

    stream = stream.map_pipeline(frozen_ml_nlp, batch_size=2)

    assert len(stream.ops) == 7

    stream = stream.shuffle(batch_by="fragment")

    assert len(stream.ops) == 7

    assert stream.reader.shuffle == "fragment"

def test_unknown_shuffle():

    items = [1, 2, 3, 4, 5]

    stream = edsnlp.data.from_iterable(items)

    stream = stream.map(lambda x: x)

    with pytest.raises(ValueError):

        stream.shuffle("unknown")

def test_int_shuffle():

    items = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

    stream = edsnlp.data.from_iterable(items)

    stream = stream.map(lambda x: x)

    stream = stream.shuffle("2 docs", seed=42)

    assert list(stream) == [2, 1, 4, 3, 5, 6, 8, 7, 10, 9]

def test_parallel_preprocess_stop(run_in_test_dir, frozen_ml_nlp):

    nlp = frozen_ml_nlp

    stream = edsnlp.data.read_parquet(

        "../resources/docs.parquet",

        "omop",

        loop=True,

    )

    stream = stream.map(edsnlp.pipes.split(regex="\n+"))

    stream = stream.map(nlp.preprocess, kwargs=dict(supervision=True))

    stream = stream.batchify("128 words")

    stream = stream.map(nlp.collate)

    stream = stream.set_processing(num_cpu_workers=1, process_start_method="spawn")

    it = iter(stream)

    total = 0

    for _ in zip(it, range(10)):

        total += 1

    assert total == 10

    del it