# '
# The file "table_health_effects.r" takes the panel data for estimation 
# from "bld/data/" and estimates the treatment effect on the hospitalization and
# death rate of different diseases. The exact diseases
# are specified in "src/code/specifications.json" as outcomes. This creates Supplementary
# Table SI5 and SI6 of the paper, respectively. Tables are stored as tex files to "bld/tables".
# This file also creates the yearly effects plots with a 95% confidence interval
# (Supplementary Figures SI3 and SI4 of the paper) and stores them to "bld/figures/"
# as png files. Figure scales are later adjusted by src/code/figure_yearly_effects_plot_adjustments.R
# 
# '

rm(list = ls())
options(scipen = 999) # Turn off scientific notation
options(warn = -0) # 0 to turn on again, -1 to turn off
source(paste0("src/code/ggplot_theme_publication.r")) # Load a user-written ggplot theme and style.

# Load libraries
library(tidyverse) # More intuitive data wrangling
library(stargazer) # easy way to make model output look more appealing (R-inline, html, or latex)
library(fixest) # Fast two-way fixed-effects estimation
library(rjson) # Read JSON files in R

# Load variables for which the treatment effect shall be estimated from the model specification.
model <- fromJSON(file = paste0("src/code/specifications.json"))
outcome_list <- model$outcomes

# Load analysis data
analysis <- readRDS(paste0("bld/data/analysis_year.RDS"))

# Drop duplicates and keep only obs at 100km bandwidth
dat <- subset(analysis, analysis$dist2cutoff <= model$distance_to_border & analysis$year > 2002)

# Loop over hospitalization and death rate variables for estimation.
for (outcome in c("cases_prevalence", "death_prevalence")){

  # Loop over each health classification and create model for each
  # Overall treatment effect
  allModels = lapply(seq_along(outcome_list)[-1], function(x){
      dat$timeTreated <- dat$time * dat$treated
      frm <- as.formula(paste0(outcome, " ~ ", paste(names(model$controls)[-1], collapse = "+"),
                               "+ timeTreated + factor(segments)*time + time*mun_lon*mun_lat| mcode + year"))
      feols(frm, data = subset(dat, dat$icd == names(outcome_list[x])) %>%
              distinct(mcode, year, .keep_all = TRUE))#, vcov = "twoway")
  })

  # Yearly treatment effects
  allYears = lapply(seq_along(outcome_list)[-1], function(x){
      frm <- as.formula(paste0(outcome, " ~ ", paste(names(model$controls)[-1], collapse = "+"),
                               "+ factor(segments)*factor(year) + factor(year)*mun_lon*mun_lat + factor(year)*treated | mcode + year"))
      feols(frm, data = subset(dat, dat$icd == names(outcome_list[x])) %>%
              distinct(mcode, year, .keep_all = TRUE))#, vcov = "twoway")
  })

  # Plot yearly effects
  lapply(seq_along(allYears), function(x){
    results <- coeftable(allYears[[x]])
    i <- grep(":treated", rownames(results), fixed=T)
    impact <- as.data.frame(results[i, 1:4])
    impact$year <- c(2004:2017)
    add <- c(0, 0, 0, 0, 2003)
    impact <- rbind(impact, add)
    impact$beta <- as.numeric(impact$Estimate)
    impact$SE <- as.numeric(impact$`Std. Error`)
    impact[is.na(impact)] <- 0
    plot <- ggplot(impact, 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) +
      ylab("Treatment effect")  +
      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 
    ggsave(plot = plot, filename = paste0("bld/figures/", names(outcome_list[x+1]), "_", outcome, ".png"), height = 9, width = 15, units = "cm")
    # Save results
    saveRDS(results, file = paste0("bld/results/", names(outcome_list[x+1]), "_", outcome, "_yearly_effects.rds"))
  })

  # Create regression table
  star <- 
    etable(
    allModels,
    tex = TRUE,
    placement = "h",
    title = "Outcome variables",
    keep = "%timeTreated",
    se.below = TRUE,
    headers =  sapply(1:length(compact(allModels)), function(i) paste(outcome_list[[i+1]])),
    dict = c(timeTreated = "Treatment effect", mcode = "Municipality", year = "Year"),
    label = paste0("table:", outcome),
    depvar = FALSE,
    notes = "Will be replaced",
    fixef_sizes = TRUE,
    # drop.section = "fixef",
    # coef.just = "c",
    style.tex = style.tex("qje"),
    digits = "r3",
    digits.stats = "r3"
    
  )

  # Add individual note to each table
  if (outcome == "cases_prevalence"){
    note.latex <- "[-1.8ex] \\end{tabular}
        \\begin{tablenotes}[para,flushleft]
        \\textit{Note:} The dependent variables are hospitalizations per 1,000 inhabitants associated with the named health problems
        according to the ICD-10 classification.
        All regressions use year and municipality fixed effects.
        Time variant control variables are discussed in the Methods section.
        Standard errors (in parentheses) are clustered at the municipality level. \\\\
        $^{*}$p$ < $0.1; $^{**}$p$ < $0.05; $^{***}$p$ < $0.01.\\\\
      \\end{tablenotes}"

    var.label.latex <- "\\tabularnewline\\midrule\\midrule\\\\[-1.8ex]
     & \\multicolumn{6}{c}{Hospitalizations per 1,000} \\\\"
  } else {
      note.latex <- "[-1.8ex] \\end{tabular}
        \\begin{tablenotes}[para,flushleft]
          \\textit{Note:} The dependent variables are deaths from hospitalizations per 1,000 inhabitants associated with the named health problems
          according to the ICD-10 classification.
          All regressions use year and individual fixed effects.
          Time variant control variables are discussed in the Methods section.
          Standard errors (in parentheses) are clustered at the municipality level. \\\\
          $^{*}$p$ < $0.1; $^{**}$p$ < $0.05; $^{***}$p$ < $0.01.\\\\
        \\end{tablenotes}"

    var.label.latex <- "\\tabularnewline\\midrule\\midrule\\\\[-1.8ex]
    & \\multicolumn{6}{c}{Deaths per 1,000} \\\\"
  }
  # Add note to the table and store as a tex file
  star[grepl("end{tabular}", star, fixed=TRUE)] <- note.latex
  star[grepl("\\tabularnewline\\midrule\\midrule", star, fixed=TRUE)] <- var.label.latex
  star <- star[c(1:14,16:length(star))]
  sink(paste0("bld/tables/", outcome_list[[1]], "_", outcome, ".tex"))
  sink(cat(star[5:(length(star)-4)], sep = "\n"))
}