# NOTE: Set wd to source file location to run

rm(list = ls())

require(Amelia)
require(ggplot2)
require(Hmisc)
require(cowplot)


######################################################################
###################### Figure 1 ######################################
######################################################################

## Source a function to make equations (to remove clutter)

make.equation <- function(dvs, ivs, interactions = c()){
  return(paste(dvs, "~", paste(c(ivs, interactions), sep = "", collapse = " + ")))
}

## Load in data 

load("data/uganda_citizen_imputed_manip.RData")
load("data/ghana_citizen_imputed_manip.RData")
var.info <- read.csv("data/mass_analysis_varinfo.csv", stringsAsFactors = T)
df.ug <- ug.imp.manip[ug.imp.manip$imp != 0, ] # imp = 0 is the original dataset
df.gh <- gh.imp.manip[gh.imp.manip$imp != 0, ]
rm(list = c("gh.imp.manip", "ug.imp.manip")) # Remove original objects

## Variable prep

var.groups <- c("action", "benefit", "misappropriation")
var.list <- sapply(var.groups, FUN = function(x){var.info$var_name[var.info$var_index == x]})
var.names <- sapply(var.groups, FUN = function(x){var.info$var_label[var.info$var_index == x]})

####################### Models

imp.results <- c() # Container to hold the imputed results

## Loop over each imputation

for(m in 1:max(df.ug$imp)){ # Note that this implies imp = 1 is an imputed dataset
  
  ## Set the data you want to use
  
  data.list <- list(df.ug[df.ug$imp == m, ], 
                    df.gh[df.gh$imp == m, ])
  manip.list <- c("yes", "yes") #  
  country.list <- c("uganda", "ghana")
    
  ## Define control and treatment groups 
  
  treat.vars <- list(c("aid", "ngo", "tax"), # Oil as base
                     c("ngo", "tax", "oil"), # Aid as base
                     c("ngo")) # Non-NGO as base
  ref.group <- c("oil", "aid", "non-ngo") # These are the references groups for each of these models
  vars.fe <- c("as.factor(enum_id)")
  dvs <- c("action_avg", "benefit_avg", "misappropriation_avg")
  dvs.group <- c("action", "benefit", "misappropriation") # Will use this to make plotting easier
  dvs.label <- c("Action Index (NMWI)", "Benefit Index (NMWI)", "Misappropriation Index (NMWI)")
  
  ## Build container object and begin loop
  
  
  for(l in 1:length(treat.vars)){ # For each set of treatment variables (e.g. each element in list)
    
    table.results <- c()
    ## Generate equations 
    
    treat.use <- treat.vars[[l]] # Grab the treatment variables to include in equation
    treat.ref <- ref.group[l] # Set the reference group
    eqs.fe <- make.equation(dvs = dvs, ivs = c(treat.use, vars.fe)) # Add FEs (defined above)
    
    ## Loop over the dvs and the datasets
    
      for(j in 1:length(data.list)){ # This list is defined above, includes full UG + GH datasets and then also with subjects who failed the manipulation check removed
        for(i in 1:length(dvs)){
          
          ## Run the model and store it
          
          tab <- c() # Create container to stack into dataframe
          mod <- glm(eqs.fe[i], data = data.list[[j]])
          tab <- data.frame(summary(mod)$coefficients[c(treat.use), ])
          if(ncol(tab) == 1){tab <- t(tab)
          rownames(tab) <- treat.use} # If there's only coefficient R will default to column vector, have to transpose back to the row
          
          ## Create summary object from estimates
          
          mod.results <- data.frame(est = tab[, 1], se = tab[, 2], t = tab[, 3], p = tab[, 4], 
                                    iv = treat.use, model = "fe", country = country.list[j], 
                                    manip = manip.list[j], var_group = dvs.group[i], 
                                    dv = dvs[i], label = dvs.label[i], 
                                    ref_cat = treat.ref, imp = m)
          table.results <- rbind(mod.results, table.results)
          colnames(table.results)[1:4] <- c("est", "se", "t", "p")
          
          
        } # Close i loop (over variables)
      } # Close j loop (over datasets)

  
  useCols <- c("label", "iv", "est", "se", "t", "p", 
               "country", "manip", "var_group", 
               "ref_cat", "dv", "model", "imp")
  lm.results <- table.results[, useCols]
  lm.results <- lm.results[with(lm.results, order(country, ref_cat, iv, dv, 
                                                  var_group, model, manip, imp)), ]
  rownames(lm.results) <- NULL
  
  ## Store the results in a list, where each element is an imputation
  
  imp.results <- rbind(imp.results, lm.results)
  
  } # Close l loop over treatment variable list
} # Close m loop over imputations

