# -*- coding: utf-8 -*-
"""
   Copyright 2015 Creare

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

Digital Elevation Model Processing Module
==========================================

This module implements Tarboton (96,
http://www.neng.usu.edu/cee/faculty/dtarb/96wr03137.pdf) Digital Elevation
Model (DEM) processing algorithms for calculating the magnitude and direction
(or aspect) of slopes given the DEM elevation data. This implementation
takes into account the supplied coordinate system, so it does not make an
assumption of a rectangular grid for improved accuracy.

It also implements a novel Upstream Contributing Area (UCA) calculation
algorithm that can operator on chunks of an input file, and accurately handle
the fluxes at edges of chunks.

Finally, it can calculate the Topographic Wetness Index (TWI) based on the UCA
and slope magnitude. Flats and no-data areas are in-painted, the maximum
value for TWI is capped, and the output is multiplied by 10 to get higher
accuracy when storing the array as an integer.

Usage Notes
-------------
This module consists of 3 classes, and 3 helper functions. General users should
only need to be concerned with the DEMProcessor class.

It has only been tested in USGS geotiff files.

It presently only supports WGS84 coordinate systems

Developer Notes
-----------------
The Edge and TileEdge classes keep track of the edge information for tiles

Development Notes
------------------
TODO: Implement this for the elevation data http://pastebin.com/umNmSrE2 .
    Generally the DEM data is stored as integers. So in areas with flat terrain
    due to rounding there may be large flats. We can create interpolated
    floating point values in those regions so that we can still get good
    answers

TODO: Replace complete file loading with partial loading from disk to RAM.
    Presently, an entire geotiff is loaded into memory, which makes the
    chunk calculation less useful because the memory restriction is still
    present

TODO: Improve general memory usage (following from previous TODO).

TODO: Cythonize magnitude and slope calculations

TODO: Implement pit-removal algorithm so that we can start from raw elevation
    tiles.

Created on Wed Jun 18 14:19:04 2014

@author: mpu
"""

import warnings
import gc
import numpy as np


import os
import subprocess
import scipy.sparse as sps
import scipy.ndimage as spndi
import numba

from reader.gdal_reader import GdalReader
from taudem import taudem
from test_pydem import get_test_data, make_file_names
from utils import mk_dx_dy_from_geotif_layer, get_fn

try:
    from cyfuncs import cyutils
    CYTHON = True
except:
    CYTHON = False
    warnings.warn("Cython functions are not compiled. UCA calculation will be,"
                  " slow. Consider compiling cython functions using: "
                  "python setup.py build_ext --inplace", RuntimeWarning)
# CYTHON = False

# A test aspect ration between dx and dy coordinates
TEST_DIV = 1/1.1  # 0.001
FLAT_ID = np.nan  # This is the fill value for flats in float arrays
FLAT_ID_INT = -1  # This is the fill value for flats in int arrays

# Used to extend flat region downstream when calculating angle
FLATS_KERNEL1 = np.ones((3, 3))
FLATS_KERNEL2 = np.ones((3, 3))  # Convolution used to find edges to the flats
# This is the only choice for FLATS_KERNEL3 because otherwise it won't play
# nicely with the pit-filling algorithm
FLATS_KERNEL3 = np.ones((3, 3))  # Kernel used to connect flats and edges
FILL_VALUE = -9999  # This is the integer fill value for no-data values


class Edge(object):
    """
    Small helper class that keeps track of data on an edge. It doesn't care
    if it's a top, bottom, left, or right edge.
    """
    todo = None
    done = None
    slice = None
    data = None

    def __init__(self, size, slice_):
        self.todo = np.ones(size, bool)
        self.done = np.zeros(size, bool)
        self.slice = slice_
        self.data = np.zeros(size, float)

    @property
    def n_done(self):
        return (self.coulddo & self.done).sum()

    @property
    def n_coulddo(self):
        return self.coulddo.sum()

    @property
    def percent_done(self):
        return 1.0 * (self.coulddo & self.done).sum() \
            / (self.coulddo.sum() + 1e-16)

    @property
    def coulddo(self):
        return self.todo & (self.data > 0)


class TileEdge(object):
    """
    Class that combines 4 edges per tile, and keeps track of all the edges
    in all the tiles on an image. This is for a single image file.
    """
    left = None
    right = None
    top = None
    bottom = None
    keys = None
    coords = None
    n_chunks = None
    n_cols = None
    x_axis = None
    y_axis = None
    max_elev = None  # Used to figure out which tile first in case of ties
    n_done = None
    percent_done = None
    n_todo = None

    def __init__(self, top_edge, bottom_edge, left_edge, right_edge, overlap,
                 x_axis, y_axis, elev):
        """
        Parameters
        ----------
        top_edge : 1d array
            Array of indices giving the top edge corners
        bottom_edge : 1d array
            Array of indices giving the bottom edge corners
        left_edge : 1d array
            Array of indices giving the left edge corners
        right_edge : 1d array
            Array of indices giving the right edge corners
        overlap : int
            The number of pixels that tiles overlab
        x_axis : 1d array
            Coordinates of the x-axis
        y_axis : 1d array
            Coordinates of the y-axis
        elev : 2d array
            The elevation data on the tile
        Notes
        ------
        To create the *_edge arrays,
        see :py:func:`DEMProcessor._get_chunk_edges`
        """
        self.n_chunks = top_edge.size * left_edge.size
        self.n_cols = left_edge.size
        self.x_axis = x_axis
        self.y_axis = y_axis
        keys = {}
        coords = []
        i = 0
        left = np.empty((left_edge.size, top_edge.size), Edge)
        right = np.empty((left_edge.size, top_edge.size), Edge)
        top = np.empty((left_edge.size, top_edge.size), Edge)
        bottom = np.empty((left_edge.size, top_edge.size), Edge)
        n_done = np.zeros(left.shape, int)
        self.percent_done = np.zeros(left.shape, float)
        self.n_todo = np.zeros(left.shape, int)
        max_elev = np.zeros(left.shape, float)
        for tb in xrange(top_edge.size):
            for lr in xrange(left_edge.size):
                te = top_edge[tb]
                be = bottom_edge[tb]
                le = left_edge[lr]
                re = right_edge[lr]
                # create the key with the overlaps
                keys[(te, be, le, re)] = i
                coords.append([te, be, le, re])
                # Remove the overlaps
#                if re != right_edge.max():
#                    re -= overlap
#                if le != 0:
#                    le += overlap
#                if te != 0:
#                    te += overlap
#                if be != bottom_edge.max():
#                    be -= overlap

                left.ravel()[i] = Edge(be-te, [slice(te, be), slice(le, le+1)])
                right.ravel()[i] = Edge(be-te, [slice(te, be), slice(re-1, re)])
                top.ravel()[i] = Edge(re-le, [slice(te, te+1), slice(le, re)])
                bottom.ravel()[i] = Edge(re-le, [slice(be-1, be), slice(le, re)])
                max_elev.ravel()[i] = elev[te:be, le:re].max()
                i += 1
        self.left = left
        self.right = right
        self.top = top
        self.bottom = bottom
        self.keys = keys
        self.coords = coords
        self.n_done = n_done
        self.max_elev = max_elev

    def get(self, key, side):
        """
        Returns an edge given a particular key
        Parmeters
        ----------
        key : tuple
            (te, be, le, re) tuple that identifies a tile
        side : str
            top, bottom, left, or right, which edge to return
        """
        return getattr(self, side).ravel()[self.keys[key]]

    def get_i(self, i, side):
        """ Returns the i'th tile's 'side' edge """
        return getattr(self, side).ravel()[i]

    def set(self, key, data, field, side, local=False):
        edge = self.get(key, side)
        if local:
            if side == 'left':
                dt = data[:, 0:1]
            elif side == 'right':
                dt = data[:, -1:]
            elif side == 'top':
                dt = data[0:1, :]
            elif side == 'bottom':
                dt = data[-1:, :]
            setattr(edge, field, dt)
        else:
            setattr(edge, field, data[edge.slice])

    def set_i(self, i, data, field, side):
        """ Assigns data on the i'th tile to the data 'field' of the 'side'
        edge of that tile
        """
        edge = self.get_i(i, side)
        setattr(edge, field, data[edge.slice])

    def set_sides(self, key, data, field, local=False):
        """
        Assign data on the 'key' tile to all the edges
        """
        for side in ['left', 'right', 'top', 'bottom']:
            self.set(key, data, field, side, local)

    def set_neighbor_data(self, neighbor_side, data, key, field):
        """
        Assign data from the 'key' tile to the edge on the
        neighboring tile which is on the 'neighbor_side' of the 'key' tile.
        The data is assigned to the 'field' attribute of the neihboring tile's
        edge.
        """
        i = self.keys[key]
        found = False
        sides = []
        if 'left' in neighbor_side:
            if i % self.n_cols == 0:
                return None
            i -= 1
            sides.append('right')
            found = True
        if 'right' in neighbor_side:
            if i % self.n_cols == self.n_cols - 1:
                return None
            i += 1
            sides.append('left')
            found = True
        if 'top' in neighbor_side:
            sides.append('bottom')
            i -= self.n_cols
            found = True
        if 'bottom' in neighbor_side:
            sides.append('top')
            i += self.n_cols
            found = True
        if not found:
            print "Side '%s' not found" % neighbor_side
        # Check if i is in range
        if i < 0 or i >= self.n_chunks:
            return None
        # Otherwise, set the data
        for side in sides:
            self.set_i(i, data, field, side)

    def set_all_neighbors_data(self, data, done,  key):
        """
        Given they 'key' tile's data, assigns this information to all
        neighboring tiles
        """
        # The order of this for loop is important because the topleft gets
        # it's data from the left neighbor, which should have already been
        # updated...
        for side in ['left', 'right', 'top', 'bottom', 'topleft',
                     'topright', 'bottomleft', 'bottomright']:
            self.set_neighbor_data(side, data, key, 'data')
#            self.set_neighbor_data(side, todo, key, 'todo')
            self.set_neighbor_data(side, done, key, 'done')

    def fill_n_todo(self):
        """
        Calculate and record the number of edge pixels left to do on each tile
        """
        left = self.left
        right = self.right
        top = self.top
        bottom = self.bottom
        for i in xrange(self.n_chunks):
            self.n_todo.ravel()[i] = np.sum([left.ravel()[i].n_todo,
                                            right.ravel()[i].n_todo,
                                            top.ravel()[i].n_todo,
                                            bottom.ravel()[i].n_todo])

    def fill_n_done(self):
        """
        Calculate and record the number of edge pixels that are done one each
        tile.
        """
        left = self.left
        right = self.right
        top = self.top
        bottom = self.bottom
        for i in xrange(self.n_chunks):
            self.n_done.ravel()[i] = np.sum([left.ravel()[i].n_done,
                                            right.ravel()[i].n_done,
                                            top.ravel()[i].n_done,
                                            bottom.ravel()[i].n_done])

    def fill_percent_done(self):
        """
        Calculate the percentage of edge pixels that would be done if the tile
        was reprocessed. This is done for each tile.
        """
        left = self.left
        right = self.right
        top = self.top
        bottom = self.bottom
        for i in xrange(self.n_chunks):
            self.percent_done.ravel()[i] = \
                np.sum([left.ravel()[i].percent_done,
                        right.ravel()[i].percent_done,
                        top.ravel()[i].percent_done,
                        bottom.ravel()[i].percent_done])
            self.percent_done.ravel()[i] /= \
                np.sum([left.ravel()[i].percent_done > 0,
                        right.ravel()[i].percent_done > 0,
                        top.ravel()[i].percent_done > 0,
                        bottom.ravel()[i].percent_done > 0, 1e-16])

    def fill_array(self, array, field, add=False, maximize=False):
        """
        Given a full array (for the while image), fill it with the data on
        the edges.
        """
        self.fix_shapes()
        for i in xrange(self.n_chunks):
            for side in ['left', 'right', 'top', 'bottom']:
                edge = getattr(self, side).ravel()[i]
                if add:
                    array[edge.slice] += getattr(edge, field)
                elif maximize:
                    array[edge.slice] = np.maximum(array[edge.slice],
                                                   getattr(edge, field))
                else:
                    array[edge.slice] = getattr(edge, field)
        return array

    def fix_shapes(self):
        """
        Fixes the shape of the data fields on edges. Left edges should be
        column vectors, and top edges should be row vectors, for example.
        """
        for i in xrange(self.n_chunks):
            for side in ['left', 'right', 'top', 'bottom']:
                edge = getattr(self, side).ravel()[i]
                if side in ['left', 'right']:
                    shp = [edge.todo.size, 1]
                else:
                    shp = [1, edge.todo.size]
                edge.done = edge.done.reshape(shp)
                edge.data = edge.data.reshape(shp)
                edge.todo = edge.todo.reshape(shp)

    def find_best_candidate(self):
        """
        Determine which tile, when processed, would complete the largest
        percentage of unresolved edge pixels. This is a heuristic function
        and does not give the optimal tile.
        """
        self.fill_percent_done()
        i_b = np.argmax(self.percent_done.ravel())
        if self.percent_done.ravel()[i_b] <= 0:
            return None

        # check for ties
        I = self.percent_done.ravel() == self.percent_done.ravel()[i_b]
        if I.sum() == 1:
            return i_b
        else:
            I2 = np.argmax(self.max_elev.ravel()[I])
            return I.nonzero()[0][I2]
