TSR-Proj / Figures_Manuscript.r
Figures_Manuscript.r
Raw
#############################################################################
# Script to create the images present in the supplementary and main text
#############################################################################

# Figures for main text ----
## Loading all necessary libraries ----
rm(list = ls())
library(cowplot)
library(patchwork)
library(scales)
library(ggh4x)
library(latex2exp)
library(viridis)
library(pomp)
library(tidyverse)

## Figure 1: Causal diagram for the transmission model ----
# This was done via LaTeX
## Figure 2: Workflow figure ----
### Figure 2.1: Sampled parameter space ----
#### First sample parameter space ----
mu_year <- 1 / 80
parameters <- sobol_design(lower = c(A = 1, alpha_year = 1/20), upper = c(A = 10, alpha_year = 1), nseq = 50)
parameters <- parameters |>
  dplyr::mutate(R0 = 1 + (1 + 1 / (alpha_year + mu_year)) / (A * (mu_year + 1))) 
parameters <- parameters |> dplyr::mutate(alpha = alpha_year / 52) |>
  dplyr::mutate(eps = rep(1,nrow(parameters))) |>
  dplyr::mutate(.id = seq_len(nrow(parameters)))
final_pars <- as.data.frame(parameters)

#### Plot parameter space ----
pars_plot <- final_pars |> 
      ggplot() + 
      geom_point(aes(x = A, y = alpha, color = R0), shape = 16, size = 6) +
      scale_color_continuous(type = "viridis") +
      theme_classic() +
      labs(
        x = "Mean age at first infection",
        y = "Duration of immunity"
      ) +
      scale_y_reverse() +
      theme(
        legend.position = 'none',
        axis.title = element_text(size = 24, face = "bold"),
        text = element_text(family = "Helvetica"),
        axis.ticks.x = element_blank(),
        axis.ticks.y = element_blank(),
        axis.text = element_blank()
      )


ggsave(
      filename = "Figures/Figure2/Parameters_Workflow.pdf",
      plot = pars_plot,
      width = 15,
      height = 15,
      units = "cm",
      device = cairo_pdf
    )

### Figure 2.2: Climate time series ----
#### Import climate data for a location ----
climate <- readRDS("weather_data/picked_Locations_weather.rds") |>
  dplyr::filter(loc == "Rome, Italy") |>
  dplyr::mutate(max_year = max(yr)) |>
  dplyr::filter(yr >= (max_year - 10 + 1)) |>
  dplyr::select(-max_year) |> dplyr::group_by(loc, yr) |> dplyr::mutate(week_in_year = row_number()) |> dplyr::ungroup() |> 
  dplyr::mutate(week_no = row_number() - 1) 

#### Plot Temperature and Relative Humidity 10-year time series ----
p1 <- ggplot(climate, aes(x = week_no, y = Te_norm)) +
  geom_line() +
  labs(y = "Temperature", x = "Time") +
  theme_classic() +
  theme(
    legend.position = 'none',
        axis.title = element_text(size = 24, face = "bold"),
        text = element_text(family = "Helvetica"),
        axis.ticks.x = element_blank(),
        axis.ticks.y = element_blank(),
    axis.text = element_blank(),
        axis.title.x = element_blank(),  # Remove only x-axis title

  )

p2 <- ggplot(climate, aes(x = week_no, y = RH_norm)) +
  geom_line() +
  labs(y = "Relative humidity", x = "Time") +
  theme_classic() +
  theme(
    legend.position = 'none',
        axis.title = element_text(size = 24, face = "bold"),
        text = element_text(family = "Helvetica"),
        axis.ticks.x = element_blank(),
        axis.ticks.y = element_blank(),
        axis.text = element_blank()
  )

combined_plot_climate <- p1 / p2

ggsave(
  filename = "Figures/Figure2/Climate_TS_Workflow.pdf",
  plot = combined_plot_climate,
  width = 25,
  height = 18,
  units = "cm",
  device = cairo_pdf
)

### Figure 2.3: Example of simulations ----
df <- readRDS(file = "Results_Aggregated/R_sims_Te-0.04_RH-0.04_sigma0_rho_mean0.5_Rome__rho_k0-1_import_1e-05/sims_tot.rds")

for(replicate in c(38, 49)){
  pl <- df |> dplyr::filter(.id == 17, rep == replicate) |> ggplot() +
    geom_line(aes(x=week_no,y=CC_obs), linewidth = 1) + 
    labs(y = "Observed incidence", x = "Time") + theme_classic() +
    theme(
        axis.title = element_text(size = 28, face = "bold"),
        axis.text.x = element_blank(),  
        axis.text.y = element_blank(),  
        axis.ticks.x = element_blank(), 
        axis.ticks.y = element_blank(),
        text = element_text(family = "Helvetica")
    )
  
  ggsave(
    filename = sprintf("Figures/Figure2/Simulation%s_Workflow.pdf", replicate),
    plot = pl,
    width = 18,
    height = 16,
    units = "cm",
    device = cairo_pdf
  )
}

### Figure 2.4: Dummy performance evaluation figure ----
df <- readRDS(file = "Results_Aggregated/R_sims_Te-0.04_RH-0.04_sigma0_rho_mean0.5_Rio_de_Janeiro__rho_k0-1_import_1e-05/sims_reg_tot.rds")
final_pars <- readRDS(file = "Results_Aggregated/R_sims_Te-0.04_RH-0.04_sigma0_rho_mean0.5_Rio_de_Janeiro__rho_k0-1_import_1e-05/final_pars.rds")

filtered_df <- df |> dplyr::filter(type == "CC_stand_smooth_AC",
                                .id == 41,
                                rep %in% c(38, 48, 49, 50, 65, 75, 81, 83, 87, 94)) |> 
  dplyr::mutate(y_position = seq(1, 10))

##### Plot for e_Te
plot_Te <- ggplot(filtered_df, aes(x = e_Te, y = y_position)) +
  geom_point(size = 3) +
  geom_errorbar(aes(xmin = e_Te_low_CI, xmax = e_Te_high_CI), width = 0.1, orientation = "y") +
  geom_vline(xintercept = final_pars[1,"True.e_Te"], linetype = "dashed", color = "black", linewidth = 1.2) +
  labs(x = latex2exp::TeX(r'($\delta_{{Te}}$)'), y = NULL, title = "Temperature") +
  theme_classic() + 
  theme(
        legend.position = "none",
        plot.title = element_text(size = 24, hjust = 0.5, margin=margin(0,0,20,0)),
        axis.title = element_text(size = 22, face = "bold"),
        axis.text.x = element_blank(),  
        axis.text.y = element_blank(),  
        axis.ticks.x = element_blank(), 
        axis.ticks.y = element_blank(),
        text = element_text(family = "Helvetica"))

##### Plot for e_RH
plot_RH <- ggplot(filtered_df, aes(x = e_RH, y = y_position)) +
  geom_point(size = 3) +
  geom_errorbar(aes(xmin = e_RH_low_CI, xmax = e_RH_high_CI), width = 0.1, orientation = "y") +
  geom_vline(xintercept = final_pars[1,"True.e_RH"], linetype = "dashed", color = "black", linewidth = 1.2) +
  labs(x = latex2exp::TeX(r'($\delta_{{RH}}$)'), y = NULL, title = "Relative humidity") +
  theme_classic() + 
  theme(
    legend.position = "none",
        plot.title = element_text(size = 24, hjust = 0.5, margin=margin(0,0,20,0)),
        axis.title = element_text(size = 22, face = "bold"),
        axis.text.x = element_blank(),  
        axis.text.y = element_blank(),  
        axis.ticks.x = element_blank(), 
        axis.ticks.y = element_blank(),
        text = element_text(family = "Helvetica"))

##### Combine the plots into a single figure
combined_plot <- plot_grid(plot_Te, plot_RH, nrow = 1, ncol = 2, align = "h")

ggsave(
  filename = "Figures/Figure2/Performance_Dummy_Workflow.pdf",
  plot = combined_plot,                  
  width = 15,
  height = 10,
  units = "cm",
  device = cairo_pdf
)

## Figure 3: Sample of simulations across locations and scenarios ----
to_use <- c("Dubai", "Rio_de_Janeiro", "Rome", "Toronto")
True_e_Te <- -0.04
True_e_RH <- -0.04
rho_mean <- 0.5
School.Terms <- F
years <- 10
rho_K <- 0.1
import <- 1e-05
sims_total <- NULL
sigma_val <- 0

for(roiLocation in to_use){
  data_Path <- paste("Results_Aggregated/R_sims_Te",
                        True_e_Te, "_RH", True_e_RH,
                        "_",
                        "sigma", sigma_val,
                        "_",
                        "rho_mean", rho_mean,
                        "_",
                        roiLocation, "_",
                        ifelse(School.Terms == T, "_School-Forced", ""),
                        sprintf("_rho_k%s", gsub("\\.", "-", as.character(rho_K))),
                        sprintf("_import_%s", as.character(import)),
                        sep = "")
  final_pars <- readRDS(paste(data_Path, "/final_pars.rds", sep = "")) |> dplyr::select(A,R0,alpha,.id, alpha_year)
  sims_tot <- readRDS(paste(data_Path, "/sims_tot.rds", sep = "")) |>  
    dplyr::filter(rep == 1) |> 
    dplyr::select(.id, week_no, CC, RH_norm, Te_norm) |> 
    dplyr::mutate(location = roiLocation)
  
  sims_total <- sims_total |>
    bind_rows(sims_tot)
}

location_labels <- c(
  "Rome" = "Rome, Italy",
  "Rio_de_Janeiro" = "Rio de Janeiro, Brazil",
  "Toronto" = "Toronto, Canada",
  "Dubai" = "Dubai, UAE"
)

periods_plot <- sims_total |>
  dplyr::mutate(
    .id = factor(.id)
  ) |>
  ggplot(aes(x = week_no, y = CC / 10000)) +
  geom_line(
    aes(group = .id),
    colour = "grey80",
    alpha = 0.6,
    linewidth = 0.4
  ) +
  geom_line(
    data = sims_total |>
      dplyr::filter(.id %in% c(32, 34, 27)) |>
      dplyr::mutate(.id = factor(.id, levels = c(32, 34, 27))),
    aes(colour = .id, group = .id),
    alpha = 0.8,
    linewidth = .8
  ) +
  facet_wrap(
    ~ location,
    labeller = labeller(
      location = as_labeller(location_labels)
    ),
    scales = "fixed"
  ) +
  scale_colour_manual(
    values = c(
      "32" = "#f7ae02",
      "34" = "#16c970",
      "27" = "#ae02f7"
    )
  ) +
  theme_bw() +
  labs(
    x = "Time (weeks)",
    y = bquote("Total incidence rates (per " ~ 10^4 ~ " individuals)")
  ) + 
  theme(
    strip.background = element_rect(fill = "white", colour = NA),
    strip.text.x = element_text(
      size = 14, face = "bold",
      family = "Helvetica", colour = "black"
    ),
    axis.title = element_text(
      size = 12, family = "Helvetica"
    ),
    axis.text = element_text(
      size = 11, family = "Helvetica"
    ),
    panel.grid.major = element_line(colour = "grey90", linewidth = 0.5),
    panel.grid.minor = element_blank(),
    legend.position = "none",
    plot.title = element_text(hjust = 0.5)
  )

ggsave(
  filename = "Figures/Fig3.pdf",
  plot = periods_plot,
  width = 18,
  height = 15,
  units = "cm" ,
  scale = 1, device = cairo_pdf
)

