#code for Chapter 6 "Propensity Score Methods for Multiple Treatments" of book:
#Leite, W. L. (2017). Practical propensity score methods using R. 
#Thousand Oaks, CA: Sage. 
# 
#PART 2 - OBTAIN IPTW AND CHECK COVARIATE BALANCE PAIRWISE

#this is the code that was used to generate the example results in the book
#As the R software and the R packages used in this example are updated frequently
#some incompatibilities between the current code and new R versions or package versions
#may appear
#Any updates to the code will be posted at:
# http://www.practicalpropensityscore.com

#this code estimates  the effects of assigning mentors 
#of different areas to new teachers on the probability 
#that they will continue in the teaching profession 
#in the following year
#using data from the 1999--2000 School and Staffing Survey (SASS) 
#and 2000--2001 Teacher Follow-up Survey (TFS)

#load data with generalized propensity scores
load("chapter6_generalized_propensity_score_estimation_results.Rdata")

#Obtain Inverse Probability Treatment Weights for Propensity score weighting
imputedData$IPTW <- ifelse(imputedData$Treat=="noMentor", 1/ps$noMentor, 
                          ifelse(imputedData$Treat=="sameArea", 1/ps$sameArea, 1/ps$otherArea))
with(imputedData, by(IPTW,Treat,summary))


#Obtain the final weight that is the multiplication of IPTW and sampling weight TFNLWGT
imputedData$IPTW.TFNLWGT <- with(imputedData,IPTW*TFNLWGT)
with(imputedData, by(IPTW.TFNLWGT,Treat,summary))

#normalize weights (make them sum to sample size)
imputedData$finalWeightATE <- with(imputedData,IPTW.TFNLWGT/
                                    mean(IPTW.TFNLWGT))
#check distribution of weights
with(imputedData, by(finalWeightATE,Treat,summary))



#===============================================
#EVALUATE COVARIATE BALANCE PAIRWISE

#evaluate covariate balance with IPTW obtained with multinomial logist cregresson


#function to evaluate covariate balance pairwise
pairwise.balance <- function(condition1, condition2, data, iptw){
require(twang)
data <- subset(data, Treat == condition1 | Treat == condition2)  
data$Treat <- as.numeric(data$Treat==condition1)
  balance.iptw <- bal.stat(data=data,#dataset is defined
                           vars=covariateNames, #list of the covariates that are used in the treatment assignment model
                           treat.var ="Treat" , #treatment assignment indicator
                           w.all = data[,iptw],   #This time, PS weight times Sampling weight is defined
                           get.ks=F,            #Avoid estimating KS statistic
                           sampw = data$TFNLWGT,   #indicate the sampling weight
                           estimand="ATE",     #Define proper estimand
                           multinom=F)         #Indicate that we make pairwise comparisons,
  
  return( balance.iptw$results) } #close function
#========================================

#obtain balance for all groups pairwise
balance.iptw <- list(pair12 <- pairwise.balance("noMentor", "sameArea", imputedData, "finalWeightATE"), 
                    pair13 <- pairwise.balance("noMentor", "otherArea", imputedData, "finalWeightATE"),
                    pair23 <- pairwise.balance("sameArea", "otherArea", imputedData, "finalWeightATE"))


#obtain a dataset with just standardized effect sizes
std.eff.iptw <- data.frame(abs(balance.iptw[[1]][5]),
                            abs(balance.iptw[[2]][5]),
                            abs(balance.iptw[[3]][5]))

summary(std.eff.iptw)

#identify covariates/categories with standardized mean difference above 0.25, which will be 
#used in covariance adjustment
unbalanced.covariates <- row.names(std.eff.iptw)[as.logical(apply(std.eff.iptw>0.25,1,max))]
unbalanced.covariates <- strsplit(unbalanced.covariates, ":")
unbalanced.covariates <- unique(sapply(unbalanced.covariates, "[", 1))
print(unbalanced.covariates)