#        i_b = self.percent_done.ravel() > self.percent_done.mean()
#        i = np.argmax(self.max_elev.ravel()[i_b])
#        I2 = self.max_elev.ravel()[i_b] == self.max_elev.ravel()[i_b][i]
#        if I2.sum() == 1:
#            return i_b.nonzero()[0][i]
#        else: #tie-break is highest percent done
#            i2 = np.nonzero(I2)[0]
#            i3 = np.argmax(self.percent_done.ravel()[i_b][i2])
#            return i_b.nonzero()[0][i2][i3]


class DEMProcessor(object):
    """
    This class processes elevation data, and returns the magnitude of slopes,
    the direction of slopes, the upstream contributing area, and the
    topographic wetness index.

    """
    # Flag that when true will ensure edge UCA is continuous across chunks
    resolve_edges = True
#    apply_uca_limit_edges = True
#    apply_twi_limits = True
#    apply_twi_limits_on_uca = True

    # When resolving drainage across edges, if maximum UCA is reached, should
    # edge be marked as completed?
    apply_uca_limit_edges = False
    # When calculating TWI, should TWI be limited to max value?
    apply_twi_limits = False
    # When calculating TWI, should UCA be limited to max value?
    apply_twi_limits_on_uca = False

    save_projection = 'EPSG:4326'

    direction = None  # Direction of slope in radians
    mag = None  # magnitude of slopes m/m
    uca = None  # upstream contributing area
    twi = None  # topographic wetness index
    elev = None  # elevation data
    A = None  # connectivity matrix

    # Gives the quadrant used for determining the d_infty mag/direction
    section = None  # save for debugging purposes, not useful output otherwise
    # Give the proportion of the area to drain to the first pixel in quadrant
    proportion = None  # Also saved for debugging
    done = None  # Marks if edges are done

    plotflag = False  # Debug plots
    # Use uniform values for dx/dy or obtain from geotiff
    dx_dy_from_file = True
    file_name = None  # Elevation data filename

    dX = None  # delta x
    dY = None  # delta y

    flats = None  # Boolean array indicating location of flats

    uca_saturation_limit = 32  # units of area
    twi_min_slope = 1e-3  # Used for TWI max limiting
    twi_min_area = np.inf  # Finds min area in tile
    chunk_size_slp_dir = 512  # Size of chunk (without overlaps)
    # This has to be > 1 to avoid edge effects for flats
    chunk_overlap_slp_dir = 4  # Overlap when calculating magnitude/directions
    chunk_size_uca = 512  # Size of chunks when calculating UCA
    chunk_overlap_uca = 32  # Number of overlapping pixels for UCA calculation
    # Mostly deprecated, but maximum number of iterations used to try and
    # resolve circular drainage patterns (which should never occur)
    circular_ref_maxcount = 50

    # The pixel coordinates for the different facets used to calculate the
    # D_infty magnitude and direction (from Tarboton)
    facets = [
             [(0, 0), (0, 1), (-1, 1)],
             [(0, 0), (-1, 0), (-1, 1)],
             [(0, 0), (-1, 0), (-1, -1)],
             [(0, 0), (0, -1), (-1, -1)],
             [(0, 0), (0, -1), (1, -1)],
             [(0, 0), (1, 0), (1, -1)],
             [(0, 0), (1, 0), (1, 1)],
             [(0, 0), (0, 1), (1, 1)],
    ]
    # Helper for magnitude/direction calculation (modified from Tarboton)
    ang_adj = np.array([
                       [0, 1],
                       [1, -1],
                       [1, 1],
                       [2, -1],
                       [2, 1],
                       [3, -1],
                       [3, 1],
                       [4, -1]
                       ])