## Figure 4: Main Figure of biases for all models and all locations, for both variables ----
rm(list = ls())

Locations <- c("Dubai", "Rio_de_Janeiro", "Rome", "Toronto")
delta_climate <- -0.04
rho_k <- 0.1
sigma_beta_val <- 0
rho_mean <- 0.5
School.Terms <- F
years <- 10 
import <- 1e-5

# Locations full names
location_labels <- c(
  "Rome" = "Rome, Italy",
  "Rio_de_Janeiro" = "Rio de Janeiro, Brazil",
  "Toronto" = "Toronto, Canada",
  "Dubai" = "Dubai, UAE"
)

# Variables full names
variable_labels <- c(
    "Te_est" = "Temperature",
    "RH_est" = "Relative humidity"
)

# Models full names
model_labels <- c(CC_stand_smooth_AC = "Test",
                  CC_true = "Control")

df_main_fig <- NULL
for(roiLoc in Locations){
    data_Path <- paste("Results_Aggregated/R_sims_Te",
                      delta_climate, "_RH", delta_climate, 
                      "_",
                      "sigma", sigma_beta_val, "_",
                      paste("rho_mean", rho_mean, "_", sep=""),
                      roiLoc, "_",
                      sprintf("_rho_k%s", gsub("\\.", "-", as.character(rho_k))),
                      sprintf("_import_%s", as.character(import)),
                      ifelse(School.Terms == T, "_School-Forced", ""),
                      sep = "")
    final_pars <- readRDS(paste(data_Path, "/final_pars.rds", sep = ""))
    sims_reg_tot <- readRDS(paste(data_Path, "/sims_reg_tot.rds", sep = "")) |> 
      dplyr::filter(type %in% c("CC_true", "CC_stand_smooth_AC")) |> 
      dplyr::mutate(location = roiLoc) |>
      dplyr::select(-c(e_RH_low_CI, e_RH_high_CI,
                      log_SI_lag_VAR, mean_log_SI_lag,
                      R2, e_Te_low_CI, e_Te_high_CI,
                      TP_Te, TP_RH)) |> 
      dplyr::group_by(type, .id, location) |> 
      dplyr::summarise(
        MAB_Te = mean(abs(e_Te - unique(final_pars[,"True.e_Te"]))), # Computing the mean absolute bias for temperature
        MAB_RH = mean(abs(e_RH - unique(final_pars[,"True.e_RH"]))),
        sd_AB_Te = sd(abs(e_Te - unique(final_pars[,"True.e_Te"]))), # Computing the standard deviation in absolute biases for temperature
        sd_AB_RH = sd(abs(e_RH - unique(final_pars[,"True.e_RH"]))),
        MSE_Te = mean(e_Te_se),
        MSE_RH = mean(e_RH_se),
        MA_e_Te = mean(abs(e_Te)),
        MA_e_RH = mean(abs(e_RH)),
        SDSE_Te = sd(e_Te_se),
        SDSE_RH = sd(e_RH_se),
        MPOW_Te = mean(e_Te_pow),
        MPOW_RH = mean(e_RH_pow)
      ) |> dplyr::ungroup() |> 
      dplyr::mutate(true_val = if(unique(final_pars[,"True.e_Te"]) == unique(final_pars[,"True.e_RH"])){
                                unique(final_pars[,"True.e_Te"])      
                              })

    df_main_fig <- df_main_fig |>
      bind_rows(sims_reg_tot); 
    rm(sims_reg_tot)
}

## Plot for temperature
# Here MSE is divided by the mean absolute e_Te
### Make general colour scale filling
control_colour <- "grey50"
control_shape <- 1
global_limits <- range(
  c(df_main_fig$MPOW_Te, df_main_fig$MPOW_RH),
  na.rm = TRUE
)

scale_mean_power <- viridis::scale_color_viridis(
  option = "plasma",
  direction = 1,
  begin = 0,
  end = .95,
  limits = global_limits,
  oob = scales::squish,  # important, see below
  name = "Mean power"
)

# Plot for temperature
Te_main_plot <- df_main_fig |>
  ggplot() + 
  geom_point(
    data = subset(df_main_fig, type != "CC_true"),
    aes(
      x = MAB_Te / abs(true_val) * 100,
      y = MSE_Te / abs(true_val) * 100,
      colour = MPOW_Te,
      shape = type,
      size = type
    ),
    alpha = .7,
    stroke = 1
  ) + 
  geom_point(
    data = subset(df_main_fig, type == "CC_true"),
    aes(
      x = MAB_Te / abs(true_val) * 100,
      y = MSE_Te / abs(true_val) * 100,
      shape = type,
      size  = type
    ),
    colour = control_colour,
    alpha = .6,
    stroke = .5
  ) +
  facet_wrap(
    ~location,
    labeller = labeller(location = location_labels)
  ) +
  theme_bw() + 
  scale_mean_power +
  scale_shape_manual(
    values = c(
      CC_stand_smooth_AC = 16,
      CC_true = control_shape
    ),
    labels = model_labels,
    name = "Model"
  ) +
  scale_size_manual(
    values = c(
      CC_stand_smooth_AC = 4,
      CC_true = 2
    )
  ) +
  xlab("RMAB (%)") +
  ylab("RMSE (%)") +
  theme(
      panel.grid.minor = element_blank(),
    strip.background = element_blank(),
    strip.text.x = element_text(
      size = 18, face = "bold",
      family = "Helvetica", colour = "black"
    ),
    strip.text.y = element_text(
      size = 18, face = "bold",
      family = "Helvetica", colour = "black"
    ),
    axis.title = element_text(
      size = 16, family = "Helvetica"
    ),
    axis.text = element_text(
      size = 12, family = "Helvetica"
    ),
    axis.ticks.x = element_line(),
    panel.grid.major.x = element_line(colour = "grey90"),
    panel.grid.major.y = element_line(colour = "grey90"),
    plot.title = element_blank(),
    legend.position = "bottom",
    legend.spacing.x = unit(5, "cm"),
    panel.spacing.x = unit(0.75, "cm"),
    legend.text = element_text(
      size = 14, family = "Helvetica"
    ),
    legend.title = element_text(
      size = 16, family = "Helvetica"
    )
  ) +
  scale_x_continuous(
    breaks = seq(
      0,
      max(df_main_fig$MAB_Te / abs(df_main_fig$true_val) * 100, na.rm = TRUE),
      by = 20
    )
  ) +
  scale_y_continuous(
    breaks = seq(
      0,
      max(df_main_fig$MSE_Te / abs(df_main_fig$true_val) * 100, na.rm = TRUE),
      by = 20
    )
  ) +
guides(
  colour = guide_colourbar(
    barwidth  = unit(7.5, "cm"),
    barheight = unit(0.4, "cm"),
    title.position = "left",
    title.theme = element_text(
      margin = margin(r = 20)
    )
  ),
  size = "none",
  shape = guide_legend(
    override.aes = list(
      shape  = c(16, control_shape),
      size   = c(4, 2),
      colour = c("black", control_colour),
      alpha  = 1,
      stroke = c(1, 1)
    )
  )
)

ggsave(
  filename = "Figures/Fig4.pdf",
  plot = Te_main_plot,
  width = 35,
  height = 22,
  units = "cm",
  device = cairo_pdf
)

## Figure 5: Main figure for raltive humidity ----
RH_main_plot <- df_main_fig |>
  ggplot() + 
  geom_point(
    data = subset(df_main_fig, type != "CC_true"),
    aes(
      x = MAB_RH / abs(true_val) * 100,
      y = MSE_RH / abs(true_val) * 100,
      colour = MPOW_RH,
      shape = type,
      size = type
    ),
    alpha = .7,
    stroke = 1
  ) + 
  geom_point(
    data = subset(df_main_fig, type == "CC_true"),
    aes(
      x = MAB_RH / abs(true_val) * 100,
      y = MSE_RH / abs(true_val) * 100,
      shape = type,
      size = type
    ),
    colour = control_colour,
    alpha = .6,
    stroke = .5
  ) +
  facet_wrap(
    ~location,
    labeller = labeller(location = location_labels)
  ) +
  theme_bw() + 
  scale_mean_power +
  scale_shape_manual(
    values = c(
      CC_stand_smooth_AC = 16,
      CC_true = control_shape
    ),
    labels = model_labels,
    name = "Model"
  ) +
  scale_size_manual(
    values = c(
      CC_stand_smooth_AC = 4,
      CC_true = 2
    )
  ) +
  xlab("RMAB (%)") +
  ylab("RMSE (%)") +
  theme(
      panel.grid.minor = element_blank(),
    strip.background = element_blank(),
    strip.text.x = element_text(
      size = 18, face = "bold",
      family = "Helvetica", colour = "black"
    ),
    strip.text.y = element_text(
      size = 18, face = "bold",
      family = "Helvetica", colour = "black"
    ),
    axis.title = element_text(
      size = 16, family = "Helvetica"
    ),
    axis.text = element_text(
      size = 12, family = "Helvetica"
    ),
    axis.ticks.x = element_line(),
    panel.grid.major.x = element_line(colour = "grey90"),
    panel.grid.major.y = element_line(colour = "grey90"),
    plot.title = element_blank(),
    legend.position = "bottom",
    legend.spacing.x = unit(5, "cm"),
    panel.spacing.x = unit(0.75, "cm"),
    legend.text = element_text(
      size = 14, family = "Helvetica"
    ),
    legend.title = element_text(
      size = 16, family = "Helvetica"
    )
  ) +
  scale_x_continuous(
    breaks = seq(
      0,
      max(df_main_fig$MAB_RH / abs(df_main_fig$true_val) * 100, na.rm = TRUE),
      by = 20
    )
  ) +
  scale_y_continuous(
    breaks = seq(
      0,
      max(df_main_fig$MSE_RH / abs(df_main_fig$true_val) * 100, na.rm = TRUE),
      by = 20
    )
  ) +
  guides(
  colour = guide_colourbar(
    barwidth  = unit(7.5, "cm"),
    barheight = unit(0.4, "cm"),
    title.position = "left",
    title.theme = element_text(
      margin = margin(r = 20)
    )
  ),
  size = "none",
  shape = guide_legend(
    override.aes = list(
      shape = c(16, control_shape),
      size = c(4, 2),
      colour = c("black", control_colour),
      alpha = 1,
      stroke = c(1, 1)
    )
  )
)

ggsave(
      filename = "Figures/Fig5.pdf",
      plot = RH_main_plot,
      width = 35,
      height = 22,
      units = "cm",
      device = cairo_pdf
    )

# Figures for supplementary material ----
## Figure A: Plot of all weather time series for the 4 locations ----
# Importing only TS from 2010 to 2020
climate <- readRDS("weather_data/picked_Locations_weather.rds") |>
  dplyr::filter(between(yr, 2010, 2020)) |>
  dplyr::group_by(loc, yr) |> dplyr::mutate(week_in_year = row_number()) |> dplyr::ungroup() |> 
  dplyr::group_by(loc) |> 
  dplyr::mutate(week_no = row_number() - 1)

df_plot <- climate |>
  pivot_longer(
    cols = c(Te_norm, RH_norm),
    names_to = "variable_name",
    values_to = "variable_value"
  ) |>
  mutate(
    variable_name = case_when(
      variable_name == "Te_norm" ~ "Normalized temperature (°C)",
      variable_name == "RH_norm" ~ "Normalized relative humidity (%)",
      TRUE ~ variable_name
    ),
    loc = recode(
      loc,
      "Rio_de_Janeiro, Brazil" = "Rio de Janeiro, Brazil",
      "Dubai, United Arab Emirates" = "Dubai, UAE")) |>
  mutate(variable_name = fct_relevel(variable_name, "Normalized temperature (°C)", "Normalized relative humidity (%)"))

