# '
# The file "figure_yearly_effects_plot_adjustments.r" takes the yearly treatment effect results 
# from "bld/results/" and creates the yearly effects plots with a 95% confidence interval
# (Supplementary Figures  SI2, SI3 and SI4 of the paper) and stores them to "bld/figures/"
# as pdf files. The y-axis scale is adjusted to be the same for each outcome group for comparability,
# i.e. the same scale for the mechanism variables, the same scale for hospitalizations and
# the same scale for mortality.
# 
# '

rm(list = ls())
options(scipen = 999)
options(warn = -0) # 0 to turn on again, -1 to turn off
source(paste0("src/code/ggplot_theme_publication.r"))

# Load libraries
library(tidyverse)
library(Cairo)

# Load model specification
model <- fromJSON(file = paste0("src/code/specifications.json"))

for (outcome in c("cases_prevalence", "death_prevalence")){#
  
  # Get and read-in all files.
  results = lapply(seq_along(model$outcomes)[-1], function(i){
    readRDS(paste0("bld/results/", names(model$outcomes[i]), "_", outcome, "_yearly_effects.rds"))
  })
  
  impact <- lapply(seq_along(results), function(i){
    ii <- grep(":treated", rownames(results[[i]]), fixed=T)
    as.data.frame(cbind(
      beta = results[[i]][ii,1],
      SE = results[[i]][ii,2],
      year = c(2004:2017)
    ))
  })
  impact <- lapply(impact, function(i){
    rbind(i, c(0, 0, 2003))
  })
  
  max <- ceiling(max(unlist(lapply(impact, function(i){#
    max(i$beta + i$SE * 1.96)
  })))*100)/100
  min <- floor(min(unlist(lapply(impact, function(i){ #
    min(i$beta - i$SE * 1.96)
  })))*100)/100
  
  p <- lapply(impact, function(i){
    ggplot(i, aes(year, beta)) +
    geom_point(size = .75) +
    geom_line() +
    geom_ribbon(aes(ymin = beta - SE * 1.96, ymax = beta + SE * 1.96), alpha = 0.3) +
    geom_vline(xintercept = 2006) +
    geom_hline(yintercept = 0) +
    scale_y_continuous(
      name = "Treatment effect (per 1,000)",
      limits = c(min, max),
      breaks = round(seq(min, max, (max+abs(min))/10),2))  +
    scale_x_continuous(
      name = "Year",
      breaks = 2003:2017) +
    scale_fill_Publication() +
    scale_colour_Publication() +
    theme_Publication() +
    theme(
      #legend.title = element_blank(),
      axis.text.x = element_text(angle = 45, hjust = 1))
  # print(plot)
  })
  
  # Save it 
  for (i in seq_along(p)){
    ggsave(
      plot = p[[i]], 
      filename = paste0("bld/figures/", names(model$outcomes[i+1]), "_", outcome, ".pdf"), 
      device = cairo_pdf,
      height = 9, 
      width = 15, 
      units = "cm"
    )
  }
}


# Remove these elements from the list
my_list <- model$mechanisms[-grep("_over_limit_relative", names(model$mechanisms))]

# Get and read-in all files.
results = lapply(seq_along(my_list)[-1], function(i){
  readRDS(paste0("bld/results/", names(my_list[i]), "_yearly_effects.rds")) 
})

impact <- lapply(seq_along(results), function(i){
  ii <- grep(":treated", rownames(results[[i]]), fixed=T)
  as.data.frame(cbind(
    beta = results[[i]][ii,1],
    SE = results[[i]][ii,2],
    year = c(2004:2017)
  ))
})
impact <- lapply(impact, function(i){
  rbind(i, c(0, 0, 2003))
})

max <- ceiling(max(unlist(lapply(1:6, function(i){#
  max(impact[[i]]$beta + impact[[i]]$SE * 1.96)
})))*100)/100#
min <- floor(min(unlist(lapply(1:6, function(i){ #
  min(impact[[i]]$beta - impact[[i]]$SE * 1.96)
})))*100)/100

p <- lapply(impact, function(i){
  ggplot(i, aes(year, beta)) +
    geom_point(size = .75) +
    geom_line() +
    geom_ribbon(aes(ymin = beta - SE * 1.96, ymax = beta + SE * 1.96), alpha = 0.3) +
    geom_vline(xintercept = 2006) +
    geom_hline(yintercept = 0) +
    scale_y_continuous(
      name = "Treatment effect (%)",
      limits = c(min, max),
      breaks = round(seq(min, max, (max+abs(min))/10), 2))  +
    scale_x_continuous(
      name = "Year",
      breaks = 2003:2017) +
    scale_fill_Publication() +
    scale_colour_Publication() +
    theme_Publication() +
    theme(
      #legend.title = element_blank(),
      axis.text.x = element_text(angle = 45, hjust = 1))
  # print(plot)
})

# Save it 
for (i in seq_along(p)){
  ggsave(
    plot = p[[i]], 
    filename = paste0("bld/figures/", names(my_list[i+1]), ".pdf"), 
    device = cairo_pdf,
    height = 9, 
    width = 15, 
    units = "cm"
  )
}