# same function using variables
_NFACES = 13

# these need to be set at some higher point
xdim='i3'
ydim='i2'

def _orthogonal_dim(dim):
    if dim==xdim:
        return ydim
    elif dim==ydim:
        return xdim
    else:
        raise ValueError("Invalid dim %s" % dim)

def _edge_dim_and_start(edge):
    if edge in ['N', 'S']:
        dim =  ydim
    elif edge in ['E', 'W']:
        dim = xdim
    else:
        raise ValueError('%s is not a valid edge' % edge)
    if edge in ['S', 'W']:
        start = 0
    elif edge in ['N', 'E']:
        start = -1
    return dim, start

def exch_get_t_neighbors(var, edge, count=1, facedim='face'):
    """Return the neighbor points of a face at cell centers.

    Paramters
    ---------
    var : xarray.Variable
        Must have a face dimension
    edge : {'N','S','E','W}
        The edge to fetch
    count : int
        Number of edge points to fetch
    facedim : str
        The name of the face dimension

    Returns
    -------
    edge : numpy.ndarray
        The edge points. (Note NOT an xarray object.)
    """

    face_axis = var.get_axis_num(facedim)
    this_edge_dim, _ = _edge_dim_and_start(edge)
    this_edge_axis = var.get_axis_num(this_edge_dim)

    def edge_data(fnum):
        link = face_edge_link[fnum][edge]
        if link is None:
            # TODO: actually return the right data, NaN or similar
            return None
        else:
            other_fnum, other_edge, other_orientation = link

        # Build up a slice that selects the correct edge region for a given face.
        # Start by getting everything
        face_slice = [slice(None),] * var.ndim

        # Specify the face
        # (If this is an int rather than slice, the whole axis gets dropped.)
        face_slice[face_axis] = slice(other_fnum, other_fnum+1)

        other_edge_dim, start = _edge_dim_and_start(other_edge)
        other_edge_axis = var.get_axis_num(other_edge_dim)
        if start==0:
            other_edge_slice = slice(None,count)
        elif start==-1:
            other_edge_slice = slice(-count,None)
        face_slice[other_edge_axis] = other_edge_slice

        # the orthogonal dimension might need to be reoriented
        ortho_axis = var.get_axis_num(_orthogonal_dim(other_edge_dim))
        ortho_slice = slice(None,None,other_orientation)
        face_slice[ortho_axis] = ortho_slice

        # now get the data
        data = var[tuple(face_slice)].data

        # the axis of the edge on THIS face is
        # not necessarily the same as the axis on the OTHER face
        if this_edge_axis != other_edge_axis:
            # unfortunately dask.Array has no swapaxes method
            axes = range(data.ndim)
            axes[this_edge_axis] = other_edge_axis
            axes[other_edge_axis] = this_edge_axis
            data = data.transpose(axes)
        return data

    arrays = [edge_data(fnum) for fnum in range(_NFACES)]
    data = xr.core.ops.concatenate(arrays, face_axis)
    return data

# modeled after xarray/xarray/core/variable.py

# specifcy which variables are rollable across a face
# this is valid for 2D (time, lat, lon) or 3D (depth, lat, lon)
_HORIZ_DIMS = {'i2': ('N','S'), 'i3': ('E','W')}
# instead, 4D data would have
#_HORIZ_DIMS = ['i3': ('N','S'), 'i4': ('E','W')]

# works on a variable
def _roll_one_dim_face(var, dim, count):
    if dim not in _HORIZ_DIMS:
        raise KeyError("Can't roll dim '%s' across face. "
                         "Can only roll dims %s" % (dim, repr(_HORIZ_DIMS)))

    # set to N, S, E, or W, depending on sign of count
    # a positive roll (count>0) shifts everything LEFT=S=W
    # a negative roll (count<0) shifts everything RIGHT=N=E
    roll_left = count>0
    edge = _HORIZ_DIMS[dim][1 if roll_left else 0]
    data_neighbor = exch_get_t_neighbors(var, edge, abs(count))

    axis = var.get_axis_num(dim)

    count %= var.shape[axis]

    if roll_left:
        indices = slice(None,-count)
    else:
        indices = slice(-count, None)

    slices = (slice(None),) * axis + (indices,)

    data_this_face = var[slices].data
    if roll_left:
        arrays = [data_neighbor, data_this_face]
    else:
        arrays = [data_this_face, data_neighbor]
    print([a.shape for a in arrays])

    data = xr.core.ops.concatenate(arrays, axis)

    if isinstance(data, xr.core.pycompat.dask_array_type):
        # chunked data should come out with the same chunks; this makes
        # it feasible to combine shifted and unshifted data
        # TODO: remove this once dask.array automatically aligns chunks
        data = data.rechunk(var.data.chunks)

    return type(var)(var.dims, data, var._attrs, fastpath=True)

def roll_face(result, **shifts):
    """
    Return a new Variable with rolled data.
    Parameters
    ----------
    **shifts : keyword arguments of the form {dim: offset}
        Integer offset to roll along each of the given dimensions.
        Positive offsets roll to the right; negative offsets roll to the
        left.
    Returns
    -------
    shifted : Variable
        Variable with the same dimensions and attributes but rolled data.
    """
    for dim, count in shifts.items():
        result = _roll_one_dim_face(result, dim, count)
    return result