from edsnlp.utils.collections import (

    batch_compress_dict,

    batchify,

    decompress_dict,

    dl_to_ld,

    flatten,

    flatten_once,

    get_deep_attr,

    ld_to_dl,

    multi_tee,

    set_deep_attr,

)

def test_multi_tee():

    gen = (i for i in range(10))

    tee = multi_tee(gen)

    items1 = [value for i, value in zip(tee, range(5))]

    items2 = [value for i, value in zip(tee, range(5))]

    assert items1 == items2 == [0, 1, 2, 3, 4]

    # not the behavior I'd like (continue from 5 would be nice) but at least

    # the generator is not exhausted

    assert next(gen) == 6

    assert multi_tee(items1) is items1

def test_flatten():

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

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

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

def test_dict_of_lists():

    items = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]

    assert ld_to_dl(items) == {"a": [1, 3], "b": [2, 4]}

    assert list(dl_to_ld({"a": [1, 3], "b": [2, 4]})) == items

def test_dict_compression():

    list1 = [1, 2, 3]

    mapping1 = {"a": list1, "b": {"c": list1, "d": 4}}

    list2 = [1, 2, 3]

    mapping2 = {"a": list2, "b": {"c": list2, "d": 4}}

    samples = [mapping1, mapping2]

    assert list(batch_compress_dict(samples)) == [

        {"a|b/c": [1, 2, 3], "b/d": 4},

        {"a|b/c": [1, 2, 3], "b/d": 4},

    ]

    res = decompress_dict(

        [

            {"a|b/c": [1, 2, 3], "b/d": 4},

            {"a|b/c": [1, 2, 3], "b/d": 4},

        ]

    )

    assert res == {

        "a": [[1, 2, 3], [1, 2, 3]],

        "b": {"c": [[1, 2, 3], [1, 2, 3]], "d": [4, 4]},

    }

    assert res["a"] is res["b"]["c"]

def test_batchify():

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

    batches = list(batchify(items, 3))

    assert len(batches) == 4

    batches = list(batches)

    assert batches == [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]

def test_deep_path():

    class custom:

        def __init__(self, values):

            self.values = values

    item = {"a": {"b": (0, 1), "other": custom((1, 2))}}

    assert get_deep_attr(item, "a.b.0") == 0

    assert get_deep_attr(item, "a.other.values.0") == 1

    set_deep_attr(item, "a.b.1", 2)

    set_deep_attr(item, "a.other.values.0", 1000)

    assert item["a"]["b"] == (0, 2)

    assert item["a"]["other"].values == (1000, 2)