Data_analysis_SPINZ.Rmd

---
title: "SPINZ genomic analysis"
author: "Laura T Phillips"
date: "`r format(Sys.time(), '%d %B, %Y')`"
output: 
  prettydoc::html_pretty:
    theme: cayman
    toc: true
---
## Set working directory and load packages
```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)

# set working directory to the location of the script
knitr::opts_knit$set(root.dir = getwd())
setwd(dirname(rstudioapi::getSourceEditorContext()$path))

# load packages
library(knitr)
library(tidyverse)
library(dplyr)
library(ggplot2)
library(ggrepel)
library(patchwork)
library(forcats)
library(ggtree)
library(ape)
library(readxl)
library(coda)
library(cowplot)
library(phytools)
library(lubridate)
library(waffle)
library(RColorBrewer)
library(dendextend)
library(janitor)
library(gtsummary)
library(rcompanion)
library(gtsummary)
library(ComplexUpset)
library(genoPlotR)

```
## Print session info (package versions)
```{r session_info}

# print versions
sessionInfo()

```
## Set colour palettes
```{r}

st_colours <- c('ST307' = '#a6cee3', 'ST2004' = '#fb9a99', 'ST983' = '#b2df8a','ST5856' = '#33a02c','ST101' = '#1f78b4' ,'ST985' = '#e31a1c', 'ST147' = '#fdbf6f','ST15' = '#ff7f00','ST1119' = '#cab2d6', 'Other' = '#fccde5')

born_loc <- c("Inborn" = "#d94701", "Unknown" = "#bdbdbd", "Outborn" = "#fdae6b")

clusters_cols <- c("1" = "#7fc97f", "2" = "#756bb1", "3" = "#fdc086", "4" = "#386cb0", "5" = "#f0027f", "6" = "#e41a1c", "Missing data" = "#999999", "NC" = "black")

res_colours <- c("Bla ESBL" =  '#fa9fb5', "Aminoglycosides" = '#66c2a4', "Macrolides" = '#807dba', "Sulfonamide" = '#fc8d59', "Trimethoprim" = '#0570b0', "Yersiniabactin" = '#cccccc', 'NA' = 'white')

all_res_colours <- c("Bla ESBL" =  '#fa9fb5', "Bla" = "#a6cee3", "Carbapenemase" = "#b2df8a", "Quinolone" = "#e31a1c", "Phenicol" = "#cab2d6", "Tetracycline" = "#ffff99", "Rifampicin" = "#b15928", "Aminoglycosides" = '#66c2a4', "Macrolides" = '#807dba', "Sulfonamide" = '#fc8d59', "Trimethoprim" = '#0570b0', "Yersiniabactin" = '#cccccc', 'NA' = 'white')

```
## Load data on included samples - Table S2
```{r}

kleb_clin <- read_tsv("tables/Supplementary_table_S2.tsv")

# Create ST grouped column 
kleb_clin <- kleb_clin %>%
  add_count(`Sequence type`, name="ST_count") %>%
  mutate(`Sequence type category` = ifelse(ST_count >=4, `Sequence type`, ifelse(`Sequence type` == "ST307-1LV", "ST307", 'Other'))) %>%
  mutate(days_to_onset_cat = case_when(
    is.na(`Days to onset`) ~ "Unknown",
    `Days to onset` < 0  ~ "Unknown",   
    `Days to onset` >= 8  ~ "8+ days",  
    TRUE ~ as.character(`Days to onset`))) %>%
  mutate(onset_group = case_when(
    days_to_onset_cat %in% c("0", "1", "2") ~ "Rapid",
    days_to_onset_cat == "Unknown" ~ "Unknown",
    TRUE ~ "Delayed"
  )) %>%
  select(-ST_count)

# Create onset category data - rapid and delayed onset 
kleb_clin <- kleb_clin %>%
  mutate(days_to_onset_cat = case_when(
    is.na(`Days to onset`) ~ "Unknown",
    `Days to onset` < 0  ~ "Unknown",   
    `Days to onset` >= 8  ~ "8+ days",  
    TRUE ~ as.character(`Days to onset`))) %>%
  mutate(onset_group = case_when(
    days_to_onset_cat %in% c("0", "1", "2") ~ "Rapid",
    days_to_onset_cat == "Unknown" ~ "Unknown",
    TRUE ~ "Delayed"
  ))

```
## Resistance data cleaning
```{r}

## Subset of resistance markers, for individual ST trees
resistance_subset <- subset(kleb_clin, select = c(`Read accession`, bla_esbl_acquired, bla_carb_acquired, a_gly_acquired, mls_acquired, sul_acquired, tmt_acquired, yersiniabactin, `Sequence type category`)) %>% mutate_all(~ gsub("\\.v[12]", "", gsub("[\\^\\*]", "", sub("[\\^\\*]+$", "", .)))) ##Removing any of the * or ^ and .v1 and .v2 suffixes, for the purpose of plotting

## All resistance markers for tree of all isolates
all_res <- subset(kleb_clin, select = c(`Read accession`, bla_esbl_acquired, bla_carb_acquired, bla_acquired, flq_acquired, phe_acquired, tet_acquired, rif_acquired, a_gly_acquired, mls_acquired, sul_acquired, tmt_acquired, yersiniabactin, `Sequence type category`)) %>% mutate_all(~ gsub("\\.v[12]", "", gsub("[\\^\\*]", "", sub("[\\^\\*]+$", "", .)))) ##Removing any of the * or ^ and .v1 and .v2 suffixes, for the purpose of plotting

# List column names which need splitting (by ;)
cols <- c("bla_esbl_acquired","bla_carb_acquired","a_gly_acquired","mls_acquired","sul_acquired","tmt_acquired")
all_cols <- c("bla_esbl_acquired","bla_carb_acquired","bla_acquired", "flq_acquired", "phe_acquired", "tet_acquired", "rif_acquired", "a_gly_acquired","mls_acquired","sul_acquired","tmt_acquired")

#Function to split columns by ; and give them their own column in a new dataframe
split_cols_precisely <- function(df, cols) { 
  unique_vals <- unique(unlist(strsplit(as.vector(as.matrix(df[, cols])), ";"))) # Trim white spaces and ensure unique values are correct 
  unique_vals <- trimws(unique_vals) 
  for (val in unique_vals) { 
    val <- trimws(val) # Create columns and ensure exact matching 
    df[[val]] <- apply(df[, cols], 1, function(x) ifelse(val %in% unlist(strsplit(as.character(x), ";")), val, NA)) # 1 0
  } 
  return(df) 
}

# Run split function and remove original unsplit columns
res_split_subset <- split_cols_precisely(resistance_subset, cols) %>% select(-one_of(cols)) %>%
  rename_with(~ ifelse(. %in% c("CTX-M-14", "CTX-M-15", "TEM-116", "NDM-5", "OXA-1", "TEM-1", "TEM-1D", "LAP-2"), paste0("bla-", .), .), everything())  # Add bla- prefix

res_split_all <- split_cols_precisely(all_res, all_cols) %>% select(-one_of(all_cols)) %>%
  rename_with(~ ifelse(. %in% c("CTX-M-14", "CTX-M-15", "TEM-116", "NDM-5", "OXA-1", "TEM-1", "TEM-1D", "LAP-2"), paste0("bla-", .), .), everything())  # Add bla- prefix
                                                                                                
```
## Table 1 data
```{r}

table_1 <- kleb_clin %>%
  mutate(`Sequence type` = ifelse(`Sequence type` == "ST307-1LV", "ST307", `Sequence type`)) %>%
  group_by(`Sequence type`, `K locus`, `K type`, `O locus`, `O type`) %>%
  summarise(count = n(), .groups = "drop") %>%
  group_by(`Sequence type`) %>%
  mutate(total_count = sum(count)) %>%
  arrange(desc(total_count), desc(count)) %>%
  select(-total_count)  # optional: remove the helper column

table_1
```
## Table 2 
```{r}

summary_tab <- kleb_clin %>%
  mutate(culture_date_status = if_else(!is.na(`Days to onset`), "Present", NA)) %>% 
  mutate(admission_date_status = if_else(!is.na(week_number), "Present", NA)) %>%
  mutate(`Birth location` = na_if(`Birth location`, "Unknown")) %>%
  mutate(days_to_onset_cat = case_when(
    is.na(`Days to onset`) ~ "Unknown",
    `Days to onset` < 0  ~ "Unknown",   
    `Days to onset` >= 8  ~ "8+ days",  
    TRUE ~ as.character(`Days to onset`)))

table_2 <- summary_tab %>%
  mutate(`Neonate gender` = recode(`Neonate gender`,
                                       `1` = "Female",
                                       `2` = "Male")) %>%
  mutate(`Outcome` = recode(`Outcome`,
                                       `0` = "Discharged",
                                       `1` = "Died")) %>%
  mutate(days_to_onset_cat = factor(as.character(days_to_onset_cat), levels = c("0", "1", "2", "3", "4", "5", "6", "7", "8+ days"))) %>%
  select(`Neonate gender`, admission_date_status, culture_date_status, `Outcome`, `Days to onset`, `Birth location`, `Study timepoint`, days_to_onset_cat, `Days to outcome`, `Days to onset`) %>%
  tbl_summary(by = `Study timepoint`, 
              include = c(`Neonate gender`,
    admission_date_status,
    culture_date_status,
    `Outcome`,
    `Birth location`,
    days_to_onset_cat,
    `Study timepoint`),
              sort = list(
                `Neonate gender` ~ "frequency",
                culture_date_status ~ "frequency",
                `Outcome` ~ "frequency",
                `Birth location` ~ "frequency"),
              #sort = all_categorical() ~ "frequency",
              type = all_continuous() ~ "continuous2",
              statistic = all_continuous() ~ c(
      "{N_nonmiss}",
      "{median} ({p25}, {p75})",
      "{min}, {max}"
    ),
  label = list(
      `Neonate gender` ~ "Neonate sex",
       admission_date_status ~ "Admission Date",
      culture_date_status ~ "Culture Date",
      `Outcome` ~ "Clinical outcome",
      `Birth location` ~ "Birth location",
      days_to_onset_cat ~ "Days to onset from admission"
    )) %>%
  add_n(statistic = "{N_nonmiss}", footnote = T) %>% ## Only show count of non-missing values
   bold_labels()

table_2
```
## Figure 1
```{r, warning=FALSE}

# Load pathogenwatch tree of all genomes
all_st_tree <- read.tree("output_trees/pathogenwatch_allst_nj_tree.nwk") %>% midpoint.root() 

# Change negative branch lengths to 0
all_st_tree$edge.length <- pmax(all_st_tree$edge.length, 0.0)

# Extract tree tipnames 
all_tips <- as.data.frame(get_taxa_name(ggtree(all_st_tree))) %>% setNames('all_tips')

# Create acquired resistance gene heatmap dataframe just for st307
all_res_data <- left_join(all_tips, res_split_all, by = c("all_tips" = "Read accession"))

# Clean up resistance data
all_res_data <- all_res_data %>%
  mutate_if(is.factor, as.character) %>%
  mutate(across(everything(), ~na_if(., "-"))) %>% #Replace all "-" entries with NA
  select_if(~ !all(is.na(.))) %>% #Remove columns that are entirely NA.
  arrange(factor(all_tips, levels = all_tips)) %>% #Reorder rows based on the all_tips vector
  mutate_all(as.factor)

# Assign each resistance gene to its drug class
allst_group_info <- data.frame(
  column = colnames(all_res_data), 
  group = c("None", "Yersiniabactin", "None", "Bla ESBL", "Bla ESBL", "Bla ESBL", "Carbapenemase", "Bla", "Bla", "Bla", "Bla", "Quinolone", "Quinolone", "Quinolone", "Quinolone", "Phenicol", "Phenicol", "Tetracycline", "Tetracycline", "Rifampicin", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Macrolides", "Macrolides", "Sulfonamide", "Sulfonamide", "Trimethoprim", "Trimethoprim", "Trimethoprim", "Trimethoprim", "Trimethoprim"))

allst_res_data_long <- all_res_data %>% pivot_longer(cols = -c(`Sequence type category`, all_tips), names_to = "column", values_to = "Genes") %>% left_join(allst_group_info, by = "column") %>%
   mutate(group = ifelse(!is.na(Genes), group, NA))

allst_res_data_long$all_tips <- factor(allst_res_data_long$all_tips, levels = rev(unique(all_tips$all_tips))) ## Needed to flip the y axis of the heatmap to align to tree tips 

# Heatmap
hm_allst <- allst_res_data_long %>%
  filter(!is.na(group)) %>%
  ggplot(aes(x = column, y = all_tips, fill = group)) + 
               geom_tile() +
facet_grid(. ~ group, scales = "free_x", space = "free_x") + 
scale_fill_manual(name = "", values = all_res_colours, na.value = 'white', na.translate = FALSE) + 
scale_x_discrete() + 
  scale_y_discrete() +  
  theme_minimal() +
  labs(title = "", x = "") +
  theme(axis.text.y = element_blank(), axis.text.x = element_text(angle = 60, hjust = 1, vjust = 1, size = 11), axis.title.y = element_blank(), legend.position = "top", strip.text = element_blank(), legend.text = element_text(size = 12)) 

all_t <- ggtree(all_st_tree) %<+% allst_res_data_long 
all_t_tree <- all_t + geom_tippoint(aes(colour=`Sequence type category`), shape = 19, size = 3) +
  #geom_tiplab(size = 1) +
  scale_colour_manual(values = st_colours) + theme(legend.position = "top", legend.justification.top = "right", legend.text = element_text(size = 12), legend.title = element_blank())

# Change order of ST data
all_t_tree$data$`Sequence type category` <- factor(all_t_tree$data$`Sequence type category`, levels=c("ST307", "ST2004", "ST983", "ST5856", "ST101", "ST985", "ST147", "ST15", "ST1119", "Other"))

# Plot
fig1 <- plot_grid(all_t_tree, hm_allst, align = "h", axis = "bt", ncol = 2, nrow = 1, rel_widths = c(1, 1.5))

fig1

#ggsave("figures/Figure_1.png", width=12, height=10)

```
## Figure 2
```{r}

# Figure 2A timeline of infections for all STs
timevcase_all_st <- kleb_clin %>% 
  filter(!is.na(week_number)) %>% 
  ggplot(aes(x = week_number, y = `Sequence type category`)) + 
  geom_count(aes(color = as.factor(`Sequence type category`))) +
  scale_colour_manual(values = st_colours, guide = "none") + 
  labs(title = '',
       x = 'Study week',
       y = '') +
  scale_y_discrete(name = "") + 
  geom_vline(xintercept = 31, colour = 'black') +
  #theme_minimal(base_size = 15) +
  theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1), legend.title = element_blank(), text = element_text(),  axis.title.y = element_blank()) +
  theme_minimal()

# Count the number of observations each study month per ST 
monthly_st_counts <- kleb_clin %>%
  filter(!is.na(month_number)) %>%
  group_by(month_number, `Sequence type category`) %>% 
  summarise(total_count = n(), .groups = "drop") 

# Calculate the total number of isolates per month across all ST categories
monthly_counts <- monthly_st_counts %>%
  group_by(month_number) %>%
  summarise(month_total = sum(total_count))

# Plot a bar chart of the number of infections by ST per study month
bar_chart_st <- ggplot(monthly_st_counts, aes(x = month_number, y = total_count, fill = `Sequence type category`)) +
  geom_bar(stat = "identity", alpha = 0.8) +
  scale_x_continuous(breaks = seq(min(monthly_st_counts$month_number, na.rm = TRUE), 
                                  max(monthly_st_counts$month_number, na.rm = TRUE), 
                                  by = 2)) +  # Show ticks every 2 months
  labs(fill = "") +
   ylim(0, 70) +
  scale_fill_manual(values = st_colours) +
  labs(title = "",
       x = "Study month", y = "") +
  geom_vline(xintercept = 7.5, colour = 'black') +
  theme_minimal() + #base_size = 15
theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1), legend.title = element_blank(),  axis.title.y = element_blank(), text = element_text())

# Add value labels for total number of infections per study month for all STs combined
bar_chart_st2 <- bar_chart_st +
  geom_text(data = monthly_counts,
            aes(x = month_number, y = month_total, label = month_total),
            vjust = -0.5, size = 3,
            inherit.aes = FALSE)



fig2 <- timevcase_all_st / bar_chart_st2 + plot_annotation(tag_levels = "A")

fig2
#ggsave("figures/Figure_2.png", width=8, height=8)

```
## Phylogenetic trees ST307, ST2004 and ST101
```{r, warning=FALSE}
st307 <- read.tree("output_trees/final_st307_bootstrapped_tree.tre") %>% midpoint.root() 

st307_tips <- as.data.frame(get_taxa_name(ggtree(st307))) %>% setNames('st307_tips')
st307_t <- ggtree(st307) #+ geom_tiplab(size = 1, align=F) 

st101 <- read.tree("output_trees/final_st101_bootstrapped_tree.tre") %>% midpoint.root() 

st101_tips <- as.data.frame(get_taxa_name(ggtree(st101))) %>% setNames('tips')
st101_t <- ggtree(st101) #+ geom_tiplab(size = 3, align = F)

st2004 <- read.tree("output_trees/final_st2004_bootstrapped_tree.tre") %>% midpoint.root() 

st2004_tips <- as.data.frame(get_taxa_name(ggtree(st2004))) %>% setNames('tips')
st2004_t <- ggtree(st2004) #+ geom_treescale(x = 0, y = 0, width = 5, fontsize = 3, linesize = 1)

```
## Figure 4 
```{r, warning=FALSE}

######### ST307 #########

# Filter data to just ST307
st307_kleb <- kleb_clin %>% 
  filter(`Sequence type category` == 'ST307') 

# Load cluster data ST307
st307_clusters <- read_csv("output_clusters/clusters_data_st307.csv") 

# Join cluster data and metadata, label non-clustered isolates and adjust cluster numbers
ordered_data307 <- st307_kleb %>%
  left_join(st307_clusters %>% select(id, Cluster), by = c(`Read accession` = "id")) %>% ########
  mutate(Cluster = case_when(
    grepl("ERR15165515", `Read accession`) ~ "NC", # Label only non-clustered sequence
    Cluster == 4 ~ "3",
    Cluster == 3 ~ "4",
    is.na(Cluster) ~ "Missing data",  # Assign "Missing data" to NA values
    TRUE ~ as.character(Cluster) # Keep other cluster values unchanged
  )) %>%
  mutate(`Read accession` = factor(`Read accession`, levels = st307_tips$st307_tips)) %>%
  arrange(`Read accession`) %>%
  mutate(`Read accession` = factor(`Read accession`, levels = rev(levels(`Read accession`)))) # flip the y axis to match tree labels

# Plot of infections over time, by cluster
plot307 <- ordered_data307 %>%
  ggplot(aes(x = week_number, y = `Read accession`)) +
  geom_point(aes(colour = as.factor(Cluster)), size = 2, shape = 21, fill = "white", stroke = 2) + ##shape = onset_group, size = 2.5
  scale_colour_manual(values = clusters_cols, name = "Cluster number") +
  #scale_shape_manual(values = c("Rapid" = 1, "Delayed" = 17, "Unknown" = 4), na.translate = TRUE, name = "Onset category") +
  labs(title = '',
       x = '') +
  scale_x_continuous(limits = c(0, 75), breaks = seq(0, 75, by = 8)) +
  theme_minimal() +
  geom_vline(xintercept = 31, colour = 'black') +
  theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1), axis.text.y = element_blank(), axis.title.y = element_blank(), panel.grid.minor.x = element_blank(),plot.title = element_text(hjust = 1), legend.position = "right", axis.title.x = element_blank())


# Create acquired resistance gene heatmap dataframe just for st307
st307_res_data <- left_join(st307_tips, res_split_subset, by = c("st307_tips" = "Read accession"))
st307_res_data[] <- lapply(st307_res_data, as.factor)
st307_res_data <- st307_res_data %>%
  mutate_if(is.factor, as.character) %>%
  select_if(~ !all(is.na(.))) %>%
  arrange(factor(st307_tips, levels = st307_tips)) %>%
  select(-all_of(c("-", "strA", "strB"))) %>% ## These genes are not relevant to resistance, and also columns are removed if cases of * or ^ have the same profile as the gene without it
  mutate_all(~ na_if(., '-')) %>% 
  mutate_all(as.factor)

# Assign each resistance gene to its drug class
st307_group_info <- data.frame(
  column = colnames(st307_res_data), 
  group = c("None", "Yersiniabactin", "None", "Bla ESBL", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Macrolides", "Macrolides", "Sulfonamide", "Sulfonamide", "Trimethoprim", "Trimethoprim", "Trimethoprim" ))

# Add additional metadata for a second heatmap
add_data_307 <- ordered_data307 %>%
  select(`Read accession`, `Birth location`, onset_group) %>%
  pivot_longer(cols = c(`Birth location`, onset_group),
               names_to = "column_2",
               values_to = "onset_location_groups")


st307_res_data_long <- st307_res_data %>% pivot_longer(cols = -c(`Sequence type category`, st307_tips), names_to = "column", values_to = "Genes") %>% left_join(st307_group_info, by = "column") %>%
  mutate(st307_tips = factor(st307_tips, levels = st307$tip.label)) %>% ## This is necessary in order to colour the heatmap by genes present but also the group they belong to
  mutate(group = ifelse(!is.na(Genes), group, NA)) 

st307_res_add_data_long <- left_join(st307_res_data_long, add_data_307, by = c("st307_tips" = "Read accession"))

st307_res_add_data_long$st307_tips <- factor(st307_res_add_data_long$st307_tips, levels = rev(unique(st307_tips$st307_tips))) ## Needed to flip the y axis of the heatmap to match tree tips


# resistance data heatmap
hm_plot <- ggplot(st307_res_add_data_long, aes(x = column, y = st307_tips, fill = group)) + 
               geom_tile() +
scale_fill_manual(name = "Drug class", values = res_colours, na.value = 'white', na.translate = FALSE) +
scale_x_discrete() + 
  scale_y_discrete() +  
  theme_minimal() +
  coord_fixed(ratio=0.1) +
  labs(title = "", x = "") +
  theme(axis.text.y = element_blank(), axis.text.x = element_text(angle = 45, hjust = 1, vjust = 1), axis.title.y = element_blank(), legend.position = "right", axis.title.x = element_blank()) 

# Reorder metadata legend order 
st307_res_add_data_long$onset_location_groups <- factor(
  st307_res_add_data_long$onset_location_groups,
  levels = c("Inborn", "Outborn", "Rapid", "Delayed", "Unknown")
)

# metadata heatmap
hm_plot2 <- ggplot(st307_res_add_data_long, aes(x = column_2, y = st307_tips, fill = onset_location_groups)) + 
               geom_tile() +
scale_fill_manual(values = c(
  "Inborn" = "black", 
  "Outborn" = "grey", 
  "Rapid" = "black", 
  "Delayed" = "grey",  
  "Unknown" = "white"), 
  labels = c(
    "Inborn" = "Inborn",
    "Outborn" = "Outborn",
    "Rapid" = "Rapid-onset",
    "Delayed" = "Delayed-onset",
    "Unknown" = "Unknown"
  ), name = "Metadata", na.value = 'white', na.translate = FALSE) +
scale_x_discrete(labels = c("onset_group" = "Onset category")) + 
  scale_y_discrete() +  
  coord_fixed(ratio=0.1) +
  theme_minimal() +
  labs(title = "", x = "") +
  theme(axis.text.y = element_blank(), axis.text.x = element_text(angle = 45, hjust = 1, vjust = 1), axis.title.y = element_blank(), legend.position = "right", axis.title.x = element_blank())

# Label tree tips with cluster number
st307_t$data <- left_join(st307_t$data, ordered_data307, by = c("label" = "Read accession")) 
st307_t_labs <- st307_t + geom_tippoint(aes(color=Cluster), size=2, shape = 21, fill = "white", stroke = 2) + scale_colour_manual(values = clusters_cols, name = "") + labs(title = 'ST307') + guides(color = "none") + theme_tree2() 

# Plot
fig4a <- (st307_t_labs + hm_plot + hm_plot2 + plot307) +
  plot_layout(ncol = 4) &
  theme(legend.position = "right")

######### ST101 #########

st101_kleb <- kleb_clin %>% 
  filter(`Sequence type category` == 'ST101')

### Load cluster data - this is based on SNP dists from coresnpfilter set to 80%
st101_clusters <- read_csv("output_clusters/clusters_data_st101.csv") 

st101_res_data <- left_join(st101_tips, res_split_subset, by = c("tips" = "Read accession"))
st101_res_data <- st101_res_data %>%
  select_if(~ any(. != 0)) %>%
  select(-all_of(c("-", "strB", "strA"))) %>%
  mutate_all(~ na_if(., '-')) %>% 
  mutate_all(as.factor)

# Assign each resistance gene to its drug class
st101_group_info <- data.frame(
  column = colnames(st101_res_data), 
  group = c("None", "Yersiniabactin", "None", "Bla ESBL", "Bla ESBL", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Macrolides", "Macrolides", "Sulfonamide", "Sulfonamide", "Trimethoprim", "Trimethoprim"))

ordered_data101 <- st101_kleb %>%
  left_join(st101_clusters %>% select(id, Cluster), by = c(`Read accession` = "id")) %>%
  mutate(Cluster = case_when(
    grepl("ERR6350131|ERR6350206|ERR6350791", `Read accession`) ~ "NC",   ## This removes the one sequence that isnt clustered
    Cluster == 2 ~ "1",
    Cluster == 1 ~ "2",
                          is.na(Cluster) ~ "Missing data", 
                          TRUE ~ as.character(Cluster))) %>%
  mutate(`Read accession` = factor(`Read accession`, levels = st101_tips$tips)) %>%
  arrange(`Read accession`) %>%
  mutate(`Read accession` = factor(`Read accession`, levels = rev(levels(`Read accession`)))) #this is essential to flip the y axis to match tree tip labels

# additional metadata
add_data_101 <- ordered_data101 %>%
  select(`Read accession`, `Birth location`, onset_group) %>%
  pivot_longer(cols = c(`Birth location`, onset_group),
               names_to = "column_2",
               values_to = "onset_location_groups")

st101_res_data_long <- st101_res_data %>% pivot_longer(cols = -c(`Sequence type category`, tips), names_to = "column", values_to = "Genes") %>% left_join(st101_group_info, by = "column") %>%
  mutate(group = ifelse(!is.na(Genes), group, NA))

st101_res_add_data_long <- left_join(st101_res_data_long, add_data_101, by = c("tips" = "Read accession"))

st101_res_add_data_long$tips <- factor(st101_res_add_data_long$tips, levels = rev(unique(st101_tips$tips))) #this is essential to flip the y axis to match tree tip labels

# Plot
plot101 <- ordered_data101 %>%
  mutate(onset_group = factor(`Birth location`, levels = c("Inborn", "Outborn", "unknown"))) %>%
  ggplot(aes(x = week_number, y = `Read accession`)) +
  geom_point(aes(colour = as.factor(Cluster)), size = 2, shape = 21, fill = "white", stroke = 2) + ## shape = onset_group, 
  scale_colour_manual(values = clusters_cols, name = "Cluster number") +
  labs(title = '',
       x = '',
       y = '') +
  scale_x_continuous(limits = c(0, 75), breaks = seq(0, 75, by = 8)) +
  theme_minimal() +
  geom_vline(xintercept = 31, colour = 'black') +
   theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1), axis.text.y = element_blank(), axis.title.y = element_blank(), panel.grid.minor.x = element_blank(),plot.title = element_text(hjust = 1), legend.position = "none", axis.title.x = element_blank()) 

# heatmap
hm101 <- ggplot(st101_res_data_long, aes(x = column, y = tips, fill = group)) + 
               geom_tile() +
scale_fill_manual(name = "Drug class", values = res_colours, na.value = 'white', na.translate = FALSE) + 
scale_x_discrete() + 
  scale_y_discrete() + 
  coord_fixed(ratio=1.0) +
  theme_minimal() + 
  labs(title = "", x = "") +
  theme(axis.text.y = element_blank(), axis.text.x = element_text(angle = 45, hjust = 1, vjust = 1), axis.title.y = element_blank(), legend.position = "none", axis.title.x = element_blank())

# Reorder metadata legend order 
st101_res_add_data_long$onset_location_groups <- factor(
  st101_res_add_data_long$onset_location_groups,
  levels = c("Inborn", "Outborn", "Rapid", "Delayed", "Unknown")
)

# metadata heatmap
hm101_2 <- ggplot(st101_res_add_data_long, aes(x = column_2, y = tips, fill = onset_location_groups)) + 
               geom_tile() +
scale_fill_manual(values = c(
  "Inborn" = "black", 
  "Outborn" = "grey", 
  "Rapid" = "black", 
  "Delayed" = "grey",  
  "Unknown" = "white"), 
  labels = c(
    "Inborn" = "Inborn",
    "Outborn" = "Outborn",
    "Rapid" = "Rapid-onset",
    "Delayed" = "Delayed-onset",
    "Unknown" = "Unknown"
  ), name = "Metadata", na.value = 'white', na.translate = FALSE) +
scale_x_discrete(labels = c("onset_group" = "Onset category")) + 
  scale_y_discrete() +  
  coord_fixed(ratio=1.0) +
  theme_minimal() +
  labs(title = "", x = "") +
  theme(axis.text.y = element_blank(), axis.text.x = element_text(angle = 45, hjust = 1, vjust = 1), axis.title.y = element_blank(), legend.position = "none", axis.title.x = element_blank())

st101_t$data <- left_join(st101_t$data, ordered_data101, by = c("label" = "Read accession")) 

# Label tree tips with cluster number
st101_t_labs <- st101_t + geom_tippoint(aes(color=Cluster), size=2, shape = 21, fill = "white", stroke = 2) + scale_colour_manual(values = clusters_cols, name = "") + labs(title = "ST101") +
  guides(color = "none") + theme_tree2() 

fig4b <- (st101_t_labs + hm101 + hm101_2 + plot101) +
  plot_layout(ncol = 4, ) &
  theme(legend.position = "none")

######### ST2004 #########

st2004_kleb <- kleb_clin %>% 
  filter(`Sequence type category` == 'ST2004')

st2004_clusters <- read_csv("output_clusters/clusters_data_st2004.csv") 

st2004_res_data <- left_join(st2004_tips, res_split_subset, by = c("tips" = "Read accession"))
st2004_res_data <- st2004_res_data %>%
  select_if(~ any(. != 0)) %>%
  select(-all_of(c("-", "strB", "strA"))) %>%
  mutate_all(~ na_if(., '-')) %>% 
  select_if(~ !all(is.na(.))) %>% 
  mutate_all(as.factor)


st2004_group_info <- data.frame(
  column = colnames(st2004_res_data), 
  group = c("None", "None", "Bla ESBL", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Sulfonamide", "Trimethoprim", "Trimethoprim"))

st2004_res_data_long <- st2004_res_data %>% pivot_longer(cols = -c(`Sequence type category`, tips), names_to = "column", values_to = "Genes") %>% left_join(st2004_group_info, by = "column") %>%
  mutate(group = ifelse(!is.na(Genes), group, NA))

ordered_data2004 <- st2004_kleb %>%
  left_join(st2004_clusters %>% select(id, Cluster), by = c(`Read accession` = "id")) %>%
  mutate(Cluster = case_when(Cluster == 2 ~ "1",  # Keep existing 2 values unchanged
                             Cluster == 1 ~ "2",
                           is.na(week_number) ~ "Missing data", ## remaining st2004 NAs should just be unclustered 
                           is.na(Cluster) ~ "NC",
                           TRUE ~ "1")) %>%               ## all other ST's are in 1 cluster per st
  mutate(`Read accession` = factor(`Read accession`, levels = st2004_tips$tips)) %>%
  arrange(`Read accession`) %>%
  mutate(`Read accession` = factor(`Read accession`, levels = rev(levels(`Read accession`)))) #this is essential to flip the y axis as atm its the opposite way round to the tree labels

# additional metadata
add_data_2004 <- ordered_data2004 %>%
  select(`Read accession`, `Birth location`, onset_group) %>%
  pivot_longer(cols = c(`Birth location`, onset_group),
               names_to = "column_2",
               values_to = "onset_location_groups")


st2004_res_add_data_long <- left_join(st2004_res_data_long, add_data_2004, by = c("tips" = "Read accession"))


st2004_res_add_data_long$tips <- factor(st2004_res_add_data_long$tips, levels = rev(unique(st2004_tips$tips))) #this is essential to flip the y axis to match tree tip labels

# Plot
plot2004 <- ordered_data2004 %>%
  ggplot(aes(x = week_number, y = `Read accession`)) +
  geom_point(aes(colour = as.factor(Cluster)), size = 2, shape = 21, fill = "white", stroke = 2) + 
  scale_colour_manual(values = clusters_cols, name = "Cluster number") +
  labs(title = '',
       x = 'Study week',
       y = '') +
  scale_x_continuous(limits = c(0, 75), breaks = seq(0, 75, by = 8)) +
  theme_minimal() +
  geom_vline(xintercept = 31, colour = 'black') +
   theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1), axis.text.y = element_blank(), axis.title.y = element_blank(), panel.grid.minor.x = element_blank(),plot.title = element_text(hjust = 1), legend.position = "none") 

# heatmap
hm2004 <- ggplot(st2004_res_add_data_long, aes(x = column, y = tips, fill = group)) + 
               geom_tile() +
#facet_wrap(~ group,  scales = "free_x", nrow = 1, ncol = 10) + 
scale_fill_manual(name = "Drug class", values = res_colours, na.value = 'white', na.translate = FALSE) + 
scale_x_discrete() + 
  scale_y_discrete() +  
  coord_fixed(ratio=0.1) +
  theme_minimal() + 
  labs(title = "", x = "Acquired resistance genes") +
  theme(axis.text.y = element_blank(), axis.text.x = element_text(angle = 45, hjust = 1, vjust = 1), axis.title.y = element_blank(), legend.position = "none")

# Reorder metadata legend order 
st2004_res_add_data_long$onset_location_groups <- factor(
  st2004_res_add_data_long$onset_location_groups,
  levels = c("Inborn", "Outborn", "Rapid", "Delayed", "Unknown")
)


# metadata heatmap
hm2004_2 <- ggplot(st2004_res_add_data_long, aes(x = column_2, y = tips, fill = onset_location_groups)) + 
               geom_tile() +
scale_fill_manual(values = c(
  "Inborn" = "black", 
  "Outborn" = "grey", 
  "Rapid" = "black", 
  "Delayed" = "grey",  
  "Unknown" = "white"), 
  labels = c(
    "Inborn" = "Inborn",
    "Outborn" = "Outborn",
    "Rapid" = "Rapid-onset",
    "Delayed" = "Delayed-onset",
    "Unknown" = "Unknown"
  ), name = "Metadata", na.value = 'white', na.translate = FALSE) +
scale_x_discrete(labels = c("onset_group" = "Onset category")) + 
  scale_y_discrete() +  
  coord_fixed(ratio=0.1) +
  theme_minimal() +
  labs(title = "", x = "Metadata") +
  theme(axis.text.y = element_blank(), axis.text.x = element_text(angle = 45, hjust = 1, vjust = 1), axis.title.y = element_blank(), legend.position = "none")

st2004_t$data <- left_join(st2004_t$data, ordered_data2004, by = c("label" = "Read accession")) 

# Label tree tips with cluster number
st2004_t_labs <- st2004_t + geom_tippoint(aes(color=Cluster), size = 2, shape = 21, fill = "white", stroke = 2) + scale_colour_manual(values = clusters_cols, name = "") + labs(title = "ST2004") + guides(color = "none") + theme_tree2()

# Plot
fig4c <- (st2004_t_labs + hm2004 + hm2004_2 + plot2004) +
  plot_layout(ncol = 4) &
  theme(legend.position = "none")

# Final joint plot
fig4_all <- fig4a / fig4b /fig4c + plot_layout(guides = "collect", heights = c(2, 2, 1))

fig4_all
#ggsave("figures/Figure_4.png", width=10, height=12)

```
## Supplementary figure 4
```{r, warning=FALSE}

######### 'Other' sequence types #########

# Add cluster data
other_st_kleb <- kleb_clin %>%
  filter(!`Sequence type category` %in% c('ST307', 'ST101', 'ST2004')) %>%
  filter(!is.na(week_number)) %>%
  mutate(Cluster = case_when(`Sequence type category` == 'Other' ~ "NC", 
                           TRUE ~ "1"))                

# Sort isolates by cluster 
other_st_kleb <- other_st_kleb %>%
  arrange(Cluster, `Read accession`) %>%  # Sort isolates by Cluster first
  mutate(`Read accession` = factor(`Read accession`, levels = unique(`Read accession`))) 

# Extract relevant IDs
other_st_kleb_ids <- other_st_kleb %>%
  select(`Read accession`, Cluster)

# Join resistance data
other_res_data <- left_join(other_st_kleb_ids, res_split_subset) %>% mutate_all(as.factor)

# Clean and filter resistance data
other_res_data <- other_res_data %>%
  mutate_if(is.factor, as.character) %>%
  select_if(~ any(. != 0)) %>%
  select(-all_of(c("-", "strB", "strA"))) %>%
  mutate_all(~ na_if(., '-'))

group_info_all_genes <- data.frame(
  column = colnames(other_res_data), 
  group = c("None", "None", "Yersiniabactin", "None", "Bla ESBL", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Macrolides", "Macrolides", "Sulfonamide", "Sulfonamide", "Trimethoprim", "Trimethoprim", "Trimethoprim", "Trimethoprim"))
  
other_res_data_long <- other_res_data %>% pivot_longer(cols = -c(`Sequence type category`, `Read accession`, Cluster), names_to = "column", values_to = "Genes") %>% left_join(group_info_all_genes, by = "column") %>%
  mutate(group = ifelse(!is.na(Genes), group, NA)) %>%
  filter(!is.na(`Read accession`)) 

# additional metadata
add_data_other <- other_st_kleb %>%
  select(`Read accession`, `Birth location`, onset_group) %>%
  pivot_longer(cols = c(`Birth location`, onset_group),
               names_to = "column_2",
               values_to = "onset_location_groups")

other_res_add_data_long <- left_join(other_res_data_long, add_data_other) %>% mutate_all(as.factor)
  
# Plot infections over time
plot_others <- other_st_kleb %>%
  ggplot(aes(x = week_number, y = `Read accession`)) +
  geom_point(aes(colour = as.factor(Cluster)), size = 2.5, shape = 21, fill = "white", stroke = 2) + ##shape = onset_group, 
  scale_colour_manual(values = clusters_cols, name = "Cluster number") +
  labs(title = '',
       x = 'Study week',
       y = '') +
  scale_x_continuous(limits = c(0, 75), breaks = seq(0, 75, by = 4)) +
  theme_minimal() +
  geom_vline(xintercept = 31, colour = 'black') +
  guides(colour = guide_legend(title = "Cluster number")) +
  theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1), axis.text.y = element_blank(), plot.title = element_text(hjust = 1, vjust = -1), legend.position = "none") 

# Facet plot by ST
plot_others <- plot_others +
  facet_grid(rows = vars(`Sequence type category`), scales = "free_y", space = "free_y") +
  theme(strip.text.y = element_text(size = 17, angle = 0), strip.placement = "outside")

other_res_add_data_long_sort <- other_res_add_data_long %>%
  arrange(Cluster, `Read accession`) %>%  # Sort isolates by Cluster first
  mutate(`Read accession` = factor(`Read accession`, levels = unique(`Read accession`))) # Convert to ordered factor  

# Plot AMR genes
other_st_amr <- ggplot(other_res_add_data_long_sort, aes(x = column, y = `Read accession`, fill = group)) + 
               geom_tile() +
scale_fill_manual(name = "Drug class", values = res_colours, na.value = 'white', na.translate = FALSE) + 
scale_x_discrete() + 
  scale_y_discrete() +
  #coord_fixed(ratio=0.1) +  
  facet_grid(rows = vars(`Sequence type category`), scales = "free_y", space = "free_y") +
  theme_minimal() + 
  labs(title = "", x = "Acquired resistance genes") +
  theme(axis.text.y = element_blank(), axis.text.x = element_text(angle = 75, hjust = 1, vjust = 1), axis.title.y = element_blank(), legend.position = "none", strip.text.y = element_blank()) 

# You need to subset other_res_add_data_long_sort so that it only includes rows relevant to the metadata columns you're plotting
other_hm2_data <- other_res_add_data_long_sort %>%
  filter(column_2 %in% c("onset_group", "Birth location")) 

# Reorder metadata legend order 
other_hm2_data$onset_location_groups <- factor(
  other_hm2_data$onset_location_groups,
  levels = c("Inborn", "Outborn", "Rapid", "Delayed", "Unknown")
)

# metadata heatmap
other_hm2 <- ggplot(other_hm2_data, aes(x = column_2, y = `Read accession`, fill = onset_location_groups)) + 
               geom_tile() +
scale_fill_manual(values = c(
  "Inborn" = "black", 
  "Outborn" = "grey", 
  "Rapid" = "black", 
  "Delayed" = "grey",  
  "Unknown" = "white"), 
  labels = c(
    "Inborn" = "Inborn",
    "Outborn" = "Outborn",
    "Rapid" = "Rapid-onset",
    "Delayed" = "Delayed-onset",
    "Unknown" = "Unknown"
  ),
  name = "Metadata", na.value = 'white', na.translate = FALSE) +
scale_x_discrete(labels = c("onset_group" = "Onset category")) + 
  scale_y_discrete() +  
  #coord_fixed(ratio=0.1) +  
  facet_grid(rows = vars(`Sequence type category`), scales = "free_y", space = "free_y") +
  theme_minimal() +
  labs(title = "", x = "Metadata") +
  theme(axis.text.y = element_blank(), axis.text.x = element_text(angle = 75, hjust = 1, vjust = 1), axis.title.y = element_blank(), strip.text.y = element_blank(), legend.position = "right")


fig_s4 <- (other_st_amr + other_hm2 + plot_others ) +
  plot_layout(ncol = 3, guides = "collect", widths = c(3, 0.5, 3)) &
  theme(legend.position = "right")

fig_s4
#ggsave("figures/Supplementary_figure_4.png", width=10, height=10)

```
## Stats and summary numbers 
```{r}
# Final number of high-quality K. pneumoniae genomes after all filtering
nrow(kleb_clin)

# Number of STs, with one case of ST307-1LV grouped into ST307 
kleb_clin %>% mutate(`Sequence type` = ifelse(`Sequence type` == "ST307-1LV", "ST307", `Sequence type`)) %>% group_by(`Sequence type`) %>% count() %>% nrow() 


# Proportions of each ST (of ≥4 genomes)
ST_prop <- kleb_clin %>%
  count(`Sequence type category`) %>% # counts occurrences of each ST
  mutate(proportion = n / sum(n)*100) # calculate proportion for each ST

# Count of K locus & type and O locus & type combinations, per ST
ko_tab <- kleb_clin %>%
  group_by(`Sequence type category`, `O locus`, `O type`, `K locus`, `K type`) %>%
  summarise(count = n(), .groups = "drop") %>%
  group_by(`Sequence type category`) %>%
  mutate(total_count = sum(count)) %>%
  arrange(desc(total_count), desc(count)) %>%
  select(-total_count)  # optional: remove the helper column


# Columns to select from dataframe of acquired resistance and virulence genes
acq_gene_columns <- c("a_gly_acquired", "flq_acquired", "mls_acquired", "phe_acquired", "rif_acquired", "sul_acquired", "tet_acquired", "tmt_acquired", "bla_acquired", "bla_esbl_acquired", "bla_carb_acquired", "yersiniabactin")

# Summary counts and proportions of cases with acquired resistance for each Abx, and virulence genes 
resistance_count <- kleb_clin %>%
  summarise(across(
    all_of(acq_gene_columns),
    ~sum(. != "-"),
    .names = "Genomes_with_acquired_resistance_genes_in_{.col}"
  )) %>%
  pivot_longer(everything(),
               names_to = "Column",
               values_to = "RowCount") %>%
  mutate(Proportion = round(RowCount / nrow(kleb_clin)*100, 1))

# Number of isolates with associated clinical metadata
kleb_clin %>% filter(!is.na(week_number)) %>% nrow() 

# Number of isolates with associated clinical metadata including a date of admission
kleb_clin %>% filter(!is.na(`Days to onset`)) %>% nrow() 

# Create a filtered dataframe of only isolates collected at study baseline 
stats_tab_subset_b <- kleb_clin %>%
  filter(!is.na(`Days to onset`)) %>%
  filter(`Study timepoint` == "Baseline") 
# Create a filtered dataframe of only isolates collected post-implementation 
stats_tab_subset_pi <- kleb_clin %>%
  filter(!is.na(`Days to onset`)) %>%
  filter(`Study timepoint` == "Post-implementation") 

# Days to onset summary stats for baseline cases
summary(stats_tab_subset_b$`Days to onset`)

# Days to onset summary stats for post-implementation cases
summary(stats_tab_subset_pi$`Days to onset`)

# Contingency table of days to onset for each timepoint
stats_tab_filt <- kleb_clin %>% filter(!days_to_onset_cat == "Unknown")

# Contingency table of Rapid or late onset infection for both timepoints
t_el <- table(stats_tab_filt$`Study timepoint`, stats_tab_filt$onset_group)
# Number of Rapid onset infections in each timepoint
onset_early <- c(t_el["Baseline", "Rapid"], t_el["Post-implementation", "Rapid"])

total_births <- rowSums(t_el) # total births for each timepoint

# Test whether the proportion of Inborn neonates changes between timepoints. 
prop.test(x = onset_early, n = total_births)

# Number of neonates with a recorded location of birth
kleb_clin %>% filter(!`Birth location` == "Unknown") %>% nrow() 

# Total neonates Inborn and outborn 
birth_loc_prop <- kleb_clin %>%
  filter(`Birth location` != "Unknown") %>% 
  count(`Birth location`) %>%                  
  mutate(proportion = n / sum(n)*100)

birth_loc <- kleb_clin %>% filter(!`Birth location` == "Unknown") 
# Contingency table of birth location for both timepoints
t1 <- table(birth_loc$`Study timepoint`, birth_loc$`Birth location`)

# Number of Inborn neonates in each timepoint
birth_loc_UTH <- c(t1["Baseline", "Inborn"], t1["Post-implementation", "Inborn"])
total_births <- rowSums(t1) # total births for each timepoint

# Test whether the proportion of Inborn neonates changes between timepoints. 
prop.test(x = birth_loc_UTH, n = total_births)

# Number of samples taken within the baseline and post-implementation periods
table(kleb_clin$`Study timepoint`)

##########################################

# Calculating the total, mean and median number of infections per ST, per study timepoint
# Baseline
weekly_infection_stats_baseline <- function(data, sequence_list, max_week = 31) {
  results <- lapply(sequence_list, function(st) {
    
    # Filter and summarize baseline data
    weekly_stats <- data %>%
      filter(`Sequence type category` == st) %>%
      filter(week_number <= max_week) %>%
      group_by(week_number, `Sequence type category`) %>%
      summarise(total_count = n(), .groups = "drop")
    
    # Create full week sequence
    min_week <- min(weekly_stats$week_number, na.rm = TRUE)
    max_week <- max(weekly_stats$week_number, na.rm = TRUE)
    full_weeks <- data.frame(week_number = seq(min_week, max_week))
  
    
    # Join and fill missing weeks with 0
    weekly_stats <- full_weeks %>%
      mutate(`Sequence type category` = st) %>%
      left_join(weekly_stats, by = c("week_number", "Sequence type category")) %>%
      mutate(total_count = ifelse(is.na(total_count), 0, total_count))
    
    # 🔍 Print the weekly_stats for inspection
    cat("\n--- Weekly Stats for:", st, "---\n")
    print(weekly_stats)
    
    # Return summary stats
   data.frame(
      Sequence = st,
      Sum = sum(weekly_stats$total_count),
      Median = median(weekly_stats$total_count),
      Mean = round(mean(weekly_stats$total_count),2),
      Week_Range = paste(min_week, max_week, sep = "–")
   )
  })
  
  #names(results) <- sequence_list
  #return(results)
  
  do.call(rbind, results)
}

sequence_list_baseline <- c('ST307', 'ST101', 'ST2004', 'ST985', 'ST147')
baseline_stats <- weekly_infection_stats_baseline(kleb_clin, sequence_list_baseline)

## Post-implementation
weekly_infection_stats_pi <- function(data, sequence_list, max_week = 31) {
  results <- lapply(sequence_list, function(st) {
    
    # Filter and summarize baseline data
    weekly_stats <- data %>%
      filter(`Sequence type category` == st) %>%
      filter(week_number > max_week) %>%
      group_by(week_number, `Sequence type category`) %>%
      summarise(total_count = n(), .groups = "drop")
    
    # Create full week sequence
    min_week <- min(weekly_stats$week_number, na.rm = TRUE)
    max_week <- max(weekly_stats$week_number, na.rm = TRUE)
    full_weeks <- data.frame(week_number = seq(min_week, max_week))
    
    # Join and fill missing weeks with 0
    weekly_stats <- full_weeks %>%
      mutate(`Sequence type category` = st) %>%
      left_join(weekly_stats, by = c("week_number", "Sequence type category")) %>%
      mutate(total_count = ifelse(is.na(total_count), 0, total_count))
    
    # Return summary stats
    data.frame(
      Sequence = st,
      Sum = sum(weekly_stats$total_count),
      Median = median(weekly_stats$total_count),
      Mean = round(mean(weekly_stats$total_count),2),
      Week_Range = paste(min_week, max_week, sep = "–")
    )
  })
  
  do.call(rbind, results)
}

sequence_list_pi <- c('ST307', 'ST101', 'ST2004', 'ST983', 'ST15')
pi_stats <- weekly_infection_stats_pi(kleb_clin, sequence_list_pi)

```
## Clustering summary numbers
```{r, warning=FALSE}

# Combine cluster data for each relevant ST 
clusters_df <- rbind(ordered_data307, ordered_data101, ordered_data2004, other_st_kleb) 

clusters_stats <- clusters_df %>%
  group_by(`Sequence type category`, Cluster) %>%
  filter(!Cluster %in% c('NC', 'Missing data')) %>%
  summarise(
    earliest_date = min(week_number, na.rm = TRUE),
    latest_date = max(week_number, na.rm = TRUE),
    count = n(),
    .groups = "drop"
  ) %>%
  mutate(cluster_duration = as.numeric(latest_date - earliest_date)) 

summary(clusters_stats) 

# Proportion of clustered sequences in the rapid/late onset groups
onset_cat_prop <- clusters_df %>%
  filter(!Cluster %in% c("Missing data", "NC")) %>%
  count(onset_group) %>%                  
  mutate(proportion = n / sum(n)*100)


# Proportion of cases belonging to clusters, or attributable to clusters in the baseline period and the post-implementation period` 