## Create unique ID string

imp.results$result_id <- as.numeric(factor(
  with(imp.results, paste(iv, country, manip, ref_cat, dv, model, sep = "-"))))

## Subset to the imputed data only and combine them

res.id <- unique(imp.results$result_id)
mi.results <- c()
for(i in 1:length(res.id)){
  df.tmp <- imp.results[imp.results$result_id == i & imp.results$imp != 0, ]
  est.mi <- unlist(mi.meld(q = cbind(df.tmp[, "est"], df.tmp[, "est"]), 
                           se = cbind(df.tmp[, "se"], df.tmp[, "se"])))[c(1, 3)] # element 1-2 is beta, 3-4 is se (because of repetition)
  # Note, need to duplicate here because mi.meld function won't handle vectors 
  all.mi <- c(est.mi[1], est.mi[2], est.mi[1]/est.mi[2], 2*(1 - pnorm(abs(est.mi[1]/est.mi[2]))))
  df.use <- df.tmp[1, ]
  df.use[, c("est", "se", "t", "p")] <- all.mi
  
  mi.results <- rbind(mi.results, df.use)}

mi.results[, c("dv", "country", "iv", "ref_cat", "est", "se", "t", "p")]
save(mi.results, file = "results/mass_lm_results_maintext.RData")

## Load in results and variable information for plot

load("results/lm_results_maintext.RData") # Mass LM Results
varInfo <- read.csv("data/mass_analysis_varinfo.csv", stringsAsFactors = F) # Variable Descriptions

## MI Estimates (Confidence Intervals)

alpha <- 0.05
mi.results$uci_95 <- mi.results$est + qnorm(alpha/2)*mi.results$se
mi.results$lci_95 <- mi.results$est - qnorm(alpha/2)*mi.results$se
alpha <- 0.10
mi.results$uci_90 <- mi.results$est + qnorm(alpha/2)*mi.results$se
mi.results$lci_90 <- mi.results$est - qnorm(alpha/2)*mi.results$se

##Relabel Text and Clean Up Strings for ggplot

mi.results$iv <- gsub("ngo", "NGO", mi.results$iv)
mi.results$iv <- gsub("Ngo", "NGO", mi.results$iv)
mi.results$ref_cat <- gsub("non-ngo", "non-NGO", mi.results$ref_cat)
mi.results$ref_cat <- gsub("ngo", "NGO", mi.results$ref_cat)
mi.results$comparison <- paste(capitalize(mi.results$iv), "Treatment,",
                                  capitalize(mi.results$ref_cat), "Control")
mi.results$comparison <- gsub("Treatment", "Treat.", mi.results$comparison)
mi.results$comparison <- gsub("Control", "Ctr.", mi.results$comparison)
mi.results$comparison <- gsub("Ngo", "NGO", mi.results$comparison)
mi.results$label_plot <- factor(ifelse(mi.results$var_group == "action", "Action", 
                                ifelse(mi.results$var_group == "benefit", "Benefit", "Misapp.")), 
                                levels = c("Misapp.", "Benefit", "Action"))

