# Module 6.3: Runge-Kutta 2 (Euler's predictor-corrector)
# Run the sections in order while coding along with the Quarto page.

# One large step for dP/dt = 0.1 P, starting at P(0) = 100.
old_population <- 100
growth_fraction <- 0.1       # per year
delta_t <- 8                 # years

start_rate <- growth_fraction * old_population
euler_prediction <- old_population + start_rate * delta_t
predicted_end_rate <- growth_fraction * euler_prediction
average_rate <- (start_rate + predicted_end_rate) / 2
rk2_estimate <- old_population + average_rate * delta_t

exact_value <- old_population * exp(growth_fraction * delta_t)
data.frame(method = c("Euler", "RK2", "Exact"),
           year_8_population = c(euler_prediction, rk2_estimate, exact_value))

# Plot the exact curve and the two one-step estimates.
smooth_time <- seq(0, 8, by = 0.1)
exact_curve <- 100 * exp(growth_fraction * smooth_time)

plot(smooth_time, exact_curve, type = "l", lwd = 2, col = "steelblue",
     xlab = "Years", ylab = "Population", ylim = c(90, 230))
segments(0, 100, 8, euler_prediction, col = "firebrick", lwd = 2)
segments(0, 100, 8, rk2_estimate, col = "darkgreen", lwd = 2)
points(c(8, 8), c(euler_prediction, rk2_estimate),
       col = c("firebrick", "darkgreen"), pch = 19)
legend("topleft", c("Exact curve", "Euler", "RK2"),
       col = c("steelblue", "firebrick", "darkgreen"),
       lty = 1, lwd = 2, bty = "n")

# Repeat RK2 with two-year steps.
initial_population <- 100
end_time <- 8
delta_t <- 2
time <- seq(0, end_time, by = delta_t)
population_rk2 <- numeric(length(time))
population_rk2[1] <- initial_population

for (i in 2:length(time)) {
  old_population <- population_rk2[i - 1]
  start_rate <- growth_fraction * old_population
  euler_prediction <- old_population + start_rate * delta_t
  predicted_end_rate <- growth_fraction * euler_prediction
  average_rate <- (start_rate + predicted_end_rate) / 2
  population_rk2[i] <- old_population + average_rate * delta_t
}

data.frame(time, population_rk2)

# Compare with Euler at the same step size.
population_euler <- numeric(length(time))
population_euler[1] <- initial_population

for (i in 2:length(time)) {
  old_population <- population_euler[i - 1]
  rate_now <- growth_fraction * old_population
  population_euler[i] <- old_population + rate_now * delta_t
}

plot(time, initial_population * exp(growth_fraction * time),
     type = "l", lwd = 2, col = "steelblue",
     xlab = "Years", ylab = "Population")
lines(time, population_euler, type = "b", col = "firebrick", pch = 19)
lines(time, population_rk2, type = "b", col = "darkgreen", pch = 17)
legend("topleft", c("Exact", "Euler", "RK2"),
       col = c("steelblue", "firebrick", "darkgreen"),
       lty = 1, pch = c(NA, 19, 17), lwd = 2, bty = "n")

data.frame(method = c("Euler", "RK2", "Exact"),
           year_8_population = c(population_euler[length(time)],
                                 population_rk2[length(time)],
                                 initial_population * exp(growth_fraction * end_time)))

# Book Quick Review Question 1: dP/dt = 10 + P/5.
old_population <- 500
delta_t <- 0.1

start_rate <- 10 + old_population / 5
euler_prediction <- old_population + start_rate * delta_t
predicted_end_rate <- 10 + euler_prediction / 5
corrected_population <- old_population +
  (start_rate + predicted_end_rate) / 2 * delta_t

data.frame(start_rate, euler_prediction,
           predicted_end_rate, corrected_population)