# '
# The file "figure_main_results.r" takes the estimation results created by "table_mechanism_controls.r" 
# and "table_health_effects.r", stored in "bld/results/" and summarizes and visualizes
# three main results in a graph.
# This creates Figure 2 of the paper, stored as a pdf to "bld/figures/".
# 
# '

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

# Load libraries
library(tidyverse)
library(Cairo)

# Load stored results
fire <- readRDS("bld/results/fires_yearly_effects.rds") 
pm25 <- readRDS("bld/results/pm25_ugm3_yearly_effects.rds") 
resp_deaths <- readRDS("bld/results/J_death_prevalence_yearly_effects.rds") 

i <- grep(":treated", rownames(fire), fixed=T)
ii <- grep(":treated", rownames(pm25), fixed=T)
iii <- grep(":treated", rownames(resp_deaths), fixed=T)

# Prepare the impact data frame
impact <- as.data.frame(cbind(
  beta.fire = fire[i,1],
  SE.fire = fire[i,2],
  beta.pm25 = pm25[ii,1],
  SE.pm25 = pm25[ii,2],
  beta.resp = resp_deaths[iii,1],
  SE.resp = resp_deaths[iii,2]
))
impact$year <- c(2004:2017)
impact <- rbind(impact, c(0, 0, 0, 0, 0, 0, 2003))

# Define plot limits
max <- ceiling(max(impact$beta.fire + impact$`SE.fire` * 1.96) * 100) / 100
min <- floor(min(impact$beta.fire - impact$`SE.fire` * 1.96) * 100) / 100

# Convert data to long format for easier ggplot2 handling of aesthetics
impact_long <- data.frame(
  year = rep(impact$year, 3),
  beta = c(impact$beta.fire, impact$beta.pm25, impact$beta.resp),
  SE = c(impact$SE.fire, impact$SE.pm25, impact$SE.resp),
  variable = factor(rep(c("Fires", "PM2.5", "Mortality from resp. diseases"), each = nrow(impact)))
)

# Define colors and point shapes for each variable
cols <- c("Fires" = "#386cb0", "PM2.5" = "#f87f01", "Mortality from resp. diseases" = "#7fc97f")
shapes <- c("Fires" = 16, "PM2.5" = 17, "Mortality from resp. diseases" = 15)
linetypes <- c("Fires" = "solid", "PM2.5" = "dashed", "Mortality from resp. diseases" = "dotted")

# Create the plot
p <- ggplot(impact_long) +
  # Add lines with color mapped to 'variable'
  geom_line(aes(year, beta, color = variable, linetype = variable), linewidth = 1) +
  
  # Add points with color and shape mapped to 'variable'
  geom_point(aes(year, beta, color = variable, shape = variable), size = 2) +
  
  # Add ribbons for confidence intervals
  geom_ribbon(aes(x = year, ymin = beta - SE * 1.96, ymax = beta + SE * 1.96, fill = variable), alpha = 0.3) +
  
  # Reference lines
  geom_vline(xintercept = 2006) +
  geom_hline(yintercept = 0) +
  
  # X and Y axis settings
  scale_x_continuous(breaks = 2003:2017, name = "Year") +
  scale_y_continuous(
    name = "Treatment effect",
    limits = c(min, max),
    breaks = round(seq(min, max, (max + abs(min)) / 10), 2)
  ) +
  
  # Color, shape, and fill scales mapped to 'variable'
  scale_colour_manual(values = cols) +
  scale_shape_manual(values = shapes) +
  scale_fill_manual(values = cols) +
  scale_linetype_manual(values = linetypes) +
  
  # Theme and legend settings
  theme_Publication() +
  theme(
    legend.title = element_blank(),
    axis.text.x = element_text(angle = 45, hjust = 1)
  )

# Display the plot
print(p)

# Save it 
ggsave(plot = p, filename = "bld/figures/figure2.pdf", device = cairo_pdf, height = 9, width = 15, units = "cm")