using GR, Random, Memoize

xmin = ymin = -3
xmax = ymax = 3
resx = resy = 400

setviewport(0, 1, 0, 1)
setwindow(xmin, xmax, ymin, ymax)
setcolormap(3)

# Hodge podge parameters
n = 100  # Number of states
k1 = 3   # Contagion for infected cells
k2 = 3   # Contagion factor for ill cells
g = 25   # Infection evolution speed

# Initialise grids - 2 for alternation
cells1 = Array{UInt32}(undef, resx, resy)
rand!(cells1, 0:n)
cells2 = copy(cells1)
cells = [cells1, cells2]

# Evolution functions
@inline isHealthy(t, i, j) = cells[t][i,j] == 0
@inline isInfected(t, i, j) = cells[t][i,j] > 0 && cells[t][i,j] < n
@inline isIll(t, i, j) = cells[t][i,j] >= n

@inline cx(x) = x <= 0 ? resx + x : x > resx ? x - resx : x
@inline cy(y) = y <= 0 ? resy + y : y > resy ? y - resy : y

@memoize neighbours(i, j) = [[cx(i - 1), cy(j - 1)], [cx(i), cy(j - 1)], [cx(i + 1), cy(j - 1)],
                [cx(i - 1), cy(j)], [cx(i + 1), cy(j)],
                [cx(i - 1), cy(j + 1)], [cx(i), cy(j + 1)], [cx(i + 1), cy(j + 1)]]

@inline a(t, i, j) = count(c -> isInfected(t, c[1], c[2]), neighbours(i,j))
@inline b(t, i, j) = count(c -> isIll(t, c[1], c[2]), neighbours(i,j))
@inline s(t, i, j) = sum(c -> cells[t][c[1], c[2]], neighbours(i,j)) + cells[t][i, j]

t = 1
while true

    # Display grid with scaled colour map (this slows down sumulation considerably)
    clearws()
    cellarray(xmin, xmax, ymin, ymax, resx, resy, floor.(8 .+ 64*cells[t]/n) )
    updatews()

    # Update next time step, alternating between in-memory grids
    tnext = (t == 1) ? 2 : 1
    for i in 1:resx
        for j in 1:resy
            # (a) If the cell is healthy (i.e., in state 0) then its new state is [a/k1] + [b/k2],
            # where a is the number of infected cells among its eight neighbors, b is the number of
            # ill cells among its neighbors, and k1 and k2 are constants. Here “[]” means the
            # integer part of the number enclosed, so that, for example, [7/3] = [2+1/3] = 2.
            if(isHealthy(t,i,j))
                cells[tnext][i,j] = floor(a(t,i,j) / k1) + floor(b(t,i,j) / k2)

            # (b) If the cell is ill (i.e., in state n) then it miraculously becomes healthy (i.e.,
            # its state becomes 0).
            elseif(isIll(t,i,j))
                cells[tnext][i,j] = 0

            # (c) If the cell is infected (i.e., in a state other than 0 and n) then its new state
            # is [s/(a+b+1)] + g, where a and b are as above, s is the sum of the states of the cell
            # and of its neighbors and g is a constant.
            else
                cells[tnext][i,j] = floor(s(t,i,j) / (a(t,i,j) + b(t,i,j) + 1)) + g
            end
        end
    end

    # Alternate grids
    global t
    t = tnext
end