from __future__ import print_function, division, absolute_import


import pytest
pytest.importorskip('numpy')
pytest.importorskip('pandas')

import dask
import dask.dataframe as dd
import dask.bag as db
from distributed import Client
from distributed.client import wait
from distributed.utils_test import cluster, gen_cluster
from distributed.utils_test import loop # flake8: noqa
import numpy as np
import pandas as pd
import pandas.util.testing as tm


dfs = [pd.DataFrame({'x': [1, 2, 3]}, index=[0, 10, 20]),
       pd.DataFrame({'x': [4, 5, 6]}, index=[30, 40, 50]),
       pd.DataFrame({'x': [7, 8, 9]}, index=[60, 70, 80])]


def assert_equal(a, b):
    assert type(a) == type(b)
    if isinstance(a, pd.DataFrame):
        tm.assert_frame_equal(a, b)
    elif isinstance(a, pd.Series):
        tm.assert_series_equal(a, b)
    elif isinstance(a, pd.Index):
        tm.assert_index_equal(a, b)
    else:
        assert a == b


@gen_cluster(timeout=240, client=True)
def test_dataframes(c, s, a, b):
    df = pd.DataFrame({'x': np.random.random(1000),
                       'y': np.random.random(1000)},
                      index=np.arange(1000))
    ldf = dd.from_pandas(df, npartitions=10)

    rdf = c.persist(ldf)

    assert rdf.divisions == ldf.divisions

    remote = c.compute(rdf)
    result = yield remote

    tm.assert_frame_equal(result,
                          ldf.compute(get=dask.get))

    exprs = [lambda df: df.x.mean(),
             lambda df: df.y.std(),
             lambda df: df.assign(z=df.x + df.y).drop_duplicates(),
             lambda df: df.index,
             lambda df: df.x,
             lambda df: df.x.cumsum(),
             lambda df: df.groupby(['x', 'y']).count(),
             lambda df: df.loc[50:75]]
    for f in exprs:
        local = f(ldf).compute(get=dask.get)
        remote = c.compute(f(rdf))
        remote = yield remote
        assert_equal(local, remote)


@gen_cluster(client=True)
def test__dask_array_collections(c, s, a, b):
    import dask.array as da

    x_dsk = {('x', i, j): np.random.random((3, 3)) for i in range(3)
             for j in range(2)}
    y_dsk = {('y', i, j): np.random.random((3, 3)) for i in range(2)
             for j in range(3)}
    x_futures = yield c._scatter(x_dsk)
    y_futures = yield c._scatter(y_dsk)

    dt = np.random.random(0).dtype
    x_local = da.Array(x_dsk, 'x', ((3, 3, 3), (3, 3)), dt)
    y_local = da.Array(y_dsk, 'y', ((3, 3), (3, 3, 3)), dt)

    x_remote = da.Array(x_futures, 'x', ((3, 3, 3), (3, 3)), dt)
    y_remote = da.Array(y_futures, 'y', ((3, 3), (3, 3, 3)), dt)

    exprs = [lambda x, y: x.T + y,
             lambda x, y: x.mean() + y.mean(),
             lambda x, y: x.dot(y).std(axis=0),
             lambda x, y: x - x.mean(axis=1)[:, None]]

    for expr in exprs:
        local = expr(x_local, y_local).compute(get=dask.get)

        remote = c.compute(expr(x_remote, y_remote))
        remote = yield remote

        assert np.all(local == remote)


@gen_cluster(client=True)
def test_bag_groupby_tasks_default(c, s, a, b):
    with dask.set_options(get=c.get):
        b = db.range(100, npartitions=10)
        b2 = b.groupby(lambda x: x % 13)
        assert not any('partd' in k[0] for k in b2.dask)


@pytest.mark.parametrize('wait', [wait, lambda x: None])
def test_dataframe_set_index_sync(loop, wait):
    with cluster() as (s, [a, b]):
        with Client(s['address'], loop=loop) as c:
            with dask.set_options(get=c.get):
                df = dd.demo.make_timeseries('2000', '2001',
                                             {'value': float, 'name': str, 'id': int},
                                             freq='2H', partition_freq='1M', seed=1)
                df = c.persist(df)
                wait(df)

                df2 = df.set_index('name', shuffle='tasks')
                df2 = c.persist(df2)

                assert len(df2)


def test_loc_sync(loop):
    with cluster() as (s, [a, b]):
        with Client(s['address'], loop=loop) as c:
            df = pd.util.testing.makeTimeDataFrame()
            ddf = dd.from_pandas(df, npartitions=10)
            ddf.loc['2000-01-17':'2000-01-24'].compute(get=c.get)


def test_rolling_sync(loop):
    with cluster() as (s, [a, b]):
        with Client(s['address'], loop=loop) as c:
            df = pd.util.testing.makeTimeDataFrame()
            ddf = dd.from_pandas(df, npartitions=10)
            dd.rolling_mean(ddf.A, 2).compute(get=c.get)


@gen_cluster(client=True)
def test_loc(c, s, a, b):
    df = pd.util.testing.makeTimeDataFrame()
    ddf = dd.from_pandas(df, npartitions=10)
    future = c.compute(ddf.loc['2000-01-17':'2000-01-24'])
    yield future


def test_dataframe_groupby_tasks(loop):
    df = pd.util.testing.makeTimeDataFrame()
    df['A'] = df.A // 0.1
    df['B'] = df.B // 0.1
    ddf = dd.from_pandas(df, npartitions=10)
    with cluster() as (s, [a, b]):
        with Client(s['address'], loop=loop) as c:
            with dask.set_options(get=c.get):
                for ind in [lambda x: 'A', lambda x: x.A]:
                    a = df.groupby(ind(df)).apply(len)
                    b = ddf.groupby(ind(ddf)).apply(len, meta=int)
                    assert_equal(a, b.compute(get=dask.get).sort_index())
                    assert not any('partd' in k[0] for k in b.dask)

                    a = df.groupby(ind(df)).B.apply(len)
                    b = ddf.groupby(ind(ddf)).B.apply(len, meta=('B', int))
                    assert_equal(a, b.compute(get=dask.get).sort_index())
                    assert not any('partd' in k[0] for k in b.dask)

                with pytest.raises(NotImplementedError):
                    ddf.groupby(ddf[['A', 'B']]).apply(len, meta=int)

                a = df.groupby(['A', 'B']).apply(len)
                b = ddf.groupby(['A', 'B']).apply(len, meta=int)

                assert_equal(a, b.compute(get=dask.get).sort_index())


@gen_cluster(client=True)
def test_sparse_arrays(c, s, a, b):
    sparse = pytest.importorskip('sparse')
    da = pytest.importorskip('dask.array')

    x = da.random.random((100, 10), chunks=(10, 10))
    x[x < 0.95] = 0
    s = x.map_blocks(sparse.COO)
    future = c.compute(s.sum(axis=0)[:10])

    yield future