#    def __del__(self):
#        self.elev_file = None #Close the elevation file

    def __init__(self, file_name, dx_dy_from_file=True, plotflag=False):
        """
        Parameters
        -----------
        file_name : str
            File name of elevation data
        dx_dy_from_file : bool, optional
            Default True. If true, will extract coordinates from geotiff file
            and use those to calculate the magnitude/direction of slopes.
            Otherwise assumes rectangular (uniform) coordinates.
        plotflag : bool, optional
            Default False: If True, will plot debug image. For a large
            file this is not advised.
        """
        # %%
        if os.path.exists(file_name):
            elev_file = GdalReader(file_name=file_name)
            elev, = elev_file.raster_layers
            data = elev.raster_data
            self.elev = elev
            self.data = data
            try:  # if masked array
                self.data.mask[(np.isnan(self.data))
                               | (self.data < -9998)] = True
                self.data[data.mask] = FILL_VALUE
            except:
                self.data = np.ma.masked_array(self.data,
                                               mask=(np.isnan(self.data))
                                               | (self.data < -9998))
            del elev_file  # close the file
            self.file_name = file_name

        shp = np.array(self.data.shape) - 1
        if dx_dy_from_file == 'test':  # This is a hidden option for dev/test
            # dX = np.linspace(0.9, 1.1, data.shape[0] - 1)
            dY = np.linspace(0.9, 0.9, data.shape[0] - 1)
            dX = np.ones((data.shape[0] - 1)) / TEST_DIV
        elif dx_dy_from_file:
            dX, dY = mk_dx_dy_from_geotif_layer(elev)

            if plotflag:
                from matplotlib.pyplot import plot, gca
                x = np.dot(np.arange(1, 10)[:, None], dX[None, :])
                y = np.dot(np.arange(1, 10)[:, None], dY[None, :])
                plot(x.ravel(), y.ravel(), '.')
                gca().invert_yaxis()
        else:
            dX = np.ones((data.shape[0]-1)) / (shp[1])
            dY = np.ones((data.shape[0]-1)) / (shp[0])
        if dx_dy_from_file == 'hack':
            dx = np.mean([dX.mean(), dY.mean()])
            dX[:] = dx
            dY[:] = dx

        self.dX = dX
        self.dY = dY

    def get_fn(self, name=None):
        return get_fn(self.elev, name)

    def get_full_fn(self, name, rootpath='.'):
        return os.path.join(rootpath, name, self.get_fn(name))

    def save_array(self, array, name=None, partname=None, rootpath='.',
                   raw=False, as_int=True):
        """
        Standard array saving routine

        Parameters
        -----------
        array : array
            Array to save to file
        name : str, optional
            Default 'array.tif'. Filename of array to save. Over-writes
            partname.
        partname : str, optional
            Part of the filename to save (with the coordinates appended)
        rootpath : str, optional
            Default '.'. Which directory to save file
        raw : bool, optional
            Default False. If true will save a .npz of the array. If false,
            will save a geotiff
        as_int : bool, optional
            Default True. If true will save array as an integer array (
            excellent compression). If false will save as float array.
        """
        if name is None and partname is not None:
            fnl_file = self.get_full_fn(partname, rootpath)
            tmp_file = os.path.join(rootpath, partname,
                                    self.get_fn(partname + '_tmp'))
        elif name is not None:
            fnl_file = name
            tmp_file = fnl_file + '_tmp.tiff'
        else:
            fnl_file = 'array.tif'
        if not raw:
            s_file = self.elev.clone_traits()
            s_file.raster_data = np.ma.masked_array(array)
            count = 10
            while count > 0 and (s_file.raster_data.mask.sum() > 0 \
                    or np.isnan(s_file.raster_data).sum() > 0):
                s_file.inpaint()
                count -= 1

            s_file.export_to_geotiff(tmp_file)

            if as_int:
                cmd = "gdalwarp -multi -wm 2000 -co BIGTIFF=YES -of GTiff -co compress=lzw -ot Int16 -co TILED=YES -wo OPTIMIZE_SIZE=YES -r near -t_srs %s %s %s" \
                    % (self.save_projection, tmp_file, fnl_file)
            else:
                cmd = "gdalwarp -multi -wm 2000 -co BIGTIFF=YES -of GTiff -co compress=lzw -co TILED=YES -wo OPTIMIZE_SIZE=YES -r near -t_srs %s %s %s" \
                    % (self.save_projection, tmp_file, fnl_file)
            print "<<"*4, cmd, ">>"*4
            subprocess.call(cmd)
            os.remove(tmp_file)
        else:
            np.savez_compressed(fnl_file, array)

    def save_uca(self, rootpath, raw=False, as_int=False):
        """ Saves the upstream contributing area to a file
        """
        self.save_array(self.uca, None, 'uca', rootpath, raw, as_int=as_int)

    def save_twi(self, rootpath, raw=False, as_int=True):
        """ Saves the topographic wetness index to a file
        """
        self.twi = np.ma.masked_array(self.twi, mask=self.twi <= 0,
                                      fill_value=-9999)
        #  self.twi = self.twi.filled()
        self.twi[self.flats] = 0
        self.twi.mask[self.flats] = True
        # self.twi = self.flats
        self.save_array(self.twi, None, 'twi', rootpath, raw, as_int=as_int)

    def save_slope(self, rootpath, raw=False, as_int=False):
        """ Saves the magnitude of the slope to a file
        """
        self.save_array(self.mag, None, 'mag', rootpath, raw, as_int=as_int)

    def save_direction(self, rootpath, raw=False, as_int=False):
        """ Saves the direction of the slope to a file
        """
        self.save_array(self.direction, None, 'ang', rootpath, raw, as_int=as_int)

    def save_outputs(self, rootpath='.', raw=False):
        """Saves TWI, UCA, magnitude and direction of slope to files.
        """
        self.save_twi(rootpath, raw)
        self.save_uca(rootpath, raw)
        self.save_slope(rootpath, raw)
        self.save_direction(rootpath, raw)

    def load_array(self, fn, name):
        """
        Can only load files that were saved in the 'raw' format.
        Loads previously computed field 'name' from file
        Valid names are 'mag', 'direction', 'uca', 'twi'
        """

        if os.path.exists(fn + '.npz'):
            array = np.load(fn + '.npz')
            try:
                setattr(self, name, array['arr_0'])
            except Exception, e:
                print e
            finally:
                array.close()

        else:
            raise RuntimeError("File %s does not exist." % (fn + '.npz'))

    def load_slope(self, fn):
        """Loads pre-computed slope magnitude from file
        """
        self.load_array(fn, 'mag')

    def load_direction(self, fn):
        """Loads pre-computed slope direction from file
        """
        self.load_array(fn, 'direction')

    def load_uca(self, fn):
        """Loads pre-computed uca from file
        """
        self.load_array(fn, 'uca')

    def _get_chunk_edges(self, NN, chunk_size, chunk_overlap):
        """
        Given the size of the array, calculate and array that gives the
        edges of chunks of nominal size, with specified overlap
        Parameters
        ----------
        NN : int
            Size of array
        chunk_size : int
            Nominal size of chunks (chunk_size < NN)
        chunk_overlap : int
            Number of pixels chunks will overlap
        Returns
        -------
        start_id : array
            The starting id of a chunk. start_id[i] gives the starting id of
            the i'th chunk
        end_id : array
            The ending id of a chunk. end_id[i] gives the ending id of
            the i'th chunk
        """
        left_edge = np.arange(0, NN - chunk_overlap, chunk_size)
        left_edge[1:] -= chunk_overlap
        right_edge = np.arange(0, NN - chunk_overlap, chunk_size)
        right_edge[:-1] = right_edge[1:] + chunk_overlap
        right_edge[-1] = NN
        right_edge = np.minimum(right_edge, NN)
        return left_edge, right_edge

    def _assign_chunk(self, data, arr1, arr2, te, be, le, re, ovr, add=False):
        """
        Assign data from a chunk to the full array. The data in overlap regions
        will not be assigned to the full array

        Parameters
        -----------
        data : array
            Unused array (except for shape) that has size of full tile
        arr1 : array
            Full size array to which data will be assigned
        arr2 : array
            Chunk-sized array from which data will be assigned
        te : int
            Top edge id
        be : int
            Bottom edge id
        le : int
            Left edge id
        re : int
            Right edge id
        ovr : int
            The number of pixels in the overlap
        add : bool, optional
            Default False. If true, the data in arr2 will be added to arr1,
            otherwise data in arr2 will overwrite data in arr1
        """
        if te == 0:
            i1 = 0
        else:
            i1 = ovr
        if be == data.shape[0]:
            i2 = 0
            i2b = None
        else:
            i2 = -ovr
            i2b = -ovr
        if le == 0:
            j1 = 0
        else:
            j1 = ovr
        if re == data.shape[1]:
            j2 = 0
            j2b = None
        else:
            j2 = -ovr
            j2b = -ovr
        if add:
            arr1[te+i1:be+i2, le+j1:re+j2] += arr2[i1:i2b, j1:j2b]
        else:
            arr1[te+i1:be+i2, le+j1:re+j2] = arr2[i1:i2b, j1:j2b]

    def find_flats(self):
        flats = self._find_flats_edges(self.data, self.dX, self.dY,
                                       self.mag, self.direction)
        self.direction[flats] = FLAT_ID
        self.mag[flats] = FLAT_ID
        self.flats = flats

    def calc_slopes_directions(self, plotflag=False):
        """
        Calculates the magnitude and direction of slopes and fills
        self.mag, self.direction
        """
        # %% Calculate the slopes and directions based on the 8 sections from
        # Tarboton http://www.neng.usu.edu/cee/faculty/dtarb/96wr03137.pdf
        if self.data.shape[0] <= self.chunk_size_slp_dir and \
                self.data.shape[1] <= self.chunk_size_slp_dir:
            print "starting slope/direction calculation"
            self.mag, self.direction = self._slopes_directions(
                self.data, self.dX, self.dY, 'tarboton')
            # Find the flat regions. This is mostly simple (look for mag < 0),
            # but the downstream pixel at the edge of a flat will have a
            # calcuable angle which will not be accurate. We have to also find
            # these edges and set their magnitude to -1 (that is, the flat_id)

            self.find_flats()
        else:
            self.direction = np.ones(self.data.shape) * FLAT_ID_INT
            self.mag = np.ones_like(self.direction) * FLAT_ID_INT
            self.flats = np.zeros(self.data.shape, bool)
            top_edge, bottom_edge = \
                self._get_chunk_edges(self.data.shape[0],
                                      self.chunk_size_slp_dir,
                                      self.chunk_overlap_slp_dir)
            left_edge, right_edge = \
                self._get_chunk_edges(self.data.shape[1],
                                      self.chunk_size_slp_dir,
                                      self.chunk_overlap_slp_dir)
            ovr = self.chunk_overlap_slp_dir
            count = 1
            for te, be in zip(top_edge, bottom_edge):
                for le, re in zip(left_edge, right_edge):
                    print "starting slope/direction calculation for chunk", \
                        count, "[%d:%d, %d:%d]" % (te, be, le, re)
                    count += 1
                    mag, direction = \
                        self._slopes_directions(self.data[te:be, le:re],
                                                self.dX[te:be-1],
                                                self.dY[te:be-1])

                    flats = self._find_flats_edges(self.data[te:be, le:re],
                                                   self.dX[te:be-1],
                                                   self.dY[te:be-1], mag,
                                                   direction)
                    direction[flats] = FLAT_ID
                    mag[flats] = FLAT_ID
                    self._assign_chunk(self.data, self.mag, mag,
                                       te, be, le, re, ovr)
                    self._assign_chunk(self.data, self.direction, direction,
                                       te, be, le, re, ovr)
                    self._assign_chunk(self.data, self.flats, flats,
                                       te, be, le, re, ovr)

        if plotflag:
            self._plot_debug_slopes_directions()

        gc.collect()  # Just in case
        return self.mag, self.direction

    def _slopes_directions(self, data, dX, dY, method='tarboton'):
        """ Wrapper to pick between various algorithms
        """
        # %%
        if method == 'tarboton':
            return self._tarboton_slopes_directions(data, dX, dY)
        elif method == 'central':
            return self._central_slopes_directions(data, dX, dY)

    def _tarboton_slopes_directions(self, data, dX, dY):
        """
        Calculate the slopes and directions based on the 8 sections from
        Tarboton http://www.neng.usu.edu/cee/faculty/dtarb/96wr03137.pdf
        """

        return _tarboton_slopes_directions(data, dX, dY,
                                           self.facets, self.ang_adj)

    def _central_slopes_directions(self, data, dX, dY):
        """
        Calculates magnitude/direction of slopes using central difference
        """
        shp = np.array(data.shape) - 1

        direction = np.ones(data.shape) * FLAT_ID_INT
        mag = np.ones_like(direction) * FLAT_ID_INT

        ind = 0
        d1, d2, theta = _get_d1_d2(dX, dY, ind, [0, 1], [1, 1], shp)
        s2 = (data[0:-2, 1:-1] - data[2:, 1:-1]) / d2
        s1 = -(data[1:-1, 0:-2] - data[1:-1, 2:]) / d1
        direction[1:-1, 1:-1] = np.arctan2(s2, s1) + np.pi
        mag = np.sqrt(s1**2 + s2**2)

        return mag, direction

    def _find_flats_edges(self, data, dX, dY, mag, direction):
        """
        Extend flats 1 square downstream
        Flats on the downstream side of the flat might find a valid angle,
        but that doesn't mean that it's a correct angle. We have to find
        these and then set them equal to a flat
        """

        flats = mag == FLAT_ID_INT
        assigned, n_flats = spndi.label(flats, FLATS_KERNEL3)
        # Perhaps the code below will be faster?
#        edges = ndimage.convolve(flats, FLATS_KERNEL1) - flats
#        edges = ndimage.convolve(edges, FLATS_KERNEL1) & flats
#        assigned, n_flats = spndi.label(edges, FLATS_KERNEL3)
        nn, mm = flats.shape
        flat_ids, flat_coords, flat_labelsf = _get_flat_ids(assigned)
        for ii in xrange(n_flats):
            ids_flats = flat_ids[flat_coords[ii]:flat_coords[ii+1]]
            elev_flat = data.ravel()[flat_ids[flat_coords[ii]]]
            if elev_flat is np.ma.masked:
                continue
            j = ids_flats % mm
            i = ids_flats // mm
            for iii in [-1, 0, 1]:
                for jjj in [-1, 0, 1]:
                    i_2 = i + iii
                    j_2 = j + jjj

                    ids_tmp = (i_2 >= 0) & (j_2 >= 0) & (i_2 < nn) & (j_2 < mm)
                    ids_tmp2 = data[i_2[ids_tmp], j_2[ids_tmp]] == elev_flat
                    flats[i_2[ids_tmp][ids_tmp2], j_2[ids_tmp][ids_tmp2]]\
                        += FLATS_KERNEL3[iii+1, jjj+1]

        return flats

    def calc_uca(self, plotflag=False, edge_init_data=None, uca_init=None):
        """Calculates the upstream contributing area.

        Parameters
        ----------
        plotflag : bool, optional
            Default False. If true will plot debugging plots. For large files,
            this will be very slow
        edge_init_data : list, optional
            edge_init_data = [uca_data, done_data, todo_data]
            uca_data : dict
                Dictionary with 'left', 'right', 'top', 'bottom' keys that
                gives the arrays filled with uca data on the edge corresponding
                to the key
            done_data : dict
                As uca_data, but bool array indicating if neighboring tiles
                have computed a finished value for that edge pixel
            todo_data : dict
                As uca_data, but bool array indicating if edges on tile still
                have to be computed
        uca_init : array, optional
            Array with pre-computed upstream contributing area
            (without edge contributions)

        Notes
        -------
        if edge_init_data is given, then the initialized area will be modified
        such that the edges are equal to the edge_init_data.

        If uca_init is given, then the interior of the upstream area will not
        be calculated. Only the information from the edges will be updated.
        Unless the tile is too large so that the calculation is chunked. In
        that case, the whole tile is re-computed.
        """
        if self.direction is None:
            self.calc_slopes_directions()

        # Initialize the upstream area
        uca_edge_init = np.zeros(self.data.shape)
        uca_edge_done = np.zeros(self.data.shape, dtype=bool)
        uca_edge_todo = np.zeros(self.data.shape, dtype=bool)
        edge_init_done, edge_init_todo = None, None
        if edge_init_data is not None:
            edge_init_data, edge_init_done, edge_init_todo = edge_init_data
            slices = {'left': [slice(None), slice(0, 1)],
                      'right': [slice(None), slice(-1, None)],
                      'top': [slice(0, 1), slice(None)],
                      'bottom': [slice(-1, None), slice(None)]}
            for key, val in slices.iteritems():
                # To initialize and edge it needs to have data and be finished
                uca_edge_done[val] += \
                    edge_init_done[key].reshape(uca_edge_init[val].shape)
                uca_edge_init[val] = \
                    edge_init_data[key].reshape(uca_edge_init[val].shape)
                uca_edge_init[val][~uca_edge_done[val]] = 0
                uca_edge_todo[val] += \
                    edge_init_todo[key].reshape(uca_edge_init[val].shape)

        if uca_init is None:
            self.uca = np.ones(self.data.shape) * FLAT_ID_INT
        else:
            self.uca = uca_init

        if self.data.shape[0] <= self.chunk_size_uca and \
                self.data.shape[1] <= self.chunk_size_uca:
            if uca_init is None:
                print "Starting uca calculation"
                res = self._calc_uca_chunk(self.data, self.dX, self.dY,
                                           self.direction, self.mag,
                                           self.flats,
                                           area_edges=uca_edge_init,
                                           plotflag=plotflag)
                self.edge_todo = res[1]
                self.edge_done = res[2]
                self.uca = res[0]
            else:
                print "Starting edge resolution round: ",
                # last return value will be None: edge_
                area, e2doi, edone, _ = \
                    self._calc_uca_chunk_update(self.data, self.dX, self.dY,
                                                self.direction, self.mag,
                                                self.flats,
                                                area_edges=uca_edge_init,
                                                edge_todo=uca_edge_todo,
                                                edge_done=uca_edge_done)
                self.uca += area
                self.edge_todo = e2doi
                self.edge_done = edone

        else:

            top_edge, bottom_edge = \
                self._get_chunk_edges(self.data.shape[0], self.chunk_size_uca,
                                      self.chunk_overlap_uca)
            left_edge, right_edge = \
                self._get_chunk_edges(self.data.shape[1], self.chunk_size_uca,
                                      self.chunk_overlap_uca)
            ovr = self.chunk_overlap_uca

            # Initialize the edge_todo and done arrays
            edge_todo = np.zeros(self.data.shape, bool)
            edge_todo_tile = np.zeros(self.data.shape, bool)
            edge_not_done_tile = np.zeros(self.data.shape, bool)
            edge_done = np.zeros(self.data.shape, bool)

            tile_edge = TileEdge(top_edge, bottom_edge, left_edge,
                                 right_edge, ovr,
                                 self.elev.grid_coordinates.x_axis,
                                 self.elev.grid_coordinates.y_axis, self.data)
            count = 1

            # Mask out the edges because we're just trying to resolve the
            # internal edge conflicts
            self.data.mask[:, 0] = True
            self.data.mask[:, -1] = True
            self.data.mask[0, :] = True
            self.data.mask[-1, :] = True

            # if 1:  # uca_init == None:
            print "Starting uca calculation for chunk: ",
            # %%
            for te, be in zip(top_edge, bottom_edge):
                for le, re in zip(left_edge, right_edge):
                    print count, "[%d:%d, %d:%d]" % (te, be, le, re),
                    count += 1
                    area, e2doi, edone, e2doi_no_mask, e2o_no_mask = \
                        self._calc_uca_chunk(self.data[te:be, le:re],
                                             self.dX[te:be-1],
                                             self.dY[te:be-1],
                                             self.direction[te:be, le:re],
                                             self.mag[te:be, le:re],
                                             self.flats[te:be, le:re],
                                             area_edges=uca_edge_init[te:be, le:re],
                                             plotflag=plotflag,
                                             edge_todo_i_no_mask=uca_edge_todo[te:be, le:re])
                    self._assign_chunk(self.data, self.uca, area,
                                       te, be, le, re, ovr)
                    edge_todo[te:be, le:re] += e2doi
                    edge_not_done_tile[te:be, le:re] += e2o_no_mask
                    # if this tile is on the edge of the domain, we actually
                    # want to keep the edge information
                    # UPDATE: I don't think we actually need this here as it
                    # will be handled by chunk update ???
                    self._assign_chunk(self.data, edge_todo_tile, e2doi_no_mask,
                                       te, be, le, re, ovr)
