# Purpose:  Produce summary figure (Figure 1)
# Project:  McMahon, Baylis, Sweeney, and Funk (2025)
# Author:   Katie McMahon

#-------------------Base Setup--------------------------------------------------

#--clear environment
rm(list=ls())

#--load libraries
library(tidyverse)
library(janitor)
library(sf)
library(rworldxtra)

#--define directories
inputdir <- "YOUR_FOLDER/" # this inputdir should now be your `outdir` from the 03 script
outdir <- "YOUR_FOLDER/" # this is where you will store all tables and figures

#-------------------Read data---------------------------------------------------

#--load clean linked heat-DHS dataset
kid <- readRDS(paste0(inputdir, "merged_sa_dhs_clean.Rds"))

#--load shapefile of DHS clusters
gps <- read_sf(paste0(inputdir, "YOUR_SHAPEFILE.shp")) %>% 
  janitor::clean_names() %>% 
  filter(longnum != 0)

#-------------------Produce Figure 1--------------------------------------------

#--get country outlines and define map extent
data(countriesHigh) 
countriesHigh <- st_as_sf(countriesHigh) %>% 
  st_transform(st_crs(gps)) # make sure coordinate systems match
bbox <- st_bbox(gps)
new_bb <- c(66, 7, 97, 36) # define plot window in degrees
names(new_bb) <- c("xmin", "ymin", "xmax", "ymax")
attr(new_bb, "class") <- "bbox"
borders <- countriesHigh %>% 
  st_crop(new_bb + c(-2, -2, 4, 2))

#--plot panel A
both <- kid2 %>% 
  group_by(dhsid) %>% 
  summarize(
    com_ninemo_both = mean(ninemo_both) # aggregate to the cluster level
  ) %>% 
  ungroup() %>% 
  inner_join(gps,., by = "dhsid") %>%
  select(dhsid, geometry, com_ninemo_both, latnum) %>%
  drop_na(com_ninemo_both) %>% 
  mutate(z_both = (com_ninemo_both - mean(com_ninemo_both))/sd(com_ninemo_both)) %>% 
  filter(z_both < abs(3)) %>% # remove extreme outliers (more than 3 standard deviations from the mean)
  mutate(percent_both = (com_ninemo_both/275)*100) %>% # convert to percentage
  ggplot() + 
  geom_sf(aes(color = percent_both, alpha = 0.5)) +
  geom_sf(
    data = borders,
    fill = NA, colour = "dark gray",
    linewidth = 0.5,
    inherit.aes = FALSE
  ) +
  coord_sf(
    crs = st_crs(gps),
    xlim = new_bb[c(1,3)],
    ylim = new_bb[c(2,4)],
    expand = 1
  ) +
  scale_color_distiller(palette = "YlOrBr", direction = 1, limits = c(0, 50)) + 
  scale_alpha(guide = 'none') +
  theme_classic() +
  labs(color = "% of Days\nIn Utero",
       title = NULL,
       size = NULL,
       alpha = NULL) +
  theme(legend.key.size = unit(1.2, "cm"),
        legend.title = element_text(size = 18),
        legend.text = element_text(size = 14),
        axis.text = element_text(size = 14))

#--plot panel B
none <- kid2 %>% 
  group_by(dhsid) %>% 
  summarize(
    com_ninemo_none = mean(ninemo_none)
  ) %>% 
  ungroup() %>% 
  inner_join(gps,., by = "dhsid") %>%
  select(dhsid, geometry, com_ninemo_none, latnum) %>%
  drop_na(com_ninemo_none) %>% 
  mutate(z_none = (com_ninemo_none - mean(com_ninemo_none))/sd(com_ninemo_none)) %>% 
  filter(z_none < abs(3)) %>% 
  mutate(percent_none = (com_ninemo_none/275)*100) %>% 
  ggplot() + 
  geom_sf(aes(color = percent_none, alpha = 0.5)) +
  geom_sf(
    data = borders,
    fill = NA, colour = "dark gray",
    linewidth = 0.5,
    inherit.aes = FALSE
  ) +
  coord_sf(
    crs = st_crs(gps),
    xlim = new_bb[c(1,3)],
    ylim = new_bb[c(2,4)],
    expand = 1
  ) +
  scale_color_distiller(palette = "GnBu", direction = 1) + 
  scale_alpha(guide = 'none') +
  theme_classic() +
  labs(color = "% of Days\nIn Utero",
       title = NULL,
       size = NULL,
       alpha = NULL) +
  theme(legend.key.size = unit(1.2, "cm"),
        legend.title = element_text(size = 18),
        legend.text = element_text(size = 14),
        axis.text = element_text(size = 14))

