# -*- coding: utf-8 -*-
"""
Created on Wed Apr 29 15:29:34 2020

@author: neil
"""

# this is a skeleton of the core CA routine to route water on the DEM
# at this point we have the DEM in the array Z

# now the actuall cellular automata
W  = np.ones_like(Z)       # the current water positions, at the start we put one 'unit' of water
                           # on every cell of the grid W which is the same size as Z
Wn = np.zeros_like(Z)      # An array to move the water to, in the next iteration


for t in range(200):       # this is an iteration, should really have a convergence crit.
                           # but we know that it can't take more than 200 steps to get to the edge
                           # of a 200x200 grid, unless the water gets stuck
                           
    for i in range(1,rows-1):      # only interior, not boundary notes
        for j in range(1,cols-1):  # apply our automaton to all interior nodes
            # find lowest nearest neighbour, numbered 1 colckwise from right
            d = np.argmin([Z[i+1,j], Z[i,j+1], Z[i,j-1], Z[i-1,j] ])
            # argmin returns the index of the minimum in the array

            if W[i,j] > 0:        # if water is at a location W>0, move water to lowest neighbor
                if d == 1:        # note we ignore the case where i,j is a local depression!!
                    Wn[i+1,j] = 1
                elif d == 2:
                    Wn[i,j+1] = 1
                elif d == 3:
                    Wn[i,j-1] = 1
                elif d == 4:
                    Wn[i-1,j] = 1

    W[:] = Wn[:]                   # move new to old, to start next iteration
    Wn[:] = 0                      # and zero out the space for the new water to go

# BIG NOTE, we haven't written the code to show the results !!!