summary_nc_timepoint <- clusters_df %>% # non-clustered
  filter(Cluster == "NC",
         !is.na(`Study timepoint`)) %>%
  group_by(`Study timepoint`) %>%
  summarise(nc_count = n(), .groups = "drop")


summary_clustered_timepoint <- clusters_df %>% # clustered
  filter(!Cluster %in% c("Missing data", "NC"),
         !is.na(`Study timepoint`)) %>%
  group_by(`Study timepoint`) %>%
  summarise(total_count = sum(!is.na(Cluster)), .groups = "drop")

summary_2x2_timepoint <- full_join(summary_clustered_timepoint, summary_nc_timepoint, by = "Study timepoint") %>%
  mutate(
    Clustered = total_count,
    Non_clustered = nc_count,
    .keep = "unused"
  ) %>%
  select(`Study timepoint`, Clustered, Non_clustered)

# Number of successes (clustered)
tot_clustered <- summary_2x2_timepoint$Clustered
# Total cases per group
totals <- summary_2x2_timepoint$Clustered + summary_2x2_timepoint$Non_clustered

# Test whether the proportion of clustered cases differs significantly between the two study timepoints.
prop.test(tot_clustered, totals)


## For calculating attributable to a cluster, need to include cluster number
summary_clustered_timepoint_att <- clusters_df %>%
  filter(!Cluster %in% c("Missing data", "NC"),
         !is.na(`Study timepoint`)) %>%
  group_by(`Study timepoint`, Cluster, `Sequence type category`) %>%
  summarise(total_count = sum(!is.na(Cluster)), .groups = "drop")


