import functools
import operator

import dask
import dask.array as da
import dask.dataframe as dd

import numpy as np

META = dd.utils.make_meta([("foo", int)])


def test_sequential():
    image = da.zeros((2, 1), chunks=(1, -1))

    indices = list(np.ndindex(*image.numblocks))
    slices = da.core.slices_from_chunks(image.chunks)
    chunks = list(map(functools.partial(operator.getitem, image), slices))

    # Create one dataframe per chunk.
    ddf_all = np.empty(image.numblocks, dtype=object)
    for idx_chunk, chunk in zip(indices, chunks):
        ddf_delayed = dask.delayed(lambda: None)(chunk)
        ddf_all[idx_chunk] = dd.from_delayed(ddf_delayed, meta=META)

    # Merge dataframes along first axis.
    ddf_array = np.empty(image.numblocks[2:], dtype=object)
    for idx_spatial in np.ndindex(*ddf_array.shape):
        idx_all = (...,) + idx_spatial
        ddf = dd.concat(ddf_all[idx_all].flatten().tolist())
        ddf_array[idx_spatial] = ddf

    # Run a per-chunk calculation using the merged dataframes,
    # where a single merged dataframe is passed back to the all chunks that
    # went into making it.
    signal = []
    for idx_chunk, chunk in zip(indices, chunks):
        ddf = ddf_array[idx_chunk[2:]]
        ddf_signal_chunk = dask.delayed(lambda: None)(chunk, ddf)
        signal.append(ddf_signal_chunk)

    ddf_signal = dd.from_delayed(signal, meta=META)
    ddf_signal.__dask_graph__().validate()  # FAILS