# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 1.6.2
#     language: julia
#     name: julia-1.6
# ---

# ## The simulation and duck typing
#
# You might have noticed that the simulation on the first page has some dark blue cells at the beginning. These cells are immune. We did not have any immune cells in the notebook, though. Let's add this final piece to the program.
#
# The cell below will load the simulation as we left it.

include("../../epidemic_no_immunity.jl")

# Write a new cell type called where the `status` of a cell can be `immune` in addition to the states we had before.

# +
@enum InfectionStatus uninfected infected dead recovered immune

mutable struct Cell
    status::InfectionStatus
    infection_time::Int8
end
# -

# Modify the loop below to make roughly 20% of the cells immune.

# +
# About 20% of the cells should be immune
immunity_rate = 0.2

# Create an array of cells, with each cell undefined at first
cells = Array{Cell}(undef, 32, 32)

# Now loop over the positions in the array and create either an
# uninfected cell or an immune cell
for i in 1:size(cells)[1]
    for j in 1:size(cells)[2]
        if rand(1)[1] < immunity_rate
            cells[i,j] = Cell(immune, 0)
        else
            cells[i,j] = Cell(uninfected, 0)
        end
    end
end

# Infect one cell at the start
cells[16,16].status = infected
# -

# We don't need to change the `update!`, since only `uninfected` cells can get infected. We can just run the simulation and thanks to duck-typing, everything works.

parameters = Parameters(0.2, 0.001, 8, 0.1)
update!(cells, parameters)

# Plotting does not work immediately, however. Since `to_colors` does not know about immune cells, it will not return anything for them.
#
# Update to_colors function below to use a different color for immune cells.

# +
"Map the cells to colors for plotting"
function to_colors(cell)
    if cell.status == uninfected
        return RGB(0.0,0.8,0.0)
    end
    if cell.status == infected
        return RGB(0.8,0.0,0.0)
    end
    if cell.status == dead
        return RGB(0.1,0.1,0.1)
    end
    if cell.status == recovered
        return RGB(0.0,0.0,0.8)
    end
    if cell.status == immune
        return RGB(0.2,0.2,1.0)
    end
end

to_colors.(cells)

StarcoderdataJupyterScriptsDedupFiltered