summary_2x2_timepoint_attributable <- summary_clustered_timepoint_att %>%
  group_by(`Study timepoint`) %>%
  summarise(
    included = sum(total_count),
    not_included = n(),  # number of cluster labels per timepoint
    total = included - not_included,
    .groups = "drop"
  )

summary_2x2_timepoint_att <- full_join(summary_2x2_timepoint_attributable, summary_nc_timepoint, by = "Study timepoint") %>%
  mutate(
    Clustered = total,
    Non_clustered = nc_count + not_included, #The not included are the index cases for each cluster, which belong in the non-clustered category when looking at attributable to a cluster
    .keep = "unused"
  ) %>%
  select(`Study timepoint`, Clustered, Non_clustered)

# Number of successes (clustered)
tot_clustered_att <- summary_2x2_timepoint_att$Clustered
# Total cases per group
total_att <- summary_2x2_timepoint_att$Clustered + summary_2x2_timepoint_att$Non_clustered

# Test whether the proportion of clustered cases differs significantly between the two study timepoints.
prop.test(tot_clustered_att, total_att)

## logistic regression
# Univariate, association between birth location and outcome: clustered/non-clustered
# filter out any neonates with an unknown birth location
clusters_df_r_loc <- clusters_df %>%
  filter(!Cluster %in% c("Missing data"),
         !`Birth location` %in% c("Unknown")) %>%
  mutate(Cluster_r = ifelse(Cluster != "NC", 1, 0),
    `Birth location` = relevel(factor(`Birth location`), ref = "Outborn")
  ) 

