# '
# The file "table_mechanism_controls.r" takes the panel data 
# from "bld/data/" and estimates the treatment effect on the potential mechanism
# variables and the control variables as a robustness check. The exact variables
# are specified in "src/code/specifications.json". This creates Supplementary
# Tables SI3 and SI11 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 Figure SI2 of the paper) and stores it to "bld/figures/"
# as pdf file. 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 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"))
lists <- list(model$mechanisms, model$controls)

# Load the data for estimation.
analysis <- readRDS("bld/data/analysis_year.RDS")

# 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) 

# Loop over mechanism and control variables for estimation.
for (list in lists){
  if (list$title == "Mechanism"){
    
    # Select the right outcome variables
    list <- list[-grep("_over_limit_relative", names(list))]
    
    # Loop over each variable and create model for each
    # Overall treatment effect
    allModels = lapply(names(list[c(2:length(list))]), function(outcome){
      if(outcome == "fires"){
        d$timeTreated <- d$time * d$treated.fires
        frm <- as.formula(paste0("asinh(", outcome, ") ~ ", paste(names(model$controls)[-1], collapse = "+"),
                                 "+ time*factor(segments)+ time*mun_lon*mun_lat + treated.fires + timeTreated | mcode + year"))
        feols(frm, data = d)#, vcov = "twoway")
      }
      else{
        d$timeTreated <- d$time * d$treated
        frm <- as.formula(paste0("asinh(", outcome, ") ~ ", paste(names(model$controls)[-1], collapse = "+"),
                                 "+ time*factor(segments)+ time*mun_lon*mun_lat + timeTreated | mcode + year"))
        feols(frm, data = d %>% distinct(mcode, year, .keep_all = TRUE))#, vcov = "twoway")
      }
    })
    
    # Yearly treatment effects
    allYears = lapply(names(list[c(2:length(list))]), function(outcome){
      if(outcome == "fires"){
        frm <- as.formula(paste0("asinh(", outcome, ") ~ factor(segments)*factor(year) + factor(year)*mun_lon*mun_lat + 
                                       treated.fires*factor(year) +",
                                 paste(names(model$controls)[-1], collapse = "+"), "| mcode + year"))
        feols(frm, data = d)#, vcov = "twoway")
      }
      else{
        frm <- as.formula(paste0("asinh(", outcome, ") ~ factor(segments)*factor(year) + factor(year)*mun_lon*mun_lat + 
                                       treated*factor(year) +",
                                 paste(names(model$controls)[-1], collapse = "+"), "| mcode + year"))
        feols(frm, data = d %>% distinct(mcode, year, .keep_all = TRUE)) #, vcov = "twoway")
      }
    })


    # Plot yearly effects
    lapply(seq_along(allYears), function(x){
      results <- coeftable(allYears[[x]])#[["coeftable"]]
      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 <- impact$Estimate 
      impact$SE <-  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(list[x+1]), ".png"), height = 9, width = 15, units = "cm")
      # Save results
      saveRDS(results, file = paste0("bld/results/", names(list[x+1]), "_yearly_effects.rds"))
    })

  } else {
    # Loop over each variable and create model for each
    # Overall treatment effect
    allModels = lapply(names(list[c(2:length(list))]), function(x){
      d$timeTreated <- d$time * d$treated
      frm <- as.formula(paste0(x, " ~ time*factor(segments)+ time*mun_lon*mun_lat + timeTreated | mcode + year"))
      feols(frm, data = d %>% distinct(mcode, year, .keep_all = TRUE))#, vcov = "twoway")
    })
  }
  
  # Create regression table
  star <- 
    etable(
      allModels,
      tex = T,
      placement = "h",
      title = "Mechanism variables",
      keep = "%timeTreated",
      se.below = TRUE,
      headers =  sapply(2:length(list), function(i) paste(list[[i]])),
      dict = c(timeTreated = "Treatment effect", mcode = "Municipality", year = "Year"),
      label = paste0("table:", list$title),
      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 (list$title == "Mechanism"){
    note.latex <- "[-1.8ex] \\end{tabular}
      \\begin{tablenotes}[para,flushleft]
        \\textit{Note:} Coefficients are given as decimal percentage changes. 
        All regressions use year and municipality fixed effects.
        Time variant control variables are discussed in the Methods section.
        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.
        Standard errors (in parentheses) are clustered at the municipality level.\\\\
        $^{*}$p$ < $0.1; $^{**}$p$ < $0.05; $^{***}$p$ < $0.01. \\\\
      \\end{tablenotes}"
  } else {
    note.latex <- "[-1.8ex] \\end{tabular}
      \\begin{tablenotes}[para,flushleft]
        \\textit{Note:} All regressions use year and municipality fixed effects.
        Dependent variables: GDP per capita is in thousand Brazilian Reais in 2012 prices; agricultural share is the GDP share of the agricultural sector in percentage;
        population density is inhabitants per square km; doctors and health establishments are per 1,000 inhabitants;
        temperature is in degree Celsius; precipitation is in millimeters per day; humidity is in percentage; wind speed is in meters per second; and wind direction is in degrees.
        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
  star[grepl("end{tabular}", star, fixed=TRUE)] <- note.latex
  star <- star[c(1:14,16:length(star))]
  sink(paste0("bld/tables/", list[[1]], ".tex"))
  sink(cat(star[5:(length(star)-4)], sep = "\n"))
}