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

[((67, 1, 67, 54), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(67, 25, 67, 38): '"""inject_dict"""', (67, 40, 67, 53): '[True, False]'}, {}), "('inject_dict', [True, False])", False, 'import pytest\n'), ((91, 1, 91, 61), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(91, 25, 91, 45): '"""use_layer_map_task"""', (91, 47, 91, 60): '[True, False]'}, {}), "('use_layer_map_task', [True, False])", False, 'import pytest\n'), ((120, 1, 126, 1), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(121, 4, 121, 16): '"""annotation"""', (122, 4, 125, 5): "[{'worker': 'alice'}, {'block_id': annot_map_fn}]"}, {}), "('annotation', [{'worker': 'alice'}, {'block_id':\n    annot_map_fn}])", False, 'import pytest\n'), ((15, 4, 15, 35), 'pytest.importorskip', 'pytest.importorskip', ({(15, 24, 15, 34): '"""graphviz"""'}, {}), "('graphviz')", False, 'import pytest\n'), ((17, 8, 17, 32), 'dask.array.ones', 'da.ones', (), '', True, 'import dask.array as da\n'), ((22, 11, 22, 29), 'os.path.exists', 'os.path.exists', ({(22, 26, 22, 28): 'fn'}, {}), '(fn)', False, 'import os\n'), ((30, 9, 30, 45), 'dask.highlevelgraph.HighLevelGraph', 'HighLevelGraph', ({(30, 24, 30, 30): 'layers', (30, 32, 30, 44): 'dependencies'}, {}), '(layers, dependencies)', False, 'from dask.highlevelgraph import HighLevelGraph, BasicLayer, Layer\n'), ((37, 8, 37, 32), 'dask.array.ones', 'da.ones', (), '', True, 'import dask.array as da\n'), ((58, 9, 58, 45), 'dask.highlevelgraph.HighLevelGraph', 'HighLevelGraph', ({(58, 24, 58, 30): 'layers', (58, 32, 58, 44): 'dependencies'}, {}), '(layers, dependencies)', False, 'from dask.highlevelgraph import HighLevelGraph, BasicLayer, Layer\n'), ((88, 4, 88, 26), 'dask.array.utils.assert_eq', 'assert_eq', ({(88, 14, 88, 15): 'y', (88, 17, 88, 25): '([42] * 3)'}, {}), '(y, [42] * 3)', False, 'from dask.array.utils import assert_eq\n'), ((113, 4, 113, 26), 'dask.array.utils.assert_eq', 'assert_eq', ({(113, 14, 113, 15): 'y', (113, 17, 113, 25): '([42] * 3)'}, {}), '(y, [42] * 3)', False, 'from dask.array.utils import assert_eq\n'), ((71, 8, 71, 44), 'dask.array.ones', 'da.ones', (), '', True, 'import dask.array as da\n'), ((94, 8, 94, 44), 'dask.array.ones', 'da.ones', (), '', True, 'import dask.array as da\n'), ((110, 36, 110, 77), 'functools.partial', 'partial', ({(110, 44, 110, 59): 'Layer.map_tasks', (110, 61, 110, 76): 'blockwise_layer'}, {}), '(Layer.map_tasks, blockwise_layer)', False, 'from functools import partial\n'), ((128, 9, 128, 36), 'dask.annotate', 'dask.annotate', ({}, {}), '(**annotation)', False, 'import dask\n'), ((129, 12, 129, 44), 'dask.array.ones', 'da.ones', (), '', True, 'import dask.array as da\n'), ((133, 11, 133, 47), 'dask.config.get', 'dask.config.get', ({(133, 27, 133, 40): '"""annotations"""', (133, 42, 133, 46): 'None'}, {}), "('annotations', None)", False, 'import dask\n'), ((137, 9, 137, 45), 'dask.annotate', 'dask.annotate', (), '', False, 'import dask\n'), ((145, 11, 145, 47), 'dask.config.get', 'dask.config.get', ({(145, 27, 145, 40): '"""annotations"""', (145, 42, 145, 46): 'None'}, {}), "('annotations', None)", False, 'import dask\n'), ((81, 19, 81, 34), 'dask.highlevelgraph.BasicLayer', 'BasicLayer', ({(81, 30, 81, 33): 'dsk'}, {}), '(dsk)', False, 'from dask.highlevelgraph import HighLevelGraph, BasicLayer, Layer\n'), ((138, 13, 138, 42), 'dask.annotate', 'dask.annotate', (), '', False, 'import dask\n'), ((139, 16, 139, 48), 'dask.array.ones', 'da.ones', (), '', True, 'import dask.array as da\n')]