##################################################
################ PREPARATION #####################
##################################################
# Clear environment
rm(list=setdiff(ls(), "path_to"))

# Load data
gcd = readRDS(path_to("gcd"))
world = readRDS(path_to("world"))

# Alerts
if (!any(gcd$iso3 %in% world$iso3)) {message("Alert. Participant without world iso3.")}
if (any(is.na(gcd$wgt))) {message("Alert. NA wgt.")}


##################################################
################ GCD COUNTRY AVG ################
##################################################
# Variable list
vars = c(
  "risktaking", "risktaking_std", "patience", "patience_std", "altruism", "altruism_std",
  "pg_preference", "norm", "exp_hypo", "coop", "coop_belief", "education_lev2", "education_lev3"
)

# Full World Poll Sample
for (v in vars) {
  # Averages
  world[, paste0("gcd_", v)] = sapply(world$iso3, function(i) {
    data_vector = pull(gcd[gcd$iso3 %in% i,], v)
    weights = gcd$wgt[gcd$iso3 %in% i]
    weighted.mean(data_vector, weights, na.rm=T)
  })
  var_label(world[, paste0("gcd_", v)]) = paste0("Country average of: ", var_label(pull(gcd, v)), ". Weighted (wgt).")
  
  # Standard deviations
  world[, paste0("gcd_sd_", v)] = sapply(world$iso3, function(i) {
    data_vector = pull(gcd[gcd$iso3 %in% i,], v)
    weights = gcd$wgt[gcd$iso3 %in% i]
    if (all(is.na(data_vector))) {NA}
    else {Hmisc::wtd.var(data_vector, weights, normwt=T, na.rm=T) %>% sqrt}
  })
  var_label(world[, paste0("gcd_sd_", v)]) = paste0("Country-level standard deviation of: ", var_label(pull(gcd, v)), ". Weighted (wgt).")
}

# Full cooperation sample
for (v in vars) {
  # Averages
  world[, paste0("gcd_", v, "_wgtexpfull")] = sapply(world$iso3, function(i) {
    data_vector = pull(gcd[gcd$iso3 %in% i,], v)
    weights = gcd$wgt_expfull[gcd$iso3 %in% i]
    Hmisc::wtd.mean(data_vector, weights, na.rm=T, normwt=T)
  })
  var_label(world[, paste0("gcd_", v, "_wgtexpfull")]) = paste0("Country average of: ", var_label(pull(gcd, v)), ". Control and treatment. Weighted (weights for cooperation sample, control and treatment).")
  
  # Standard deviation
  world[, paste0("gcd_sd_", v, "_wgtexpfull")] = sapply(world$iso3, function(i) {
    data_vector = pull(gcd[gcd$iso3 %in% i,], v)
    weights = gcd$wgt_expfull[gcd$iso3 %in% i]
    if (all(is.na(data_vector))) {NA}
    else {Hmisc::wtd.var(data_vector, weights, normwt=T, na.rm=T) %>% sqrt}
  })
  var_label(world[, paste0("gcd_sd_", v, "_wgtexpfull")]) = paste0("Country-level standard deviation of: ", var_label(pull(gcd, v)), ". Control and treatment. Weighted (weights for cooperation sample, control and treatment).")
}

# Cooperation sample, control group only
for (v in vars) {
  # Averages
  world[, paste0("gcd_", v, "_wgtexpctrl")] = sapply(world$iso3, function(i) {
    data_vector = pull(gcd[gcd$iso3 %in% i,], v)
    weights = gcd$wgt_expctrl[gcd$iso3 %in% i]
    Hmisc::wtd.mean(data_vector, weights, na.rm=T, normwt=T)
  })
  var_label(world[, paste0("gcd_", v, "_wgtexpctrl")]) = paste0("Country average of: ", var_label(pull(gcd, v)), ". Control only. Weighted (weights for cooperation sample, control only).")
  
  # Standard deviation
  world[, paste0("gcd_sd_", v, "_wgtexpctrl")] = sapply(world$iso3, function(i) {
    data_vector = pull(gcd[gcd$iso3 %in% i,], v)
    weights = gcd$wgt_expctrl[gcd$iso3 %in% i]
    if (all(is.na(data_vector))) {NA}
    else {Hmisc::wtd.var(data_vector, weights, normwt=T, na.rm=T) %>% sqrt}
  })
  var_label(world[, paste0("gcd_sd_", v, "_wgtexpctrl")]) = paste0("Country-level standard deviation of: ", var_label(pull(gcd, v)), ". Control only. Weighted (weights for cooperation sample, control only).")
}


## MONGOLIA ADJUSTMENT ##
# Adjust pg_preference score for Mongolia
replacement = 0.896
world[world$iso3 %in% "MNG", "gcd_pg_preference"] = replacement
world[world$iso3 %in% "MNG", "gcd_pg_preference_wgtexpctrl"] = replacement
world[world$iso3 %in% "MNG", "gcd_pg_preference_wgtexpfull"] = replacement


##################################################
################ COUNTRY REG EFFECTS #############
##################################################
## Bivariate regressions
# Define variables
reg_vars = c("coop_belief", "norm", "altruism_std", "patience_std", "risktaking_std", "pg_preference",
             "female", "age20", "income_pctile_10", "education_lev2", "education_lev3")
world[, paste0("biregcoeff_", reg_vars)] = NA
world[, paste0("biregse_", reg_vars)] = NA
world[, paste0("biregp_", reg_vars)] = NA

# Regression
for (i in world$iso3) {for (v in reg_vars) {
  reg_data = gcd[gcd$iso3 %in% i & gcd$treatment %in% 0, ]
  if (v == "education_lev2") {reg_data = reg_data[!reg_data$education_lev3 %in% 1,]}
  if (v == "education_lev3") {reg_data = reg_data[!reg_data$education_lev2 %in% 1,]}
  if (length(setdiff(unique(reg_data[,v,drop=T]), NA)) < 2) {next} # Skip if not enough RHS data
  mod = lm(paste0("coop ~ ", v), reg_data, weights = reg_data$wgt_expctrl)
  test = mod %>% coeftest(vcovHC(., type="HC1"))
  world[world$iso3 %in% i, paste0("biregcoeff_", v)] = mod$coefficients[v] %>% t
  world[world$iso3 %in% i, paste0("biregse_", v)] = test[v, 2] %>% t
  world[world$iso3 %in% i, paste0("biregp_", v)] = test[v, 4] %>% t
}}

# Labels
for (v in reg_vars) {
  var_label(world[, paste0("biregcoeff_", v)]) = paste0("OLS coefficient of bivariate regression of cooperation decision on ", v," . Control only. Weighted (weights for cooperation sample, control only).")
  var_label(world[, paste0("biregse_", v)])    = paste0("OLS coefficient's standard error in bivariate regression of cooperation decision on ", v," . Control only. Weighted (weights for cooperation sample, control only).")
  var_label(world[, paste0("biregp_", v)])     = paste0("OLS coefficient's p-value in bivariate regression of cooperation decision on ", v," . Control only. Weighted (weights for cooperation sample, control only).")
}


##################################################
################ EXPORT ##########################
##################################################
saveRDS(world, path_to("world"))

