from functools import partial
import os

import pytest

import dask
import dask.array as da
from dask.utils_test import inc
from dask.highlevelgraph import HighLevelGraph, BasicLayer, Layer
from dask.blockwise import Blockwise
from dask.array.utils import assert_eq


def test_visualize(tmpdir):
    pytest.importorskip("graphviz")
    fn = str(tmpdir)
    a = da.ones(10, chunks=(5,))
    b = a + 1
    c = a + 2
    d = b + c
    d.dask.visualize(fn)
    assert os.path.exists(fn)


def test_basic():
    a = {"x": 1}
    b = {"y": (inc, "x")}
    layers = {"a": a, "b": b}
    dependencies = {"a": set(), "b": {"a"}}
    hg = HighLevelGraph(layers, dependencies)

    assert dict(hg) == {"x": 1, "y": (inc, "x")}
    assert all(isinstance(layer, Layer) for layer in hg.layers.values())


def test_keys_values_items_methods():
    a = da.ones(10, chunks=(5,))
    b = a + 1
    c = a + 2
    d = b + c
    hg = d.dask

    keys, values, items = hg.keys(), hg.values(), hg.items()
    assert all(isinstance(i, list) for i in [keys, values, items])
    assert keys == [i for i in hg]
    assert values == [hg[i] for i in hg]
    assert items == [(k, v) for k, v in zip(keys, values)]


def test_cull():
    a = {"x": 1, "y": (inc, "x")}
    layers = {
        "a": BasicLayer(
            a, dependencies={"x": set(), "y": {"x"}}, global_dependencies=set()
        )
    }
    dependencies = {"a": set()}
    hg = HighLevelGraph(layers, dependencies)

    culled_by_x = hg.cull({"x"})
    assert dict(culled_by_x) == {"x": 1}

    culled_by_y = hg.cull({"y"})
    assert dict(culled_by_y) == a


@pytest.mark.parametrize("inject_dict", [True, False])
def test_map_basic_layers(inject_dict):
    """Check map_basic_layers() by injecting an inc() call"""

    y = da.ones(3, chunks=(3,), dtype="int") + 40

    def inject_inc(dsk):
        assert isinstance(dsk, BasicLayer)
        dsk = dict(dsk)
        k = next(iter(dsk))
        dsk[k] = (inc, dsk[k])
        if inject_dict:
            return dsk  # map_basic_layers() should automatically convert it to a `BasicLayer`
        else:
            return BasicLayer(dsk)

    dsk = y.__dask_graph__()
    y.dask = dsk.map_basic_layers(inject_inc)
    layers = list(y.dask.layers.values())
    assert isinstance(layers[0], BasicLayer)
    assert isinstance(layers[1], Blockwise)
    assert_eq(y, [42] * 3)


@pytest.mark.parametrize("use_layer_map_task", [True, False])
def test_map_tasks(use_layer_map_task):
    """Check map_tasks() by injecting an +1 to the `40` literal"""
    y = da.ones(3, chunks=(3,), dtype="int") + 40

    def plus_one(tasks):
        ret = []
        for t in tasks:
            if t == 40:
                t += 1
            ret.append(t)
        return tuple(ret)

    dsk = y.__dask_graph__()

    if use_layer_map_task:
        # In order to test the default map_tasks() implementation on a Blockwise Layer,
        # we overwrite Blockwise.map_tasks with Layer.map_tasks
        blockwise_layer = list(dsk.layers.values())[1]
        blockwise_layer.map_tasks = partial(Layer.map_tasks, blockwise_layer)

    y.dask = dsk.map_tasks(plus_one)
    assert_eq(y, [42] * 3)


def annot_map_fn(key):
    return key[1:]


@pytest.mark.parametrize(
    "annotation",
    [
        {"worker": "alice"},
        {"block_id": annot_map_fn},
    ],
)
def test_single_annotation(annotation):
    with dask.annotate(**annotation):
        A = da.ones((10, 10), chunks=(5, 5))

    alayer = A.__dask_graph__().layers[A.name]
    assert alayer.annotations == annotation
    assert dask.config.get("annotations", None) is None


def test_multiple_annotations():
    with dask.annotate(block_id=annot_map_fn):
        with dask.annotate(resource="GPU"):
            A = da.ones((10, 10), chunks=(5, 5))

        B = A + 1

    C = B + 1

    assert dask.config.get("annotations", None) is None

    alayer = A.__dask_graph__().layers[A.name]
    blayer = B.__dask_graph__().layers[B.name]
    clayer = C.__dask_graph__().layers[C.name]
    assert alayer.annotations == {"resource": "GPU", "block_id": annot_map_fn}
    assert blayer.annotations == {"block_id": annot_map_fn}
    assert clayer.annotations is None