#################### Main results using facet wrap ##############
#################### Plot the 3 Main Index Variables ##############

##MI Results
manipSwitch <- c("yes")
models <- c("fe") 
plotGroups <- c("action")
plotNames <- c("Index Measures") 
cbbPalette <- c("#000000", "#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2", "#D55E00", "#CC79A7")
compsUse1 <- c("Tax Treat., Oil Ctr.",  "Tax Treat., Aid Ctr.", "Oil Treat., Aid Ctr.", "NGO Treat., Non-NGO Ctr.")
compsUses <- list(compsUse1)
titleNote1 <- c(", Enumerator Fixed Effects (Imputed)")
titleNotes <- list(titleNote1)
plotVars <- varInfo$var_name[varInfo$is_index == 1 & varInfo$include_plot == 1] #

## First obtain estimates in dataframe
estFrame <- mi.results[mi.results$model == models[1]
                            & mi.results$comparison %in% compsUses[[1]]
                            & mi.results$manip == manipSwitch[1]
                            & mi.results$dv %in% plotVars, ]
estFrame$comparison<-factor(estFrame$comparison, compsUses[[1]])
table(estFrame$label)

##Then plot
p <- ggplot(estFrame, aes(y = est, x = label_plot, colour = country)) + 
  facet_wrap(~comparison, ncol=4) + 
  geom_linerange(aes(x = label_plot, ymin = lci_90,
                     ymax = uci_90, size = I(10)),
                 lwd = 1.25, position = position_dodge(width = 1/2)) +
  geom_pointrange(aes(x = label_plot, y = est, ymin =  lci_95,
                      ymax = uci_95),
                  lwd = 0.4, position = position_dodge(width = 1/2),
                  shape = 21, fill = "WHITE") + 
  coord_flip() +  xlab("") +   theme_gray() +  
  ylab("Average Treatment Effect") + theme(legend.position=c(1, -0.3),  legend.justification = "right",
                                           legend.direction = "horizontal", legend.text = element_text(size=6), 
                                           legend.title = element_text(size=6), legend.margin=margin(0,0,0,0),
                                           legend.box.margin=margin(0,0,0,0)) +
  theme(axis.text = element_text(size = 7), 
        axis.title = element_text(size = 7), 
        plot.title = element_text(size = 7, face = "bold"), 
        strip.text.y = element_text(size=7, face="bold"),
        strip.text.x = element_text(size = 7)) + 
  ggtitle(paste(plotNames[1], titleNotes[[1]], sep = ""))  + geom_hline(yintercept = 0, lty = 2, col = "red") +
  labs(color = "Country") + scale_color_manual(labels = c("Ghana", "Uganda"), 
                                               values = cbbPalette[c(2, 3)], guide = guide_legend(reverse = TRUE))
w = (12/1.72); h = (5/2.72) ## to fit PNAS specifications
pdf(file = "figures/fig1.pdf", width = w, height = h)
plot(p)
dev.off()



######################################################################
###################### Figure 2 ######################################
######################################################################

rm(list = ls())

## Source a function to make equations (to remove clutter)

make.equation <- function(dvs, ivs, interactions = c()){
  return(paste(dvs, "~", paste(c(ivs, interactions), sep = "", collapse = " + ")))
}

## Load in Uganda and Ghana elite data

load("data/uganda_mp_full.RData")
load("data/ghana_mp_full.RData")
var.info <- read.csv("data/mp_analysis_varinfo.csv", stringsAsFactors = F)

data.list <- list(d.ug.mp, d.gh.mp)
country.list <- c("uganda", "ghana")


######################## EQUATION PREP ##########################################

treat.vars <- list(c("oil", "aid", "ngo"), # Tax as base
                   c("aid", "ngo", "tax"), # Oil as base
                   c("ngo", "tax", "oil"), # Aid as base
                   c("ngo")) # Non-NGO as base
ref.group <- c("tax", "oil", "aid", "non-ngo")
vars.fe <- c("as.factor(enum_id)")