model1 <- glm(Cluster_r ~ `Birth location`, family = binomial, data = clusters_df_r_loc)
model1_tab <- tbl_regression(model1, exponentiate = TRUE) %>%
  add_global_p()

# Univariate, association between onset category and outcome: clustered/non-clustered
# filter out any neonates with an unknown birth location
clusters_df_r_onset <- clusters_df %>%
  filter(!Cluster %in% c("Missing data"),
         !onset_group %in% c("Unknown")) %>%
  mutate(Cluster_r = ifelse(Cluster != "NC", 1, 0),
    onset_group = relevel(factor(onset_group), ref = "Delayed"),
  ) 

model2 <- glm(Cluster_r ~ onset_group, family = binomial, data = clusters_df_r_onset)
model2_tab <- tbl_regression(model2, exponentiate = TRUE) %>%
  add_global_p()



# Number of clustered (1) and non-clustered (0) neonates which were inborn/outborn 
clust_loc <- table(clusters_df_r_loc$`Birth location`, clusters_df_r_loc$Cluster_r)
# Row-wise proportions
clust_loc_prop <- prop.table(clust_loc, margin = 1)
round(clust_loc_prop, 3)

# Number of clustered (1) and non-clustered (0) neonates which were inborn/outborn and had rapid or delayed onset infection
clust_onset <- table(clusters_df_r_onset$onset_group, clusters_df_r_onset$Cluster_r)
clust_onset_prop <- prop.table(clust_onset, margin = 1)
round(clust_onset_prop, 3)


