from __future__ import division
from dask.context import config
import dask.dataframe as dd
import numpy as np
import pandas as pd
import xarray as xr

import datashader as ds
import datashader.utils as du

import pytest

config.set(scheduler='synchronous')

df = pd.DataFrame({'x': np.array(([0.] * 10 + [1] * 10)),
                   'y': np.array(([0.] * 5 + [1] * 5 + [0] * 5 + [1] * 5)),
                   'log_x': np.array(([1.] * 10 + [10] * 10)),
                   'log_y': np.array(([1.] * 5 + [10] * 5 + [1] * 5 + [10] * 5)),
                   'i32': np.arange(20, dtype='i4'),
                   'i64': np.arange(20, dtype='i8'),
                   'f32': np.arange(20, dtype='f4'),
                   'f64': np.arange(20, dtype='f8'),
                   'empty_bin': np.array([0.] * 15 + [np.nan] * 5),
                   'cat': ['a']*5 + ['b']*5 + ['c']*5 + ['d']*5})
df.cat = df.cat.astype('category')
df.f32[2] = np.nan
df.f64[2] = np.nan

ddf = dd.from_pandas(df, npartitions=3)

c = ds.Canvas(plot_width=2, plot_height=2, x_range=(0, 1), y_range=(0, 1))
c_logx = ds.Canvas(plot_width=2, plot_height=2, x_range=(1, 10),
                   y_range=(0, 1), x_axis_type='log')
c_logy = ds.Canvas(plot_width=2, plot_height=2, x_range=(0, 1),
                   y_range=(1, 10), y_axis_type='log')
c_logxy = ds.Canvas(plot_width=2, plot_height=2, x_range=(1, 10),
                    y_range=(1, 10), x_axis_type='log', y_axis_type='log')

axis = ds.core.LinearAxis()
lincoords = axis.compute_index(axis.compute_scale_and_translate((0, 1), 2), 2)
coords = [lincoords, lincoords]
dims = ['y', 'x']


def assert_eq(agg, b):
    assert agg.equals(b)


def floats(n):
    """Returns contiguous list of floats from initial point"""
    while True:
        yield n
        n = n + np.spacing(n)