climates_TS <- ggplot(df_plot, aes(x = week_no, y = variable_value)) +
  geom_line(linewidth = .35) +
  facet_grid(loc ~ variable_name, scales = "free_y") +
  theme_light() + 
  labs(x = "Time (week)") + 
  scale_x_continuous(expand = expansion(mult = c(0.01, 0.01))) +
  theme(
    panel.border = element_rect(color = "black", fill = NA, linewidth = 1),
    text = element_text(family = "Helvetica"),
    axis.title.x = element_text(size = 14, color = "black"),
    axis.title.y = element_blank(),
    strip.background = element_blank(),
    strip.text.x = element_text(size = 14, face = "bold", color = "black"),
    strip.text.y = element_text(size = 14, face = "bold", color = "black"),
    axis.text.y = element_text(size = 11),
    axis.text.x = element_text(size = 11),
    panel.grid.major.x = element_blank(),
    panel.grid.minor.x = element_blank()
  )

  ggsave(
  filename = "Figures/A_fig.pdf",
  plot = climates_TS,
  width = 12,
  height = 9.2,
  device = cairo_pdf
)

## Figure B: Sampled parameters with tiks ----
# alpha_year = alpha * 52
# 20 yrs immunity -> alpha = 1 / (52 * 20) -> alpha_year = 52 * 1 / (20 * 52) = 1 / 20
# 1 yr immunity -> alpha = 1 / (52 * 1) -> alpha_year = 52 * 1 / (1 * 52) = 1
mu_year = 1 / 80  
parameters = sobol_design(lower = c(A = 1, alpha_year = 1/20), upper = c(A = 10, alpha_year = 1), nseq = 50)
parameters = parameters |>
  dplyr::mutate(R0 = 1 + (1 + 1 / (alpha_year + mu_year)) / (A * (mu_year + 1))) 
parameters = parameters |> dplyr::mutate(alpha = alpha_year / 52) |>
  dplyr::mutate(eps = rep(1,nrow(parameters))) |>
  dplyr::mutate(.id = seq_len(nrow(parameters)))
final_pars = as.data.frame(parameters)

years_num <- c(1.0, 1.5, 2.0, 3.0, 5.0, 7.0, 10.0, 15.0)
alpha_breaks <- 1 / (52 * years_num)
years_labels  <- as.character(years_num)

ParamSpace <- final_pars |>
  ggplot() +
  geom_point(aes(x = A, y = alpha, color = R0), shape = 16, size = 4) +
  scale_color_continuous(type = "viridis") +
  scale_y_reverse(
    breaks = alpha_breaks,
    labels = years_labels
  ) +
  theme_light() +
  labs(
    x = "Mean age at first infection (years)",
    y = "Average duration of immunity (years)",
    color = expression(R[bold(0)])
  ) +
  theme(
    panel.border = element_rect(color = "black", fill = NA, linewidth = 1),
    legend.position = "right",
    axis.title = element_text(size = 16),
    text = element_text(family = "Helvetica"), axis.text.x = element_text(size=12),
    axis.text.y = element_text(size=12),
    legend.title= element_text(size=16),
    legend.text = element_text(size = 12)
   ) 

  ggsave(
  filename = "Figures/B_fig.pdf",
  plot = ParamSpace,
  width = 8,
  height = 6,
  device = cairo_pdf
)

## Figure C -----
# Prepare data to be plotted
rm(list = ls())

Locations <- c("Dubai", "Rio_de_Janeiro", "Rome", "Toronto")
delta_climate = -0.04
rho_k <- 0.1
Conds <- expand.grid(loc = Locations)
sigma_beta_val = 0
rho_mean = 0.5
School.Terms = F
years = 10 
import = 1e-5

## Locations full names
location_labels <- c(
  "Rome" = "Rome, Italy",
  "Rio_de_Janeiro" = "Rio de Janeiro, Brazil",
  "Toronto" = "Toronto, Canada",
  "Dubai" = "Dubai, UAE"
)

## Variables full names
variable_labels <- c(
    "Te_est" = "Temperature (°C)",
    "RH_est" = "Relative humidity (%)"
)

## Models full names
model_labels <- c(CC_AC = "Autocorrelation",
                  CC_stand_smooth = "Time-smooth",
                  CC_stand_smooth_AC = "Time-smooth + \n Autocorrelation",
                  CC_true = "Control")

df_rel_bias = NULL
for(i in 1:nrow(Conds)){
    data_Path = paste("Results_Aggregated/R_sims_Te",
                      delta_climate, "_RH", delta_climate, 
                      "_",
                      "sigma", sigma_beta_val, "_",
                      paste("rho_mean", rho_mean, "_", sep=""),
                      Conds[i, "loc"], "_",
                      ifelse(School.Terms == T, "_School-Forced", ""),
                      sprintf("_rho_k%s", gsub("\\.", "-", as.character(rho_k))),
                      sprintf("_import_%s", as.character(import)),
                      sep = "")
    R0s = readRDS(paste(data_Path, "/final_pars.rds", sep = "")) |> dplyr::select(.id, R0) |> distinct()
    final_pars = readRDS(paste(data_Path, "/final_pars.rds", sep = ""))
    sims_reg_tot = readRDS(paste(data_Path, "/sims_reg_tot.rds", sep = "")) |> 
      dplyr::mutate(location = Conds[i, "loc"]) |>
      dplyr::select(-c(e_RH_low_CI, e_RH_high_CI,
                      log_SI_lag_VAR, mean_log_SI_lag,
                      R2, e_Te_low_CI, e_Te_high_CI,
                      TP_Te, e_Te_pow, TP_RH, e_RH_pow)) |> 
      dplyr::group_by(type, .id, location) |> 
      dplyr::summarise(Te_est = mean(abs(e_Te - final_pars[1,"True.e_Te"]) / abs(final_pars[1,"True.e_Te"]) * 100),
        RH_est = mean(abs(e_RH - final_pars[1,"True.e_RH"]) / abs(final_pars[1,"True.e_RH"]) * 100),
        Te_se = mean(e_Te_se) / abs(final_pars[1,"True.e_Te"]) * 100,
        RH_se = mean(e_RH_se) / abs(final_pars[1,"True.e_RH"]) * 100) |>
      dplyr::full_join(
                      y = R0s, 
                      by = c(".id")
                    )

    df_rel_bias = df_rel_bias |>
      bind_rows(sims_reg_tot); 
    rm(sims_reg_tot)
}

df_longer <- df_rel_bias |> tidyr::pivot_longer(
    cols = tidyselect::starts_with(c("Te_", "RH_")),
    names_to = c("variable", ".value"),
    names_pattern = "(.*)_(est|se)") |> 
    dplyr::mutate(variable = factor(variable, levels = c("Te", "RH")))

# Plot for temperature
plot_data <- df_longer |> 
  full_join(y = final_pars |> select(-R0), by = c(".id")) |>
  filter(type == "CC_stand_smooth_AC", variable == "Te")

## Compute breaks for fill
est_values <- plot_data$est
est_min <- min(est_values, na.rm = TRUE)
est_max <- max(est_values, na.rm = TRUE)
est_median <- median(est_values, na.rm = TRUE)

## Round to nearest multiple of 10
round_to_10 <- function(x) round(x / 10) * 10

fill_breaks <- sort(unique(c(
  round_to_10(est_min),
  round_to_10(est_median),
  round_to_10(est_max)
)))

## Compute breaks for size
se_values <- plot_data$se
se_min <- min(se_values, na.rm = TRUE)
se_max <- max(se_values, na.rm = TRUE)
se_median <- median(se_values, na.rm = TRUE)

# Round to nearest multiple of 10
size_breaks <- sort(unique(c(
  round_to_10(se_min),
  round_to_10(se_median),
  round_to_10(se_max)
)))

## Actual plotting
Bubble_Te <- plot_data |>
  ggplot(aes(
    x = A,
    y = alpha_year,
    fill = est,
    size = se
  )) +
  geom_point(shape = 21, color = "black", alpha = 0.8) +
  scale_fill_viridis_c(
  option = "C",
  name = "RMAB (%)",
  breaks = fill_breaks,
  limits = range(fill_breaks),
  oob = scales::squish
  ) +
  scale_size_continuous(
    name = "RMSE (%)",
    breaks = size_breaks,
    range = c(.7, 6.5)
  ) + 
  scale_y_reverse(
    name = "Years of immunity (years)",
    trans = "reverse",
    breaks = c(1, 1/2, 1/5, 1/10, 1/20),
    labels = c(1, 2, 5, 10, 20)
  ) +
  facet_wrap(~location, scales = "free_x", labeller = labeller(
    location = as_labeller(location_labels))) +
  theme_bw() +
  guides(
    fill = guide_colorbar(order = 1, barheight = 6, barwidth = 1),
    size = guide_legend(
      order = 2,
      override.aes = list(fill = "grey50")
    )
  ) +
  theme(
    text = element_text(family = "Helvetica"),
    strip.background = element_blank(),
    strip.text = element_text(size = 26, face = "bold"),
    axis.title = element_text(size = 22),
    axis.text = element_text(size = 22),
    legend.text = element_text(size = 22), 
    legend.title = element_text(size = 22, margin = margin(b = 10)),
    legend.spacing.y = unit(1, "cm"),
    legend.key.size = unit(0.8, "cm")
  ) +
  labs(
    x = "Mean age at first infection (years)"
  )

ggsave(
  filename = "Figures/C_fig.pdf",
  plot = Bubble_Te,
  width = 40,
  height = 35,
  units = "cm",
  device = cairo_pdf
)

## Figure D ----
# FOR RELATIVE HUMIDITY
plot_data <- df_longer |> 
  full_join(
    y = final_pars |> select(-R0),
    by = ".id"
  ) |>
  filter(
    type == "CC_stand_smooth_AC",
    variable == "RH"
  )

# Round min down to nearest multiple of 5, max up to nearest multiple of 5
round_down_to_5 <- function(x) floor(x / 5) * 5
round_up_to_5   <- function(x) ceiling(x / 5) * 5
round_to_5      <- function(x) round(x / 5) * 5

# Fill (RMAB)
est_values <- plot_data$est
fill_breaks <- sort(unique(c(
  round_down_to_5(min(est_values, na.rm = TRUE)),
  round_to_5(median(est_values, na.rm = TRUE)),
  round_up_to_5(max(est_values, na.rm = TRUE))
)))
fill_limits <- range(fill_breaks, na.rm = TRUE)

# Size (RMSE)
se_values <- plot_data$se
size_breaks <- sort(unique(c(
  round_down_to_5(min(se_values, na.rm = TRUE)),
  round_to_5(median(se_values, na.rm = TRUE)),
  round_up_to_5(max(se_values, na.rm = TRUE))
)))
size_limits <- range(size_breaks, na.rm = TRUE)

