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

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

# Set seed
set.seed(220214)


##################################################
################ WEIGHTS #########################
##################################################
# Define raking variables
vars = c("raking_female", "raking_education", "raking_urban_rural", "raking_age", "raking_income")
gcd$raking_female = paste0("raking_female", gcd$female)
gcd$raking_education = paste0("raking_education", gcd$education)
gcd$raking_urban_rural = paste0("raking_urban_rural", gcd$urban_rural)
gcd$raking_age = paste0("raking_age", cut(gcd$age, breaks=c(14, 25, 35, 50, 70, 100)))
gcd$raking_income = paste0("raking_income", gcd$INCOME_5)

# Replace NAs with random other level
for (i in world$iso3) {for (v in vars) {
  nas = grepl("NA$", gcd[gcd$iso3 %in% i, v, drop=T])
  gcd[gcd$iso3 %in% i, v][nas,] = ifelse(
    all(nas),
    "allNA",
    sample(gcd[gcd$iso3 %in% i, v, drop=T][!nas], sum(nas))
  )
}}

# For each country
rakectrl = rakefull = list()
for (i in world$iso3) {
  # Extract country data
  df = as.data.frame(gcd[gcd$iso3 %in% i,])
  
  # Factorize variables
  for (v in vars) {df[,v] = factor(df[,v])}
  
  # Derive target
  target = list()
  for (v in vars) {
    target[[v]] = wtd.table(df[,v], weights=df$wgt)$sum.of.weights/sum(df$wgt)
    names(target[[v]]) = wtd.table(df[,v], weights=df$wgt)$x
  }
  
  # Restrict to those who participated in experiment
  dfctrl = df[!is.na(df$coop) & df$treatment %in% 0,]
  dffull = df[!is.na(df$coop), ]
  
  # Rake!
  rakectrl[[i]] = suppressWarnings(anesrake(
    inputter = target, dataframe = dfctrl,
    caseid = dfctrl$wpid, weightvec = dfctrl$wgt,
    cap = 5, verbose = F, type = "nolim"
  ))
  rakefull[[i]] = suppressWarnings(anesrake(
    inputter = target, dataframe = dffull,
    caseid = dffull$wpid, weightvec = dffull$wgt,
    cap = 5, verbose = F, type = "nolim"
  ))
  
  # Normalize weights
  rakectrl[[i]]$weightvec = rakectrl[[i]]$weightvec / sum(rakectrl[[i]]$weightvec) * nrow(dfctrl)
  rakefull[[i]]$weightvec = rakefull[[i]]$weightvec / sum(rakefull[[i]]$weightvec) * nrow(dffull)
}

# Extract weights and attach to GCD data
weights = unlist(lapply(rakectrl, function(x) {x$weightvec}))
names(weights) = gsub("[A-Z]|\\.", "", names(weights))
gcd$wgt_expctrl = weights[as.character(gcd$wpid)] 
var_label(gcd$wgt_expctrl) = "Individual-level weight for cooperation control group. Weights results from raking algorithm to render cooperation control sample rep. of full weighted country sample in the dimensions: age, gender, urban/rural, income, education."

# Same for the full cooperation sample
weights = unlist(lapply(rakefull, function(x) {x$weightvec}))
names(weights) = gsub("[A-Z]|\\.", "", names(weights))
gcd$wgt_expfull = weights[as.character(gcd$wpid)] 
var_label(gcd$wgt_expfull) = "Individuel-level weight for cooperation data (control and treatment). Weights results from raking algorithm to render cooperation sample rep. of full weighted country sample in the dimensions: age, gender, urban/rural, income, education."

# Weight distribution
#quantile(gcd$wgt_expctrl, seq(0,1,0.1), na.rm=T)
#quantile(gcd$wgt_expfull, seq(0,1,0.1), na.rm=T)

# Population adjusted
gcd$wgt_expctrl_pop15 = gcd$wgt_expfull_pop15 = NA
for (i in world$iso3) {
  # Control group only
  n = sum(gcd$wgt_expctrl[gcd$iso3 %in% i], na.rm=T)
  gcd[gcd$iso3 %in% i, "wgt_expctrl_pop15"] = gcd[gcd$iso3 %in% i, "wgt_expctrl"] *
    world[world$iso3 %in% i, "wgt_pop15", drop=T] / n * 
    sum(!is.na(gcd$coop) & gcd$treatment == 0)
  
  # Control and treatment
  n = sum(gcd$wgt_expfull[gcd$iso3 %in% i], na.rm=T)
  gcd[gcd$iso3 %in% i, "wgt_expfull_pop15"] = gcd[gcd$iso3 %in% i, "wgt_expfull"] *
    world[world$iso3 %in% i, "wgt_pop15", drop=T] / n * 
    sum(!is.na(gcd$coop))
}
var_label(gcd$wgt_expctrl_pop15) = "Individual level weights for cooperation control group (see wgt_expctrl), population adjusted, base weight * share of country in total population among surveyed countries, aged 15+"
var_label(gcd$wgt_expfull_pop15) = "Individual level weights for cooperation data (control + treatment, see wgt_expfull), population adjusted, base weight * share of country in total population among surveyed countries, aged 15+"

# Drop raking vars
gcd[, vars] = NULL


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