import pytest
from tornado import gen

from dask.core import flatten
import dask
from dask import delayed, persist

from distributed.utils_test import gen_cluster, inc, slowinc, slowdec
from distributed import wait
from distributed.utils import tokey


@gen_cluster(client=True)
def test_submit(c, s, a, b):
    low = c.submit(inc, 1, priority=-1)
    futures = c.map(slowinc, range(10), delay=0.1)
    high = c.submit(inc, 2, priority=1)
    yield wait(high)
    assert all(s.processing.values())
    assert s.tasks[low.key].state == 'processing'


@gen_cluster(client=True)
def test_map(c, s, a, b):
    low = c.map(inc, [1, 2, 3], priority=-1)
    futures = c.map(slowinc, range(10), delay=0.1)
    high = c.map(inc, [4, 5, 6], priority=1)
    yield wait(high)
    assert all(s.processing.values())
    assert s.tasks[low[0].key].state == 'processing'


@gen_cluster(client=True)
def test_compute(c, s, a, b):
    da = pytest.importorskip('dask.array')
    x = da.random.random((10, 10), chunks=(5, 5))
    y = da.random.random((10, 10), chunks=(5, 5))

    low = c.compute(x, priority=-1)
    futures = c.map(slowinc, range(10), delay=0.1)
    high = c.compute(y, priority=1)
    yield wait(high)
    assert all(s.processing.values())
    assert s.tasks[tokey(low.key)].state in ('processing', 'waiting')


@gen_cluster(client=True)
def test_persist(c, s, a, b):
    da = pytest.importorskip('dask.array')
    x = da.random.random((10, 10), chunks=(5, 5))
    y = da.random.random((10, 10), chunks=(5, 5))

    low = x.persist(priority=-1)
    futures = c.map(slowinc, range(10), delay=0.1)
    high = y.persist(priority=1)
    yield wait(high)
    assert all(s.processing.values())
    assert all(s.tasks[tokey(k)].state in ('processing', 'waiting')
               for k in flatten(low.__dask_keys__()))


@gen_cluster(client=True)
def test_expand_compute(c, s, a, b):
    low = delayed(inc)(1)
    many = [delayed(slowinc)(i, delay=0.1) for i in range(10)]
    high = delayed(inc)(2)

    low, many, high = c.compute([low, many, high], priority={low: -1, high: 1})
    yield wait(high)
    assert s.tasks[low.key].state == 'processing'


@gen_cluster(client=True)
def test_expand_persist(c, s, a, b):
    low = delayed(inc)(1, dask_key_name='low')
    many = [delayed(slowinc)(i, delay=0.1) for i in range(4)]
    high = delayed(inc)(2, dask_key_name='high')

    low, high, x, y, z, w = persist(low, high, *many, priority={low: -1, high: 1})
    yield wait(high)
    assert s.tasks[low.key].state == 'processing'


@gen_cluster(client=True, ncores=[('127.0.0.1', 1)])
def test_repeated_persists_same_priority(c, s, w):
    xs = [delayed(slowinc)(i, delay=0.05, dask_key_name='x-%d' % i) for i in range(10)]
    ys = [delayed(slowinc)(x, delay=0.05, dask_key_name='y-%d' % i) for i, x in enumerate(xs)]
    zs = [delayed(slowdec)(x, delay=0.05, dask_key_name='z-%d' % i) for i, x in enumerate(xs)]

    ys = dask.persist(*ys)
    zs = dask.persist(*zs)

    while sum(t.state == 'memory' for t in s.tasks.values()) < 5:  # TODO: reduce this number
        yield gen.sleep(0.01)

    assert any(s.tasks[y.key].state == 'memory' for y in ys)
    assert any(s.tasks[z.key].state == 'memory' for z in zs)


@gen_cluster(client=True, ncores=[('127.0.0.1', 1)])
def test_last_in_first_out(c, s, w):
    xs = [c.submit(slowinc, i, delay=0.05) for i in range(5)]
    ys = [c.submit(slowinc, x, delay=0.05) for x in xs]
    zs = [c.submit(slowinc, y, delay=0.05) for y in ys]

    while len(s.tasks) < 15 or not any(s.tasks[z.key].state == 'memory' for z in zs):
        yield gen.sleep(0.01)

    assert not all(s.tasks[x.key].state == 'memory' for x in xs)