# Interaction between onset category and birth location and association with being clustered/non-clustered
clusters_df_r <- clusters_df %>%
  filter(!Cluster %in% c("Missing data"),
         !onset_group %in% c("Unknown"),
         !`Birth location` %in% c("Unknown")) %>%
  mutate(Cluster_r = ifelse(Cluster != "NC", 1, 0),
    onset_group = relevel(factor(onset_group), ref = "Delayed"),
    `Birth location` = relevel(factor(`Birth location`), ref = "Outborn")
  ) 

model3 <- glm(Cluster_r ~ onset_group*`Birth location`, family = binomial, data = clusters_df_r)
model3_tab <- tbl_regression(model3, exponentiate = TRUE) %>% add_global_p() 


```
## Figure 3 
```{r}

# Add a column which captures ST and cluster number combinations - then add cluster_timepoint column which captures the starting timepoint of the overall cluster (rather than some clusters having both baseline and post-implementation cases)
clusters_df <- clusters_df %>% mutate(ST_cluster=
                                        case_when(Cluster!="NC" & Cluster !="Missing data" ~ paste(`Sequence type`, Cluster),
                                                  TRUE ~ Cluster)) %>%
  group_by(ST_cluster) %>%
  mutate(
    cluster_timepoint = case_when(
      n_distinct(`Study timepoint`) > 1 & any(`Study timepoint` == "Baseline") ~ "Baseline",
      TRUE ~ `Study timepoint`
    )
  ) %>%
  ungroup()