## Note: what the treat.vars list object is doing is setting up a set of dummies
## to be included in the estimating equation such that the omitted variable
## is the base category. The ``ref group'' object is the omitted feature
## and is used below to specify the reference group in the output table

####################### BEGIN LOOP ###############################################
table.results <- c() ## Container for results

for(l in 1:length(treat.vars)){ # For each set of treatment variables (e.g. each element in list)
  
  ## GENERATE EQUATIONS 
  
  treat.use <- treat.vars[[l]] # Grab the treatment variables to include in equation
  treat.ref <- ref.group[l] # Set the reference group
  eqs.fe <- make.equation(dvs = var.info$var_name, ivs = c(treat.use, vars.fe)) # Add FEs (defined above)  
  eqs.type <- c("fe") # Name types models to run, will loop over this below
  eqs.list <- list(eqs.fe) # Gather corresponding equations into lists
      # Note that if you want to use other types of models, just specify them using make.equation
      # name them in eqs.type and include them in eqs.list
  
  ## NOTE: CAN REMOVE CERTAIN MODELS BY COMMENTING OUT AND ADJUST eqs.type and eqs.list objects
  
  
  
  ############################ ESTIMATION ######################  
  
  model.obj <- list()
  for(k in 1:length(eqs.type)){
    for(j in 1:length(data.list)){ # This list is defined above, includes full UG + GH datasets 
      for(i in 1:length(var.info$var_name)){
        
        ## Run models and store objects depending on variable type
        
        tab <- c()
        
        mod <- glm(eqs.list[[k]][i], data = data.list[[j]])
        
        if(length(treat.use) == 1){ # if it's just a single dummy, have to transpose
          tab <- data.frame(t(as.data.frame(summary(mod)$coefficients[c(treat.use), ])))
          rownames(tab) <- treat.use}
        else{ # if multiple treatment dummies, can take summary frame as-is
          tab <- data.frame(summary(mod)$coefficients[c(treat.use), ])}
        
        
        ## Create summary object from estimates
        
        colnames(tab) <- c("est", "se", "t", "p")
        tab$iv <- rownames(tab)
        tab$model <- eqs.type[k]
        tab$country <- country.list[j]
        tab$var_group <- var.info$var_index[i]
        tab$type <- var.info$var_type[i]
        tab$sig <- ifelse(tab$p >= 0.05, "no", "yes")
        tab$dv <- var.info$var_name[i]
        tab$label <- var.info$var_label[i]
        tab$reference_category <- treat.ref
        table.results <- rbind(tab, table.results)
      } # Close i loop (over variables)
    } # Close j loop (over datasets)
  } # Close model loop (over different equations)
} # Close l loop (over different base categories)

lm.results <- table.results[, c("label", "iv", "est", "se", "t", "p", "sig",
                                "country", "var_group", 
                                "reference_category", "dv", "model")]
lm.results <- lm.results[with(lm.results, order(country, reference_category, iv, dv, 
                                                var_group, model)), ]
rownames(lm.results) <- NULL
save(lm.results, file = "results/mp_lm_results_maintext.RData")

##### Prep data for figure

## Create necessary QOIs for ggplot
final.results <- lm.results
alpha <- 0.05
final.results$uci_95 <- final.results$est + qnorm(alpha/2)*final.results$se
final.results$lci_95 <- final.results$est - qnorm(alpha/2)*final.results$se
alpha <- 0.10
final.results$uci_90 <- final.results$est + qnorm(alpha/2)*final.results$se
final.results$lci_90 <- final.results$est - qnorm(alpha/2)*final.results$se