#                    if te == top_edge[0] or be == bottom_edge[-1] \
#                            or le == left_edge[0] or re == right_edge[-1]:
#                        edge_todo_tile[te:be, le:re] = e2doi
                    self._assign_chunk(self.data, edge_done, edone,
                                       te, be, le, re, ovr)
                    tile_edge.set_all_neighbors_data(self.uca,
                                                     edge_done,
                                                     (te, be, le, re))
                    tile_edge.set_sides((te, be, le, re), e2doi, 'todo',
                                        local=True)
            # %%
            print '..Done'
            # This needs to be much more sophisticated because we have to
            # follow the tile's edge value through the interior.
            # Since we have to do that anyway, we might as well recompute
            # the UCA from scratch. So the above branch does that. The branch
            # below would be more efficient if we can get it working.
            if 0:  # else:
                # need to populate tile_edge somehow
                edge_todo_tile = uca_edge_todo & ~uca_edge_done
                edge_not_done_tile = edge_todo_tile.copy()
                for te, be in zip(top_edge, bottom_edge):
                    for le, re in zip(
                            [left_edge[0], left_edge[-1]],
                            [right_edge[0], right_edge[-1]]):
                        e2doi = uca_edge_todo[te:be, le:re]
                        tiledata = uca_edge_init[te:be, le:re]
                        tile_edge.set_sides((te, be, le, re), tiledata, 'data',
                                            local=True)
                        tiledone = uca_edge_done[te:be, le:re]
                        tile_edge.set_sides((te, be, le, re), tiledone, 'done',
                                            local=True)
                for te, be in zip([top_edge[0], top_edge[-1]],
                                  [bottom_edge[0], bottom_edge[-1]]):
                    for le, re in zip(left_edge, right_edge):
                        e2doi = uca_edge_todo[te:be, le:re]
                        tiledata = uca_edge_init[te:be, le:re]
                        tile_edge.set_sides((te, be, le, re), tiledata, 'data',
                                            local=True)
                        tiledone = uca_edge_done[te:be, le:re]
                        tile_edge.set_sides((te, be, le, re), tiledone, 'done',
                                            local=True)

            if not self.resolve_edges:
                # This branch is probably horribly broken (but it might have
                # always been that way)
                self.tile_edge = tile_edge
                self.edge_todo = edge_todo
                self.edge_done = edge_done
                return self.uca

            # ## RESOLVING EDGES ## #

            # Get a good starting tile for the iteration
            i = tile_edge.find_best_candidate()
            tile_edge.fix_shapes()
#            dbug = np.zeros_like(self.uca)
            print "Starting edge resolution round: ",
            count = 0
            i_old = -1
            while i is not None and i != i_old:
                count += 1
                print count, '(%d) .' % i,
                # %%
                te, be, le, re = tile_edge.coords[i]
                data, dX, dY, direction, mag, flats = \
                    [self.data[te:be, le:re],
                     self.dX[te:be-1], self.dY[te:be-1],
                     self.direction[te:be, le:re],
                     self.mag[te:be, le:re], self.flats[te:be, le:re]]
                area, e2doi, edone, e2doi_tile = self._calc_uca_chunk_update(
                    data, dX, dY, direction, mag, flats, tile_edge, i,
                    edge_todo=edge_not_done_tile[te:be, le:re])
                self._assign_chunk(self.data, self.uca, area,
                                   te, be, le, re, ovr, add=True)
                self._assign_chunk(self.data, edge_done, edone,
                                   te, be, le, re, ovr)
                tile_edge.set_all_neighbors_data(self.uca,
                                                 edge_done, (te, be, le, re))

                edge_not_done_tile[te:be, le:re] += e2doi_tile
                tile_edge.set_sides((te, be, le, re), e2doi, 'todo',
                                    local=True)
                i_old = i
                i = tile_edge.find_best_candidate()
                # Debugging plots below. Feel free to uncomment for debugging