# Actual plot
Bubble_RH <- plot_data |>
  mutate(
    se_clamped = pmin(pmax(se, size_limits[1]), size_limits[2])
  ) |> 
  ggplot(aes(
    x = A,
    y = alpha_year,
    fill = est,
    size = se_clamped
  )) +
  geom_point(
    shape = 21,
    color = "black",
    alpha = 0.8
  ) +
  scale_fill_viridis_c(
    option = "C",
    name = "RMAB (%)",
    breaks = fill_breaks,
    labels = fill_breaks,
    limits = fill_limits
  ) +
  scale_size_continuous(
    name = "RMSE (%)",
    breaks = size_breaks,
    labels = size_breaks,
    limits = size_limits,
    range = c(.7, 6.5)
  ) +
  scale_y_reverse(
    name = "Years of immunity (years)",
    breaks = c(1, 1/2, 1/5, 1/10, 1/20),
    labels = c(1, 2, 5, 10, 20)
  ) +
  facet_wrap(
    ~ location,
    scales = "free_x",
    labeller = labeller(location = as_labeller(location_labels))
  ) +
  theme_bw() +
  guides(
    fill = guide_colorbar(
      order = 1,
      barheight = 6,
      barwidth = 1
    ),
    size = guide_legend(
      order = 2,
      override.aes = list(
        size = c(2.5, 4.5, 6.5),
        fill = "grey50"
      )
    )
  ) +
  theme(
    text = element_text(family = "Helvetica"),
    strip.background = element_blank(),
    strip.text = element_text(size = 26, face = "bold"),
    axis.title = element_text(size = 22),
    axis.text = element_text(size = 22),
    legend.text = element_text(size = 22), 
    legend.title = element_text(size = 22, margin = margin(b = 10)),
    legend.spacing.y = unit(1, "cm"),
    legend.key.size = unit(0.8, "cm")
  ) +
  labs(
    x = "Mean age at first infection (years)"
  )

ggsave(
  filename = "Figures/D_fig.pdf",
  plot = Bubble_RH,
  width = 40,
  height = 35,
  units = "cm",
  device = cairo_pdf
)

## Figure E: transmissibility component of temperature -----
rm(list=ls())

# Now import all regression results to get the .id and relative estimated e_RH ----
Locations <- c("Dubai", "Rio_de_Janeiro", "Rome", "Toronto")
delta_climate = -0.04
rho_k <- 0.1
sigma_beta_val = 0
rho_mean = 0.5
School.Terms = F
years = 10 
import = 1e-5

# Locations full names
location_labels <- c(
  "Rome" = "Rome, Italy",
  "Rio_de_Janeiro" = "Rio de Janeiro, Brazil",
  "Toronto" = "Toronto, Canada",
  "Dubai" = "Dubai, UAE"
)

sims_total <- NULL
df_regs <- NULL
df_perf_tot <- NULL
for(roiLoc in Locations){
    data_Path <- paste("Results_Aggregated/R_sims_Te",
                      delta_climate, "_RH", delta_climate, 
                      "_",
                      "sigma", sigma_beta_val, "_",
                      paste("rho_mean", rho_mean, "_", sep=""),
                      roiLoc, "_",
                      sprintf("_rho_k%s", gsub("\\.", "-", as.character(rho_k))),
                      sprintf("_import_%s", as.character(import)),
                      ifelse(School.Terms == T, "_School-Forced", ""),
                      sep = "")

    # Get parameters
    final_pars <- readRDS(paste(data_Path, "/final_pars.rds", sep = ""))

    # Get regression estimates
    sims_reg_tot <- readRDS(paste(data_Path, "/sims_reg_tot.rds", sep = "")) |> 
      dplyr::filter(type == "CC_stand_smooth_AC") |> 
      dplyr::mutate(location = roiLoc) |>
      dplyr::select(-c(e_RH_low_CI, e_RH_high_CI,
                      log_SI_lag_VAR, mean_log_SI_lag,
                      R2, e_Te_low_CI, e_Te_high_CI,
                      TP_Te, TP_RH)) |> 
      dplyr::ungroup() |> 
      dplyr::group_by(type, .id, location) |> 
      dplyr::summarise(
        M_Te = mean(e_Te),
        M_RH = mean(e_RH)) |>
      dplyr::ungroup() |> dplyr::select(-type)

    # Get RMABs
    df_perf <- readRDS(paste(data_Path, "/sims_reg_tot.rds", sep = "")) |> 
      dplyr::filter(type == "CC_stand_smooth_AC") |> dplyr::ungroup() |> dplyr::select(-rep) |> 
      dplyr::mutate(location = roiLoc) |>
      dplyr::select(-c(e_RH_low_CI, e_RH_high_CI,
                      log_SI_lag_VAR, mean_log_SI_lag,
                      R2, e_Te_low_CI, e_Te_high_CI,
                      TP_Te, TP_RH)) |> 
      dplyr::group_by(.id, location) |> 
      dplyr::summarise(
        RMAB_Te = mean(abs(e_Te - unique(final_pars[,"True.e_Te"]))) / abs(unique(final_pars[,"True.e_Te"])) * 100,
        RMAB_RH = mean(abs(e_RH - unique(final_pars[,"True.e_RH"]))) / abs(unique(final_pars[,"True.e_RH"])) * 100
      ) |> dplyr::ungroup() |> 
      dplyr::mutate(true_val = if(unique(final_pars[,"True.e_Te"]) == unique(final_pars[,"True.e_RH"])){
                                unique(final_pars[,"True.e_Te"])      
                              })

    # Get climate time series
    sims_tot <- readRDS(paste(data_Path, "/sims_tot.rds", sep = "")) |>  
        dplyr::filter(rep == 1) |> 
        dplyr::select(.id, week_no, CC, RH_norm, Te_norm) |> 
        dplyr::mutate(location = roiLoc) |> dplyr::ungroup() |> dplyr::select(-rep)
  
    # Bind all needed
    df_regs <- df_regs |>
      bind_rows(sims_reg_tot); 
    sims_total <- sims_total |>
        bind_rows(sims_tot)
    df_perf_tot <- df_perf_tot |> 
        bind_rows(df_perf)
    rm(sims_reg_tot, sims_tot, df_perf)
}

# Now put these together to get our plotting dataframe
extremes <- df_perf_tot |> 
  dplyr::group_by(location) |> 
  dplyr::summarise(
    id_max_RMAB_Te = .id[which.max(RMAB_Te)],
    id_min_RMAB_Te = .id[which.min(RMAB_Te)],
    id_max_RMAB_RH = .id[which.max(RMAB_RH)],
    id_min_RMAB_RH = .id[which.min(RMAB_RH)],
  )

# Create dataframe to be used for the plotting
plot_df <- sims_total |>
  dplyr::left_join(df_regs, by = c(".id", "location")) |>
  dplyr::left_join(extremes, by = "location") |>
  dplyr::mutate(
    beta_seas_Te = exp(M_Te * Te_norm),
    beta_seas_RH = exp(M_RH * RH_norm)
  )

# Now to the actual plotting
transm_Te <- plot_df |> ggplot() +
  geom_line(
    data = ~ dplyr::filter(.x,
      .id != id_max_RMAB_Te & .id != id_min_RMAB_Te
    ),
    aes(
      x = Te_norm,
      y = beta_seas_Te,
      group = as.factor(.id)
    ),
    linewidth = .65,
    alpha = .5,
    colour = "grey70"
  ) +
  geom_line(
    data = ~ dplyr::filter(.x,
      .id == id_max_RMAB_Te | .id == id_min_RMAB_Te
    ),
    aes(
      x = Te_norm,
      y = beta_seas_Te,
      group = as.factor(.id),
      colour = case_when(
        .id == id_max_RMAB_Te ~ "max",
        .id == id_min_RMAB_Te ~ "min"
      )
    ),
    linewidth = .8,
    alpha = .9,
    linetype = "solid"
  ) +
  scale_colour_manual(values = c(max = "#FF4613", min = "#2E95FB")) +
  geom_line(
    data = plot_df |>
        dplyr::distinct(location, Te_norm) |>
        dplyr::mutate(beta_true_Te = exp(final_pars[1, "True.e_Te"] * Te_norm)),
    aes(x = Te_norm, y = beta_true_Te, group = location),
    linewidth = .8,
    colour = "black",
    alpha = .9,
    linetype = "dashed"
  ) +
  labs(x = "Normalized temperature (°C)", y = expression(beta[Te])) +
  facet_wrap(~location, scales = "free_y",
    labeller = labeller(location = location_labels)) +
  theme_classic() +
  theme(
      panel.grid.minor = element_blank(),
    strip.background = element_blank(),
    strip.text.x = element_text(
      size = 20, face = "bold",
      family = "Helvetica", colour = "black"
    ),
    strip.text.y = element_text(
      size = 20,
      family = "Helvetica", colour = "black"
    ),
    axis.title = element_text(
      size = 20, family = "Helvetica"
    ),
    axis.text = element_text(
      size = 14, family = "Helvetica"
    ),
    axis.ticks.x = element_line(),
    panel.grid.major.x = element_line(colour = "grey90"),
    panel.grid.major.y = element_line(colour = "grey90"),
    plot.title = element_blank(),
    legend.position = "none",
    legend.spacing.x = unit(5, "cm"),
    panel.spacing.x = unit(0.75, "cm"),
    legend.text = element_text(
      size = 20, family = "Helvetica"
    ),
    legend.title = element_text(
      size = 22, family = "Helvetica", face = "bold"
    )
  )

  ggsave(
      filename = "Figures/E_fig.pdf",
      plot = transm_Te,
      width = 35,
      height = 25,
      units = "cm",
      device = cairo_pdf
    )

## Figure F: transmissibility component of relative humidity -----
# Using the same cumulative dataframe obtained before
transm_RH <- plot_df |> ggplot() +
  geom_line(
    data = ~ dplyr::filter(.x,
      .id != id_max_RMAB_RH & .id != id_min_RMAB_RH
    ),
    aes(
      x = RH_norm,
      y = beta_seas_RH,
      group = as.factor(.id)
    ),
    linewidth = .65,
    alpha = .5,
    colour = "grey70"
  ) +
  geom_line(
    data = plot_df |>
      dplyr::distinct(location, RH_norm) |>
      dplyr::mutate(beta_true_RH = exp(final_pars[1, "True.e_RH"] * RH_norm)),
    aes(x = RH_norm, y = beta_true_RH, group = location),
    linewidth = .8,
    colour = "black",
    alpha = .9,
    linetype = "dashed"
  ) +
  geom_line(
    data = ~ dplyr::filter(.x,
      .id == id_max_RMAB_RH | .id == id_min_RMAB_RH
    ),
    aes(
      x = RH_norm,
      y = beta_seas_RH,
      group = as.factor(.id),
      colour = case_when(
        .id == id_max_RMAB_RH ~ "max",
        .id == id_min_RMAB_RH ~ "min"
      )
    ),
    linewidth = .9,
    alpha = 1
  ) +
  scale_colour_manual(values = c(max = "#FF4613", min = "#2E95FB")) +
  labs(x = "Normalized relative humidity (%)", y = expression(beta[RH])) +
  facet_wrap(~location, scales = "free_y",
    labeller = labeller(location = location_labels)) +
  theme_classic() +
  theme(
      panel.grid.minor = element_blank(),
    strip.background = element_blank(),
    strip.text.x = element_text(
      size = 20, face = "bold",
      family = "Helvetica", colour = "black"
    ),
    strip.text.y = element_text(
      size = 20,
      family = "Helvetica", colour = "black"
    ),
    axis.title = element_text(
      size = 20, family = "Helvetica"
    ),
    axis.text = element_text(
      size = 14, family = "Helvetica"
    ),
    axis.ticks.x = element_line(),
    panel.grid.major.x = element_line(colour = "grey90"),
    panel.grid.major.y = element_line(colour = "grey90"),
    plot.title = element_blank(),
    legend.position = "none",
    legend.spacing.x = unit(5, "cm"),
    panel.spacing.x = unit(0.75, "cm"),
    legend.text = element_text(
      size = 20, family = "Helvetica"
    ),
    legend.title = element_text(
      size = 22, family = "Helvetica", face = "bold"
    )
  )

