from __future__ import division

from collections import defaultdict

import numpy
from scipy import ndimage

DISTANCE = numpy.sqrt([
    2., 1., 2.,
    1., 1., 1.,
    2., 1., 2.
])
DIR_MAP = dict(zip(range(9), [32, 64, 128, 16, -1, 1, 8, 4, 2]))
FLOWS_IN = numpy.array([2, 4, 8, 1, numpy.nan, 16, 128, 64, 32])


def _stack_neighbors(topo, radius=1, **padkwargs):
    """ Create a MxNx9 array of neighbors for each element of a MxN
    array.

    Creates a MxNx9 array where each layer represents all of the
    adjacent values at a give row/col. Input array is edge padded to
    handle the blocks on edges and corners.

    Parameters
    ----------
    topo : numpy array (MxN)
        An array represeting a digital elevation model (DEM)
    radius : int
        Search radius for stacking the neighbors. A radius = 1 implies
        all eight immediately adjacent cells (plus the actual cell)
    **padkwargs : optional parameters
        Positional and keyword arguments padded to numpy.pad

    Returns
    -------
    blocked : numpy array (MxNx9)
        Array of neighboros where `blocked[row, col, :].reshape(3, 3)`
        returns the block cells adjacent to and including
        `topo[row, col]`

    See Also
    --------
    numpy.pad

    References
    ----------
    `Stack Overflow <http://goo.gl/Y3h8ti>`_

    """

    mode = padkwargs.pop('mode')
    if mode == 'constant':
        pad_width = ((radius, radius), (radius, radius))
    elif mode == 'edge':
        pad_width = radius
    else:
        raise NotImplementedError("only 'edge' and 'constant' modes are supported")

    # pad the edges
    padded = numpy.pad(topo, pad_width=pad_width, mode=mode, **padkwargs)

    # new rows and cols count
    M, N = padded.shape

    # width is typically 3 -- length of each of
    # block that defines the neighbors
    width = radius * 2 + 1
    row_length = N - width + 1
    col_length = M - width + 1

    # Linear indices for the starting width-by-width block
    idx1 = numpy.arange(width)[:, None] * N + numpy.arange(width)

    # Offset (from the starting block indices) linear indices for all the blocks
    idx2 = numpy.arange(col_length)[:, None] * N + numpy.arange(row_length)

    # Finally, get the linear indices for all blocks
    all_idx = idx1.ravel()[None, None, :] + idx2[:, :, None]

    # Index into padded for the final output
    blocked = padded.ravel()[all_idx]

    return blocked


def _adjacent_slopes(topo):
    """ Compute the slope from each to cell to all of its neighbors.

    Parameters
    ----------
    topo : numpy array (MxN)
        Elevation data

    Returns
    -------
    slope : numpy array (MxNx9)
        3-D array where the z-axis is the unraveled slope array of
        each cell's neighbors.

    Notes
    -----
    Downward slopes are represented with positive numbers.

    See Also
    --------
    watershed.algo._stack_neighbors

    """

    # initial array shape
    rows, cols = topo.shape

    # make 3-D array of each cell's neighbors
    blocks = _stack_neighbors(topo, radius=1, mode='edge')

    # change in elevation (dz/dx)
    drop = topo.reshape(rows, cols, 1) - blocks

    # Slope (dz/dx) masked to exclude uphill directions
    slope = numpy.ma.masked_less_equal(drop / DISTANCE, 0)

    return slope


def _mark_sinks(topo):
    """ Marks sink areas in a DEM.

    Parameters
    ----------
    topo : numpy array
        Eelvation data

    Returns
    -------
    sink : numpy array
        Bool array. True value indicate cell is (in) a sink.

    """

    # compute the slopes in every direction at each cell
    slope = _adjacent_slopes(topo)

    # find where this is no downward slope
    sinks = slope.data.max(axis=2) == 0

    # remove 'sinks' on the edges of the array
    sinks[(0, -1), :] = False
    sinks[:, (0, -1)] = False

    return sinks


def fill_sinks(topo, copy=True):
    """ Fills sink areas in a DEM with the lowest adjacent elevation.

    Parameters
    ----------
    topo : numpy array
        Elevation data
    copy : bool, optional
        When True, operates on a copy of the `topo` array. Set to
        False if memory is a concern.

    Returns
    -------
    filled : numpy array
        Numpy array with all the sinks filled

    See Also
    --------
    watershed.algo.flow_direction_d8
    watershed.algo.trace_upstream
    watershed.algo.flow_accumulation