## File Name: iptw.R
## File version: 1.0

#' Compute inverse probability of treatement weights (IPTW) for use with PCM
#'
#' This function computes IPTW weights for causal inference with the pcm function
#'
#' @param df data.frame containing the data
#' @param Y string containing the name of the column where the dependent variable is stored in df
#' @param X vector of strings containing the names of the columns where the independent variables are stored in df
#' @param target string containing the target causal effect of interest. Either "ate" (average treatment effect, default), "stab-ate" (stabilized ATE), "att" (average treatment effect on the treated) or "atu" (average treatment effect on the untreated)
#' @return A vector of IPT weights
#' @export

iptw <- function(df=NULL,Y=NULL,X=NULL,target="ate") {
  if (any(!(Y %in% colnames(df)))) {
    stop("ERROR: provided Y variable name does not exist in df")
  }
  if (!("id"%in%colnames(df))) {
    stop('ERROR: no column named id provided')
  }
  if (target !="ate" & target != "att" & target !="atu" & target !="stab-ate")
  dff <- df[,c(Y,X)]
  if (length(X)==1) {
    formu <- paste0(Y,"~",X)
  } else {
    formu <- paste0(Y,"~",X[1])
    for (k in 2:length(X)) {
      xx <- X[k]
      formu <- paste0(formu,"+",xx)
    }
  }
  lr_out <- glm(formula = as.formula(formu),data=df,family = binomial(link = 'logit'))
  if (target == "ate") {
    psw <- df$TT/fitted(lr_out) + (1-df$TT)/(1-fitted(lr_out))
  } else if (target=="stab-ate") {
    pt <- sum(df$TT)/nrow(df)
    psw <- pt*(df$TT/fitted(lr_out)) + (1-pt)*(1-df$TT)/(1-fitted(lr_out))
  } else if (target=="att") {
    psw <- fitted(lr_out)/(1-fitted(lr_out))
    psw[df[,Y]==1] <- 1
  } else if (target=="atu") {
    psw <- (1-fitted(lr_out))/(fitted(lr_out))
    psw[df[,Y]==0] <- 1
  }
  return(psw)
}