# COSC/MATH 201 - Module 10.3: Spreading of Fire
# A cellular automaton: each cell is EMPTY, TREE, or BURNING, and every
# cell updates from its north, east, south, and west neighbors.
# Periodic boundaries: the grid wraps around at the edges.

EMPTY <- 0
TREE <- 1
BURNING <- 2

init_forest <- function(n, prob_tree) {
  forest <- matrix(EMPTY, nrow = n, ncol = n)
  forest[runif(n * n) < prob_tree] <- TREE  # one draw per cell
  forest
}

set.seed(1031)
small_forest <- init_forest(10, 0.7)
small_forest
sum(small_forest == TREE)

forest_colors <- c("wheat", "forestgreen", "orangered")  # EMPTY, TREE, BURNING

show_forest <- function(forest, title = "") {
  n <- nrow(forest)
  image(1:n, 1:n, t(forest[n:1, ]), col = forest_colors,
        zlim = c(EMPTY, BURNING), axes = FALSE, asp = 1,
        xlab = "", ylab = "", main = title)
}

show_forest(small_forest, "A 10 × 10 forest, density 0.7")

spread <- function(site, N, E, S, W, prob_immune, prob_lightning) {
  if (site == EMPTY) {
    return(EMPTY)
  }
  if (site == BURNING) {
    return(EMPTY)                      # a burning tree burns down
  }
  # site is a TREE
  if (N == BURNING || E == BURNING || S == BURNING || W == BURNING) {
    if (runif(1) < prob_immune) {
      return(TREE)                     # fire nearby, but immune this step
    } else {
      return(BURNING)
    }
  }
  if (runif(1) < prob_lightning * (1 - prob_immune)) {
    return(BURNING)                    # struck by lightning and not immune
  }
  TREE
}

spread(TREE, BURNING, EMPTY, TREE, TREE, 0, 0)  # burning neighbor to the north
spread(TREE, TREE, TREE, TREE, TREE, 0, 0)      # no fire nearby
spread(BURNING, TREE, TREE, TREE, TREE, 0, 0)   # burns down

extend_periodic <- function(forest) {
  n <- nrow(forest)
  forest_ns <- rbind(forest[n, ], forest, forest[1, ])    # north and south
  cbind(forest_ns[, n], forest_ns, forest_ns[, 1])        # east and west
}

tiny <- matrix(1:9, nrow = 3, byrow = TRUE)
tiny
extend_periodic(tiny)

step_forest <- function(forest, prob_immune, prob_lightning) {
  n <- nrow(forest)
  extended <- extend_periodic(forest)
  new_forest <- forest

  for (i in 1:n) {
    for (j in 1:n) {
      r <- i + 1
      c <- j + 1
      new_forest[i, j] <- spread(extended[r, c],
                                 extended[r - 1, c],  # north
                                 extended[r, c + 1],  # east
                                 extended[r + 1, c],  # south
                                 extended[r, c - 1],  # west
                                 prob_immune, prob_lightning)
    }
  }
  new_forest
}

set.seed(1032)
n <- 40
forest <- init_forest(n, 0.7)
center <- n / 2
forest[center, center] <- BURNING

number_of_steps <- 30
grids <- list(forest)              # grids[[1]] is time 0

for (t in 1:number_of_steps) {
  forest <- step_forest(forest, prob_immune = 0, prob_lightning = 0)
  grids[[t + 1]] <- forest
}

par(mfrow = c(2, 3), mar = c(1, 1, 2, 1))
for (t in c(0, 4, 8, 12, 20, 30)) {
  show_forest(grids[[t + 1]], paste("t =", t))
}

par(mfrow = c(1, 1))

fraction_burned <- function(n, prob_tree, prob_immune = 0) {
  forest <- init_forest(n, prob_tree)
  center <- ceiling(n / 2)
  forest[center, center] <- BURNING
  trees_at_start <- sum(forest != EMPTY)

  while (any(forest == BURNING)) {
    forest <- step_forest(forest, prob_immune, prob_lightning = 0)
  }
  (trees_at_start - sum(forest == TREE)) / trees_at_start
}

set.seed(1033)
densities <- seq(0.3, 0.9, by = 0.05)
runs_per_density <- 20
mean_burned <- numeric(length(densities))

for (k in seq_along(densities)) {
  results <- numeric(runs_per_density)
  for (run in 1:runs_per_density) {
    results[run] <- fraction_burned(30, densities[k])
  }
  mean_burned[k] <- mean(results)
}

data.frame(density = densities, mean_fraction_burned = round(mean_burned, 3))

plot(densities, mean_burned, type = "o", pch = 16, col = "orangered",
     ylim = c(0, 1), xlab = "Tree density (prob_tree)",
     ylab = "Mean fraction of trees burned",
     main = "One fire in a 30 × 30 forest (20 runs per density)")