from datashader.compiler import compile_components
from datashader.utils import Dispatcher
from datashader.glyphs.line import LinesXarrayCommonX
from datashader.glyphs.quadmesh import (
    QuadMeshRaster, QuadMeshRectilinear, QuadMeshCurvilinear, build_scale_translate
)
from datashader.utils import apply
import dask
import numpy as np
import xarray as xr
from dask.base import tokenize, compute
from dask.array.overlap import overlap
dask_glyph_dispatch = Dispatcher()


def dask_xarray_pipeline(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=False):
    dsk, name = dask_glyph_dispatch(
        glyph, xr_ds, schema, canvas, summary, antialias=antialias, cuda=cuda)

    # Get user configured scheduler (if any), or fall back to default
    # scheduler for dask DataFrame
    scheduler = dask.base.get_scheduler() or xr_ds.__dask_scheduler__
    keys = xr_ds.__dask_keys__()
    optimize = xr_ds.__dask_optimize__
    graph = xr_ds.__dask_graph__()

    dsk.update(optimize(graph, keys))

    return scheduler(dsk, name)


def shape_bounds_st_and_axis(xr_ds, canvas, glyph):
    if not canvas.x_range or not canvas.y_range:
        x_extents, y_extents = glyph.compute_bounds_dask(xr_ds)
    else:
        x_extents, y_extents = None, None

    x_range = canvas.x_range or x_extents
    y_range = canvas.y_range or y_extents
    x_min, x_max, y_min, y_max = bounds = compute(*(x_range + y_range))
    x_range, y_range = (x_min, x_max), (y_min, y_max)

    width = canvas.plot_width
    height = canvas.plot_height

    x_st = canvas.x_axis.compute_scale_and_translate(x_range, width)
    y_st = canvas.y_axis.compute_scale_and_translate(y_range, height)
    st = x_st + y_st
    shape = (height, width)

    x_axis = canvas.x_axis.compute_index(x_st, width)
    y_axis = canvas.y_axis.compute_index(y_st, height)
    axis = dict([(glyph.x_label, x_axis), (glyph.y_label, y_axis)])

    return shape, bounds, st, axis