#                def drawgrid():
#                    ax = gca();
#                    ax.set_xticks(np.linspace(-0.5, 63.5, 9))
#                    ax.set_yticks(np.linspace(-0.5, 63.5, 9))
#                    grid(lw=2, ls='-', c=(0.5, 0.5, 0.5))
#                figure(1);clf();imshow((self.uca), interpolation='none');colorbar(); title("uca" + str(i_old) + " " + str(i));drawgrid()
#                figure(2);clf();imshow(area, interpolation='none');colorbar(); title("local area" + str(i_old) + " " + str(i))
##                edge_todo[:] = 0
##                edge_todo = tile_edge.fill_array(edge_todo, 'todo', add=True)
#                figure(3);clf();imshow(edge_todo*1.0 + edge_done*2.0, interpolation='none');colorbar(); title("todo" + str(i_old) + " " + str(i));clim(0, 3)
#                edge_todo[:] = 0
#                edge_todo = tile_edge.fill_array(edge_todo, 'coulddo', add=True)
#                figure(3);clf();imshow(edge_todo*1.0 + edge_done*2.0, interpolation='none');colorbar(); title("todo" + str(i_old) + " " + str(i));clim(0, 3);drawgrid()
#                figure(4);clf();imshow(tile_edge.percent_done, interpolation='none');colorbar(); title("percent done" + str(i_old) + " " + str(i));clim(0, 1)
#                dbug[:] = 0
#                dbug = tile_edge.fill_array(dbug, 'coulddo', maximize=False)
#                dbug[dbug > 0] -= (self.uca - ref_area)[dbug > 0]
#                figure(5);clf();imshow(dbug, interpolation='none');colorbar(); title("data diff" + str(i_old) + " " + str(i));drawgrid()
#                dbug = (self.uca - area1)
#                figure(6);clf();imshow(np.log10(np.abs(dbug)), interpolation='none');colorbar(); title("uca diff" + str(i_old) + " " + str(i));drawgrid()
                # %%
            self.tile_edge = tile_edge
            self.edge_todo = edge_todo_tile
            self.edge_done = ~edge_not_done_tile
        print '..Done'

        # Fix the very last pixel on the edges
        self.fix_edge_pixels(edge_init_data, edge_init_done, edge_init_todo)

        gc.collect()  # Just in case
        return self.uca

    def fix_edge_pixels(self, edge_init_data, edge_init_done, edge_init_todo):
        """
        This function fixes the pixels on the very edge of the tile.
        Drainage is calculated if the edge is downstream from the interior.
        If there is data available on the edge (from edge_init_data, for eg)
        then this data is used.

        This is a bit of hack to take care of the edge-values. It could
        possibly be handled through the main algorithm, but at least here
        the treatment is explicit.
        """
        data, dX, dY, direction, flats = \
            self.data, self.dX, self.dY, self.direction, self.flats
        sides = ['left', 'right', 'top', 'bottom']
        slices_o = [[slice(None), slice(1, 2)], [slice(None), slice(-2, -1)],
                    [slice(1, 2), slice(None)], [slice(-2, -1), slice(None)]]
        slices_d = [[slice(None), slice(0, 1)], [slice(None), slice(-1, None)],
                    [slice(0, 1), slice(None)], [slice(-1, None), slice(None)]]

        # The first set of edges will have contributions from two nodes whereas
        # the second set of edges will only have contributinos from one node
        indices = {'left': [[3, 4], [2, 5]], 'right': [[0, 7], [1, 6]],
                   'top': [[1, 2], [0, 3]], 'bottom': [[5, 6], [4, 7]]}

        # Figure out which section the drainage goes towards, and what
        # proportion goes to the straight-sided (as opposed to diagonal) node.

        for side, slice_o, slice_d in zip(sides, slices_o, slices_d):
            section, proportion = \
                self._calc_uca_section_proportion(data[slice_o],
                                                  dX[slice_o[0]],
                                                  dY[slice_o[0]],
                                                  direction[slice_o],
                                                  flats[slice_o])
            # self-initialize:
            if side in ['left', 'right']:
                self.uca[slice_d] = \
                    np.concatenate(([dX[slice_d[0]][0] * dY[slice_d[0]][0]],
                                    dX[slice_d[0]] * dY[slice_d[0]]))\
                    .reshape(self.uca[slice_d].shape)
            else:
                self.uca[slice_d] = dX[slice_d[0]][0] * dY[slice_d[0]][0]
            for e in range(2):
                for i in indices[side][e]:
                    ed = self.facets[i][2]
                    ids = section == i
                    if e == 0:
                        self.uca[slice_d][ids] += self.uca[slice_o][ids] \
                            * proportion[ids]
                    self.uca[slice_d][ids] += \
                        np.roll(np.roll(self.uca[slice_o] * (1 - proportion),
                                        ed[0], 0),
                                ed[1], 1)[ids]
                    if e == 1:
                        self.uca[slice_d][ids] += \
                            np.roll(np.roll(self.uca[slice_o] * (proportion),
                                            ed[0], 0),
                                    ed[1], 1)[ids]

            # Finally, add the edge data from adjacent tiles
            if edge_init_done is not None:
                ids = edge_init_done[side]  # > 0
                if side in ['left', 'right']:
                    self.uca[slice_d][ids, :] = \
                        edge_init_data[side][ids][:, None]
                else:
                    self.uca[slice_d][:, ids] = edge_init_data[side][ids]

    def _calc_uca_chunk_update(self, data, dX, dY, direction, mag, flats,
                               tile_edge=None, i=None,
                               area_edges=None, edge_todo=None, edge_done=None,
                               plotflag=False):
        """
        Calculates the upstream contributing area due to contributions from
        the edges only.
        """
        # %%

        sides = ['left', 'right', 'top', 'bottom']
        slices = [[slice(None), slice(0, 1)], [slice(None), slice(-1, None)],
                  [slice(0, 1), slice(None)], [slice(-1, None), slice(None)]]
        # Figure out which section the drainage goes towards, and what
        # proportion goes to the straight-sided (as opposed to diagonal) node.

        section, proportion = self._calc_uca_section_proportion(
            data, dX, dY, direction, flats)

        # Build the drainage or adjacency matrix
        A = self._mk_adjacency_matrix(section, proportion, flats, data, mag)
        if CYTHON:
            B = A
            C = A.tocsr()
        if 1:  # not CYTHON:
            A = A.tocoo()

        ids = np.zeros(data.shape, bool)
        area = np.zeros(data.shape, float)
        # Set the ids to the edges that are now done, and initialize the
        # edge area
        if tile_edge is not None:
            if edge_todo is not None:
                edge_todo_tile = edge_todo
            else:
                edge_todo_tile = None
            edge_todo = np.zeros(data.shape, bool)
            for side, slice0 in zip(sides, slices):
                edge = getattr(tile_edge, side).ravel()[i]
                ids[slice0] = edge.done & edge.coulddo
                # only add area from the finished edges
                area[slice0] = edge.data * edge.done * edge.coulddo
                edge_todo[slice0] = edge.todo & ~edge.done
        elif area_edges is not None and edge_todo is not None \
                and edge_done is not None:
            area[:, 0] = area_edges[:, 0]
            area[:, -1] = area_edges[:, -1]
            area[-1, :] = area_edges[-1, :]
            area[0, :] = area_edges[0, :]

            # Initialize starting ids
            ids = edge_done & edge_todo
            edge_todo = edge_todo & ~edge_done
            edge_todo_tile = None
        else:
            raise RuntimeError("Need to specify either tile_edge or area_edges"
                               "in _calc_uca_chunk_update")

        ids = ids.ravel()
        ids0 = ids.copy()

        area[flats] = np.nan
        edge_done = ~edge_todo

        edge_todo_i = edge_todo.copy()
        ids_old = np.zeros_like(ids)

        # I need this to keep track of when I have to add the area, and when
        # I have to replace the area.
        ids_i = np.arange(ids.size)

        done = np.ones(data.shape, bool)
        done.ravel()[ids] = False

        # Now we have to advance done through the mesh to figure out which
        # contributions matter (i.e. what's done already)
        def drain_pixels_done(ids, arr, rows_A, cols_A):
            ids_old = ids.copy()
            ids_old[:] = False
            # If I use ids.sum() > 0 then I might get stuck in circular
            # references.
            while (ids - ids_old).sum() > 0:
                # %%
                print "x",
                ids_old = ids.copy()
                ids_todo = ids_i[ids.ravel()]
                ids[:] = False
                for id_todo in ids_todo:
                    rows = cols_A == id_todo
                    rows_id = rows_A[rows]
                    ids[rows_id] += arr.ravel()[rows_id] is True
                    arr.ravel()[rows_id] = False  # Set second arrival new id

            return arr
        if CYTHON:
            a = cyutils.drain_connections(done.ravel().astype(int),
                                          ids.astype(int), B.indptr, B.indices,
                                          set_to=False)
            a = a.reshape(done.shape)
            done = a.astype(bool)
        else:
            done = drain_pixels_done(ids, done, A.row, A.col)

        done[data.mask] = True  # deal with no-data values
        #
        ids = ids0.copy()
        # Set all the edges to "done" for ids0. This ensures that no edges
        # will ever be updated, whether they are done or not.
        ids0 = ids0.reshape(data.shape)
        ids0[:, 0] = True
        ids0[:, -1] = True
        ids0[0, :] = True
        ids0[-1, :] = True
        ids0 = ids0.ravel()

        ids_old[:] = 0

        # %%
        def drain_area(ids, area, done, rows_A, cols_A, data_A,
                       edge_todo_tile):
            ids_old = ids.copy()
            ids_old[:] = False
            # If I use ids.sum() > 0 then I might get stuck in
            # circular references.
            while (ids - ids_old).sum() > 0:
                # %%
                print "o",
                ids_old = ids.copy()
                done.ravel()[ids] = True
                ids_todo = ids_i[ids.ravel()]
                ids[:] = False
                for id_todo in ids_todo:
                    rows = cols_A == id_todo
                    rows_id = rows_A[rows]
                    factor = data_A[rows]
                    # not allowed to modify edge values
                    edge_filter_ids = ~ids0[rows_id]
                    factor = factor[edge_filter_ids]
                    rows_id = rows_id[edge_filter_ids]
                    area.ravel()[rows_id] += area.ravel()[id_todo] * factor
                    if edge_todo_tile is not None:
                        edge_todo_tile.ravel()[rows_id] += \
                            edge_todo_tile.ravel()[id_todo] * factor
                    # Figure out of this cell that just received a contribution
                    # should give its contribution to what's next... i.e. make
                    # sure all inputs have been added together
                    for row_id in rows_id:  # this is the 'waiting' part
                        cols = cols_A[rows_A == row_id]
                        ids[row_id] += (~(done.ravel()[cols])).sum() == 0
#                        for col in cols:
#                            print 'row', row_id, 'col', col, 'done', done.ravel()[col]
                # Follow the drainage along. New candidates are cells that
                # just changed
                #ids = (track_id_old.ravel() == -1) \
                #        & (track_id_old.ravel() != track_id.ravel())
#                done.ravel()[ids] = True
    #            figure(7);clf();imshow(ids.reshape(mag.shape), interpolation='none')
    #            figure(8);clf();imshow(area, interpolation='none');colorbar()
    #            figure(9);clf();imshow(done, interpolation='none');colorbar()
    #            figure(10);clf();imshow(a + area - b, interpolation='none');colorbar()
                #%%
    #            self._plot_connectivity(A, data=data)
            return area, done, edge_todo_tile

        if CYTHON:
            if edge_todo_tile is not None:
                a, b, c, d = cyutils.drain_area(area.ravel(),
                                                done.astype(int).ravel(),
                                                ids.astype(int),
                                                B.indptr, B.indices, B.data,
                                                C.indptr, C.indices,
                                                area.shape[0], area.shape[1],
                                                edge_todo_tile.astype(float).ravel(),
                                                skip_edge=True)
                edge_todo_tile = c.reshape(edge_todo_tile.shape)
            else:
                a, b, c, d = cyutils.drain_area(area.ravel(),
                                                done.astype(int).ravel(),
                                                ids.astype(int),
                                                B.indptr, B.indices, B.data,
                                                C.indptr, C.indices,
                                                area.shape[0], area.shape[1],
                                                skip_edge=True)
            area = a.reshape(area.shape)
            done = b.reshape(done.shape)

        else:
            area, done, edge_todo_tile = \
                drain_area(ids, area, done, A.row, A.col, A.data,
                           edge_todo_tile)

        # Rather unfortunately, we still have to follow through the boolean
        # edge_todo matrix...
        ids = edge_todo.copy().ravel()
        # %%

        def drain_pixels_todo(ids, arr, rows_A, cols_A):
            ids_old = ids.copy()
            ids_old[:] = False
            # If I use ids.sum() > 0 then I might get stuck in
            # circular references.
            while (ids - ids_old).sum() > 0:
                # %%
                print "x",
                ids_old = ids.copy()
#                edge_todo_old = arr.copy()
                ids_todo = ids_i[ids.ravel()]
                ids[:] = False
                for id_todo in ids_todo:
                    rows = cols_A == id_todo
                    rows_id = rows_A[rows]
                    ids[rows_id] += arr.ravel()[rows_id] == False
                    arr.ravel()[rows_id] = True  # Set new id of second arrival