##Relabel Text for ggplot
final.results$iv <- gsub("ngo", "NGO", final.results$iv)
final.results$iv <- gsub("Ngo", "NGO", final.results$iv)
final.results$iv <- gsub("nontax_direct", "Aid/Oil", final.results$iv)
final.results$reference_category <- gsub("non-ngo", "non-NGO", final.results$reference_category)
final.results$reference_category <- gsub("ngo", "NGO", final.results$reference_category)
final.results$reference_category <- gsub("non-tax", "non-Tax", final.results$reference_category)
final.results$comparison <- paste(capitalize(final.results$iv), "Treatment,",
                                  capitalize(final.results$reference_category), "Control")
final.results$comparison <- gsub("Treatment", "Treat.", final.results$comparison)
final.results$comparison <- gsub("Control", "Ctr.", final.results$comparison)
final.results$comparison <- gsub("Ngo", "NGO", final.results$comparison)
final.results$label <- gsub("Index \\(NMWI\\)", "", final.results$label)
trim <- function (x) gsub("^\\s+|\\s+$", "", x)
final.results$label <- trim(final.results$label)
final.results$label <- gsub("Influence", " Influence", final.results$label) # Add extra whitespace to get order with facet wrap 

final.results.2 <- final.results

## Build objects necessary for plot

varInfo <- read.csv("data/mp_analysis_varinfo.csv", stringsAsFactors = F)
models <- c("fe") 
plotGroups <- c("action")
plotNames <- c("") 
cbbPalette <- c("#000000", "#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2", "#D55E00", "#CC79A7")
compsUse1 <- c("Tax Treat., Oil Ctr.",  "Tax Treat., Aid Ctr.", "Oil Treat., Aid Ctr.", "NGO Treat., Non-NGO Ctr.")
compsUses <- list(compsUse1)
titleNote1 <- c("MP Surveys, Index")
titleNotes <- list(titleNote1)
subgroups <-c("_") # Subgroup Labels
plotVars <- varInfo$var_name[varInfo$is_index == 1 & varInfo$include_plot == 1] #

## First obtain estimates in dataframe
estFrame <- final.results.2[final.results.2$model == models[1]
                            & final.results.2$comparison %in% compsUses[[1]]
                            & final.results.2$dv %in% plotVars, ]
estFrame$comparison<-factor(estFrame$comparison, compsUses[[1]])
dvorder <- c("Benefit", "Influence")
table(estFrame$label)

##Then plot

p <- ggplot(estFrame, aes(y = est, x = label, colour = country)) + 
  facet_wrap(~comparison, ncol=4) + 
  geom_linerange(aes(x = label, ymin = lci_90,
                     ymax = uci_90, size = I(10)),
                 lwd = 1.25, position = position_dodge(width = 1/2)) +
  geom_pointrange(aes(x = label, y = est, ymin =  lci_95,
                      ymax = uci_95),
                  lwd = 0.5, position = position_dodge(width = 1/2),
                  shape = 21, fill = "WHITE") + 
  coord_flip() +  xlab("") +   theme_gray() + 
  ylab("Average Treatment Effect") + theme(legend.position=c(1, -0.3),  legend.justification = "right",
                                           legend.direction = "horizontal", legend.text = element_text(size=6), 
                                           legend.title = element_text(size=6), legend.margin=margin(0,0,0,0),
                                           legend.box.margin=margin(0,0,0,0)) +
  theme(axis.text = element_text(size = 7), 
        axis.title = element_text(size = 7), 
        plot.title = element_text(size = 7, face = "bold"), 
        strip.text.y = element_text(size=7, face="bold"),
        strip.text.x = element_text(size = 7)) + 
  ggtitle(paste(plotNames[1], titleNotes[[1]], sep = ""))  + geom_hline(yintercept = 0, lty = 2, col = "red") +
  labs(color = "Country") + scale_color_manual(labels = c("Ghana", "Uganda"), 
                                               values = cbbPalette[c(2, 3)], guide = guide_legend(reverse = TRUE))
w = (12/1.72); h = (5/2.5)
pdf(file = "figures/fig2.pdf", width = w, height = h)
plot(p)
dev.off()