#--plot panel C
tmax <- kid2 %>% 
  group_by(dhsid) %>% 
  summarize(
    com_ninemo_tmax_only = mean(ninemo_tmax_only)
  ) %>% 
  ungroup() %>% 
  inner_join(gps,., by = "dhsid") %>%
  select(dhsid, geometry, com_ninemo_tmax_only, latnum) %>%
  drop_na(com_ninemo_tmax_only) %>% 
  mutate(z_tmax_only = (com_ninemo_tmax_only - mean(com_ninemo_tmax_only))/sd(com_ninemo_tmax_only)) %>% 
  filter(z_tmax_only < abs(3)) %>% 
  mutate(percent_tmax_only = (com_ninemo_tmax_only/275)*100) %>% 
  ggplot() + 
  geom_sf(aes(color = percent_tmax_only, alpha = 0.5)) + 
  geom_sf(
    data = borders,
    fill = NA, colour="dark gray",
    linewidth = 0.5,
    inherit.aes = FALSE
  ) +
  coord_sf(
    crs = st_crs(gps),
    xlim = new_bb[c(1,3)],
    ylim = new_bb[c(2,4)],
    expand = 1
  ) +
  scale_color_distiller(palette = "OrRd", direction = 1, limits = c(0, 50)) + 
  scale_alpha(guide = 'none') +
  theme_classic() +
  labs(color = "% of Days\nIn Utero",
       title = NULL,
       size = NULL,
       alpha = NULL) +
  theme(legend.key.size = unit(1.2, "cm"),
        legend.title = element_text(size = 18),
        legend.text = element_text(size = 14),
        axis.text = element_text(size = 14))

#--plot panel D
wbgtmax <- kid2 %>% 
  group_by(dhsid) %>% 
  summarize(
    com_ninemo_wbgtmax_only = mean(ninemo_wbgtmax_only)
  ) %>% 
  ungroup() %>% 
  inner_join(gps,., by = "dhsid") %>%
  select(dhsid, geometry, com_ninemo_wbgtmax_only, latnum) %>%
  drop_na(com_ninemo_wbgtmax_only) %>% 
  mutate(z_wbgtmax_only = (com_ninemo_wbgtmax_only - mean(com_ninemo_wbgtmax_only))/sd(com_ninemo_wbgtmax_only)) %>% 
  filter(z_wbgtmax_only < abs(3)) %>% 
  mutate(percent_wbgtmax_only = (com_ninemo_wbgtmax_only/275)*100) %>% 
  ggplot() + 
  geom_sf(aes(color = percent_wbgtmax_only, alpha = 0.5)) + 
  geom_sf(
    data = borders,
    fill = NA, colour = "dark gray",
    linewidth = 0.5,
    inherit.aes = FALSE
  ) +
  coord_sf(
    crs = st_crs(gps),
    xlim = new_bb[c(1,3)],
    ylim = new_bb[c(2,4)],
    expand = 1
  ) +
  scale_color_distiller(palette = "BuPu", direction = 1, limits = c(0, 50)) + 
  scale_alpha(guide = 'none') +
  theme_classic() +
  labs(color = "% of Days\nIn Utero",
       title = NULL,
       size = NULL,
       alpha = NULL) +
  theme(legend.key.size = unit(1.2, "cm"),
        legend.title = element_text(size = 18),
        legend.text = element_text(size = 14),
        axis.text = element_text(size = 14))

#--combine panels
(none + both) /
  plot_spacer() /
  (tmax + wbgtmax) /
  plot_spacer() +
  plot_layout(heights = c(7, 1, 7, 1))

#--save output
ggsave(
  "Figure1.pdf",
  path = outdir,
  width = 18,
  height = 16,
  units = "in",
  dpi = 300
)