ggsave(
  filename = "Figures/F_fig.pdf",
  plot = transm_RH,
  width = 35,
  height = 25,
  units = "cm",
  device = cairo_pdf
)

## Figure G: Differences in overdispersion ----
rm(list=ls())

Locations <- c("Dubai", "Rio_de_Janeiro", "Rome", "Toronto")
delta_climate = -0.04
Rho_ks <- c(0.1, 0.16)
sigma_beta_val = 0
Conds <- expand.grid(loc = Locations, rho_k = Rho_ks)
rho_mean = 0.5
School.Terms = F
years = 10
import = 1e-5

# Locations full names
location_labels <- c(
  "Rome" = "Rome, Italy",
  "Rio_de_Janeiro" = "Rio de Janeiro, Brazil",
  "Toronto" = "Toronto, Canada",
  "Dubai" = "Dubai, UAE"
)

# Variables full names
variable_labels <- c(
    "Te" = "Temperature",
    "RH" = "Relative humidity"
)

# Initialise the dataframe to store all
df_main = NULL
# For loop to generate big cumulative dataframe
for(i in 1:nrow(Conds)){
    data_Path = paste("Results_Aggregated/R_sims_Te",
                      delta_climate, "_RH", delta_climate,
                      "_",
                      "sigma", sigma_beta_val, "_",
                      paste("rho_mean", rho_mean, "_", sep=""),
                      Conds[i, "loc"], "_",
                      ifelse(School.Terms == T, "_School-Forced", ""),
                      sprintf("_rho_k%s", gsub("\\.", "-", as.character(Conds[i, "rho_k"]))),
                      sprintf("_import_%s", as.character(import)),
                      sep = "")
    R0s = readRDS(paste(data_Path, "/final_pars.rds", sep = "")) |> dplyr::select(.id, R0) |> distinct()
    sims_reg_tot = readRDS(paste(data_Path, "/sims_reg_tot.rds", sep = "")) |> 
      dplyr::filter(type == "CC_stand_smooth_AC") |> 
      dplyr::mutate(location = Conds[i, "loc"],
                    rho_k = Conds[i, "rho_k"]) |>
      dplyr::select(-c(e_RH_low_CI, e_RH_high_CI,
                      log_SI_lag_VAR, mean_log_SI_lag,
                      R2, e_Te_low_CI, e_Te_high_CI,
                      TP_Te, e_Te_pow, TP_RH, e_RH_pow)) |> 
      dplyr::group_by(.id, location, rho_k, type) |> 
      dplyr::full_join(
                      y = R0s, 
                      by = c(".id")
                      )

    df_main = df_main |>
      bind_rows(sims_reg_tot); 
    rm(sims_reg_tot)
}

df_diff <- df_main |> mutate(bias_Te = abs(e_Te - delta_climate), 
                             bias_RH = abs(e_RH - delta_climate)) |> 
                  dplyr::group_by(location, .id, rho_k, type, .groups = "keep") |> 
                  dplyr::summarise(MAB_id_Te = mean(bias_Te), 
                            MAB_id_RH = mean(bias_RH),
                            MSE_id_Te = mean(e_Te_se),
                            MSE_id_RH = mean(e_RH_se)) |> 
                            ungroup() |> 
                  mutate(RMAB_id_Te = MAB_id_Te / abs(delta_climate) * 100,
                         RMAB_id_RH = MAB_id_RH / abs(delta_climate) * 100,
                         RMSE_id_Te = MSE_id_Te / abs(delta_climate) * 100,
                         RMSE_id_RH = MSE_id_RH / abs(delta_climate) * 100) |> 
                  select(-c(MAB_id_Te, MAB_id_RH, MSE_id_Te, MSE_id_RH)) |> 
        pivot_wider(
          names_from = c(rho_k),
          values_from = c(RMAB_id_Te, RMAB_id_RH, RMSE_id_Te, RMSE_id_RH)
  ) |> 
  mutate( # New_Sensitivity - reference -> diff > 0 -> ref is better
    diff_RMAB_id_RH = `RMAB_id_RH_0.16` - `RMAB_id_RH_0.1`,
    diff_RMAB_id_Te = `RMAB_id_Te_0.16` - `RMAB_id_Te_0.1`,
    diff_RMSE_id_RH = `RMSE_id_RH_0.16` - `RMSE_id_RH_0.1`,
    diff_RMSE_id_Te = `RMSE_id_Te_0.16` - `RMSE_id_Te_0.1`
  ) |> 
  select(location, type, .id, diff_RMAB_id_RH, diff_RMAB_id_Te, diff_RMSE_id_RH, diff_RMSE_id_Te) |> 
  pivot_longer(cols = c(diff_RMAB_id_RH, diff_RMAB_id_Te, diff_RMSE_id_RH, diff_RMSE_id_Te),
               names_to= "variable", values_to= 'vals')

df_scatter <- df_diff |> dplyr::filter(type == "CC_stand_smooth_AC") |> 
  separate(variable, into = c("metric", "var_type"), sep = "_id_") |>
  pivot_wider(
    names_from = metric,
    values_from = vals
  ) |> mutate(var_type = factor(var_type, levels = c("Te", "RH")))

df_means <- df_scatter |> 
  dplyr::group_by(location, var_type) |> 
  dplyr::summarise(
    diff_RMAB = mean(diff_RMAB, na.rm = TRUE),
    diff_RMSE = mean(diff_RMSE, na.rm = TRUE)
  )

diff_rho_k_slim <- ggplot(df_scatter,
  aes(x = diff_RMAB, y = diff_RMSE)) +
  geom_point(aes(colour = as.factor(location)), size = 2, alpha = .7) +
  geom_point(data = df_means, colour = "black", 
    shape = 8, size = 3, stroke = 1, alpha = .6) + 
    geom_vline(xintercept = 0, colour = "black", linetype = "dashed") +
    geom_hline(yintercept = 0, colour = "black", linetype = "dashed") +
  facet_grid2(
    location ~ var_type,
    labeller = labeller(
      location = as_labeller(location_labels),
      var_type = as_labeller(variable_labels)
    ),
    scales = "fixed"#, independent = "all"
  ) + 
  labs(
  x = "Difference in RMAB (%)",
  y = "Difference in RMSE (%)"
) +
  theme_bw() +
  theme(
    strip.background = element_blank(),
    strip.text.x = element_text(size = 18, face = "bold", family = "Helvetica", colour = "black"),
    strip.text.y = element_text(size = 18, face = "bold", family = "Helvetica", colour = "black"),
    axis.title = element_text(size = 16, family = "Helvetica"),
    axis.text = element_text(size = 14, family = "Helvetica"),
    axis.ticks.x = element_line(),
    panel.grid.major = element_line(colour = "grey90", linewidth = 0.5),
    panel.grid.minor = element_blank(),
    legend.position = "none",
    title = element_blank(),
    plot.title = element_text(hjust = 0.5)
  ) +
  scale_colour_brewer(palette = "Set2")


ggsave(
  filename = "Figures/G_fig.pdf",
  plot = diff_rho_k_slim,
  width = 38,
  height = 38,
  units = "cm" ,
  scale = 1, device = cairo_pdf
)

## Figure H: Differneces in effect of climate ----
rm(list=ls())

Locations <- c("Dubai", "Rio_de_Janeiro", "Rome", "Toronto")
Delta_climates = c(-0.04, -.02)
Conds <- expand.grid(loc = Locations, delta_climate = Delta_climates)
sigma_beta_val = 0
rho_mean = 0.5
School.Terms = F
years = 10 
Pois <- F
rho_k <- 0.1
import = 1e-5
roi_model <- "CC_stand_smooth_AC"
# Locations full names
location_labels <- c(
  "Rome" = "Rome, Italy",
  "Rio_de_Janeiro" = "Rio de Janeiro, Brazil",
  "Toronto" = "Toronto, Canada",
  "Dubai" = "Dubai, UAE"
)

# Variables full names
variable_labels <- c(
    "Te" = "Temperature",
    "RH" = "Relative humidity"
)

df_effects = NULL
# For loop to generate big cumulative dataframe 
for(i in 1:nrow(Conds)){
    data_Path = paste("Results_Aggregated/R_sims_Te",
                      Conds[i, "delta_climate"], "_RH", Conds[i, "delta_climate"], 
                      "_",
                      "sigma", sigma_beta_val, "_",
                      paste("rho_mean", rho_mean, "_", sep=""),
                      Conds[i, "loc"], "_",
                      ifelse(School.Terms == T, "_School-Forced", ""),
                      sprintf("_rho_k%s", gsub("\\.", "-", as.character(rho_k))),
                      sprintf("_import_%s", as.character(import)),
                      sep = "")
    climate_effects <- readRDS(paste(data_Path, "/final_pars.rds", sep = "")) |> dplyr::select(True.e_Te, True.e_RH) |> distinct()
    sims_reg_tot <- readRDS(paste(data_Path, "/sims_reg_tot.rds", sep = "")) |>  filter(type == roi_model) |> 
      dplyr::mutate(location = Conds[i, "loc"],
             True_e_Te = climate_effects$True.e_Te,
             True_e_RH = climate_effects$True.e_RH) |>
      dplyr::ungroup() |> 
      dplyr::select(c(location, rep, .id, True_e_RH, True_e_Te, e_Te, e_RH, e_Te_se, e_RH_se))

    df_effects <- df_effects |>
      bind_rows(sims_reg_tot); 
    rm(sims_reg_tot)
}

df_diff <- df_effects |> mutate(bias_Te = abs(e_Te - True_e_Te), 
                                bias_RH = abs(e_RH - True_e_RH)) |> 
                  group_by(location, .id, True_e_Te, True_e_RH) |> 
                  dplyr::summarise(MAB_id_Te = mean(bias_Te), 
                            MAB_id_RH = mean(bias_RH),
                            MSE_id_Te = mean(e_Te_se),
                            MSE_id_RH = mean(e_RH_se)) |> 
                            ungroup() |> 
                  mutate(RMAB_id_Te = MAB_id_Te / abs(True_e_Te) * 100,
                         RMAB_id_RH = MAB_id_RH / abs(True_e_RH) * 100,
                         RMSE_id_Te = MSE_id_Te / abs(True_e_Te) * 100,
                         RMSE_id_RH = MSE_id_RH / abs(True_e_RH) * 100) |> 
                  select(-c(MAB_id_Te, MAB_id_RH, MSE_id_Te, MSE_id_RH)) |> 
        pivot_wider(
          names_from = c(True_e_Te, True_e_RH),
          values_from = c(RMAB_id_Te, RMAB_id_RH, RMSE_id_Te, RMSE_id_RH)
  ) |> 
  mutate( # New_Sensitivity - reference -> diff > 0 -> ref is better
    diff_RMAB_id_RH = `RMAB_id_RH_-0.02_-0.02` - `RMAB_id_RH_-0.04_-0.04`,
    diff_RMAB_id_Te = `RMAB_id_Te_-0.02_-0.02` - `RMAB_id_Te_-0.04_-0.04`,
    diff_RMSE_id_RH = `RMSE_id_RH_-0.02_-0.02` - `RMSE_id_RH_-0.04_-0.04`,
    diff_RMSE_id_Te = `RMSE_id_Te_-0.02_-0.02` - `RMSE_id_Te_-0.04_-0.04`
  ) |> select(location, .id, diff_RMAB_id_RH, diff_RMAB_id_Te, diff_RMSE_id_RH, diff_RMSE_id_Te) |> 
  pivot_longer(cols = c(diff_RMAB_id_RH, diff_RMAB_id_Te, diff_RMSE_id_RH, diff_RMSE_id_Te),
               names_to = "variable", values_to= 'vals')

  df_scatter <- df_diff |>
  separate(variable, into = c("metric", "var_type"), sep = "_id_") |>
  pivot_wider(
    names_from = metric,
    values_from = vals
  ) |>   mutate(var_type = factor(var_type, levels = c("Te", "RH")))