def dask_rectilinear(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=False):
    shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)

    # Compile functions
    create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
        compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=True)
    x_mapper = canvas.x_axis.mapper
    y_mapper = canvas.y_axis.mapper
    extend = glyph._build_extend(
        x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
    x_range = bounds[:2]
    y_range = bounds[2:]

    # Build chunk indices for coordinates
    chunk_inds = {}
    for k, chunks in xr_ds.chunks.items():
        chunk_inds[k] = [0] + list(np.cumsum(chunks))

    x_name = glyph.x
    y_name = glyph.y
    coords = xr_ds[glyph.name].coords
    coord_dims = list(coords.dims)
    xdim_ind = coord_dims.index(x_name)
    ydim_ind = coord_dims.index(y_name)

    var_name = list(xr_ds.data_vars.keys())[0]

    # Compute interval breaks
    xs = xr_ds[x_name].values
    ys = xr_ds[y_name].values
    x_breaks = glyph.infer_interval_breaks(xs)
    y_breaks = glyph.infer_interval_breaks(ys)

    def chunk(np_arr, *inds):
        # Reconstruct dataset for chunk from numpy array and chunk indices
        chunk_coords_list = []
        # for i, (coord_name, coord_vals) in enumerate(coords.items()):
        for i, coord_name in enumerate(coords.dims):
            chunk_number = inds[i]
            coord_slice = slice(
                chunk_inds[coord_name][chunk_number],
                chunk_inds[coord_name][chunk_number + 1]
            )
            chunk_coords_list.append([coord_name, coords[coord_name][coord_slice]])

        chunk_coords = dict(chunk_coords_list)
        chunk_ds = xr.DataArray(
            np_arr, coords=chunk_coords, dims=coord_dims, name=var_name
        ).to_dataset()

        # Compute chunk x/y breaks
        x_chunk_number = inds[xdim_ind]
        x_breaks_slice = slice(
            chunk_inds[x_name][x_chunk_number],
            chunk_inds[x_name][x_chunk_number + 1] + 1
        )
        x_breaks_chunk = x_breaks[x_breaks_slice]

        y_chunk_number = inds[ydim_ind]
        y_breaks_slice = slice(
            chunk_inds[y_name][y_chunk_number],
            chunk_inds[y_name][y_chunk_number + 1] + 1
        )
        y_breaks_chunk = y_breaks[y_breaks_slice]

        # Initialize aggregation buffers
        aggs = create(shape)

        # Perform aggregation
        extend(aggs, chunk_ds, st, bounds,
               x_breaks=x_breaks_chunk, y_breaks=y_breaks_chunk)
        return aggs

    name = tokenize(xr_ds.__dask_tokenize__(), canvas, glyph, summary)
    keys = [k for row in xr_ds.__dask_keys__()[0] for k in row]
    keys2 = [(name, i) for i in range(len(keys))]
    dsk = dict((k2, (chunk, k, k[1], k[2])) for (k2, k) in zip(keys2, keys))
    dsk[name] = (apply, finalize, [(combine, keys2)],
                 dict(cuda=cuda, coords=axis, dims=[glyph.y_label, glyph.x_label],
                      attrs=dict(x_range=x_range, y_range=y_range)))
    return dsk, name


def dask_raster(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=False):
    shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)

    # Compile functions
    create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
        compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=True)
    x_mapper = canvas.x_axis.mapper
    y_mapper = canvas.y_axis.mapper
    extend = glyph._build_extend(
        x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
    x_range = bounds[:2]
    y_range = bounds[2:]

    # Build chunk indices for coordinates
    chunk_inds = {}
    for k, chunks in xr_ds.chunks.items():
        chunk_inds[k] = [0] + list(np.cumsum(chunks))

    x_name = glyph.x
    y_name = glyph.y

    coords = xr_ds[glyph.name].coords

    coord_dims = list(coords.dims)
    xdim_ind = coord_dims.index(x_name)
    ydim_ind = coord_dims.index(y_name)
    var_name = list(xr_ds.data_vars.keys())[0]

    # Pre-compute bin sizes. We do this here to handle length-1 chunks
    src_x0, src_x1 = glyph. _compute_bounds_from_1d_centers(
        xr_ds, x_name, maybe_expand=False, orient=False
    )
    src_y0, src_y1 = glyph._compute_bounds_from_1d_centers(
        xr_ds, y_name, maybe_expand=False, orient=False
    )
    xbinsize = abs(float(xr_ds[x_name][1] - xr_ds[x_name][0]))
    ybinsize = abs(float(xr_ds[y_name][1] - xr_ds[y_name][0]))

    # Compute scale/translate
    out_h, out_w = shape
    src_h, src_w = [xr_ds[glyph.name].shape[i] for i in [ydim_ind, xdim_ind]]
    out_x0, out_x1, out_y0, out_y1 = bounds
    scale_y, translate_y = build_scale_translate(
        out_h, out_y0, out_y1, src_h, src_y0, src_y1
    )

    scale_x, translate_x = build_scale_translate(
        out_w, out_x0, out_x1, src_w, src_x0, src_x1
    )

    def chunk(np_arr, *inds):
        # Reconstruct dataset for chunk from numpy array and chunk indices
        chunk_coords_list = []
        for i, coord_name in enumerate(coords.dims):
            chunk_number = inds[i]
            coord_slice = slice(
                chunk_inds[coord_name][chunk_number],
                chunk_inds[coord_name][chunk_number + 1]
            )
            chunk_coords_list.append([coord_name, coords[coord_name][coord_slice]])

        chunk_coords = dict(chunk_coords_list)
        chunk_ds = xr.DataArray(
            np_arr, coords=chunk_coords, dims=coord_dims, name=var_name
        ).to_dataset()

        # Compute offsets
        x_chunk_number = inds[xdim_ind]
        offset_x = chunk_inds[x_name][x_chunk_number]

        y_chunk_number = inds[ydim_ind]
        offset_y = chunk_inds[y_name][y_chunk_number]

        # Initialize aggregation buffers
        aggs = create(shape)

        # Perform aggregation
        extend(aggs, chunk_ds, st, bounds,
               scale_x=scale_x, scale_y=scale_y,
               translate_x=translate_x, translate_y=translate_y,
               offset_x=offset_x, offset_y=offset_y,
               src_xbinsize=xbinsize, src_ybinsize=ybinsize)

        return aggs

    name = tokenize(xr_ds.__dask_tokenize__(), canvas, glyph, summary)
    keys = [k for row in xr_ds.__dask_keys__()[0] for k in row]
    keys2 = [(name, i) for i in range(len(keys))]
    dsk = dict((k2, (chunk, k, k[1], k[2])) for (k2, k) in zip(keys2, keys))
    dsk[name] = (apply, finalize, [(combine, keys2)],
                 dict(cuda=cuda, coords=axis, dims=[glyph.y_label, glyph.x_label],
                      attrs=dict(x_range=x_range, y_range=y_range)))
    return dsk, name


def dask_curvilinear(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=False):
    shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)

    # Compile functions
    create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
        compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=True)
    x_mapper = canvas.x_axis.mapper
    y_mapper = canvas.y_axis.mapper
    extend = glyph._build_extend(
        x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
    x_range = bounds[:2]
    y_range = bounds[2:]

    x_coord_name = glyph.x
    y_coord_name = glyph.y
    z_name = glyph.name

    data_dim_names = list(xr_ds[z_name].dims)
    x_coord_dim_names = list(xr_ds[x_coord_name].dims)
    y_coord_dim_names = list(xr_ds[y_coord_name].dims)
    zs = xr_ds[z_name].data
    x_centers = xr_ds[glyph.x].data
    y_centers = xr_ds[glyph.y].data

    var_name = list(xr_ds.data_vars.keys())[0]

    # Validate coordinates
    err_msg = (
        "DataArray {name} is backed by a Dask array, \n"
        "but coordinate {coord} is not backed by a Dask array with identical \n"
        "dimension order and chunks"
    )
    if (not isinstance(x_centers, dask.array.Array) or
            xr_ds[glyph.name].dims != xr_ds[glyph.x].dims or
            xr_ds[glyph.name].chunks != xr_ds[glyph.x].chunks):
        raise ValueError(err_msg.format(name=glyph.name, coord=glyph.x))

    if (not isinstance(y_centers, dask.array.Array) or
            xr_ds[glyph.name].dims != xr_ds[glyph.y].dims or
            xr_ds[glyph.name].chunks != xr_ds[glyph.y].chunks):
        raise ValueError(err_msg.format(name=glyph.name, coord=glyph.y))

    # Make sure coordinates are floats so that overlap with nan will behave properly
    if x_centers.dtype.kind != 'f':
        x_centers = x_centers.astype(np.float64)
    if y_centers.dtype.kind != 'f':
        y_centers = y_centers.astype(np.float64)

    x_overlapped_centers = overlap(x_centers, depth=1, boundary=np.nan)
    y_overlapped_centers = overlap(y_centers, depth=1, boundary=np.nan)

    def chunk(np_zs, np_x_centers, np_y_centers):

        # Handle boundaries that have nothing to overlap with
        for centers in [np_x_centers, np_y_centers]:
            if np.isnan(centers[0, :]).all():
                centers[0, :] = centers[1, :] - (centers[2, :] - centers[1, :])
            if np.isnan(centers[-1, :]).all():
                centers[-1, :] = centers[-2, :] + (centers[-2, :] - centers[-3, :])
            if np.isnan(centers[:, 0]).all():
                centers[:, 0] = centers[:, 1] - (centers[:, 2] - centers[:, 1])
            if np.isnan(centers[:, -1]).all():
                centers[:, -1] = centers[:, -2] + (centers[:, -2] - centers[:, -3])

        # compute interval breaks
        x_breaks_chunk = glyph.infer_interval_breaks(np_x_centers)
        y_breaks_chunk = glyph.infer_interval_breaks(np_y_centers)

        # trim breaks
        x_breaks_chunk = x_breaks_chunk[1:-1, 1:-1]
        y_breaks_chunk = y_breaks_chunk[1:-1, 1:-1]

        # Reconstruct dataset for chunk from numpy array and chunk indices
        chunk_coords = {
            x_coord_name: (x_coord_dim_names, np_x_centers[1:-1, 1:-1]),
            y_coord_name: (y_coord_dim_names, np_y_centers[1:-1, 1:-1]),
        }
        chunk_ds = xr.DataArray(
            np_zs, coords=chunk_coords, dims=data_dim_names, name=var_name
        ).to_dataset()

        # Initialize aggregation buffers
        aggs = create(shape)

        # Perform aggregation
        extend(aggs, chunk_ds, st, bounds,
               x_breaks=x_breaks_chunk, y_breaks=y_breaks_chunk)
        return aggs

    result_name = tokenize(xr_ds.__dask_tokenize__(), canvas, glyph, summary)

    z_keys = [k for row in zs.__dask_keys__() for k in row]
    x_overlap_keys = [k for row in x_overlapped_centers.__dask_keys__() for k in row]
    y_overlap_keys = [k for row in y_overlapped_centers.__dask_keys__() for k in row]

    result_keys = [(result_name, i) for i in range(len(z_keys))]

    dsk = dict(
        (res_k, (chunk, z_k, x_k, y_k))
        for (res_k, z_k, x_k, y_k) in zip(
            result_keys, z_keys, x_overlap_keys, y_overlap_keys
        )
    )

    dsk[result_name] = (
        apply, finalize, [(combine, result_keys)],
        dict(cuda=cuda, coords=axis, dims=[glyph.y_label, glyph.x_label],
             attrs=dict(x_range=x_range, y_range=y_range))
    )

    # Add x/y coord tasks to task graph
    dsk.update(x_overlapped_centers.dask)
    dsk.update(y_overlapped_centers.dask)

    return dsk, result_name


def dask_xarray_lines(
    glyph: LinesXarrayCommonX, xr_ds: xr.Dataset, schema, canvas, summary,
    *, antialias=False, cuda=False,
):
    shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)

    # Compile functions
    create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, \
        column_names = compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda,
                                          partitioned=True)
    x_mapper = canvas.x_axis.mapper
    y_mapper = canvas.y_axis.mapper
    extend = glyph._build_extend(
        x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
    x_range = bounds[:2]
    y_range = bounds[2:]

    x_name = glyph.x
    x_dim_index = glyph.x_dim_index
    other_dim_index = 1 - x_dim_index
    other_dim_name = xr_ds[glyph.y].coords.dims[other_dim_index]
    xs = xr_ds[x_name]

    # Build chunk offsets for coordinates
    chunk_offsets = {}
    for k, chunks in xr_ds.chunks.items():
        chunk_offsets[k] = [0] + list(np.cumsum(chunks))

    partitioned = True
    uses_row_index = summary.uses_row_index(cuda, partitioned)

    def chunk(np_array, *chunk_indices):
        aggs = create(shape)

        start_x_index = chunk_offsets[x_name][chunk_indices[x_dim_index]]
        end_x_index = start_x_index + np_array.shape[x_dim_index]
        x = xs[start_x_index:end_x_index].values

        start_other_index = chunk_offsets[other_dim_name][chunk_indices[other_dim_index]]
        end_other_index = start_other_index + np_array.shape[other_dim_index]

        data_vars = dict(
            name=(("x", other_dim_name) if x_dim_index == 0 else (other_dim_name, "x"), np_array),
        )
        # Other required columns are chunked in the other_dim
        for column_name in column_names:
            values = xr_ds[column_name][start_other_index:end_other_index].values
            data_vars[column_name] = (other_dim_name, values)

        chunk_ds = xr.Dataset(
            data_vars=data_vars,
            coords=dict(
                x=("x", x),
                other_dim_name=(other_dim_name, np.arange(start_other_index, end_other_index)),
            ),
        )

        if uses_row_index:
            row_offset = start_other_index
            chunk_ds.attrs["_datashader_row_offset"] = row_offset
            chunk_ds.attrs["_datashader_row_length"] = end_other_index - start_other_index

        extend(aggs, chunk_ds, st, bounds)
        return aggs

    name = tokenize(xr_ds.__dask_tokenize__(), canvas, glyph, summary)
    keys = [k for row in xr_ds.__dask_keys__()[0] for k in row]
    keys2 = [(name, i) for i in range(len(keys))]
    dsk = dict((k2, (chunk, k, k[1], k[2])) for (k2, k) in zip(keys2, keys))
    dsk[name] = (apply, finalize, [(combine, keys2)],
                 dict(cuda=cuda, coords=axis, dims=[glyph.y_label, glyph.x_label],
                      attrs=dict(x_range=x_range, y_range=y_range)))
    return dsk, name


dask_glyph_dispatch.register(QuadMeshRectilinear)(dask_rectilinear)
dask_glyph_dispatch.register(QuadMeshRaster)(dask_raster)
dask_glyph_dispatch.register(QuadMeshCurvilinear)(dask_curvilinear)
dask_glyph_dispatch.register(LinesXarrayCommonX)(dask_xarray_lines)