def test_count():
    out = xr.DataArray(np.array([[5, 5], [5, 5]], dtype='i4'),
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.count('i32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.count('i64')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.count()), out)
    out = xr.DataArray(np.array([[4, 5], [5, 5]], dtype='i4'),
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.count('f32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.count('f64')), out)


def test_any():
    out = xr.DataArray(np.array([[True, True], [True, True]]),
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.any('i64')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.any('f64')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.any()), out)
    out = xr.DataArray(np.array([[True, True], [True, False]]),
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.any('empty_bin')), out)


def test_sum():
    out = xr.DataArray(df.i32.values.reshape((2, 2, 5)).sum(axis=2, dtype='f8').T,
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.sum('i32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.sum('i64')), out)
    out = xr.DataArray(np.nansum(df.f64.values.reshape((2, 2, 5)), axis=2).T,
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.sum('f32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.sum('f64')), out)


def test_min():
    out = xr.DataArray(df.i64.values.reshape((2, 2, 5)).min(axis=2).astype('f8').T,
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.min('i32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.min('i64')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.min('f32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.min('f64')), out)


def test_max():
    out = xr.DataArray(df.i64.values.reshape((2, 2, 5)).max(axis=2).astype('f8').T,
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.max('i32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.max('i64')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.max('f32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.max('f64')), out)


def test_mean():
    out = xr.DataArray(df.i32.values.reshape((2, 2, 5)).mean(axis=2, dtype='f8').T,
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.mean('i32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.mean('i64')), out)
    out = xr.DataArray(np.nanmean(df.f64.values.reshape((2, 2, 5)), axis=2).T,
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.mean('f32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.mean('f64')), out)


def test_var():
    out = xr.DataArray(df.i32.values.reshape((2, 2, 5)).var(axis=2, dtype='f8').T,
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.var('i32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.var('i64')), out)
    out = xr.DataArray(np.nanvar(df.f64.values.reshape((2, 2, 5)), axis=2).T,
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.var('f32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.var('f64')), out)


def test_std():
    out = xr.DataArray(df.i32.values.reshape((2, 2, 5)).std(axis=2, dtype='f8').T,
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.std('i32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.std('i64')), out)
    out = xr.DataArray(np.nanstd(df.f64.values.reshape((2, 2, 5)), axis=2).T,
                       coords=coords, dims=dims)
    assert_eq(c.points(ddf, 'x', 'y', ds.std('f32')), out)
    assert_eq(c.points(ddf, 'x', 'y', ds.std('f64')), out)


def test_count_cat():
    sol = np.array([[[5, 0, 0, 0],
                     [0, 0, 5, 0]],
                    [[0, 5, 0, 0],
                     [0, 0, 0, 5]]])
    out = xr.DataArray(sol, coords=(coords + [['a', 'b', 'c', 'd']]),
                       dims=(dims + ['cat']))
    agg = c.points(ddf, 'x', 'y', ds.count_cat('cat'))
    assert_eq(agg, out)


def test_multiple_aggregates():
    agg = c.points(ddf, 'x', 'y',
                   ds.summary(f64_std=ds.std('f64'),
                              f64_mean=ds.mean('f64'),
                              i32_sum=ds.sum('i32'),
                              i32_count=ds.count('i32')))

    f = lambda x: xr.DataArray(x, coords=coords, dims=dims)
    assert_eq(agg.f64_std, f(np.nanstd(df.f64.values.reshape((2, 2, 5)), axis=2).T))
    assert_eq(agg.f64_mean, f(np.nanmean(df.f64.values.reshape((2, 2, 5)), axis=2).T))
    assert_eq(agg.i32_sum, f(df.i32.values.reshape((2, 2, 5)).sum(axis=2, dtype='f8').T))
    assert_eq(agg.i32_count, f(np.array([[5, 5], [5, 5]], dtype='i4')))


def test_auto_range_points():
    n = 10
    data = np.arange(n, dtype='i4')
    df = pd.DataFrame({'time': np.arange(n),
                       'x': data,
                       'y': data})
    ddf = dd.from_pandas(df, npartitions=2)

    cvs = ds.Canvas(plot_width=n, plot_height=n)
    agg = cvs.points(ddf, 'x', 'y', ds.count('time'))
    sol = np.zeros((n, n), int)
    np.fill_diagonal(sol, 1)
    np.testing.assert_equal(agg.data, sol)

    cvs = ds.Canvas(plot_width=n+1, plot_height=n+1)
    agg = cvs.points(ddf, 'x', 'y', ds.count('time'))
    sol = np.zeros((n+1, n+1), int)
    np.fill_diagonal(sol, 1)
    sol[5, 5] = 0
    np.testing.assert_equal(agg.data, sol)

    n = 4
    data = np.arange(n, dtype='i4')
    df = pd.DataFrame({'time': np.arange(n),
                       'x': data,
                       'y': data})
    ddf = dd.from_pandas(df, npartitions=2)

    cvs = ds.Canvas(plot_width=2*n, plot_height=2*n)
    agg = cvs.points(ddf, 'x', 'y', ds.count('time'))
    sol = np.zeros((2*n, 2*n), int)
    np.fill_diagonal(sol, 1)
    sol[[range(1, 4, 2)]] = 0
    sol[[range(4, 8, 2)]] = 0
    np.testing.assert_equal(agg.data, sol)

    cvs = ds.Canvas(plot_width=2*n+1, plot_height=2*n+1)
    agg = cvs.points(ddf, 'x', 'y', ds.count('time'))
    sol = np.zeros((2*n+1, 2*n+1), int)
    sol[0, 0] = 1
    sol[3, 3] = 1
    sol[6, 6] = 1
    sol[8, 8] = 1
    np.testing.assert_equal(agg.data, sol)


def test_uniform_points():
    n = 101
    df = pd.DataFrame({'time': np.ones(2*n, dtype='i4'),
                       'x': np.concatenate((np.arange(n, dtype='f8'),
                                            np.arange(n, dtype='f8'))),
                       'y': np.concatenate(([0.] * n, [1.] * n))})

    cvs = ds.Canvas(plot_width=10, plot_height=2, y_range=(0, 1))
    agg = cvs.points(df, 'x', 'y', ds.count('time'))
    sol = np.array([[10] * 9 + [11], [10] * 9 + [11]], dtype='i4')
    np.testing.assert_equal(agg.data, sol)


@pytest.mark.parametrize('high', [9, 10, 99, 100])
@pytest.mark.parametrize('low', [0])
def test_uniform_diagonal_points(low, high):
    bounds = (low, high)
    x_range, y_range = bounds, bounds

    width = x_range[1] - x_range[0]
    height = y_range[1] - y_range[0]
    n = width * height
    df = pd.DataFrame({'time': np.ones(n, dtype='i4'),
                       'x': np.array([np.arange(*x_range, dtype='f8')] * width).flatten(),
                       'y': np.array([np.arange(*y_range, dtype='f8')] * height).flatten()})

    cvs = ds.Canvas(plot_width=2, plot_height=2, x_range=x_range, y_range=y_range)
    agg = cvs.points(df, 'x', 'y', ds.count('time'))

    diagonal = agg.data.diagonal(0)
    assert sum(diagonal) == n
    assert abs(bounds[1] - bounds[0]) % 2 == abs(diagonal[1] / high - diagonal[0] / high)


def test_log_axis_points():
    axis = ds.core.LogAxis()
    logcoords = axis.compute_index(axis.compute_scale_and_translate((1, 10), 2), 2)

    axis = ds.core.LinearAxis()
    lincoords = axis.compute_index(axis.compute_scale_and_translate((0, 1), 2), 2)

    sol = np.array([[5, 5], [5, 5]], dtype='i4')
    out = xr.DataArray(sol, coords=[lincoords, logcoords],
                       dims=['y', 'log_x'])
    assert_eq(c_logx.points(ddf, 'log_x', 'y', ds.count('i32')), out)
    out = xr.DataArray(sol, coords=[logcoords, lincoords],
                       dims=['log_y', 'x'])
    assert_eq(c_logy.points(ddf, 'x', 'log_y', ds.count('i32')), out)
    out = xr.DataArray(sol, coords=[logcoords, logcoords],
                       dims=['log_y', 'log_x'])
    assert_eq(c_logxy.points(ddf, 'log_x', 'log_y', ds.count('i32')), out)


def test_line():
    axis = ds.core.LinearAxis()
    lincoords = axis.compute_index(axis.compute_scale_and_translate((-3., 3.), 7), 7)

    df = pd.DataFrame({'x': [4, 0, -4, -3, -2, -1.9, 0, 10, 10, 0, 4],
                       'y': [0, -4, 0, 1, 2, 2.1, 4, 20, 30, 4, 0]})
    ddf = dd.from_pandas(df, npartitions=2)
    cvs = ds.Canvas(plot_width=7, plot_height=7,
                    x_range=(-3, 3), y_range=(-3, 3))
    agg = cvs.line(ddf, 'x', 'y', ds.count())
    sol = np.array([[0, 0, 1, 0, 1, 0, 0],
                    [0, 1, 0, 0, 0, 1, 0],
                    [1, 0, 0, 0, 0, 0, 1],
                    [0, 0, 0, 0, 0, 0, 0],
                    [1, 0, 0, 0, 0, 0, 1],
                    [0, 2, 0, 0, 0, 1, 0],
                    [0, 0, 1, 0, 1, 0, 0]], dtype='i4')
    out = xr.DataArray(sol, coords=[lincoords, lincoords],
                       dims=['y', 'x'])
    assert_eq(agg, out)


# # Line tests
@pytest.mark.parametrize('df,x,y,ax', [
    # axis1 none constant
    (pd.DataFrame({
        'x0': [4, -4, 4],
        'x1': [0,  0, 0],
        'x2': [-4, 4, -4],
        'y0': [0,  0, 0],
        'y1': [-4, 4, 0],
        'y2': [0,  0, 0]
    }), ['x0', 'x1', 'x2'], ['y0', 'y1', 'y2'], 1),

    # axis1 x constant
    (pd.DataFrame({
        'y0': [0, 0,  0],
        'y1': [0, 4, -4],
        'y2': [0, 0,  0]
    }), np.array([-4, 0, 4]), ['y0', 'y1', 'y2'], 1),

    # axis0 single
    (pd.DataFrame({
        'x': [4, 0, -4, np.nan, -4, 0, 4, np.nan, 4, 0, -4],
        'y': [0, -4, 0, np.nan, 0, 4, 0, np.nan, 0, 0, 0],
    }), 'x', 'y', 0),

    # axis0 multi
    (pd.DataFrame({
        'x0': [4,  0, -4],
        'x1': [-4, 0,  4],
        'x2': [4,  0, -4],
        'y0': [0, -4,  0],
        'y1': [0,  4,  0],
        'y2': [0,  0,  0]
    }), ['x0', 'x1', 'x2'], ['y0', 'y1', 'y2'], 0),

    # axis0 multi with string
    (pd.DataFrame({
        'x0': [-4,  0, 4],
        'y0': [0, -4,  0],
        'y1': [0,  4,  0],
        'y2': [0,  0,  0]
    }), 'x0', ['y0', 'y1', 'y2'], 0),

    # axis1 RaggedArray
    (pd.DataFrame({
        'x': [[4, 0, -4], [-4, 0, 4, 4, 0, -4]],
        'y': [[0, -4, 0], [0, 4, 0, 0, 0, 0]],
    }, dtype='Ragged[int64]'), 'x', 'y', 1),
])
def test_line_manual_range(df, x, y, ax):
    axis = ds.core.LinearAxis()
    lincoords = axis.compute_index(axis.compute_scale_and_translate((-3., 3.), 7), 7)

    ddf = dd.from_pandas(df, npartitions=2)

    cvs = ds.Canvas(plot_width=7, plot_height=7,
                    x_range=(-3, 3), y_range=(-3, 3))

    agg = cvs.line(ddf, x, y, ds.count(), axis=ax)

    sol = np.array([[0, 0, 1, 0, 1, 0, 0],
                    [0, 1, 0, 0, 0, 1, 0],
                    [1, 0, 0, 0, 0, 0, 1],
                    [1, 1, 1, 1, 1, 1, 1],
                    [1, 0, 0, 0, 0, 0, 1],
                    [0, 1, 0, 0, 0, 1, 0],
                    [0, 0, 1, 0, 1, 0, 0]], dtype='i4')

    out = xr.DataArray(sol, coords=[lincoords, lincoords],
                       dims=['y', 'x'])
    assert_eq(agg, out)


@pytest.mark.parametrize('df,x,y,ax', [
    # axis1 none constant
    (pd.DataFrame({
        'x0': [0,  0, 0],
        'x1': [-4, 0, 4],
        'x2': [0,  0, 0],
        'y0': [-4, 4, -4],
        'y1': [0,  0,  0],
        'y2': [4, -4,  4]
    }), ['x0', 'x1', 'x2'], ['y0', 'y1', 'y2'], 1),

    # axis1 y constant
    (pd.DataFrame({
        'x0': [0,  0, 0],
        'x1': [-4, 0, 4],
        'x2': [0,  0, 0],
    }), ['x0', 'x1', 'x2'], np.array([-4, 0, 4]), 1),

    # axis0 single
    (pd.DataFrame({
        'x': [0, -4, 0, np.nan, 0, 0, 0, np.nan, 0, 4, 0],
        'y': [-4, 0, 4, np.nan, 4, 0, -4, np.nan, -4, 0, 4],
    }), 'x', 'y', 0),

    # axis0 multi
    (pd.DataFrame({
        'x0': [0, -4,  0],
        'x1': [0,  0,  0],
        'x2': [0,  4,  0],
        'y0': [-4, 0,  4],
        'y1': [4,  0, -4],
        'y2': [-4, 0,  4]
    }), ['x0', 'x1', 'x2'], ['y0', 'y1', 'y2'], 0),

    # axis0 multi with string
    (pd.DataFrame({
        'x0': [0, -4,  0],
        'x1': [0,  0,  0],
        'x2': [0,  4,  0],
        'y0': [-4, 0,  4]
    }), ['x0', 'x1', 'x2'], 'y0', 0),

    # axis1 RaggedArray
    (pd.DataFrame({
        'x': [[0, -4, 0], [0, 0, 0], [0, 4, 0]],
        'y': [[-4, 0, 4], [4, 0, -4], [-4, 0, 4]],
    }, dtype='Ragged[int64]'), 'x', 'y', 1),

])
def test_line_autorange(df, x, y, ax):
    axis = ds.core.LinearAxis()
    lincoords = axis.compute_index(
        axis.compute_scale_and_translate((-4., 4.), 9), 9)

    ddf = dd.from_pandas(df, npartitions=2)

    cvs = ds.Canvas(plot_width=9, plot_height=9)

    agg = cvs.line(ddf, x, y, ds.count(), axis=ax)

    sol = np.array([[0, 0, 0, 0, 3, 0, 0, 0, 0],
                    [0, 0, 0, 1, 1, 1, 0, 0, 0],
                    [0, 0, 1, 0, 1, 0, 1, 0, 0],
                    [0, 1, 0, 0, 1, 0, 0, 1, 0],
                    [1, 0, 0, 0, 1, 0, 0, 0, 1],
                    [0, 1, 0, 0, 1, 0, 0, 1, 0],
                    [0, 0, 1, 0, 1, 0, 1, 0, 0],
                    [0, 0, 0, 1, 1, 1, 0, 0, 0],
                    [0, 0, 0, 0, 3, 0, 0, 0, 0]], dtype='i4')

    out = xr.DataArray(sol, coords=[lincoords, lincoords],
                       dims=['y', 'x'])
    assert_eq(agg, out)


def test_line_x_constant_autorange():
    # axis1 y constant
    df = pd.DataFrame({
        'y0': [0, 0, 0],
        'y1': [-4, 0, 4],
        'y2': [0, 0, 0],
    })

    x = np.array([-4, 0, 4])
    y = ['y0', 'y1', 'y2']
    ax = 1

    axis = ds.core.LinearAxis()
    lincoords = axis.compute_index(
        axis.compute_scale_and_translate((-4., 4.), 9), 9)

    ddf = dd.from_pandas(df, npartitions=2)

    cvs = ds.Canvas(plot_width=9, plot_height=9)

    agg = cvs.line(ddf, x, y, ds.count(), axis=ax)

    sol = np.array([[0, 0, 0, 0, 1, 0, 0, 0, 0],
                    [0, 0, 0, 1, 0, 1, 0, 0, 0],
                    [0, 0, 1, 0, 0, 0, 1, 0, 0],
                    [0, 1, 0, 0, 0, 0, 0, 1, 0],
                    [3, 1, 1, 1, 1, 1, 1, 1, 3],
                    [0, 1, 0, 0, 0, 0, 0, 1, 0],
                    [0, 0, 1, 0, 0, 0, 1, 0, 0],
                    [0, 0, 0, 1, 0, 1, 0, 0, 0],
                    [0, 0, 0, 0, 1, 0, 0, 0, 0]], dtype='i4')

    out = xr.DataArray(sol, coords=[lincoords, lincoords],
                       dims=['y', 'x'])
    assert_eq(agg, out)


def test_log_axis_line():
    axis = ds.core.LogAxis()
    logcoords = axis.compute_index(axis.compute_scale_and_translate((1, 10), 2), 2)

    axis = ds.core.LinearAxis()
    lincoords = axis.compute_index(axis.compute_scale_and_translate((0, 1), 2), 2)

    sol = np.array([[5, 5], [5, 5]], dtype='i4')
    out = xr.DataArray(sol, coords=[lincoords, logcoords],
                       dims=['y', 'log_x'])
    assert_eq(c_logx.line(ddf, 'log_x', 'y', ds.count('i32')), out)
    out = xr.DataArray(sol, coords=[logcoords, lincoords],
                       dims=['log_y', 'x'])
    assert_eq(c_logy.line(ddf, 'x', 'log_y', ds.count('i32')), out)
    out = xr.DataArray(sol, coords=[logcoords, logcoords],
                       dims=['log_y', 'log_x'])
    assert_eq(c_logxy.line(ddf, 'log_x', 'log_y', ds.count('i32')), out)


def test_auto_range_line():
    axis = ds.core.LinearAxis()
    lincoords = axis.compute_index(axis.compute_scale_and_translate((-10., 10.), 5), 5)

    df = pd.DataFrame({'x': [-10,  0, 10,   0, -10],
                       'y': [  0, 10,  0, -10,   0]})
    ddf = dd.from_pandas(df, npartitions=2)
    cvs = ds.Canvas(plot_width=5, plot_height=5)
    agg = cvs.line(ddf, 'x', 'y', ds.count())
    sol = np.array([[0, 0, 1, 0, 0],
                    [0, 1, 0, 1, 0],
                    [2, 0, 0, 0, 1],
                    [0, 1, 0, 1, 0],
                    [0, 0, 1, 0, 0]], dtype='i4')
    out = xr.DataArray(sol, coords=[lincoords, lincoords],
                       dims=['y', 'x'])
    assert_eq(agg, out)

def test_trimesh_no_double_edge():
    """Assert that when two triangles share an edge that would normally get
    double-drawn, the edge is only drawn for the rightmost (or bottommost)
    triangle.
    """
    import multiprocessing as mp
    # Test left/right edge shared
    verts = dd.from_pandas(pd.DataFrame({'x': [4, 1, 5, 5, 5, 4],
                                         'y': [4, 5, 5, 5, 4, 4]}), npartitions=mp.cpu_count())
    tris = dd.from_pandas(pd.DataFrame({'v0': [0, 3], 'v1': [1, 4], 'v2': [2, 5], 'val': [1, 2]}), npartitions=mp.cpu_count())
    # Plot dims and x/y ranges need to be set such that the edge is drawn twice:
    cvs = ds.Canvas(plot_width=20, plot_height=20, x_range=(0, 5), y_range=(0, 5))
    agg = cvs.trimesh(verts, tris)
    sol = np.array([
        [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    ], dtype='i4')
    np.testing.assert_array_equal(np.flipud(agg.fillna(0).astype('i4').values)[:5], sol)


@pytest.mark.parametrize('npartitions', list(range(1, 6)))
def test_trimesh_dask_partitions(npartitions):
    """Assert that when two triangles share an edge that would normally get
    double-drawn, the edge is only drawn for the rightmost (or bottommost)
    triangle.
    """
    # Test left/right edge shared
    verts = dd.from_pandas(pd.DataFrame({'x': [4, 1, 5, 5, 5, 4],
                                         'y': [4, 5, 5, 5, 4, 4]}),
                           npartitions=npartitions)
    tris = dd.from_pandas(
        pd.DataFrame(
            {'v0': [0, 3], 'v1': [1, 4], 'v2': [2, 5], 'val': [1, 2]}),
        npartitions=npartitions)

    cvs = ds.Canvas(plot_width=20, plot_height=20,
                    x_range=(0, 5), y_range=(0, 5))

    # Precompute mesh with dask dataframes
    mesh = du.mesh(verts, tris)

    # Make sure mesh is a dask DataFrame
    assert isinstance(mesh, dd.DataFrame)

    # Check mesh length
    n = len(mesh)
    assert n == 6

    # Make sure we have expected number of partitions
    expected_chunksize = int(np.ceil(len(mesh) / (3*npartitions)) * 3)
    expected_npartitions = int(np.ceil(n / expected_chunksize))
    assert expected_npartitions == mesh.npartitions

    # Make sure triangles don't straddle partitions
    partitions_lens = mesh.map_partitions(len).compute()
    for partitions_len in partitions_lens:
        assert partitions_len % 3 == 0

    agg = cvs.trimesh(verts, tris, mesh)
    sol = np.array([
        [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    ], dtype='i4')
    np.testing.assert_array_equal(
        np.flipud(agg.fillna(0).astype('i4').values)[:5], sol)