df_means <- df_scatter |> 
  dplyr::group_by(location, var_type) |> 
  dplyr::summarise(
    diff_RMAB = mean(diff_RMAB, na.rm = TRUE),
    diff_RMSE = mean(diff_RMSE, na.rm = TRUE)
  )

diff_delta_slim <- ggplot(df_scatter,
  aes(x = diff_RMAB, y = diff_RMSE)) +
  geom_point(aes(colour = as.factor(location)), size = 2, alpha = .7) +
  geom_point(data = df_means, colour = "black", 
    shape = 8, size = 3, stroke = 1, alpha = .6) + 
    geom_vline(xintercept = 0, colour = "black", linetype = "dashed") +
    geom_hline(yintercept = 0, colour = "black", linetype = "dashed") +
  facet_grid2(
    location ~ var_type,
    labeller = labeller(
      location = as_labeller(location_labels),
      var_type = as_labeller(variable_labels)
    ),
    scales = "fixed"
  ) + 
  labs(
    x = "Difference in RMAB (%)",
    y = "Difference in RMSE (%)") +
  theme_bw() +
  theme(
    strip.background = element_blank(),
    strip.text.x = element_text(size = 18, face = "bold", family = "Helvetica", colour = "black"),
    strip.text.y = element_text(size = 18, face = "bold", family = "Helvetica", colour = "black"),
    axis.title = element_text(size = 16, family = "Helvetica"),
    axis.text = element_text(size = 14, family = "Helvetica"),
    axis.ticks.x = element_line(),
    panel.grid.major = element_line(colour = "grey90", linewidth = 0.5),
    panel.grid.minor = element_blank(),
    legend.position = "none",
    title = element_blank(),
    plot.title = element_text(hjust = 0.5)
  ) +
  scale_colour_brewer(palette = "Set2")

ggsave(
  filename = "Figures/H_fig.pdf",
  plot = diff_delta_slim,
  width = 38,
  height = 38,
  units = "cm" ,
  scale = 1, device = cairo_pdf
)

## Figure I: Differences in process noise ----
rm(list=ls())

Locations <- c("Dubai", "Rio_de_Janeiro", "Rome", "Toronto")
delta_climate = -0.04
rho_k <- 0.1
Sigma_beta_vals = c(0, 0.05)
Conds <- expand.grid(loc = Locations, sigma_beta = Sigma_beta_vals)
rho_mean = 0.5
School.Terms = F
years = 10 
import = 1e-5

# Locations full names
location_labels <- c(
  "Rome" = "Rome, Italy",
  "Rio_de_Janeiro" = "Rio de Janeiro, Brazil",
  "Toronto" = "Toronto, Canada",
  "Dubai" = "Dubai, UAE"
)

# Variables full names
variable_labels <- c(
    "Te" = "Temperature",
    "RH" = "Relative humidity"
)

# Initialise the dataframe to store all
df_main = NULL
# For loop to generate big cumulative dataframe 
for(i in 1:nrow(Conds)){
    data_Path = paste("Results_Aggregated/R_sims_Te",
                      delta_climate, "_RH", delta_climate, 
                      "_",
                      "sigma", Conds[i, "sigma_beta"], "_",
                      paste("rho_mean", rho_mean, "_", sep=""),
                      Conds[i, "loc"], "_",
                      ifelse(School.Terms == T, "_School-Forced", ""),
                      sprintf("_rho_k%s", gsub("\\.", "-", as.character(rho_k))),
                      sprintf("_import_%s", as.character(import)),
                      sep = "")
    R0s = readRDS(paste(data_Path, "/final_pars.rds", sep = "")) |> dplyr::select(.id, R0) |> distinct()
    sims_reg_tot = readRDS(paste(data_Path, "/sims_reg_tot.rds", sep = "")) |> 
    dplyr::filter(type == "CC_stand_smooth_AC") |> 
      dplyr::mutate(location = Conds[i, "loc"],
                    sigma_beta = Conds[i, "sigma_beta"]) |>
      dplyr::select(-c(e_RH_low_CI, e_RH_high_CI,
                      log_SI_lag_VAR, mean_log_SI_lag,
                      R2, e_Te_low_CI, e_Te_high_CI,
                      TP_Te, e_Te_pow, TP_RH, e_RH_pow)) |> 
      dplyr::group_by(.id, location, sigma_beta) |> 
      dplyr::full_join(
                      y = R0s, 
                      by = c(".id")
                      )

    df_main = df_main |>
      bind_rows(sims_reg_tot); 
    rm(sims_reg_tot)
}

df_diff <- df_main |> mutate(bias_Te = abs(e_Te - delta_climate), 
                             bias_RH = abs(e_RH - delta_climate)) |> 
                  dplyr::group_by(location, .id, sigma_beta, .groups = "keep") |> 
                  dplyr::summarise(MAB_id_Te = mean(bias_Te), 
                            MAB_id_RH = mean(bias_RH),
                            MSE_id_Te = mean(e_Te_se),
                            MSE_id_RH = mean(e_RH_se)) |> 
                            ungroup() |> 
                  mutate(RMAB_id_Te = MAB_id_Te / abs(delta_climate) * 100,
                         RMAB_id_RH = MAB_id_RH / abs(delta_climate) * 100,
                         RMSE_id_Te = MSE_id_Te / abs(delta_climate) * 100,
                         RMSE_id_RH = MSE_id_RH / abs(delta_climate) * 100) |> 
                  select(-c(MAB_id_Te, MAB_id_RH, MSE_id_Te, MSE_id_RH)) |> 
        pivot_wider(
          names_from = c(sigma_beta),
          values_from = c(RMAB_id_Te, RMAB_id_RH, RMSE_id_Te, RMSE_id_RH)
  ) |> 
  mutate( # New_Sensitivity - reference -> diff > 0 -> ref is better
    diff_RMAB_id_RH = `RMAB_id_RH_0.05` - `RMAB_id_RH_0`,
    diff_RMAB_id_Te = `RMAB_id_Te_0.05` - `RMAB_id_Te_0`,
    diff_RMSE_id_RH = `RMSE_id_RH_0.05` - `RMSE_id_RH_0`,
    diff_RMSE_id_Te = `RMSE_id_Te_0.05` - `RMSE_id_Te_0`
  ) |> select(location, .id, diff_RMAB_id_RH, diff_RMAB_id_Te, diff_RMSE_id_RH, diff_RMSE_id_Te) |> 
  pivot_longer(cols = c(diff_RMAB_id_RH, diff_RMAB_id_Te, diff_RMSE_id_RH, diff_RMSE_id_Te),
               names_to= "variable", values_to= 'vals')

df_scatter <- df_diff |>
  tidyr::separate(variable, into = c("metric", "var_type"), sep = "_id_") |>
  tidyr::pivot_wider(
    names_from = metric,
    values_from = vals
  ) |> dplyr::mutate(var_type = factor(var_type, levels = c("Te", "RH")))

df_means <- df_scatter |> 
  dplyr::group_by(location, var_type) |> 
  dplyr::summarise(
    diff_RMAB = mean(diff_RMAB, na.rm = TRUE),
    diff_RMSE = mean(diff_RMSE, na.rm = TRUE)
  )

diff_sigma_slim <- ggplot(df_scatter,
  aes(x = diff_RMAB, y = diff_RMSE)) +
  geom_point(aes(colour = as.factor(location)), size = 2, alpha = .7) +
  geom_point(data = df_means, colour = "black", 
    shape = 8, size = 3, stroke = 1, alpha = .6) + 
    geom_vline(xintercept = 0, colour = "black", linetype = "dashed") +
    geom_hline(yintercept = 0, colour = "black", linetype = "dashed") +
  facet_grid2(
    location ~ var_type,
    labeller = labeller(
      location = as_labeller(location_labels),
      var_type = as_labeller(variable_labels)
    ),
    scales = "fixed"#, independent = "all"
  ) + 
  labs(
  x = "Difference in RMAB (%)",
  y = "Difference in RMSE (%)"
) + 
  theme_bw() +
  theme(
    strip.background = element_blank(),
    strip.text.x = element_text(size = 18, face = "bold", family = "Helvetica", colour = "black"),
    strip.text.y = element_text(size = 18, face = "bold", family = "Helvetica", colour = "black"),
    axis.title = element_text(size = 16, family = "Helvetica"),
    axis.text = element_text(size = 14, family = "Helvetica"),
    axis.ticks.x = element_line(),
    panel.grid.major = element_line(colour = "grey90", linewidth = 0.5),
    panel.grid.minor = element_blank(),
    legend.position = "none",
    title = element_blank(),
    plot.title = element_text(hjust = 0.5)
  ) +
  scale_colour_brewer(palette = "Set2")

ggsave(
  filename = "Figures/I_fig.pdf",
  plot = diff_sigma_slim,
  width = 38,
  height = 38,
  units = "cm" ,
  scale = 1, device = cairo_pdf
)

## Figure J: Differences including term-time forcing ----
rm(list=ls())

Locations <- c("Dubai", "Rio_de_Janeiro", "Rome", "Toronto")
delta_climate = -0.04
rho_k <- 0.1
sigma_beta_val = 0
School.Terms = c(T,F)
Conds <- expand.grid(loc = Locations, term_time = School.Terms)
rho_mean = 0.5
years = 10 
import = 1e-5

# Locations full names
location_labels <- c(
  "Rome" = "Rome, Italy",
  "Rio_de_Janeiro" = "Rio de Janeiro, Brazil",
  "Toronto" = "Toronto, Canada",
  "Dubai" = "Dubai, UAE"
)

# Variables full names
variable_labels <- c(
    "Te" = "Temperature",
    "RH" = "Relative humidity"
)

# Initialise the dataframe to store all
df_main = NULL
# For loop to generate big cumulative dataframe 
for(i in 1:nrow(Conds)){
    data_Path = paste("Results_Aggregated/R_sims_Te",
                      delta_climate, "_RH", delta_climate, 
                      "_",
                      "sigma", sigma_beta_val, "_",
                      paste("rho_mean", rho_mean, "_", sep=""),
                      Conds[i, "loc"], "_",
                      sprintf("_rho_k%s", gsub("\\.", "-", as.character(rho_k))),
                      sprintf("_import_%s", as.character(import)),
                      ifelse(Conds[i, "term_time"] == T, "_School-Forced", ""),
                      sep = "")
    sims_reg_tot = readRDS(paste(data_Path, "/sims_reg_tot.rds", sep = "")) |> 
    dplyr::filter(type == "CC_stand_smooth_AC") |> 
      dplyr::mutate(location = Conds[i, "loc"],
                    term_time = Conds[i, "term_time"]) |>
      dplyr::ungroup() |> 
      dplyr::select(-c(e_RH_low_CI, e_RH_high_CI, type,
                      log_SI_lag_VAR, mean_log_SI_lag,
                      R2, e_Te_low_CI, e_Te_high_CI,
                      TP_Te, e_Te_pow, TP_RH, e_RH_pow))

    df_main = df_main |>
      bind_rows(sims_reg_tot); 
    rm(sims_reg_tot)
}