# Summarise ST-clusters by onset group and birth location - baseline and post-implementation
clusters_stats_onset_b <- clusters_df %>%
  group_by(cluster_timepoint, ST_cluster, `Birth location`, onset_group) %>%
  filter(!Cluster %in% c('NC', 'Missing data')) %>%
  filter(cluster_timepoint %in% c('Baseline')) %>%
  summarise(count = n(), .groups = "drop_last") %>%
  mutate(prop = count / sum(count)*100) %>%
  ungroup()

clusters_stats_onset_pi <- clusters_df %>%
  group_by(ST_cluster, onset_group, `Birth location`) %>%
  filter(!Cluster %in% c('NC', 'Missing data')) %>%
  filter(cluster_timepoint %in% c('Post-implementation')) %>%
  summarise(count = n(), .groups = "drop_last") %>%
  mutate(prop = count / sum(count)*100) %>%
  ungroup()


# Plot ST-cluster combinations, with number of inborn/outborn and number of Rapid/late onset infections per cluster
p_baseline <- clusters_df %>%
  filter(cluster_timepoint %in% c("Baseline")) %>%
  filter(!ST_cluster %in% c("ST101 1", "Missing data", "NC")) %>% # Remove this cluster as data unknown for birth loc and onset cat - added to legend
  mutate(`Birth location`= fct_relevel(`Birth location`, "Inborn", "Outborn", "Unknown")) %>%
  mutate(onset_group=fct_relevel(onset_group, "Unknown", "Rapid", "Delayed")) %>%
  ggplot(aes(x = `Birth location`, fill = onset_group)) +
  #geom_histogram(position = "identity", alpha = 0.5, bins = 30, binwidth = 1, color = "black") +
  geom_bar(position = "stack") +
  scale_fill_manual(values = c("Rapid" = "#c51b8a", "Delayed" = "#08306b", "Unknown" = "grey"), name = "Onset category") +
  labs(x = "", y = "Count", title = "Baseline") +
  #scale_y_continuous(breaks = function(x) floor(seq(0, max(x), by = 1))) +
  #scale_y_continuous(labels = function(x) floor(x)) + #forces all labels to display as whole numbers
  facet_wrap(~ ST_cluster, nrow = 4, ncol = 3, scales = "free_y") +
  theme_bw() +
  theme(axis.text.x = element_text(angle=60, hjust=1)) 

  
p_pi <- clusters_df %>%
  filter(cluster_timepoint %in% c("Post-implementation")) %>%
  filter(!ST_cluster %in% c("ST101 1", "Missing data", "NC")) %>% # Remove this cluster as data unknown for birth loc and onset cat - added to legend
  mutate(`Birth location`= fct_relevel(`Birth location`, "Inborn", "Outborn", "Unknown")) %>%
  mutate(onset_group=fct_relevel(onset_group, "Unknown", "Rapid", "Delayed")) %>%
  ggplot(aes(x = `Birth location`, fill = onset_group)) +
  #geom_histogram(position = "identity", alpha = 0.5, bins = 30, binwidth = 1, color = "black") +
  geom_bar(position = "stack") +
  scale_fill_manual(values = c("Rapid" = "#c51b8a", "Delayed" = "#08306b", "Unknown" = "grey"), name = "Onset category") +
  labs(x = "Birth location", y = "Count", title = "Post-implementation") +
  #scale_y_continuous(breaks = function(x) floor(seq(0, max(x), by = 1))) +
  #scale_y_continuous(labels = function(x) floor(x)) + #forces all labels to display as whole numbers
  facet_wrap(~ ST_cluster, nrow = 4, ncol = 3, scales = "free_y") +
  theme_bw() +
  theme(legend.position = "none", axis.text.x = element_text(angle=60, hjust=1)) 


fig_3_new <- p_baseline / p_pi +
  plot_annotation(tag_levels = 'A') +  # Adds 'A', 'B' labels
  plot_layout(guides = "collect") &
  theme(legend.position = "bottom")    # Optional: shared legend position

fig_3_new
#ggsave("figures/Figure_3.png", width=8, height=8)


```
## Supplementary figure 3
```{r, warning=FALSE}

# Plot of ST307 clusters based on date of culture if available, and date of admission if not
plot307_a <- ordered_data307 %>%
  ggplot(aes(x = week_number, y = `Read accession`)) +
  geom_point(aes(colour = as.factor(Cluster)), size = 2, shape = 16) + ##shape = onset_group, size = 2.5
  scale_colour_manual(values = clusters_cols, name = "Cluster number") +
  #scale_shape_manual(values = c("Rapid" = 1, "Delayed" = 17, "Unknown" = 4), na.translate = TRUE, name = "Onset category") +
  labs(title = '',
       x = 'Study week') +
  scale_x_continuous(limits = c(0, 75), breaks = seq(0, 75, by = 4)) +
  theme_bw() +
  geom_vline(xintercept = 31, colour = 'black') +
  theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1), axis.text.y = element_blank(), axis.title.y = element_blank(), panel.grid.minor.x = element_blank(),plot.title = element_text(hjust = 1), legend.position = "right")

# Clusters based on isolates with a culture date available only
st307_clusters_cult <- read_csv("output_clusters/clusters_data_st307_culturedate_only.csv")

ordered_data307_cult <- st307_kleb %>%
  left_join(st307_clusters_cult %>% select(id, Cluster), by = c(`Read accession` = "id")) %>% ########
  mutate(Cluster = case_when(
    grepl("ERR15165515", `Read accession`) ~ "NC", # Removes non-clustered sequence
    Cluster == 2 ~ "6", # Relabel cluster names chronologically
    Cluster == 4 ~ "3",
    Cluster == 3 ~ "2",
    Cluster == 5 ~ "4",
    Cluster == 6 ~ "5",
    is.na(Cluster) ~ "Missing data",  # Assign "Missing data" to NA values
    TRUE ~ as.character(Cluster)                    # Keep other cluster values unchanged
  )) %>%
  mutate(`Read accession` = factor(`Read accession`, levels = st307_tips$st307_tips)) %>%
  arrange(`Read accession`) %>%
  mutate(`Read accession` = factor(`Read accession`, levels = rev(levels(`Read accession`)))) #this is essential to flip the y axis as atm its the opposite way round to the tree labels

plot307_b <- ordered_data307_cult %>%
  #mutate(onset_group = factor(onset_group, levels = c("UTH born", "Outborn", "unknown"))) %>%
  ggplot(aes(x = week_number, y = `Read accession`)) +
  geom_point(aes(colour = as.factor(Cluster)), size = 2, shape = 16) +
  scale_colour_manual(values = clusters_cols, name = "Cluster number") +
  labs(title = '',
       x = 'Study week') +
  scale_x_continuous(limits = c(0, 75), breaks = seq(0, 75, by = 4)) +
  theme_bw() +
  geom_vline(xintercept = 31, colour = 'black') +
  theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1), axis.text.y = element_blank(), axis.title.y = element_blank(), panel.grid.minor.x = element_blank(),plot.title = element_text(hjust = 1), legend.position = "right")

## supplementary figure 3
Figure_s3b <- plot307_b / plot307_a + plot_annotation(tag_levels = "A")

Figure_s3b

#ggsave("figures/Supplementary_figure_3.png", width=8, height=8)

```
## Genetic comparison of baseline and post-implementation strains - ST307, ST101, ST2004
```{r}
## Comparison of clustered cases across baseline and post-implementation period for ST307, ST2004 and ST101
st_clusters <- clusters_df %>% ## change ST of interest  - summarises all clusters for that ST
  filter(`Sequence type category` == c("ST307")) %>% ## change ST of interest  - summarizes all clusters for that ST
  group_by(Cluster, week_number) %>%
  summarise(weekly_count = n(), .groups = "drop") %>%
  group_by(Cluster) %>%
  summarise(
    median_weekly_count = median(weekly_count),
    week_range = paste0(min(week_number), "–", max(week_number)),
    cluster_total_count = sum(weekly_count),
    .groups = "drop"
  )


## Load SNV distance data for each ST - renamed to new ENA labels, and extract SNV distance summaries 
dist_mat_307 <- read.csv("output_snv_distances/st307_snp_dists.csv", row.names = 1, check.names = FALSE) %>% as.matrix()

cluster_post_imp_st307 <- clusters_df %>%
  filter(`Sequence type category` == "ST307", Cluster %in% c("4", "5")) %>%
  pull(`Read accession`)

cluster_baseline_st307 <- clusters_df %>%
  filter(`Sequence type category` == "ST307", Cluster %in% c("1", "2", "3")) %>%
  pull(`Read accession`)

# extract pairwise distances
dist_values_st307 <- dist_mat_307[as.character(cluster_post_imp_st307), as.character(cluster_baseline_st307)]
dist_vector_st307 <- as.vector(dist_values_st307)

# calculate stats for the SNP distances between baseline and post-implementation clustered sequences
summary_stats_st307 <- data.frame(
  min = min(dist_vector_st307, na.rm = TRUE),
  max = max(dist_vector_st307, na.rm = TRUE),
  median = median(dist_vector_st307, na.rm = TRUE),
  mean = mean(dist_vector_st307, na.rm = TRUE)
)

# snp-dists matrix ST101
dist_mat_101 <- read.csv("output_snv_distances/st101_snp_dists.csv", row.names = 1, check.names = FALSE) %>% as.matrix()

cluster_post_imp_st101 <- clusters_df %>%
  filter(`Sequence type category` == "ST101", Cluster %in% c("2")) %>%
  pull(`Read accession`)

cluster_baseline_st101 <- clusters_df %>%
  filter(`Sequence type category` == "ST101", Cluster %in% c("1")) %>%
  pull(`Read accession`)

# extract pairwise distances
dist_values_st101 <- dist_mat_101[as.character(cluster_post_imp_st101), as.character(cluster_baseline_st101)]
dist_vector_st101 <- as.vector(dist_values_st101)

