# -*- coding: utf-8 -*-
"""
Created on Sun Apr 11 17:20:16 2021
@author: neil
"""
# quick and dirty cellular automata for water flow
# The basic procedure is to place one unit of water on every cell of the DEM
# and with each iteration, move the water downhill to the lowest local cell
# Illustrates some nice graphics
# And this basic structure can be expanded to do landscape modelling, erosion etc.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D

print("Please wait, it takes a minute to produce the animation")
with open('manaslu.txt') as f:  # open a file 'descriptor'
    Z = np.loadtxt(f)           # read in an ASCII array as a DEM
# the 'with' statement opens and closes the file, and more importantly handles
# any errors that occur dealing with the operating system (example of good coding)

rows = np.size(Z,0)          # get the shape of the DEM
cols = np.size(Z,1)

fig = plt.figure()        # this next section allows us to see (plot) the DEM
ax = fig.add_subplot(111, projection='3d') # psuedo 3D plots require the Axes3D import
Y = np.linspace(0, 19000, rows)  # since we only have the Z (elevations) on a
X = np.linspace(0, 18000, cols)  # regular grid, we need to produce X,Y values
Xv,Yv = np.meshgrid(X, Y)        # 18000x19000 is the approx DEM size in meters
ax.plot_surface(-Xv,Yv,Z,cmap='terrain')  #  X is east west, Y is north south
plt.show()

fig2 = plt.figure()          # creates figure to plot the water flow
ax2  = fig2.add_subplot(111) # this will be an animation, which takes time to record

# now the actuall cellular automata
W  = np.ones_like(Z)       # the current locations of water on the landscape
                           # this puts 1 unit of water on every cell
Wn = np.zeros_like(Z)      # the next water positions
S  = np.zeros_like(Z)      # place to store the amount of water that passed
ims = []                   # a place to save the animation images

for t in range(100):       # this is the iteration, (should have a convergence crit?)   
    for i in range(1,rows-1):
        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,j], 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:        # move water to lowest neighbor (should randomize?)
                if d == 0:        # If a low spot, increase elevation to let the water
                    Z[i,j] += 2   #    flow out of the depression
                    Wn[i,j] = 1
                    S[i,j] += 1   # keep track of how many units of water passby
                elif d == 1:
                    Wn[i+1,j] = 1 # otherwise send the water units to the lowest
                    S[i+1,j] += 1 # nearby cell
                elif d == 2:
                    Wn[i,j+1] = 1
                    S[i,j+1] += 1
                elif d == 3:
                    Wn[i,j-1] = 1
                    S[i,j-1] += 1
                elif d == 4:
                    Wn[i-1,j] = 1
                    S[i-1,j] += 1
    im = ax2.spy(W,animated=True)  # make a single frame of the animation             
    ims.append([im])               # and add it to the list
    W[:] = Wn[:]                   # move new to old, to start next iteration
    Wn[:] = 0                      # clear the water units that have gone downhill

fig3 = plt.figure()                # creates separate figure to plot the water flow
ax3  = fig3.add_subplot(111, projection='3d')
ax.plot_surface(-Xv,Yv,S*5,cmap='Blues')
ax3.plot_surface(-Xv,Yv,S,cmap='Blues')
ax3.set_zlim([0,1000])

# now show the animation on figure 2, this takes a while to process  
im_ani = animation.ArtistAnimation(fig2, ims, interval=500, repeat_delay=1000,
                                   blit=True)