final_pars <- readRDS(paste(data_Path, "/final_pars.rds", sep = ""))
True_e_Te <- unique(final_pars$True.e_Te)
True_e_RH <- unique(final_pars$True.e_RH)

df_diff <- df_main |> mutate(bias_Te = abs(e_Te - delta_climate), 
                  bias_RH = abs(e_RH - delta_climate)) |> 
                  select(-c(e_Te, e_RH)) |> 
                  group_by(location, .id, term_time) |> 
                  dplyr::summarise(MAB_id_Te = mean(bias_Te), 
                            MAB_id_RH = mean(bias_RH),
                            MSE_id_Te = mean(e_Te_se),
                            MSE_id_RH = mean(e_RH_se)) |> 
                            ungroup() |> 
                  mutate(RMAB_id_Te = MAB_id_Te / abs(True_e_Te) * 100,
                         RMAB_id_RH = MAB_id_RH / abs(True_e_RH) * 100,
                         RMSE_id_Te = MSE_id_Te / abs(True_e_Te) * 100,
                         RMSE_id_RH = MSE_id_RH / abs(True_e_RH) * 100) |> 
                  select(-c(MAB_id_Te, MAB_id_RH, MSE_id_Te, MSE_id_RH)) |> 
        pivot_wider(
          names_from = c(term_time),
          values_from = c(RMAB_id_Te, RMAB_id_RH, RMSE_id_Te, RMSE_id_RH)) |> 
  mutate(
    diff_RMAB_id_RH = `RMAB_id_RH_TRUE` - `RMAB_id_RH_FALSE`,
    diff_RMAB_id_Te = `RMAB_id_Te_TRUE` - `RMAB_id_Te_FALSE`,
    diff_RMSE_id_RH = `RMSE_id_RH_TRUE` - `RMSE_id_RH_FALSE`,
    diff_RMSE_id_Te = `RMSE_id_Te_TRUE` - `RMSE_id_Te_FALSE`
  ) |> select(location, .id, diff_RMAB_id_RH, diff_RMAB_id_Te, diff_RMSE_id_RH, diff_RMSE_id_Te) |> 
  pivot_longer(cols = c(diff_RMAB_id_RH, diff_RMAB_id_Te, diff_RMSE_id_RH, diff_RMSE_id_Te),
               names_to= "variable", values_to= 'vals')

  df_scatter <- df_diff |>
  separate(variable, into = c("metric", "var_type"), sep = "_id_") |>
  pivot_wider(
    names_from = metric,
    values_from = vals
  ) |>   mutate(var_type = factor(var_type, levels = c("Te", "RH"))) 

df_means <- df_scatter |> 
  dplyr::group_by(location, var_type) |> 
  dplyr::summarise(
    diff_RMAB = mean(diff_RMAB, na.rm = TRUE),
    diff_RMSE = mean(diff_RMSE, na.rm = TRUE)
  )

diff_Term_Time_slim <- ggplot(df_scatter,
  aes(x = diff_RMAB, y = diff_RMSE)) +
  geom_point(aes(colour = as.factor(location)), size = 2, alpha = .7) +
  geom_point(data = df_means, colour = "black", 
    shape = 8, size = 3, stroke = 1, alpha = .6) + 
    geom_vline(xintercept = 0, colour = "black", linetype = "dashed") +
    geom_hline(yintercept = 0, colour = "black", linetype = "dashed") +
  facet_grid2(
    location ~ var_type,
    labeller = labeller(
      location = as_labeller(location_labels),
      var_type = as_labeller(variable_labels)
    ),
    scales = "fixed"
  ) + 
  labs(
    x = "Difference in RMAB (%)",
    y = "Difference in RMSE (%)"
  ) +
  theme_bw() +
  theme(
    strip.background = element_blank(),
    strip.text.x = element_text(size = 18, face = "bold", family = "Helvetica", colour = "black"),
    strip.text.y = element_text(size = 18, face = "bold", family = "Helvetica", colour = "black"),
    axis.title = element_text(size = 16, family = "Helvetica"),
    axis.text = element_text(size = 14, family = "Helvetica"),
    axis.ticks.x = element_line(),
    panel.grid.major = element_line(colour = "grey90", linewidth = 0.5),
    panel.grid.minor = element_blank(),
    legend.position = "none",
    title = element_blank(),
    plot.title = element_text(hjust = 0.5)
  ) +
  scale_colour_brewer(palette = "Set2")

ggsave(
  filename = "Figures/J_fig.pdf",
  plot = diff_Term_Time_slim,
  width = 38,
  height = 38,
  units = "cm" ,
  scale = 1, device = cairo_pdf
)

## Figure K: difference in dimension of spline -----
rm(list=ls())

Locations <- c("Dubai", "Rio_de_Janeiro", "Rome", "Toronto")
delta_climate = -0.04
Rho_ks <- 0.1
sigma_beta_val = 0
Knots = 52
Conds <- expand.grid(loc = Locations, rho_k = Rho_ks, k_max = Knots)
rho_mean = 0.5
School.Terms = F
years = 10
import = 1e-5

# Locations full names
location_labels <- c(
  "Rome" = "Rome, Italy",
  "Rio_de_Janeiro" = "Rio de Janeiro, Brazil",
  "Toronto" = "Toronto, Canada",
  "Dubai" = "Dubai, UAE"
)

# Variables full names
variable_labels <- c(
  diff_RMAB_id_Te = "Temperature",
  diff_RMAB_id_RH = "Relative humidity"
)

# Initialise the dataframe to store all
df_52 = NULL
# For loop to generate big cumulative dataframe
for(i in 1:nrow(Conds)){
        data_Path = paste("Results_Aggregated/R_sims_Te",
                      delta_climate, "_RH", delta_climate,
                      "_",
                      "sigma", sigma_beta_val, "_",
                      paste("rho_mean", rho_mean, "_", sep=""),
                      Conds[i, "loc"], "_",
                      sprintf("_rho_k%s", gsub("\\.", "-", as.character(Conds[i, "rho_k"]))),
                      sprintf("_import_%s", as.character(import)),
                      ifelse(School.Terms == T, "_School-Forced", ""),
                      sep = "")
    sims_reg_tot <- readRDS(paste(data_Path, "/sims_reg_tot.rds", sep = "")) |> 
      dplyr::filter(type == "CC_stand_smooth_AC") |> 
      dplyr::mutate(loc = Conds[i, "loc"],
                    rho_k = Conds[i, "rho_k"],
                    k_max = Conds[i, "k_max"]) |>
        dplyr::ungroup() |> 
      dplyr::select(loc, rep, .id, k_max, e_Te, e_RH, e_Te_se, e_RH_se)
    df_52 = df_52 |>
      bind_rows(sims_reg_tot); 
    rm(sims_reg_tot)
}

df_75 <- readRDS("Results_Aggregated/Regressions_75Knots.rds") |> dplyr::select(-c(k_index, edf))

df_main <- rbind(df_75,df_52)

df_diff <- df_main |> mutate(bias_Te = abs(e_Te - delta_climate), 
                             bias_RH = abs(e_RH - delta_climate)) |> 
                  dplyr::group_by(loc, .id, k_max, .groups = "keep") |> 
                  dplyr::summarise(MAB_id_Te = mean(bias_Te), 
                            MAB_id_RH = mean(bias_RH),
                            MSE_id_Te = mean(e_Te_se),
                            MSE_id_RH = mean(e_RH_se)) |> 
                            dplyr::ungroup() |> 
                  dplyr::mutate(RMAB_id_Te = MAB_id_Te / abs(delta_climate) * 100,
                         RMAB_id_RH = MAB_id_RH / abs(delta_climate) * 100,
                         RMSE_id_Te = MSE_id_Te / abs(delta_climate) * 100,
                         RMSE_id_RH = MSE_id_RH / abs(delta_climate) * 100) |> 
                  dplyr::select(-c(MAB_id_Te, MAB_id_RH, MSE_id_Te, MSE_id_RH)) |> 
        tidyr::pivot_wider(
          names_from = c(k_max),
          values_from = c(RMAB_id_Te, RMAB_id_RH, RMSE_id_Te, RMSE_id_RH)
  ) |> 
  dplyr::mutate( # New_Sensitivity - reference -> diff > 0 -> ref is better
    diff_RMAB_id_RH = `RMAB_id_RH_75` - `RMAB_id_RH_52`,
    diff_RMAB_id_Te = `RMAB_id_Te_75` - `RMAB_id_Te_52`,
    diff_RMSE_id_RH = `RMSE_id_RH_75` - `RMSE_id_RH_52`,
    diff_RMSE_id_Te = `RMSE_id_Te_75` - `RMSE_id_Te_52`
  ) |> 
  dplyr::select(loc, .id, diff_RMAB_id_RH, diff_RMAB_id_Te, diff_RMSE_id_Te, diff_RMSE_id_RH) |> 
  tidyr::pivot_longer(cols = c(diff_RMAB_id_RH, diff_RMAB_id_Te, diff_RMSE_id_RH, diff_RMSE_id_Te),
               names_to= "variable", values_to= 'vals')

 variable_labels <- c(
    "Te" = "Temperature",
    "RH" = "Relative humidity"
)

 df_scatter <- df_diff |>
  tidyr::separate(variable, into = c("metric", "var_type"), sep = "_id_") |>
  tidyr::pivot_wider(
    names_from = metric,
    values_from = vals
  ) |> dplyr::mutate(var_type = factor(var_type, levels = c("Te", "RH")))

df_means <- df_scatter |> 
  dplyr::group_by(loc, var_type) |> 
  dplyr::summarise(
    diff_RMAB = mean(diff_RMAB, na.rm = TRUE),
    diff_RMSE = mean(diff_RMSE, na.rm = TRUE)
  )

  diff_knots_slim <- ggplot(df_scatter,
  aes(x = diff_RMAB, y = diff_RMSE)) +
  geom_point(aes(colour = as.factor(loc)), size = 2, alpha = .7) +
  geom_point(data = df_means, colour = "black", 
    shape = 8, size = 3, stroke = 1, alpha = .6) + 
    geom_vline(xintercept = 0, colour = "black", linetype = "dashed") +
    geom_hline(yintercept = 0, colour = "black", linetype = "dashed") +
  facet_grid2(
    loc ~ var_type,
    labeller = labeller(
      loc = as_labeller(location_labels),
      var_type = as_labeller(variable_labels)
    ),
    scales = "fixed"#, independent = "all"
  ) + 
  labs(
    x = "Difference in RMAB (%)",
    y = "Difference in RMSE (%)"
  ) +
  theme_bw() +
  theme(
    strip.background = element_blank(),
    strip.text.x = element_text(size = 18, face = "bold", family = "Helvetica", colour = "black"),
    strip.text.y = element_text(size = 18, face = "bold", family = "Helvetica", colour = "black"),
    axis.title = element_text(size = 16, family = "Helvetica"),
    axis.text = element_text(size = 14, family = "Helvetica"),
    axis.ticks.x = element_line(),
    panel.grid.major = element_line(colour = "grey90", linewidth = 0.5),
    panel.grid.minor = element_blank(),
    legend.position = "none",
    title = element_blank(),
    plot.title = element_text(hjust = 0.5)
  ) +
  scale_colour_brewer(palette = "Set2")

ggsave(
  filename = "Figures/K_fig.pdf",
  plot = diff_knots_slim,
  width = 38,
  height = 38,
  units = "cm" ,
  scale = 1, device = cairo_pdf
)

## Figure L: performances for added complexity in lagged efffect on Toronto ----
location <- "Toronto"