# calculate stats for the SNP distances between baseline and post-implementation clustered sequences
summary_stats_st101 <- data.frame(
  min = min(dist_vector_st101, na.rm = TRUE),
  max = max(dist_vector_st101, na.rm = TRUE),
  median = median(dist_vector_st101, na.rm = TRUE),
  mean = mean(dist_vector_st101, na.rm = TRUE)
)

# snp-dists matrix ST2004
dist_mat_2004 <- read.csv("output_snv_distances/st2004_snp_dists.csv", row.names = 1, check.names = FALSE) %>% as.matrix()

cluster_post_imp_st2004 <- clusters_df %>%
  filter(`Sequence type category` == "ST2004", Cluster %in% c("1")) %>%
  pull(`Read accession`)

cluster_baseline_st2004 <- clusters_df %>%
  filter(`Sequence type category` == "ST2004", Cluster %in% c("2")) %>%
  pull(`Read accession`)

# extract pairwise distances
dist_values_st2004 <- dist_mat_2004[as.character(cluster_post_imp_st2004), as.character(cluster_baseline_st2004)]
dist_vector_st2004 <- as.vector(dist_values_st2004)

# calculate stats for the SNP distances between baseline and post-implementation clustered sequences
summary_stats_st2004 <- data.frame(
  min = min(dist_vector_st2004, na.rm = TRUE),
  max = max(dist_vector_st2004, na.rm = TRUE),
  median = median(dist_vector_st2004, na.rm = TRUE),
  mean = mean(dist_vector_st2004, na.rm = TRUE)
)

```
## Figure 5
```{r}

############### K-type ################################

# Only consider K types of those STs with >4 cases
prev_tab <- kleb_clin %>%
  filter(!(`Sequence type category` %in% c("Other"))) 

k_tab <- table(prev_tab$`K type`, prev_tab$`Study timepoint`)
fisher.test(k_tab, simulate.p.value = TRUE, B = 100000) ## p<0.001

k_prev <- prop.table(k_tab, margin = 2) # Proportion of each K type 

k_percentages_matrix <- matrix(
  c(0.87111111,0.08988764,
    0.00000000,0.00000000,
    0.01777778,0.05617978,
    0.03555556,0.43820225,
    0.02222222,0.00000000,
    0.00000000,0.04494382,
    0.01777778,0.00000000,
    0.03555556,0.24719101,
    0.00000000,0.00000000,
    0.00000000,0.12359551),
  ncol = 2,    # Two columns for Baseline and Post-implementation
  byrow = TRUE, # Fill by row
  dimnames = list(
    c("Capsule null", "K15", "K17", "K23", "K39", "K54", "K64", "unknown (KL102)", "unknown (KL110)", "unknown (KL127)"),
    c("Baseline", "Post_implementation")
  )
)

# Convert matrix to a data frame
k_df_percentages <- as.data.frame(k_percentages_matrix) 

# Remove cases of 0% prevalence at both baseline and post-implementation
k_df_percentages <- k_df_percentages %>% filter(!(Baseline == 0 & Post_implementation == 0))

# Add capsule type as a column (it's currently row names)
k_df_percentages$Capsule_Type <- rownames(k_df_percentages) 


k_df_long <- k_df_percentages %>%
  pivot_longer(
    cols = c(Baseline, Post_implementation), # Columns to pivot
    names_to = "Timepoint",                # New column for timepoint names
    values_to = "Percentage"               # New column for percentage values
)

k_df_long$Capsule_Type <- factor(k_df_long$Capsule_Type,
                               levels = k_df_percentages %>% arrange(desc(Baseline)) %>% pull(Capsule_Type))


k_df_long$Percentage_Display <- k_df_long$Percentage * 100

# To ensure 'Baseline' comes before 'Post_implementation' in the legend and grouping
k_df_long$Timepoint <- factor(k_df_long$Timepoint, levels = c("Baseline", "Post_implementation"))

k_prev_plot <- ggplot(k_df_long, aes(x = Capsule_Type, y = Percentage_Display, fill = Timepoint)) +
  geom_bar(stat = "identity", position = "dodge") + # 'dodge' creates grouped bars
  geom_text(aes(label = sprintf("%.1f", Percentage_Display)), # Add percentage labels
            position = position_dodge(width = 0.9), # Adjust position to fit dodged bars
            vjust = -0.5, size = 4) + # vjust for vertical adjustment, size for text
  labs(
    title = "K-locus",
    x = "",
    y = "Percentage (%)",
    fill = "Timepoint"
  ) +
  scale_fill_manual(values = c("Baseline" = "#67a9cf", "Post_implementation" = "#ef8a62"), labels = c("Baseline", "Post-implementation")) + # Custom colors
  scale_x_discrete(labels=c("Capsule null" = "KL102\n(Capsule null)", "K23" = "K23", "unknown (KL102)" = "KL102", "K39" = "K39", "K17" = "K17", "K64" = "K64", "K54" = "K54", "unknown (KL127)" = "KL127")) +
  theme_classic() + # A clean theme
  theme(
    axis.text.x = element_text(angle = 45,  hjust= 1), # Rotate x-axis labels
    #plot.title = element_text(hjust = 0.5, face = "bold"), # Center and bold title
    legend.position = "right"#, # Position the legend
    #text = element_text(size = 22)
  ) +
  ylim(0, max(k_df_long$Percentage_Display) * 1.15)  # Adjust y-axis limit to accommodate labels


############### O-type ################################

o_tab <- table(prev_tab$`O locus`, prev_tab$`Study timepoint`)
fisher.test(o_tab, simulate.p.value = TRUE, B = 100000) ## p<0.001

o_prev <- prop.table(o_tab, margin = 2)

percentages_matrix_o <- matrix(
  c(0.00000000,0.12359551,
    0.03555556,0.05617978,
    0.96444444,0.82022472,
    0.00000000,0.00000000,
    0.00000000,0.00000000),
  ncol = 2,    # Still two columns for Baseline and Post-implementation
  byrow = TRUE, # Fill by row
  dimnames = list(
    c("OL13", "OL2α.1", "OL2α.2", "OL3γ", "OL4"),
    c("Baseline", "Post_implementation") # Column names remain the same
  )
)


# Convert matrix to a data frame
df_percentages_o <- as.data.frame(percentages_matrix_o)

# Remove cases of 0% prevalence at both baseline and post-implementation
df_percentages_o <- df_percentages_o %>% filter(!(Baseline == 0 & Post_implementation == 0))

# Add capsule type as a column (it's currently row names)
df_percentages_o$Capsule_Type <- rownames(df_percentages_o)

df_o_long <- df_percentages_o %>%
  pivot_longer(
    cols = c(Baseline, Post_implementation), # Columns to pivot
    names_to = "Timepoint",                # New column for timepoint names
    values_to = "Percentage"               # New column for percentage values
)

df_o_long$Capsule_Type <- factor(df_o_long$Capsule_Type,
                               levels = df_percentages_o %>% arrange(desc(Baseline)) %>% pull(Capsule_Type))


df_o_long$Percentage_Display <- df_o_long$Percentage * 100

# To ensure 'Baseline' comes before 'Post_implementation' in the legend and grouping
df_o_long$Timepoint <- factor(df_o_long$Timepoint, levels = c("Baseline", "Post_implementation"))

o_prev_plot <- ggplot(df_o_long, aes(x = Capsule_Type, y = Percentage_Display, fill = Timepoint)) +
  geom_bar(stat = "identity", position = "dodge") + # 'dodge' creates grouped bars
  geom_text(aes(label = sprintf("%.1f", Percentage_Display)), # Add percentage labels
            position = position_dodge(width = 0.9), # Adjust position to fit dodged bars
            vjust = -0.5, size = 4) + # vjust for vertical adjustment, size for text
  labs(
    title = "O-locus",
    x = "",
    y = "Percentage (%)",
    fill = "Timepoint"
  ) +
  scale_fill_manual(values = c("Baseline" = "#67a9cf", "Post_implementation" = "#ef8a62"), labels = c("Baseline", "Post-implementation")) + # Custom colors
  theme_classic() + # A clean theme
  theme(
    axis.text.x = element_text(angle = 45,  hjust= 1), # Rotate x-axis labels
    #plot.title = element_text(hjust = 0.5, face = "bold"), # Center and bold title
    legend.position = "right"#, # Position the legend
    #text = element_text(size = 22)
  ) +
  ylim(0, max(df_o_long$Percentage_Display) * 1.15)  # Adjust y-axis limit to accommodate labels


##################### Acquired AMR genes #####################

# Count number of AMR genes for each Abx 
amr_genes_recoded <- kleb_clin %>%
  mutate(a_gly_acquired_r = if_else(a_gly_acquired == "-", 0, 
                                  if_else(str_count(a_gly_acquired, ";") > 0, 
                                          str_count(a_gly_acquired, ";") + 1, 1)),
  flq_acquired_r = if_else(flq_acquired == "-", 0, 
                                  if_else(str_count(flq_acquired, ";") > 0, 
                                          str_count(flq_acquired, ";") + 1, 1)),
  mls_acquired_r = if_else(mls_acquired == "-", 0, 
                                  if_else(str_count(mls_acquired, ";") > 0, 
                                          str_count(mls_acquired, ";") + 1, 1)),
  Phe_acquired_r = if_else(phe_acquired == "-", 0, 
                                  if_else(str_count(phe_acquired, ";") > 0, 
                                          str_count(phe_acquired, ";") + 1, 1)),
  Rif_acquired_r = if_else(rif_acquired == "-", 0, 
                                  if_else(str_count(rif_acquired, ";") > 0, 
                                          str_count(rif_acquired, ";") + 1, 1)),
  sul_acquired_r = if_else(sul_acquired == "-", 0, 
                                  if_else(str_count(sul_acquired, ";") > 0, 
                                          str_count(sul_acquired, ";") + 1, 1)),
  tet_acquired_r = if_else(tet_acquired == "-", 0, 
                                  if_else(str_count(tet_acquired, ";") > 0, 
                                          str_count(tet_acquired, ";") + 1, 1)),
  tmt_acquired_r = if_else(tmt_acquired == "-", 0, 
                                  if_else(str_count(tmt_acquired, ";") > 0, 
                                          str_count(tmt_acquired, ";") + 1, 1)),
  bla_acquired_r = if_else(bla_acquired == "-", 0, 
                                  if_else(str_count(bla_acquired, ";") > 0, 
                                          str_count(bla_acquired, ";") + 1, 1)),
  bla_esbl_acquired_r = if_else(bla_esbl_acquired == "-", 0, 
                                  if_else(str_count(bla_esbl_acquired, ";") > 0, 
                                          str_count(bla_esbl_acquired, ";") + 1, 1)),
  bla_carb_acquired_r = if_else(bla_carb_acquired == "-", 0, 
                                  if_else(str_count(bla_carb_acquired, ";") > 0, 
                                          str_count(bla_carb_acquired, ";") + 1, 1))
  )

# Summarise mean gene count by timepoint
amr_genes_tab <- amr_genes_recoded %>%
  filter(`Study timepoint` != "unknown") %>%
  group_by(`Study timepoint`) %>%
  summarise(across(
    .cols = c("a_gly_acquired_r", "flq_acquired_r", "mls_acquired_r", "Phe_acquired_r", "Rif_acquired_r", "sul_acquired_r", "tet_acquired_r", "tmt_acquired_r", "bla_acquired_r", "bla_esbl_acquired_r", "bla_carb_acquired_r"),  # or use c(count1, count2, ..., count10)
    .fns = list(
      mean = mean,
    count_present = ~sum(. > 0),
    prop_present = ~mean(. > 0) * 100
  ),
  .names = "{.col}_{.fn}"
  
))