#                #Follow the drainage along. New candidates are cells that just changed
#                ids = (edge_todo_old.ravel() != arr.ravel())
            return arr
        if CYTHON:
            a = cyutils.drain_connections(edge_todo.ravel().astype(int),
                                          ids.astype(int), B.indptr, B.indices,
                                          set_to=True)
            a = a.reshape(edge_todo.shape)
            edge_todo = (a.astype(bool))
        else:
            edge_todo = drain_pixels_todo(ids, edge_todo, A.row, A.col)

        area[flats] = np.nan
        edge_done = ~edge_todo

        return area, edge_todo_i, edge_done, edge_todo_tile

    def _calc_uca_chunk(self, data, dX, dY, direction, mag, flats,
                        area_edges, plotflag=False, edge_todo_i_no_mask=True):
        """
        Calculates the upstream contributing area for the interior, and
        includes edge contributions if they are provided through area_edges.
        """
        # %%
        # Figure out which section the drainage goes towards, and what
        # proportion goes to the straight-sided (as opposed to diagonal) node.
        section, proportion = self._calc_uca_section_proportion(
            data, dX, dY, direction, flats)

        # Build the drainage or adjacency matrix
        A = self._mk_adjacency_matrix(section, proportion, flats, data, mag)
        if CYTHON:
            B = A.tocsr()

        colsum = np.array(A.sum(1)).ravel()
        ids = colsum == 0  # If no one drains into me

        area = (dX * dY)
        # Record minimum area
        min_area = np.nanmin(area)
        self.twi_min_area = min(self.twi_min_area, min_area)

        area = np.concatenate((area[0:1], area)).reshape(area.size+1, 1)
        area = area.repeat(data.shape[1], 1)
        # Set the edge areas to zero, will add those contributions later
        area[:, 0] = area_edges[:, 0]
        area[:, -1] = area_edges[:, -1]
        area[-1, :] = area_edges[-1, :]
        area[0, :] = area_edges[0, :]
        # These edges are done, they have been drained already
        ids[area_edges.ravel() > 0] = True

        done = np.zeros(data.shape, bool)
        done.ravel()[ids] = True
        # deal with no-data values
        done[1:-1, 1:-1] = done[1:-1, 1:-1] | data.mask[1:-1, 1:-1]

        # Check the inlet edges
        edge_todo = np.zeros_like(done)
        ids_ed = np.arange(data.size).reshape(data.shape)
        # left
        edge_todo[:, 0] = (A[:, ids_ed[:, 0]].sum(0) > 0) \
            & (area_edges[:, 0] == 0)
        edge_todo[:, -1] = (A[:, ids_ed[:, -1]].sum(0) > 0) \
            & (area_edges[:, -1] == 0)
        edge_todo[0, :] = (A[:, ids_ed[0, :]].sum(0) > 0) \
            & (area_edges[0, :] == 0)
        edge_todo[-1, :] = (A[:, ids_ed[-1, :]].sum(0) > 0) \
            & (area_edges[-1, :] == 0)

        # Will do the tile-level doneness
        edge_todo_i_no_mask = edge_todo.copy() & edge_todo_i_no_mask
        edge_todo_no_mask = edge_todo_i_no_mask.copy()  # tile-level doneness
        edge_todo[data.mask] = False  # Don't do masked areas
        # Initialize done edges
        edge_todo_i = edge_todo.copy()
        ids_old = np.zeros_like(ids)
        # %%
        count = 1
        while (np.any(~done) and count < self.circular_ref_maxcount):
            print ".",
            count += 1
            if CYTHON:
                a, b, c, d = cyutils.drain_area(area.ravel(),
                                                done.astype(int).ravel(),
                                                ids.astype(int),
                                                A.indptr, A.indices, A.data,
                                                B.indptr, B.indices,
                                                area.shape[0], area.shape[1],
                                                edge_todo.astype(float).ravel(),
                                                edge_todo_no_mask.astype(float).ravel())
                area = a.reshape(area.shape)
                done = b.reshape(done.shape).astype(bool)
                edge_todo = c.reshape(edge_todo.shape).astype(bool)
                edge_todo_no_mask = \
                    d.reshape(edge_todo_no_mask.shape).astype(bool)
            else:
                # If I use ids.sum() > 0 then I might get stuck in
                # circular references.
                while (ids - ids_old).sum() > 0:
                    # %%
                    ids_old = ids.copy()
                    ids, area, done, edge_todo = \
                        self._drain_step(A, ids, area, done, edge_todo)
    #                figure(1);clf();imshow(area, interpolation='none');colorbar()
    #                figure(2);clf();imshow(ids.reshape(area.shape), interpolation='none');colorbar()
    #                figure(3);clf();imshow(done, interpolation='none');colorbar()
            #%%
            ids[:] = False
            max_elev = (data * (~done)).max()
            ids[((data * (~done) - max_elev) / max_elev > -0.01).ravel()] = True

        area[flats] = np.nan

        edge_done = ~edge_todo
        edge_done[data.mask] = True  # Don't do masked areas

        if self.apply_uca_limit_edges:
            # 2x because of bifurcations (maybe should be more than 2x, but
            # should be ok
            edge_done[area > self.uca_saturation_limit * 2 * min_area] = True

        # %%
        if plotflag:
            self._plot_connectivity(A, (done.astype(float) is False)
                                    + flats.astype(float) * 2, [0, 3])
        return area, edge_todo_i, edge_done, edge_todo_i_no_mask, edge_todo_no_mask

    def _drain_step(self, A, ids, area, done, edge_todo):
        """
        Does a single step of the upstream contributing area calculation.
        Here the pixels in ids are drained downstream, the areas are updated
        and the next set of pixels to drain are determined for the next round.
        """
        # Only drain to cells that have a contribution
        A_todo = A[:, ids.ravel()]
        colsum = np.array(A_todo.sum(1)).ravel()
        # Only touch cells that actually receive a contribution
        # during this stage
        ids_new = colsum != 0
        # Is it possible that I may drain twice from my own cell?
        # -- No, I don't think so...
        # Is it possible that other cells may drain into me in
        # multiple iterations -- yes
        # Then say I check for when I'm done ensures that I don't drain until
        # everyone has drained into me
        area.ravel()[ids_new] += (A_todo[ids_new, :]
                                  * (area.ravel()[ids].ravel()))
        edge_todo.ravel()[ids_new] += (A_todo[ids_new, :]
                                       * (edge_todo.ravel()[ids].ravel()))
        # Figure out what's left to do.
        done.ravel()[ids] = True

        colsum = A * (~done.ravel())
        ids = colsum == 0
        # Figure out the new-undrained ids
        ids = ids & (~done.ravel())

        return ids, area, done, edge_todo

    def _calc_uca_section_proportion(self, data, dX, dY, direction, flats):
        """
        Given the direction, figure out which nodes the drainage will go
        toward, and what proportion of the drainage goes to which node
        """
        shp = np.array(data.shape) - 1

        facets = self.facets
        adjust = self.ang_adj[:, 1]

        d1, d2, theta = _get_d1_d2(dX, dY, 0, facets[0][1], facets[0][2], shp)
        if dX.size > 1:
            theta = np.row_stack((theta[0, :], theta, theta[-1, :]))
        # Which quadrant am I in?
        section = ((direction / np.pi * 2.0) // 1).astype(int)
        # Gets me in the quadrant
        quadrant = (direction - np.pi / 2.0 * section)
        proportion = np.zeros_like(quadrant) * np.nan
        # Now which section within the quadrant
        section = section * 2 \
            + (quadrant > theta.repeat(data.shape[1], 1)) * (section % 2 == 0) \
            + (quadrant > (np.pi/2 - theta.repeat(data.shape[1], 1))) \
            * (section % 2 == 1)  # greater than because of ties resolution b4

        # %%    Calculate proportion
        # As a side note, it's crazy to me how:
        #    _get_d1_d2 needs to use indices 0, 3, 4, 7,
        #    section uses even/odd (i.e. % 2)
        #    proportion uses indices (0, 1, 4, 5) {ALl of them different! ARG!}
        I1 = (section == 0) | (section == 1) | (section == 4) | (section == 5)
    #    I1 = section % 2 == 0
        I = I1 & (quadrant <= theta.repeat(data.shape[1], 1))
        proportion[I] = quadrant[I] / theta.repeat(data.shape[1], 1)[I]
        I = I1 & (quadrant > theta.repeat(data.shape[1], 1))
        proportion[I] = (quadrant[I] - theta.repeat(data.shape[1], 1)[I]) \
            / (np.pi / 2 - theta.repeat(data.shape[1], 1)[I])
        I = (~I1) & (quadrant <= (np.pi / 2 - theta.repeat(data.shape[1], 1)))
        proportion[I] = (quadrant[I]) \
            / (np.pi / 2 - theta.repeat(data.shape[1], 1)[I])
        I = (~I1) & (quadrant > (np.pi / 2 - theta.repeat(data.shape[1], 1)))
        proportion[I] = (quadrant[I] - (np.pi / 2 - theta.repeat(data.shape[1], 1)[I])) \
            / (theta.repeat(data.shape[1], 1)[I])
        # %%Finish Proportion Calculation
        section[flats] = FLAT_ID_INT
        proportion[flats] = FLAT_ID

        section[section == 8] = 0  # Fence-post error correction
        proportion = (1 + adjust[section]) / 2.0 - adjust[section] * proportion

        return section, proportion

    def _mk_adjacency_matrix(self, section, proportion, flats, elev, mag):
        """
        Calculates the adjacency of connectivity matrix. This matrix tells
        which pixels drain to which.

        For example, the pixel i, will recieve area from np.nonzero(A[i, :])
        at the proportions given in A[i, :]. So, the row gives the pixel
        drain to, and the columns the pixels drained from.
        """
        shp = section.shape
        mat_data = np.row_stack((proportion, 1 - proportion))
        NN = np.prod(shp)
        i12 = np.arange(NN).reshape(shp)
        j1 = - np.ones_like(i12)
        j2 = - np.ones_like(i12)

        # make the connectivity for the non-flats
        j1, j2 = self._mk_connectivity(section, i12, j1, j2)

        # Now make the connectivity for the flats
        j1, j2, mat_data, flat_i, flat_j, flat_prop = \
            self._mk_connectivity_flats(i12, j1, j2, mat_data, flats,
                                        elev, mag)

        # %% Make the matrix and plot
        j = np.row_stack((j1, j2))
        i = np.row_stack((i12, i12))

        j = np.concatenate((j.ravel(), flat_j)).astype(int)
        mat_data = np.concatenate((mat_data.ravel(), flat_prop))
        i = np.concatenate((i.ravel(), flat_i)).astype(int)

        # This prevents no-data values, remove connections when not present,
        # and makes sure that floating point precision errors do not
        # create circular references where a lower elevation cell drains
        # to a higher elevation cell
        I = ~np.isnan(mat_data) & (j != -1) & (mat_data > 1e-8) \
            & (elev.ravel()[j] <= elev.ravel()[i])

        mat_data = mat_data[I]
        j = j[I]
        i = i[I]

        # %%Make the matrix and initialize
        # What is A? The row i area receives area contributions from the
        # entries in its columns. If all the entries in my columns have
        #  drained, then I can drain.
        A = sps.csc_matrix((mat_data.ravel(),
                            np.row_stack((j.ravel(), i.ravel()))),
                           shape=(NN, NN))
        normalize = np.array(A.sum(0) + 1e-16).squeeze()
        A = np.dot(A, sps.diags(1/normalize, 0))

        return A

    def _mk_connectivity(self, section, i12, j1, j2):
        """
        Helper function for _mk_adjacency_matrix. Calculates the drainage
        neighbors and proportions based on the direction. This deals with
        non-flat regions in the image. In this case, each pixel can only
        drain to either 1 or two neighbors.
        """
        shp = np.array(section.shape) - 1

        facets = self.facets

        for ii, facet in enumerate(facets):
            e1 = facet[1]
            e2 = facet[2]
            I = section[1:-1, 1:-1] == ii
            j1[1:-1, 1:-1][I] = i12[1 + e1[0]:shp[0] + e1[0],
                                    1 + e1[1]:shp[1] + e1[1]][I]
            j2[1:-1, 1:-1][I] = i12[1 + e2[0]:shp[0] + e2[0],
                                    1 + e2[1]:shp[1] + e2[1]][I]

        # Now do the edges
        # left edge
        slc0 = [slice(1, -1), slice(0, 1)]
        for ind in [0, 1, 6, 7]:
            e1 = facets[ind][1]
            e2 = facets[ind][2]
            I = section[slc0] == ind
            j1[slc0][I] = i12[1 + e1[0]:shp[0] + e1[0], e1[1]][I.ravel()]
            j2[slc0][I] = i12[1 + e2[0]:shp[0] + e2[0], e2[1]][I.ravel()]

        # right edge
        slc0 = [slice(1, -1), slice(-1, None)]
        for ind in [2, 3, 4, 5]:
            e1 = facets[ind][1]
            e2 = facets[ind][2]
            I = section[slc0] == ind
            j1[slc0][I] = i12[1 + e1[0]:shp[0] + e1[0],
                              shp[1] + e1[1]][I.ravel()]
            j2[slc0][I] = i12[1 + e2[0]:shp[0] + e2[0],
                              shp[1] + e2[1]][I.ravel()]

        # top edge
        slc0 = [slice(0, 1), slice(1, -1)]
        for ind in [4, 5, 6, 7]:
            e1 = facets[ind][1]
            e2 = facets[ind][2]
            I = section[slc0] == ind
            j1[slc0][I] = i12[e1[0], 1 + e1[1]:shp[1] + e1[1]][I.ravel()]
            j2[slc0][I] = i12[e2[0], 1 + e2[1]:shp[1] + e2[1]][I.ravel()]

        # bottom edge
        slc0 = [slice(-1, None), slice(1, -1)]
        for ind in [0, 1, 2, 3]:
            e1 = facets[ind][1]
            e2 = facets[ind][2]
            I = section[slc0] == ind
            j1[slc0][I] = i12[shp[0] + e1[0],
                              1 + e1[1]:shp[1] + e1[1]][I.ravel()]
            j2[slc0][I] = i12[shp[0] + e2[0],
                              1 + e2[1]:shp[1] + e2[1]][I.ravel()]

        # top-left corner
        slc0 = [slice(0, 1), slice(0, 1)]
        for ind in [6, 7]:
            e1 = facets[ind][1]
            e2 = facets[ind][2]
            if section[slc0] == ind:
                j1[slc0] = i12[e1[0], e1[1]]
                j2[slc0] = i12[e2[0], e2[1]]

        # top-right corner
        slc0 = [slice(0, 1), slice(-1, None)]
        for ind in [4, 5]:
            e1 = facets[ind][1]
            e2 = facets[ind][2]
            if section[slc0] == ind:
                j1[slc0] = i12[e1[0], shp[1] + e1[1]]
                j2[slc0] = i12[e2[0], shp[1] + e2[1]]

        # bottom-left corner
        slc0 = [slice(-1, None), slice(0, 1)]
        for ind in [0, 1]:
            e1 = facets[ind][1]
            e2 = facets[ind][2]
            if section[slc0] == ind:
                j1[slc0] = i12[shp[0] + e1[0], e1[1]]
                j2[slc0] = i12[shp[0] + e2[0], e2[1]]

        # bottom-right corner
        slc0 = [slice(-1, None), slice(-1, None)]
        for ind in [2, 3]:
            e1 = facets[ind][1]
            e2 = facets[ind][2]
            if section[slc0] == ind:
                j1[slc0] = i12[e1[0] + shp[0], shp[1] + e1[1]]
                j2[slc0] = i12[e2[0] + shp[0], shp[1] + e2[1]]

        return j1, j2

    def _mk_connectivity_flats(self, i12, j1, j2, mat_data, flats, elev, mag):
        """
        Helper function for _mk_adjacency_matrix. This calcualtes the
        connectivity for flat regions. Every pixel in the flat will drain
        to a random pixel in the flat. This accumulates all the area in the
        flat region to a single pixel. All that area is then drained from
        that pixel to the surroundings on the flat. If the border of the
        flat has a single pixel with a much lower elevation, all the area will
        go towards that pixel. If the border has pixels with similar elevation,
        then the area will be distributed amongst all the border pixels
        proportional to their elevation.
        """
        nn, mm = flats.shape
        NN = np.prod(flats.shape)
        # Label the flats
        assigned, n_flats = spndi.label(flats, FLATS_KERNEL3)

        flat_ids, flat_coords, flat_labelsf = _get_flat_ids(assigned)
        flat_j = [None] * n_flats
        flat_prop = [None] * n_flats
        flat_i = [None] * n_flats

        # Temporary array to find the flats
        edges = np.zeros_like(flats)
        # %% Calcute the flat drainage
        warn_flats = []
        for ii in xrange(n_flats):
            ids_flats = flat_ids[flat_coords[ii]:flat_coords[ii+1]]
            edges[:] = 0
            j = ids_flats % mm
            i = ids_flats // mm
            for iii in [-1, 0, 1]:
                for jjj in [-1, 0, 1]:
                    i_2 = i + iii
                    j_2 = j + jjj
                    ids_tmp = (i_2 >= 0) & (j_2 >= 0) & (i_2 < nn) & (j_2 < mm)
                    edges[i_2[ids_tmp], j_2[ids_tmp]] += \
                        FLATS_KERNEL3[iii+1, jjj+1]
            edges.ravel()[ids_flats] = 0
            ids_edge = np.argwhere(edges.ravel()).squeeze()

            flat_elev_loc = elev.ravel()[ids_flats]
            # It is possble for the edges to merge 2 flats, so we need to
            # take the lower elevation to avoid large circular regions
            flat_elev = flat_elev_loc.min()

            loc_elev = elev.ravel()[ids_edge]
            # Filter out any elevations larger than the flat elevation
            # TODO: Figure out if this should be <= or <
            I_filt = loc_elev < flat_elev
            try:
                loc_elev = loc_elev[I_filt]
                loc_slope = mag.ravel()[ids_edge][I_filt]
            except: # If this is fully masked out (i.e. inside a no-data area)
                loc_elev = np.array([])
                loc_slope = np.array([])

            loc_dx = self.dX.mean()

            # Now I have to figure out if I should just use the minimum or
            # distribute amongst many pixels on the flat boundary
            n = len(loc_slope)
            if n == 0:  # Flat does not have anywhere to drain
                # Let's see if the flat goes to the edge. If yes, we'll just
                # distribute the area along the edge.
                ids_flat_on_edge = ((ids_flats % mag.shape[1]) == 0) | \
                    ((ids_flats % mag.shape[1]) == (mag.shape[1] - 1)) | \
                    (ids_flats <= mag.shape[1]) | \
                    (ids_flats >= (mag.shape[1] * (mag.shape[0] - 1)))
                if ids_flat_on_edge.sum() == 0:
                    warn_flats.append(ii)
                    continue

                drain_ids = ids_flats[ids_flat_on_edge]
                loc_proportions = mag.ravel()[ids_flats[ids_flat_on_edge]]
                loc_proportions /= loc_proportions.sum()

                ids_flats = ids_flats[~ids_flat_on_edge]
                # This flat is entirely on the edge of the image
                if len(ids_flats) == 0:
                    # therefore, whatever drains into it is done.
                    continue
                flat_elev_loc = flat_elev_loc[~ids_flat_on_edge]
            else:  # Flat has a place to drain to
                min_edges = np.zeros(loc_slope.shape, bool)
                min_edges[np.argmin(loc_slope)] = True
                # Add to the min edges any edge that is within an error
                # tolerance as small as the minimum
                min_edges = (loc_slope + loc_slope * loc_dx / 2) \
                    >= loc_slope[min_edges]

                drain_ids = ids_edge[I_filt][min_edges]

                loc_proportions = loc_slope[min_edges]
                loc_proportions /= loc_proportions.sum()

            # Now distribute the connectivity amongst the chosen elevations
            # proportional to their slopes

            # First, let all the the ids in the flats drain to 1
            # flat id (for ease)
            one_id = np.zeros(ids_flats.size, bool)
            one_id[np.argmin(flat_elev_loc)] = True

            j1.ravel()[ids_flats[~one_id]] = ids_flats[one_id]
            mat_data.ravel()[ids_flats[~one_id]] = 1
            # Negative indices will be eliminated before making the matix
            j2.ravel()[ids_flats[~one_id]] = -1
            mat_data.ravel()[ids_flats[~one_id] + NN] = 0

            # Now drain the 1 flat to the drains
            j1.ravel()[ids_flats[one_id]] = drain_ids[0]
            mat_data.ravel()[ids_flats[one_id]] = loc_proportions[0]
            if len(drain_ids) > 1:
                j2.ravel()[ids_flats[one_id]] = drain_ids[1]
                mat_data.ravel()[ids_flats[one_id] + NN] = loc_proportions[1]

            if len(loc_proportions > 2):
                flat_j[ii] = drain_ids[2:]
                flat_prop[ii] = loc_proportions[2:]
                flat_i[ii] = np.ones(drain_ids[2:].size, int)\
                    * ids_flats[one_id]
        try:
            flat_j = np.concatenate([fj for fj in flat_j if fj is not None])
            flat_prop = \
                np.concatenate([fp for fp in flat_prop if fp is not None])
            flat_i = np.concatenate([fi for fi in flat_i if fi is not None])
        except:
            flat_j = np.array([], dtype=int)
            flat_prop = np.array([], dtype=float)
            flat_i = np.array([], dtype=int)

        if len(warn_flats) > 0:
            warnings.warn("Warning %d flats had no place" % len(warn_flats) +
                          " to drain to --> these are pits (check pit-remove"
                          "algorithm).")
        return j1, j2, mat_data, flat_i, flat_j, flat_prop

    def calc_twi(self):
        """
        Calculates the topographic wetness index and saves the result in
        self.twi.

        Returns
        -------
        twi : array
            Array giving the topographic wetness index at each pixel
        """
        if self.uca is None:
            self.calc_uca()
            gc.collect()  # Just in case
        min_area = self.twi_min_area
        min_slope = self.twi_min_slope
        twi = self.uca.copy()
        if self.apply_twi_limits_on_uca:
            twi[twi > self.uca_saturation_limit * min_area] = \
                self.uca_saturation_limit * min_area
        gc.collect()  # Just in case
        twi = np.log((twi) / (self.mag + min_slope))
        # apply the cap
        if self.apply_twi_limits:
            twi_sat_value = \
                np.log(self.uca_saturation_limit * min_area / min_slope)
            twi[twi > twi_sat_value] = twi_sat_value
        # multiply by 10 for better integer resolution when storing
        self.twi = twi * 10

        gc.collect()  # Just in case
        return twi

    def _plot_connectivity(self, A, data=None, lims=[None, None]):
        """
        A debug function used to plot the adjacency/connectivity matrix.
        This is really just a light wrapper around _plot_connectivity_helper
        """

        if data is None:
            data = self.data

        B = A.tocoo()
        self._plot_connectivity_helper(B.col, B.row, B.data, data, lims)

    def _plot_connectivity_helper(self, ii, ji, mat_datai, data, lims=[1, 8]):
        """
        A debug function used to plot the adjacency/connectivity matrix.
        """
        from matplotlib.pyplot import quiver, colorbar, clim,  matshow
        I = ~np.isnan(mat_datai) & (ji != -1) & (mat_datai >= 0)
        mat_data = mat_datai[I]
        j = ji[I]
        i = ii[I]
        x = i.astype(float) % data.shape[1]
        y = i.astype(float) // data.shape[1]
        x1 = (j.astype(float) % data.shape[1]).ravel()
        y1 = (j.astype(float) // data.shape[1]).ravel()
        nx = (x1 - x)
        ny = (y1 - y)
        matshow(data, cmap='gist_rainbow'); colorbar(); clim(lims)
        quiver(x, y, nx, ny, mat_data.ravel(), angles='xy', scale_units='xy',
               scale=1, cmap='bone')
        colorbar(); clim([0, 1])

    def _plot_debug_slopes_directions(self):
        """
        A debug function to plot the direction calculated in various ways.
        """
        # %%
        from matplotlib.pyplot import matshow, colorbar, clim, title

        matshow(self.direction / np.pi * 180); colorbar(); clim(0, 360)
        title('Direction')

        mag2, direction2 = self._central_slopes_directions()
        matshow(direction2 / np.pi * 180.0); colorbar(); clim(0, 360)
        title('Direction (central difference)')

        matshow(self.mag); colorbar()
        title('Magnitude')
        matshow(mag2); colorbar(); title("Magnitude (Central difference)")

        # %%
        # Compare to Taudem
        filename = self.file_name
        os.chdir('testtiff')
        try:
            os.remove('test_ang.tif')
            os.remove('test_slp.tif')
        except:
            pass
        cmd = ('dinfflowdir -fel "%s" -ang "%s" -slp "%s"' %
               (os.path.split(filename)[-1], 'test_ang.tif', 'test_slp.tif'))
        taudem._run(cmd)

        td_file = GdalReader(file_name='test_ang.tif')
        td_ang, = td_file.raster_layers
        td_file2 = GdalReader(file_name='test_slp.tif')
        td_mag, = td_file2.raster_layers
        os.chdir('..')

        matshow(td_ang.raster_data / np.pi*180); clim(0, 360); colorbar()
        title('Taudem direction')
        matshow(td_mag.raster_data); colorbar()
        title('Taudem magnitude')

        matshow(self.data); colorbar()
        title('The test data (elevation)')

        diff = (td_ang.raster_data - self.direction) / np.pi * 180.0
        diff[np.abs(diff) > 300] = np.nan
        matshow(diff); colorbar(); clim([-1, 1])
        title('Taudem direction - calculated Direction')

        # normalize magnitudes
        mag2 = td_mag.raster_data
        mag2 /= np.nanmax(mag2)
        mag = self.mag.copy()
        mag /= np.nanmax(mag)
        matshow(mag - mag2); colorbar()
        title('Taudem magnitude - calculated magnitude')
        del td_file
        del td_file2
        del td_ang
        del td_mag


def _get_flat_ids(assigned):
    """
    This is a helper function to recover the coordinates of regions that have
    been labeled within an image. This function efficiently computes the
    coordinate of all regions and returns the information in a memory-efficient
    manner.

    Parameters
    -----------
    assigned : ndarray[ndim=2, dtype=int]
        The labeled image. For example, the result of calling
        scipy.ndimage.label on a binary image

    Returns
    --------
    I : ndarray[ndim=1, dtype=int]
        Array of 1d coordinate indices of all regions in the image
    region_ids : ndarray[shape=[n_features + 1], dtype=int]
        Indexing array used to separate the coordinates of the different
        regions. For example, region k has xy coordinates of
        xy[region_ids[k]:region_ids[k+1], :]
    labels : ndarray[ndim=1, dtype=int]
        The labels of the regions in the image corresponding to the coordinates
        For example, assigned.ravel()[I[k]] == labels[k]
    """
    # MPU optimization:
    # Let's segment the regions and store in a sparse format
    # First, let's use where once to find all the information we want
    ids_labels = np.arange(len(assigned.ravel()), dtype=int)
    I = ids_labels[assigned.ravel().astype(bool)]
    labels = assigned.ravel()[I]
    # Now sort these arrays by the label to figure out where to segment
    sort_id = np.argsort(labels)
    labels = labels[sort_id]
    I = I[sort_id]
    # this should be of size n_features-1
    region_ids = np.where(labels[1:] - labels[:-1] > 0)[0] + 1
    # This should be of size n_features + 1
    region_ids = np.concatenate(([0], region_ids, [len(labels)]))

    return [I, region_ids, labels]


@numba.jit
def _tarboton_slopes_directions(data, dX, dY, facets, ang_adj):
    """
    Calculate the slopes and directions based on the 8 sections from
    Tarboton http://www.neng.usu.edu/cee/faculty/dtarb/96wr03137.pdf
    """
    shp = np.array(data.shape) - 1


    direction = np.ones(data.shape) * FLAT_ID_INT
    mag = np.ones_like(direction) * FLAT_ID_INT

    slc0 = [slice(1, -1), slice(1, -1)]
    for ind in xrange(8):

        e1 = facets[ind][1]
        e2 = facets[ind][2]
        ang = ang_adj[ind]
        slc1 = [slice(1 + e1[0], shp[0] + e1[0]),
                slice(1 + e1[1], shp[1] + e1[1])]
        slc2 = [slice(1 + e2[0], shp[0] + e2[0]),
                slice(1 + e2[1], shp[1] + e2[1])]

        d1, d2, theta = _get_d1_d2(dX, dY, ind, e1, e2, shp)
        mag, direction = _calc_direction(data, mag, direction, ang, d1, d2,
                                         theta, slc0, slc1, slc2)

    # %%Now do the edges
    # if the edge is lower than the interior, we need to copy the value
    # from the interior (as an approximation)
    ids1 = (direction[:, 1] > np.pi / 2) \
        & (direction[:, 1] < 3 * np.pi / 2)
    direction[ids1, 0] = direction[ids1, 1]
    mag[ids1, 0] = mag[ids1, 1]
    ids1 = (direction[:, -2] < np.pi / 2) \
        | (direction[:, -2] > 3 * np.pi / 2)
    direction[ids1, -1] = direction[ids1, -2]
    mag[ids1, -1] = mag[ids1, -2]
    ids1 = (direction[1, :] > 0) & (direction[1, :] < np.pi)
    direction[0, ids1] = direction[1, ids1]
    mag[0, ids1] = mag[1, ids1]
    ids1 = (direction[-2, :] > np.pi) & (direction[-2, :] < 2 * np.pi)
    direction[-1, ids1] = direction[-2, ids1]
    mag[-1, ids1] = mag[-2, ids1]

    # Now update the edges in case they are higher than the interior (i.e.
    # look at the downstream angle)

    # left edge
    slc0 = [slice(1, -1), slice(0, 1)]
    for ind in [0, 1, 6, 7]:
        e1 = facets[ind][1]
        e2 = facets[ind][2]
        ang = ang_adj[ind]
        slc1 = [slice(1 + e1[0], shp[0] + e1[0]), slice(e1[1], 1 + e1[1])]
        slc2 = [slice(1 + e2[0], shp[0] + e2[0]), slice(e2[1], 1 + e2[1])]
        d1, d2, theta = _get_d1_d2(dX, dY, ind, e1, e2, shp)
        mag, direction = _calc_direction(data, mag, direction, ang, d1, d2,
                                         theta, slc0, slc1, slc2)
    # right edge
    slc0 = [slice(1, -1), slice(-1, None)]
    for ind in [2, 3, 4, 5]:
        e1 = facets[ind][1]
        e2 = facets[ind][2]
        ang = ang_adj[ind]
        slc1 = [slice(1 + e1[0], shp[0] + e1[0]),
                slice(shp[1] + e1[1], shp[1] + 1 + e1[1])]
        slc2 = [slice(1 + e2[0], shp[0] + e2[0]),
                slice(shp[1] + e2[1], shp[1] + 1 + e2[1])]
        d1, d2, theta = _get_d1_d2(dX, dY, ind, e1, e2, shp)
        mag, direction = _calc_direction(data, mag, direction, ang, d1, d2,
                                         theta, slc0, slc1, slc2)
    # top edge
    slc0 = [slice(0, 1), slice(1, -1)]
    for ind in [4, 5, 6, 7]:
        e1 = facets[ind][1]
        e2 = facets[ind][2]
        ang = ang_adj[ind]
        slc1 = [slice(e1[0], 1 + e1[0]), slice(1 + e1[1], shp[1] + e1[1])]
        slc2 = [slice(e2[0], 1 + e2[0]), slice(1 + e2[1], shp[1] + e2[1])]
        d1, d2, theta = _get_d1_d2(dX, dY, ind, e1, e2, shp, 'top')
        mag, direction = _calc_direction(data, mag, direction, ang, d1, d2,
                                         theta, slc0, slc1, slc2)
    # bottom edge
    slc0 = [slice(-1, None), slice(1, -1)]
    for ind in [0, 1, 2, 3]:
        e1 = facets[ind][1]
        e2 = facets[ind][2]
        ang = ang_adj[ind]
        slc1 = [slice(shp[0] + e1[0], shp[0] + 1 + e1[0]),
                slice(1 + e1[1], shp[1] + e1[1])]
        slc2 = [slice(shp[0] + e2[0], shp[0] + 1 + e2[0]),
                slice(1 + e2[1], shp[1] + e2[1])]
        d1, d2, theta = _get_d1_d2(dX, dY, ind, e1, e2, shp, 'bot')
        mag, direction = _calc_direction(data, mag, direction, ang, d1, d2,
                                         theta, slc0, slc1, slc2)
    # top-left corner
    slc0 = [slice(0, 1), slice(0, 1)]
    for ind in [6, 7]:
        e1 = facets[ind][1]
        e2 = facets[ind][2]
        ang = ang_adj[ind]
        slc1 = [slice(e1[0], 1 + e1[0]), slice(e1[1], 1 + e1[1])]
        slc2 = [slice(e2[0], 1 + e2[0]), slice(e2[1], 1 + e2[1])]
        d1, d2, theta = _get_d1_d2(dX, dY, ind, e1, e2, shp, 'top')
        mag, direction = _calc_direction(data, mag, direction, ang, d1, d2,
                                         theta, slc0, slc1, slc2)
    # top-right corner
    slc0 = [slice(0, 1), slice(-1, None)]
    for ind in [4, 5]:
        e1 = facets[ind][1]
        e2 = facets[ind][2]
        ang = ang_adj[ind]
        slc1 = [slice(e1[0], 1 + e1[0]),
                slice(shp[1] + e1[1], shp[1] + 1 + e1[1])]
        slc2 = [slice(e2[0], 1 + e2[0]),
                slice(shp[1] + e2[1], shp[1] + 1 + e2[1])]
        d1, d2, theta = _get_d1_d2(dX, dY, ind, e1, e2, shp, 'top')
        mag, direction = _calc_direction(data, mag, direction, ang, d1, d2,
                                         theta, slc0, slc1, slc2)
    # bottom-left corner
    slc0 = [slice(-1, None), slice(0, 1)]
    for ind in [0, 1]:
        e1 = facets[ind][1]
        e2 = facets[ind][2]
        ang = ang_adj[ind]
        slc1 = [slice(shp[0] + e1[0], shp[0] + 1 + e1[0]),
                slice(e1[1], 1 + e1[1])]
        slc2 = [slice(shp[0] + e2[0], shp[0] + 1 + e2[0]),
                slice(e2[1], 1 + e2[1])]
        d1, d2, theta = _get_d1_d2(dX, dY, ind, e1, e2, shp, 'bot')
        mag, direction = _calc_direction(data, mag, direction, ang, d1, d2,
                                         theta, slc0, slc1, slc2)
    # bottom-right corner
    slc0 = [slice(-1, None), slice(-1, None)]
    for ind in [3, 4]:
        e1 = facets[ind][1]
        e2 = facets[ind][2]
        ang = ang_adj[ind]
        slc1 = [slice(shp[0] + e1[0], shp[0] + 1 + e1[0]),
                slice(shp[1] + e1[1], shp[1] + 1 + e1[1])]
        slc2 = [slice(shp[0] + e2[0], shp[0] + 1 + e2[0]),
                slice(shp[1] + e2[1], shp[1] + 1 + e2[1])]
        d1, d2, theta = _get_d1_d2(dX, dY, ind, e1, e2, shp, 'bot')
        mag, direction = _calc_direction(data, mag, direction, ang, d1, d2,
                                         theta, slc0, slc1, slc2)

    mag[mag > 0] = np.sqrt(mag[mag > 0])

    return mag, direction

def _get_d1_d2(dX, dY, ind, e1, e2, shp, topbot=None):
    """
    This finds the distances along the patch (within the eight neighboring
    pixels around a central pixel) given the difference in x and y coordinates
    of the real image. This is the function that allows real coordinates to be
    used when calculating the magnitude and directions of slopes.
    """
    if topbot == None:
        if ind in [0, 3, 4, 7]:
            d1 = dX[slice((e2[0] + 1) / 2, shp[0] + (e2[0] - 1) / 2)]
            d2 = dY[slice((e2[0] + 1) / 2, shp[0] + (e2[0] - 1) / 2)]
            if d1.size == 0:
                d1 = np.array([dX[0]])
                d2 = np.array([dY[0]])
        else:
            d2 = dX[slice((e1[0] + 1) / 2, shp[0] + (e1[0] - 1) / 2)]
            d1 = dY[slice((e1[0] + 1) / 2, shp[0] + (e1[0] - 1) / 2)]
            if d1.size == 0:
                d2 = dX[0]
                d1 = dY[0]
    elif topbot == 'top':
        if ind in [0, 3, 4, 7]:
            d1, d2 = dX[0], dY[0]
        else:
            d2, d1 = dX[0], dY[0]
    elif topbot == 'bot':
        if ind in [0, 3, 4, 7]:
            d1, d2 = dX[-1], dY[-1]
        else:
            d2, d1 = dX[-1], dY[-1]

    theta = np.arctan2(d2, d1)

    return d1.reshape(d1.size, 1), d2.reshape(d2.size, 1), theta.reshape(theta.size, 1)


# Let's calculate the slopes!
@numba.jit
def _calc_direction(data, mag, direction, ang, d1, d2, theta,
                    slc0, slc1, slc2):
    """
    This function gives the magnitude and direction of the slope based on
    Tarboton's D_\infty method. This is a helper-function to
    _tarboton_slopes_directions
    """
    s1 = (data[slc0] - data[slc1]) / d1
    s2 = (data[slc1] - data[slc2]) / d2
    sd = (data[slc0] - data[slc2]) / np.sqrt((d1 * d1 + d2 * d2))
    r = np.arctan2(s2, s1)
    rad2 = (s1)**2 + (s2)**2

    # Handle special cases
    # should be on diagonal
    I1 = ((s1 <= 0) & (s2 > 0)) | (r > np.arctan2(d2, d1))
    rad2[I1] = sd[I1] ** 2
    r[I1] = theta.repeat(I1.shape[1], 1)[I1]
    I2 = ((s1 > 0) & (s2 <= 0)) | (r < 0)  # should be on straight section
    rad2[I2] = s1[I2]**2
    r[I2] = 0
    I3 = (s1 <= 0) & ((s2 <= 0) | ((s2 > 0) & (sd <= 0)))  # upslope or flat
    rad2[I3] = -1

    I4 = rad2 > mag[slc0]
    mag[slc0][I4] = rad2[I4]
    direction[slc0][I4] = r[I4] * ang[1] + ang[0] * np.pi/2

    return mag, direction




