# '
# The file "table_pre_trend_test.r" takes the panel data for estimation
# from "bld/data/" and estimates the treatment effect on the hospitalization and
# death rate of different diseases, and on the potential mechanism for the before period as a 
# test for parallel pre-treatment trends.
# The exact variables are specified in "src/code/specifications.json". This creates Supplementary
# Table SI10 of the paper. Tables are stored as tex files to "bld/tables".
# 
# '

rm(list = ls())
options(scipen = 999) # Turn off scientific notation
options(warn = 0) # 0 to turn on again, -1 to turn off

# 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 regressions
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"))
mechanisms_list <- model$mechanisms[-grep("_over_limit_relative", names(model$mechanisms))]

# Load analysis data
analysis <- readRDS(paste0("bld/data/analysis_year.RDS"))
# 
# Placebo treatment
analysis$placebo_time <- 0
analysis$placebo_time[analysis$year > 2004] <- 1
# analysis$placeboTimeTreated <- analysis$treated * analysis$placebo_time

# Drop duplicates and keep only obs at 100km bandwidth
dat <- analysis %>%
  # distinct(mcode, year, .keep_all = TRUE) %>%
  filter(icd == "J") %>%
  arrange(mcode, year)

d <- subset(dat, dat$dist2cutoff <= model$distance_to_border & dat$year > 2002 & dat$year < 2007)

results <- list()
i <- 1

# Parallel pre-trend analysis
allModels <- lapply(names(mechanisms_list[c(2:length(mechanisms_list))]), function(outcome){
  if(outcome == "fires"){
    d$placeboTimeTreated <- d$treated.fires * d$placebo_time
    frm <- as.formula(paste0("asinh(", outcome, ") ~ ", paste(names(model$controls)[-1], collapse = "+"),
                             "+ placebo_time*factor(segments)+ placebo_time*mun_lon*mun_lat + treated.fires + placeboTimeTreated | mcode + year"))
    feols(frm, data = d)#, vcov = "twoway")
  }
  else{
    d$placeboTimeTreated <- d$treated * d$placebo_time
    frm <- as.formula(paste0("asinh(", outcome, ") ~ ", paste(names(model$controls)[-1], collapse = "+"),
                             "+ placebo_time*factor(segments)+ placebo_time*mun_lon*mun_lat + placeboTimeTreated | mcode + year"))
    feols(frm, data = d  %>% distinct(mcode, year, .keep_all = TRUE))#, vcov = "twoway")
  }
})

# Create regression table
star <- 
  etable(
    allModels,
    tex = TRUE,
    keep = "%placeboTimeTreated",
    se.below = TRUE,
    headers =  sapply(2:length(mechanisms_list), function(i) paste(mechanisms_list[[i]])),
    dict = c(placeboTimeTreated = "Mechanisms"),
    label = paste0("table:", mechanisms_list$title),
    title = paste0(mechanisms_list[[1]], " variables"),
    placement = "h",
    depvar = FALSE,
    notes = "Will be replaced",
    drop.section = "fixef",
    coef.just = "c",
    style.tex = style.tex("qje"),
    digits = "r3",
    digits.stats = "r3"
)

# Store regression results
results[[i]] <- c(star[7:13])

# Loop over hospitalization and death rate variables for estimation.
outcome <- "cases_prevalence"
dat <- subset(analysis, analysis$dist2cutoff <= model$distance_to_border & analysis$year > 2002 & analysis$year < 2007)

# Placebo treatment
dat$placebo_time <- 0
dat$placebo_time[dat$year > 2004] <- 1
dat$placeboTimeTreated <- dat$treated * dat$placebo_time

# Parallel pre-trend analysis
allModels = lapply(seq_along(model$outcomes)[-1], function(x){
  frm <- as.formula(paste0(outcome, " ~ ", paste(names(model$controls)[-1], collapse = "+"),
                           "+ placeboTimeTreated + factor(segments)*placebo_time + placebo_time*mun_lon*mun_lat| mcode + year"))
  feols(frm, data = subset(dat, dat$icd == names(model$outcomes[x]))%>%
          distinct(mcode, year, .keep_all = TRUE))#, vcov = "twoway"))
})

# Create regression table
star <- 
  etable(
    allModels,
    tex = TRUE,
    placement = "h",
    title = paste0(model$outcomes[[1]], " variables"),
    keep = "%placeboTimeTreated",
    se.below = TRUE,
    headers =  sapply(1:length(compact(allModels)), function(i) paste(model$outcomes[[i+1]])),
    dict = c(placeboTimeTreated = "Cases"),
    label = paste0("table:", outcome),
    depvar = FALSE,
    notes = "Will be replaced",
    drop.section = "fixef",
    coef.just = "c",
    style.tex = style.tex("qje"),
    digits = "r3",
    digits.stats = "r3"
)

# Store regression results
i <- i + 1
results[[i]] <- c(star[7:12])

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

# Create regression table
star <- 
  etable(
    allModels,
    tex = TRUE,
    placement = "h",
    title = paste0(model$outcomes[[1]], " variables"),
    keep = "%placeboTimeTreated",
    se.below = TRUE,
    headers =  sapply(1:length(compact(allModels)), function(i) paste(model$outcomes[[i+1]])),
    dict = c(placeboTimeTreated = "Deaths"),
    label = paste0("table:", outcome),
    depvar = FALSE,
    notes = "Will be replaced",
    drop.section = "fixef",
    coef.just = "c",
    style.tex = style.tex("qje"),
    digits = "r3",
    digits.stats = "r3"
)

# Store regression results
i <- i + 1
results[[i]] <- c(star[10:13])

# Append all effects to one large table
head <- c(star[1:6])
brk <- "& & & & & & \\\\"
table_out <- c(head, results[[1]], brk, results[[2]], results[[3]], star[16:length(star)]) 

# Add individual note to the table
note.latex <- "[-1.8ex] \\end{tabular}
    \\begin{tablenotes}[para,flushleft]
      \\textit{Notes:} Pre-treatment refers to the period from 2003 to 2006. All regressions use year and municipality fixed effects.
      Dependent variables: PM\\textsubscript{2.5} is particulate matter with a diameter of less than 2.5 $\\mu m$; $CO$ is carbon monoxide;
        NO\\textsubscript{2} is nitrogen dioxide; O\\textsubscript{3} is ozone; and SO\\textsubscript{2} is sulfur dioxide.
      Coefficients for fires and air pollution are given as decimal percentage changes.
    Cases are hospitalizations per 1,000 inhabitants associated with the named health problems
    according to the ICD-10 classification, and deaths are deaths resulting from these hospitalizations per 1,000 inhabitants.
    Standard errors (in parentheses) are clustered at the municipality level. \\\\
    $^{*}$p$ < $0.1; $^{**}$p$ < $0.05; $^{***}$p$ < $0.01.\\\\
    \\end{tablenotes}"

# Add note to the table and store as a tex file
table_out[grepl("end{tabular}", table_out, fixed=TRUE)] <- note.latex
table_out <- sub("[h]", "", table_out, fixed = TRUE)
sink(paste0("bld/tables/pre_trend_test.tex"))
cat(table_out[5:(length(table_out)-4)], sep = "\n")
sink()