### First import the results for the model version with additional lag ----
delta_Now <- -0.016
delta_Lag <- -0.027
rho_k <- 0.1
sigma_beta_val <- 0
rho_mean <- 0.5
School.Terms <- F
years <- 10 
import <- 1e-5

#### Aggregate results across synthetic pathogens
file_path = paste("Results/R_sims_Te",delta_Now,
                "_Te_Lag",delta_Lag,
                "_RH", delta_Now,
                "_RH_Lag",delta_Lag,
                "_sigma",sigma_beta_val, "_",
                ifelse(rho_mean != 0, paste("rho_mean", rho_mean, "_", sep="") , ""),
                location, "_",
                sprintf("_rho_k%s", gsub("\\.", "-", as.character(rho_k))),
                sprintf("_import_%s", as.character(import)),
                ifelse(School.Terms == 1, "_School-Forced", ""),
                "_LAGGED",
                sep = "")

if(!file.exists(paste0(file_path,"/sims_reg_tot.rds"))){
n <- length(list.files(
  path = file_path,
  pattern = "^results_[0-9]+\\.RData$"
))

sims_reg_tot <- NULL
for(id in 1:n){
data_Path <- paste(file_path,
                    "/results_", as.character(id), ".RData",
                    sep = "")
    load(data_Path)
    sims_reg_tot <- rbind(sims_reg_tot, sim_reg |> dplyr::mutate(.id = id ,location = location))
}
saveRDS(sims_reg_tot, file = paste0(file_path,"/sims_reg_tot.rds"))
rm(sims_reg_tot)
}

sims_reg_tot_lag <- readRDS(file = paste0(file_path,"/sims_reg_tot.rds")) |>
      dplyr::filter(type == "CC_stand_smooth_AC") |> 
      dplyr::mutate(type = "CC_stand_smooth_AC_lag") |> 
      dplyr::select(-c(e_RH_low_CI, e_RH_high_CI,
                      log_SI_lag_VAR, mean_log_SI_lag,
                      R2, e_Te_low_CI, e_Te_high_CI,
                      TP_Te, TP_RH)) |> 
      dplyr::group_by(type, .id, location) |> 
      dplyr::summarise(
        # First for the t-1 effect estimates
        MAB_Te = mean(abs(e_Te - delta_Now)),
        MAB_RH = mean(abs(e_RH - delta_Now)),
        sd_AB_Te = sd(abs(e_Te - delta_Now)),
        sd_AB_RH = sd(abs(e_RH - delta_Now)),
        MSE_Te = mean(e_Te_se),
        MSE_RH = mean(e_RH_se),
        MA_e_Te = mean(abs(e_Te)),
        MA_e_RH = mean(abs(e_RH)),
        SDSE_Te = sd(e_Te_se),
        SDSE_RH = sd(e_RH_se),
        MPOW_Te = mean(e_Te_pow),
        MPOW_RH = mean(e_RH_pow),
        # Now for the t-2 effect estimates
        MAB_Te_Lag = mean(abs(e_Te_Lag - delta_Lag)),
        MAB_RH_Lag = mean(abs(e_RH_Lag - delta_Lag)),
        sd_AB_Te_Lag = sd(abs(e_Te_Lag - delta_Lag)),
        sd_AB_RH_Lag = sd(abs(e_RH_Lag - delta_Lag)),
        MSE_Te_Lag = mean(e_Te_se_Lag),
        MSE_RH_Lag = mean(e_RH_se_Lag),
        MA_e_Te_Lag = mean(abs(e_Te_Lag)),
        MA_e_RH_Lag = mean(abs(e_RH_Lag)),
        SDSE_Te_Lag = sd(e_Te_se_Lag),
        SDSE_RH_Lag = sd(e_RH_se_Lag),
        MPOW_Te_Lag = mean(e_Te_pow_Lag),
        MPOW_RH_Lag = mean(e_RH_pow_Lag)
      ) |> dplyr::ungroup() |> 
      dplyr::mutate(true_val_Now = delta_Now,
                    true_val_Lag = delta_Lag)

### Import the baseline model results ----
data_Path <- paste0("Results_Aggregated/R_sims_Te-0.04_RH-0.04_sigma0_rho_mean0.5_",location,"__rho_k0-1_import_1e-05")
final_pars <- readRDS(paste(data_Path, "/final_pars.rds", sep = ""))
df_plot_baseline <- readRDS(paste(data_Path, "/sims_reg_tot.rds", sep = "")) |> 
dplyr::filter(type == "CC_stand_smooth_AC") |> 
dplyr::mutate(location = location) |>
dplyr::select(-c(e_RH_low_CI, e_RH_high_CI,
                log_SI_lag_VAR, mean_log_SI_lag,
                R2, e_Te_low_CI, e_Te_high_CI,
                TP_Te, TP_RH)) |> 
dplyr::group_by(type, .id, location) |> 
dplyr::summarise(
  MAB_Te = mean(abs(e_Te - unique(final_pars[,"True.e_Te"]))), # Computing the mean absolute bias for temperature
  MAB_RH = mean(abs(e_RH - unique(final_pars[,"True.e_RH"]))),
  sd_AB_Te = sd(abs(e_Te - unique(final_pars[,"True.e_Te"]))), # Computing the standard deviation in absolute biases for temperature
  sd_AB_RH = sd(abs(e_RH - unique(final_pars[,"True.e_RH"]))),
  MSE_Te = mean(e_Te_se),
  MSE_RH = mean(e_RH_se),
  MA_e_Te = mean(abs(e_Te)),
  MA_e_RH = mean(abs(e_RH)),
  SDSE_Te = sd(e_Te_se),
  SDSE_RH = sd(e_RH_se),
  MPOW_Te = mean(e_Te_pow),
  MPOW_RH = mean(e_RH_pow)
) |> dplyr::ungroup() |> 
dplyr::mutate(true_val = if(unique(final_pars[,"True.e_Te"]) == unique(final_pars[,"True.e_RH"])){
                          unique(final_pars[,"True.e_Te"])}) |> 
tidyr::pivot_longer(
  cols = c(MAB_Te, MAB_RH, MSE_Te, MSE_RH, MPOW_Te, MPOW_RH),
  names_to = c(".value", "variable"),
  names_sep = "_") |> 
dplyr::transmute(
  type,
  .id,
  RMAB = MAB / abs(true_val) * 100,
  RMSE = MSE / abs(true_val) * 100,
  MPOW,
  variable = factor(variable, levels = c("Te", "RH"))
)

# Models full names
model_labels <- c(CC_stand_smooth_AC = "Baseline",
                  CC_stand_smooth_AC_lag = "Lagged")

# Adact structure for plot ----
df_plot_lag <- sims_reg_tot_lag |>
  dplyr::rename_with(
    ~ paste0(., "_Now"),
    -c(type, .id, location, dplyr::ends_with("_Lag"), true_val_Now, true_val_Lag)
  ) |>
  tidyr::pivot_longer(
    cols = dplyr::matches("_(Now|Lag)$"),
    names_to = c(".value", "period"),
    names_pattern = "(.+)_(Now|Lag)$"
  ) |>
  dplyr::mutate(period = factor(period, levels = c("Now", "Lag"))) |> 
  pivot_longer(
    cols = c(MAB_Te, MAB_RH, MSE_Te, MSE_RH, MPOW_Te, MPOW_RH),
    names_to = c(".value", "variable"),
    names_sep = "_"
  ) %>%
  transmute(
    type,
    .id,
    period,
    RMAB = MAB / abs(true_val) * 100,
    RMSE = MSE / abs(true_val) * 100,
    MPOW,
    variable = factor(variable, levels = c("Te", "RH")),
    period = factor(period, levels = c("Now", "Lag"))
  )


# Plotting design variables
baseline_colour <- "grey30"
baseline_shape  <- 1

period_labels <- c(
  Now = "hat(delta)[0]",
  Lag = "hat(delta)[-1]"
)

variable_labels <- c(Te = "Temperature", RH = "Relative Humidity")

global_limits <- c(0,1)

scale_mean_power <- viridis::scale_color_viridis(
  option = "plasma",
  direction = 1,
  begin = 0,
  end = .95,
  limits = global_limits,
  oob = scales::squish,
  name = "Mean power"
)

# Plot object creation
lagged_perf <- ggplot() +
  # Bubbles for lagged model
  geom_point(
    data = subset(df_plot_lag, type == "CC_stand_smooth_AC_lag"),
    aes(x = RMAB, y = RMSE, colour = MPOW, shape = type, size = type),
    alpha = .7, stroke = 1
  ) +
  # Points for baseline model
  geom_point(
    data = subset(df_plot_baseline, type == "CC_stand_smooth_AC"),
    aes(x = RMAB, y = RMSE, shape = type, size = type),
    colour = baseline_colour, alpha = .6, stroke = .5
  ) +
  facet_grid(
    period ~ variable,
    labeller = labeller(
      variable = variable_labels,
      period = as_labeller(period_labels, label_parsed)
    ),
    scales = "free_x"
  ) +
  theme_bw() +
  scale_mean_power +
  scale_shape_manual(
    breaks = c("CC_stand_smooth_AC_lag", "CC_stand_smooth_AC"),
    values = c(
      CC_stand_smooth_AC_lag = 16,
      CC_stand_smooth_AC = baseline_shape
    ),
    labels = model_labels,
    name = "Model"
  ) +
  scale_size_manual(
    breaks = c("CC_stand_smooth_AC_lag", "CC_stand_smooth_AC"),
    values = c(
      CC_stand_smooth_AC_lag = 4,
      CC_stand_smooth_AC = 2
    )
  ) +
  xlab("RMAB (%)") +
  ylab("RMSE (%)") +
  theme(
    panel.grid.minor = element_blank(),
    strip.background = element_blank(),
    strip.text.x = element_text(
      size = 18, face = "bold",
      family = "Helvetica", colour = "black"
    ),
    strip.text.y = element_text(
      size = 18, face = "bold",
      family = "Helvetica", colour = "black"
    ),
    axis.title = element_text(
      size = 16, family = "Helvetica"
    ),
    axis.text = element_text(
      size = 12, family = "Helvetica"
    ),
    axis.ticks.x = element_line(),
    panel.grid.major.x = element_line(colour = "grey90"),
    panel.grid.major.y = element_line(colour = "grey90"),
    plot.title = element_blank(),
    legend.position = "bottom",
    legend.spacing.x = unit(5, "cm"),
    panel.spacing.x = unit(0.75, "cm"),
    legend.text = element_text(
      size = 14, family = "Helvetica"
    ),
    legend.title = element_text(
      size = 16, family = "Helvetica"
    )
  ) +
  scale_x_continuous(
    breaks = seq(0, max(df_plot_lag$RMAB, na.rm = TRUE), by = 20)
  ) +
  scale_y_continuous(
    breaks = seq(0, max(df_plot_lag$RMSE, na.rm = TRUE), by = 40)
  ) +
  guides(
    colour = guide_colourbar(
      barwidth  = unit(7.5, "cm"),
      barheight = unit(0.4, "cm"),
      title.position = "left",
      title.theme = element_text(margin = margin(r = 20))
    ),
    size = "none",
    shape = guide_legend(
      override.aes = list(
        shape  = c(16, baseline_shape),
        size   = c(4, 2),
        colour = c("black", baseline_colour),
        alpha  = 1,
        stroke = c(1, 1)
      )
    )
  )

ggsave(
  filename = "Figures/L_fig.pdf",
  plot = lagged_perf,
  width = 35,
  height = 22,
  units = "cm",
  device = cairo_pdf
)