# -*- coding: utf-8 -*-
"""
@author: neil 2022
"""

# perform a random walk from an initial spike of 'temperature'
# minimal cellular automata model of heat flow decay of point source
# here we do stacking of several simulations

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

# define a method to smooth the results, a 3x3 box filter
def boxfilter(S1,size):
    """Accepts a square np.array, S1, with size, and performs a 2D boxcar
    filter of size fs, returns a float array, fs must be odd """
    # note the input can be an int or float array
    fs2 = 9       # number of nodes in filter, for normalizing at end
    S2 = np.zeros([size,size])  # init the filtered array to return
    S2[:,:] = S1                # copy input to output (for edge, where box
                                # does not reach)
    for i in range(1,size-1):
        for j in range(1,size-1):     #move boxcar over the interior of the array
            v = 0                     # count of values in box
            for ii in range(-1,2):    # iterate over the 2D box
                for jj in range(-1,2):
                    v += S1[i+ii,j+jj]# add values                    
            S2[i,j] = v/fs2           # normalize
    return S2                         # return the 9 cell averages

# define the problem, a 51x51 empty square grid, with a central pile of 10000 units
# the units represent temperature, the boundary conditions will be 0.
size  = 51          # size of space (squared)
iter1 = 40          # number of random steps to take, this is our 'time' variable
iter2 = 10          # number of realizations to stack for output
particles = 10000   # this is our temperature spike

# initial arrays
S = np.zeros([size,size],dtype='int')     # S and Sn are the starting and
Sn = np.zeros([size,size],dtype='int')    # ending arrays for one randomizing time step
Sb = np.zeros([size,size],dtype='int')    # a place to stack realizations
Sn[size//2,size//2]=particles              #put a 'pile' in the center of array

# now the CA loop: loop over 'iter2' realizations, of 'iter1' time steps, over the CA grid
for m in range(iter2):                    # do 'iter2' realizations
    Sn = Sn*0                             # zero-out storage array for new realization 
    Sn[size//2,size//2]=particles         # replace pile for each new realization
    
    for k in range(iter1):                # this is the random walk iteration, 'k' is 'time'
        S[:,:] = Sn                       # after one step put new into old
        for i in range(1,size-1):         # standard looping over x,y grid
            for j in range(1,size-1):     # the actual random walk
                r = S[i,j]                # number of particles at Si,j
                while r:                  # move r particles if any exist
                    x,y = np.random.randint(0,3,2)-1
                                          # calc the cell to go to  -1,0,1
                    Sn[i,j] -= 1          # remove the particle
                    Sn[i+x,j+y] += 1      # and put in new location
                    r -= 1
        Sn[0,:] = 0                       # remove particles that make it to the boundaries
        Sn[:,0] = 0                       # this will is the boundary condition, T = 0
        Sn[size-1,:] = 0
        Sn[:,size-1] = 0
                    
    Sb[:,:] += Sn[:,:]                    # sum each new realization
Sb = Sb/iter2
### plot the results, smooth over the field of particles (and realizations)
SS = boxfilter(Sb,size)  # filter the result to smooth (not necessary)

plt.spy(Sb)  # plot the raw output

# plot both a contour plot and a 3D view
xx,zz = np.meshgrid(np.arange(size),np.arange(size))                
fig4 = plt.figure()
ax4  = fig4.add_subplot(111)
cs = ax4.contour(xx,zz,SS,colors ='k')
ax4.contourf(xx,zz,SS)
ax4.clabel(cs,inline=1,fontsize=12)
ax4.axis('equal')
ax4.set_ylabel("Depth")
ax4.set_xlabel("Width in Meters")
ax4.set_title("Temperature field as a contour")

fig3 = plt.figure()
ax3  = fig3.add_subplot(111, projection='3d')
ax3.plot_surface(xx,zz,SS, rstride = 1, cstride = 1)
ax3.set_ylabel("Depth")
ax3.set_xlabel("Width in Meters")
ax3.set_title("Temperature field as a 3D surface")


plt.show()