# Generates assets/cohort.gif for the ggpop page.
#
# Run from this directory:  Rscript make-cohort-gif.R
#
# Not a page chunk on purpose: calling gganimate::animate() inside a knitr
# chunk hangs the render (it completes in ~30s standalone), so the animation
# is generated here and committed. The code below is what the page displays.

suppressPackageStartupMessages({
  library(dplyr); library(ggplot2); library(ggpop); library(gganimate)
})

p_HS <- 0.045; p_HD <- 0.006; p_SH <- 0.02; p_SD <- 0.05
ages  <- 40:100
trace <- matrix(0, length(ages), 3,
                dimnames = list(ages, c("Healthy", "Sick", "Dead")))
trace[1, ] <- c(1, 0, 0)
for (i in 2:length(ages)) {
  h <- trace[i - 1, "Healthy"]; s <- trace[i - 1, "Sick"]; d <- trace[i - 1, "Dead"]
  trace[i, "Healthy"] <- h * (1 - p_HS - p_HD) + s * p_SH
  trace[i, "Sick"]    <- s * (1 - p_SH - p_SD) + h * p_HS
  trace[i, "Dead"]    <- d + h * p_HD + s * p_SD
}

# one fixed 100-icon circle, ordered bottom row first
grid100 <- fetch_df_coordinates() |>
  filter(size == 100) |>
  arrange(y1, x1) |>
  select(px = x1, py = y1)

# fill it from the bottom up: Dead, then Sick, then Healthy
cohort <- lapply(ages, function(a) {
  sh <- trace[as.character(a), ]
  n_dead <- round(sh["Dead"] * 100); n_sick <- round(sh["Sick"] * 100)
  grid100 |>
    mutate(age   = a,
           state = c(rep("Dead", n_dead), rep("Sick", n_sick),
                     rep("Healthy", 100 - n_dead - n_sick)),
           icon  = case_when(state == "Healthy" ~ "person",
                             state == "Sick"    ~ "person-cane",
                             state == "Dead"    ~ "skull"))
}) |> bind_rows()

p_anim <- ggplot(cohort, aes(x = px, y = py, icon = icon, color = state)) +
  geom_icon_point(size = 1.5, dpi = 72) +
  scale_color_manual(values = c(Healthy = "#81C784", Sick = "#FFB74D",
                                Dead    = "#6B7C8F"),
                     breaks = c("Healthy", "Sick", "Dead")) +
  coord_fixed() +
  theme_pop_dark(bg_color = "#000001", legend_position = "bottom") +
  labs(title = "One cohort of 100, followed from age 40",
       subtitle = "Age {closest_state}", color = NULL) +
  transition_states(age, transition_length = 0, state_length = 1)

anim <- animate(p_anim, nframes = length(ages), fps = 12, end_pause = 12,
                width = 520, height = 560,
                renderer = gifski_renderer(), bg = "#000001")
anim_save("assets/cohort.gif", anim)

cat("wrote assets/cohort.gif\n")
print(round(trace[c("40", "60", "80", "100"), ], 3))