# Reshape data for plotting
amr_genes_long <- amr_genes_tab %>%
  pivot_longer(
    cols = c("a_gly_acquired_r_mean", "flq_acquired_r_mean", "mls_acquired_r_mean", "Phe_acquired_r_mean", "Rif_acquired_r_mean", "sul_acquired_r_mean", "tet_acquired_r_mean", "tmt_acquired_r_mean", "bla_acquired_r_mean", "bla_esbl_acquired_r_mean", "bla_carb_acquired_r_mean"),
    names_to = "genes",
    values_to = "mean"
  )

amr_mean_plot <- ggplot(amr_genes_long, aes(x = genes, y = mean, fill = `Study timepoint`)) +
  geom_bar(stat = "identity", position = "dodge") +
   labs(
    title = "Mean number of acquired AMR resistance genes",
    x = "",
    y = "Mean",
    fill = "Timepoint"
  ) +
  scale_fill_manual(values = c("Baseline" = "#67a9cf", "Post-implementation" = "#ef8a62"), labels = c("Baseline", "Post-implementation")) + # Custom colors
  scale_x_discrete(labels = c("a_gly_acquired_r_mean" = "Aminoglycosides", "Col_acquired_r_mean" = "Colistin", "Fcyn_acquired_r_mean" = "Fosfomycin", "flq_acquired_r_mean" = "Fluoroquinolones", "Gly_acquired_r_mean" = "glycopeptides", "mls_acquired_r_mean" = "Macrolides", "Phe_acquired_r_mean" = "Phenicols", "Rif_acquired_r_mean" = "Rifampicin", "sul_acquired_r_mean" = "Sulfonamides", "tet_acquired_r_mean" = "Tetracyclines", "Tgc_acquired_r_mean" = "Tigecycline", "tmt_acquired_r_mean" = "Trimethoprim", "bla_acquired_r_mean" = "Penicillins", "bla_esbl_acquired_r_mean" = "3rd generation cephalosporins", "bla_carb_acquired_r_mean" = "Carbapenems")) +
  theme_classic() + # A clean theme
  theme(
    axis.text.x = element_text(angle = 60,  hjust= 1), # Rotate x-axis labels
    #plot.title = element_text(hjust = 0.5, face = "bold"), # Center and bold title
    legend.position = "right"#, # Position the legend
    #text = element_text(size = 22)
  ) 

k_o_prev_plots <- k_prev_plot / o_prev_plot / amr_mean_plot + plot_annotation(tag_levels = "A") + plot_layout(guides = "collect") & theme(legend.position = "bottom") 

k_o_prev_plots

#ggsave("figures/Figure_5.png", width=8, height=10)

```
## Supplementary figure 5
```{r, warning=FALSE}

# genbank file of reference KL102 locus, and example of a disrupted KL102 locus
seg1 <- read_dna_seg_from_genbank("data/KL102.gbk")
seg2 <- read_dna_seg_from_genbank("data/KL102_disrupted.gbk")

# colour the gene arrows
seg2$col <- "navy"
seg1$col <- "navy"

seg1$fill <- "navy"
seg2$fill <- ifelse(seg2$gene %in% c("wbaP*", "wzy*"), "maroon", "navy")

dna_segs <- list(seg2, seg1)

# annotate DNA segments
annot1 <- auto_annotate(seg1)
annot2 <- auto_annotate(seg2)

# Combine them into a list for the plotting function
annotations <- list(annot2, annot1)

comparisons <- read_comparison_from_blast("data/KL102_vs_disrupted.crunch.txt")

# colour the homology regions grey
comparisons$col <- "darkgray" 

#pdf("figures/Supplementary_figure_5.pdf", width=9,height=3)

plot_gene_map(
  dna_segs = dna_segs,
  comparisons = list(comparisons),
  gene_type = "arrows", 
  dna_seg_scale=T,
  annotations = annotations, 
  annotation_height = 2, 
  annotation_cex = 0.75,
  dna_seg_labels = list("a) Disrupted KL102", "b) Wildtype KL102")
)

dev.off()


```
## Supplementary figure 7
```{r, warning=FALSE}

# Trimmed two genetically distant and NC genomes from the tree 
st101_trimmed <- drop.tip(st101, c("ERR6350131", "ERR6350206"))

st101_t_trimmed <- ggtree(st101_trimmed) 

# Join to ST101 metadata
st101_t_trimmed$data <- left_join(st101_t_trimmed$data, ordered_data101, by = c("label" = "Read accession")) 

subset_st101_t <- st101_t_trimmed + geom_tippoint(aes(color=Cluster), size=2, shape = 21, fill = "white", stroke = 2) + scale_colour_manual(values = clusters_cols, name = "Cluster number") + theme_tree(text = element_text(size=15)) 

# Add custom labels indicating week of infection
custom_labels <- c("6", "6", "39", "40", "41", "44", "48")  # your labels
names(custom_labels) <- c("ERR6350777", "ERR6350762", "ERR15165532", "ERR15165534", "ERR15165536", "ERR6350761", "ERR6350791")  # matching tip labels

subset_st101_t$data$custom_label <- custom_labels[subset_st101_t$data$label]

# Extract SNV distances 
dist_mat_101_clean <- dist_mat_101[rownames(dist_mat_101) != c("ERR6350131","ERR6350206","Reference"), colnames(dist_mat_101) != c("ERR6350131","ERR6350206","Reference")]

dist_df <- as.data.frame(dist_mat_101_clean)

dist_long <- dist_df %>%
  rownames_to_column("ID1") %>%
  pivot_longer(-ID1, names_to = "ID2", values_to = "SNV_dist")

dist_long <- dist_long %>%
  filter(!grepl("Reference", ID1) & !grepl("Reference", ID2))

# Extract tip labels from tree
tip_order <- st101_t_trimmed$data %>%
  filter(isTip) %>%
  arrange(y) %>%
  pull(label)

# Reorder your heatmap rows and columns to match tip labels extracted from tree
dist_long$ID1 <- factor(dist_long$ID1, levels = tip_order)
dist_long$ID2 <- factor(dist_long$ID2, levels = tip_order)

# Plot
tree_hm <- subset_st101_t + geom_tiplab(aes(label = custom_label), nudge_x = 0.2, size = 3) + ggplot(dist_long, aes(x = ID2, y = ID1, fill = SNV_dist)) +
  geom_tile(width = 0.9, height = 0.9) +
  geom_text(aes(label = round(SNV_dist, 1)), size = 4, colour = "darkgrey") +
  scale_x_discrete(name = "Sample ID") +
  scale_fill_gradient(low = "white", high = "navy", name = "SNP distance") +
  theme_minimal() +
  theme(axis.text.x = element_text(angle = 90, hjust = 1),  axis.title.x = element_blank(), axis.title.y = element_blank(), axis.text.y.left = element_text(), legend.position = "none") +
  coord_cartesian(clip = "off")

tree_hm
#ggsave("figures/Supplementary_figure_7.png", width=8, height=4)

```
## Bonus AMR upset plot 
```{r}

# Load binary presence/absence of resistance genes data per genome
all_res_binary <- read_csv("data/res_subset_seperated_binary.csv") %>% select(-all_of(c("-"))) %>% select_if(~ any(. != 0)) %>% select(-all_of(c("strB.v1", "strA.v1"))) %>%
  select_if(~ !all(is.na(.))) %>% 
  mutate(across(everything(), as.character))

# Create a lookup table that maps resistance markers to an antibiotic class
abx_group_info <- data.frame(
  column = colnames(all_res_binary), 
  group = c("None", "None", "Bla ESBL", "Bla ESBL", "Bla ESBL", "Carbapenemase", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Aminoglycosides", "Macrolides", "Macrolides", "Sulfonamide", "Sulfonamide", "Trimethoprim", "Trimethoprim", "Trimethoprim", "Trimethoprim", "Trimethoprim", "Trimethoprim"))

# Reshape data to long format and join group info
all_res_binary <- all_res_binary %>%
  pivot_longer(-c('Read accession', 'Sequence type category'), names_to = "marker.type", values_to = "marker") %>%
  left_join(abx_group_info, by = c("marker.type" = "column"))


# Split the long-format data by antibiotic group
#abx_data_list <- split(all_res_binary, all_res_binary$group)

# Keep only groups with at least 2 unique markers
#abx_data_list <- abx_data_list[
#  sapply(abx_data_list, function(df) length(unique(df$marker.type)) >= 2)
#]


# generate_upset_plot <- function(marker_data, st_colours, abx_group_info, plot_title = "Group") {
#   
#   # Step 1: Transform to wide format
#   wide_data <- marker_data %>%
#     select(ID = `Read accession`, ST = `Sequence type category`, 
#            marker.type, present = marker, group) %>%
#     mutate(present = as.numeric(present)) %>%
#     pivot_wider(
#       id_cols = c(ID, ST),
#       names_from = marker.type,
#       values_from = present,
#       values_fill = 0
#     )
#   
#   # Step 2: Identify marker columns
#   marker_cols_indices <- 3:ncol(wide_data)
#   
#   # Step 3: Prepare data for ComplexUpset
#   st_data <- wide_data %>%
#     select(all_of(marker_cols_indices), ST) %>%
#     mutate(ST = factor(ST))
#   
#   set_names <- colnames(st_data)[-ncol(st_data)]  # exclude ST column
#   
#   label_colors <- abx_group_info %>%
#   filter(column %in% set_names) %>%
#   #arrange(group, column) %>%
#   mutate(color = all_res_colours[group]) %>%
#   select(column, color) %>%
#   deframe()
#   
#   # Step 4: Generate UpSet plot
#   p <- ComplexUpset::upset(
#     st_data,
#     set_names,
#     intersections = "observed",
#     min_size = 3,
#     sort_sets = FALSE,
#     
#     base_annotations = list(
#       'Percent (and count) of genomes' = list(
#         aes = aes(x = intersection, fill = ST),
#         geom = list(
#           geom_bar(stat = 'count', position = 'fill'),
#           theme(axis.title = element_text(size = 10)),
#           scale_fill_manual(values = st_colours),
#           scale_y_continuous(labels = scales::percent_format()),
#           geom_text(aes(label = ..count..), stat = "count", 
#                     position = position_fill(vjust = 0.5), angle = 60, size = 3)
#         )
#       )
#     ),
#     
#     set_sizes = FALSE,
#     name = plot_title,
#     width_ratio = 0.7
#   )
# 
# 
#   
#   # Step 6: Apply custom label colors to x-axis text
#   p <- p + theme(
#     axis.text.y = element_text(
#     #colour = label_colors[set_names],  # match order of sets
#     angle = 0,
#     hjust = 1,
#     size = 9
#     )
#   )
#  
#   set_labels <- ggplot_build(p)$data[[2]]$label
#   # Step 7: Return everything you want to inspect
#   #return(p)
#   return(list(
#     plot = p,
#     label_colors = label_colors,
#     set_names = set_labels,
#     st_data = st_data
#   ))
# 
# }
# 
# 
# upset_amr <- generate_upset_plot(all_res_binary, st_colours, abx_group_info, plot_title = "Groups")


```