[((1597, 1650), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""inject_dict"""', '[True, False]'], {}), "('inject_dict', [True, False])\n", (1620, 1650), False, 'import pytest\n'), ((2361, 2421), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""use_layer_map_task"""', '[True, False]'], {}), "('use_layer_map_task', [True, False])\n", (2384, 2421), False, 'import pytest\n'), ((3200, 3292), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""annotation"""', "[{'worker': 'alice'}, {'block_id': annot_map_fn}]"], {}), "('annotation', [{'worker': 'alice'}, {'block_id':\n    annot_map_fn}])\n", (3223, 3292), False, 'import pytest\n'), ((300, 331), 'pytest.importorskip', 'pytest.importorskip', (['"""graphviz"""'], {}), "('graphviz')\n", (319, 331), False, 'import pytest\n'), ((361, 385), 'dask.array.ones', 'da.ones', (['(10)'], {'chunks': '(5,)'}), '(10, chunks=(5,))\n', (368, 385), True, 'import dask.array as da\n'), ((464, 482), 'os.path.exists', 'os.path.exists', (['fn'], {}), '(fn)\n', (478, 482), False, 'import os\n'), ((629, 665), 'dask.highlevelgraph.HighLevelGraph', 'HighLevelGraph', (['layers', 'dependencies'], {}), '(layers, dependencies)\n', (643, 665), False, 'from dask.highlevelgraph import HighLevelGraph, BasicLayer, Layer\n'), ((837, 861), 'dask.array.ones', 'da.ones', (['(10)'], {'chunks': '(5,)'}), '(10, chunks=(5,))\n', (844, 861), True, 'import dask.array as da\n'), ((1414, 1450), 'dask.highlevelgraph.HighLevelGraph', 'HighLevelGraph', (['layers', 'dependencies'], {}), '(layers, dependencies)\n', (1428, 1450), False, 'from dask.highlevelgraph import HighLevelGraph, BasicLayer, Layer\n'), ((2335, 2357), 'dask.array.utils.assert_eq', 'assert_eq', (['y', '([42] * 3)'], {}), '(y, [42] * 3)\n', (2344, 2357), False, 'from dask.array.utils import assert_eq\n'), ((3130, 3152), 'dask.array.utils.assert_eq', 'assert_eq', (['y', '([42] * 3)'], {}), '(y, [42] * 3)\n', (3139, 3152), False, 'from dask.array.utils import assert_eq\n'), ((1762, 1798), 'dask.array.ones', 'da.ones', (['(3)'], {'chunks': '(3,)', 'dtype': '"""int"""'}), "(3, chunks=(3,), dtype='int')\n", (1769, 1798), True, 'import dask.array as da\n'), ((2537, 2573), 'dask.array.ones', 'da.ones', (['(3)'], {'chunks': '(3,)', 'dtype': '"""int"""'}), "(3, chunks=(3,), dtype='int')\n", (2544, 2573), True, 'import dask.array as da\n'), ((3046, 3087), 'functools.partial', 'partial', (['Layer.map_tasks', 'blockwise_layer'], {}), '(Layer.map_tasks, blockwise_layer)\n', (3053, 3087), False, 'from functools import partial\n'), ((3372, 3399), 'dask.annotate', 'dask.annotate', ([], {}), '(**annotation)\n', (3385, 3399), False, 'import dask\n'), ((3413, 3445), 'dask.array.ones', 'da.ones', (['(10, 10)'], {'chunks': '(5, 5)'}), '((10, 10), chunks=(5, 5))\n', (3420, 3445), True, 'import dask.array as da\n'), ((3549, 3585), 'dask.config.get', 'dask.config.get', (['"""annotations"""', 'None'], {}), "('annotations', None)\n", (3564, 3585), False, 'import dask\n'), ((3638, 3674), 'dask.annotate', 'dask.annotate', ([], {'block_id': 'annot_map_fn'}), '(block_id=annot_map_fn)\n', (3651, 3674), False, 'import dask\n'), ((3815, 3851), 'dask.config.get', 'dask.config.get', (['"""annotations"""', 'None'], {}), "('annotations', None)\n", (3830, 3851), False, 'import dask\n'), ((2108, 2123), 'dask.highlevelgraph.BasicLayer', 'BasicLayer', (['dsk'], {}), '(dsk)\n', (2118, 2123), False, 'from dask.highlevelgraph import HighLevelGraph, BasicLayer, Layer\n'), ((3689, 3718), 'dask.annotate', 'dask.annotate', ([], {'resource': '"""GPU"""'}), "(resource='GPU')\n", (3702, 3718), False, 'import dask\n'), ((3736, 3768), 'dask.array.ones', 'da.ones', (['(10, 10)'], {'chunks': '(5, 5)'}), '((10, 10), chunks=(5, 5))\n', (3743, 3768), True, 'import dask.array as da\n')]