N = 128
iterations = 50
noiseFct = .01
noiseSeed(1)

Kr, Ks, Kc, Ke = .01, .01, .01, .5 # Constants
w, m, wDiff, mDiff = [[[0 for y in xrange(N)] for x in xrange(N)] for n in xrange(4)] # Water and Material 2D arrays + corresponding differences arrays
h = [[noise(x * noiseFct, y * noiseFct) for y in xrange(N)] for x in xrange(N)] # 2D (Perlin) noised terrain

mat = ((-1, -1), (0, -1), (1, -1), 
       (-1, 0),           (1, 0), 
       (-1, 1),  (0, 1),  (1, 1)) # Moore Neighborhood

for iter in xrange(iterations):

    for i in xrange(1, N-1):
        for j in xrange(1, N-1):

            # Add precipitation
            w[i][j] += Kr

            # Remove sediment from terrain
            sediment = w[i][j] * Ks
            h[i][j] -= sediment
            m[i][j] += sediment

            # Compute differences with lower surrounding neighbors 
            a = h[i][j] + w[i][j] # Current altitude = current height + current amount of water
            aTotal = 0 
            dTotal = 0
            cells = 0
            for dx, dy in mat:
               x, y = (i + dx), (j + dy) # Coordinates of surrounding cells
                ai = h[x][y] + w[x][y] # altitude of ith neighbor
                d = a - ai # difference between current altitude and altitude of ith neighbor
                if d > 0.0: # if lower altitudes
                    dTotal += d # cumulating differences for later average computation
                    aTotal += ai # cumulating altitudes for later average computation
                    cells += 1


            if aTotal == 0: continue # Sometimes no cells with lower altitude are found

            # Water and Sediment distribution
            avgAlt = aTotal / cells # average of total heights
            deltaA = a - avgAlt
            waterAmount = min(w[i][j], deltaA)


            ### STEP 3 ###            
            for dx, dy in mat:
                x, y = (i + dx), (j + dy) # Coordinates of surrounding cells
                ai = h[x][y] + w[x][y] # altitude of ith neighbor
                di = a - ai # difference between current altitude and altitude of ith neighbor

                if di > 0: # if lower altitudes

                    deltaW = waterAmount * (di / dTotal) # water amount 

                    wDiff[i][j] += deltaW #storing water amount to be removed from central cell
                    w[x][y] += deltaW #water distribution to selected surrounding cells

                    deltaM = m[i][j] * (deltaW / w[i][j]) # material amount

                    mDiff[i][j] += deltaM #storing material amount to be removed from central cell
                    m[x][y] += deltaM #material distribution to selected surrounding cells

                    w[i][j] -= wDiff[i][j] #removing amount of water that has been distributed from central cell
                    m[i][j] -= mDiff[i][j] #removing amount of material that has been distributed from central cell


            #Evaporation
            mMax = Kc * w[i][j]
            deltaM = max(0, m[i][j] - mMax)
            m[i][j] -= deltaM
            h[i][j] += deltaM