install.packages("ggplot2")
install.packages(c("ggalluvial", "ggridges"))
install.packages("ipumsr")
install.packages("srvyr")

library(ipumsr)
library(ggalluvial)
library(ggridges)
library(tidyverse)
library(srvyr)

#Make sure the data from the Stata code is in your R working directory to read it in
wide_dat <- read.csv("alluvial_wide.csv")
  
joint_tbl <- wide_dat %>% 
  filter(!is.na(gen_cat_1) & !is.na(gen_cat_2) & !is.na(gen_cat_3)) %>%
  group_by(geog) %>% 
  summarise(
    .groups = "keep",
    cur_data() %>% 
      as_survey_design(weight = fullpanelweight, id = eaid_1, strata = strata_1) %>%
      group_by(interact(gen_cat_1, gen_cat_2, gen_cat_3)) %>% 
      summarise(joint = survey_mean(prop = TRUE, prop_method = "logit", vartype = "ci"))
  )

joint_tbl2 <- joint_tbl %>% 
  rowid_to_column("alluvium") %>% 
  pivot_longer(
    c(gen_cat_1, gen_cat_2, gen_cat_3), 
    names_to = "x", 
    values_to = "stratum"
  ) %>% 
  mutate(x = ifelse(x == "gen_cat_1", "Phase 1", ifelse(x == "gen_cat_2", "Phase 2", "Phase 3"))) %>%   
  arrange(x, alluvium)

pma_alluvial <- function(
    title = "Fertility preferences across three longitudinal phases",     # an optional title 
    subtitle = "Four Sub-Saharan countries, 2019-2021",  # an optional subtitle 
    xaxis = NULL,     # an optional label for the x-axis (displayed below)
    yaxis = "Weighted percent"      # an optional label for the y-axis (displayed left)
){
  components <- list(
    theme(
      plot.title = element_text(size = 22, color = "#541E5A", 
                                hjust = 0.5, mar = margin(b = 5)),
      plot.subtitle = element_text(size = 15, hjust = 0.5, margin = margin(b = 20)),
      axis.text.x = element_text(size = 10, margin = margin(t = 5, b = 10)),
      strip.text.x = element_text(size = 14, margin = margin(b = 5)),
      plot.margin = margin(0, 100, 0, 100),
      legend.position = "bottom",
      legend.title = element_blank(),
      legend.spacing.x = unit(10, "pt"),
      legend.text=element_text(size=14),
      panel.grid = element_blank()
    ),
    labs(
      title = title,
      subtitle = subtitle,
      x = xaxis,
      y = str_wrap(yaxis, 10),
    ),
    scale_fill_manual(values = c(
      "Don't know" = "#B4B3B3", 
      "No more children" = "#4E4F71", 
      "Wants a child" = "#EFD372"
    )),
    scale_y_continuous(expand = c(0, 0))
  )
}

status_alluvial <- joint_tbl2 %>%
  ggplot(aes(
    x = x, 
    y = joint,
    fill = stratum,
    stratum = stratum,
    alluvium = alluvium
  )) + 
  geom_flow() + 
  geom_stratum(size = 0) + 
  facet_wrap(~geog, scales = "free_x", nrow = 2) +
  pma_alluvial()

status_alluvial 

status_alluvial2 <- joint_tbl2 %>% 
  filter(country == 7 | country == 9) %>%
  ggplot(aes(
    x = x, 
    y = joint,
    fill = stratum,
    stratum = stratum,
    alluvium = alluvium
  )) + 
  geom_flow() + 
  geom_stratum(size = 0) + 
  facet_wrap(~geog, scales = "free_x", nrow = 1) +
  pma_alluvial()