# Transfer Learning for Individualized Treatment Rules Functions #

## A. Computing linear rule (eta)

phi_su <- function(s, u) {
  phis <- (s-u)**2
  phis[u >= s] <- 0 
  phis[u < s-1] <- (2*s-2*u-1)[u < s-1]
  return(phis)
}

phi_0_grad <- function(u) {
  phis <- 2*u
  phis[u >= 0] <- 0
  phis[u < -1] <- -2
  return(phis)
}

smooth.ramp.loss <- function(eta, betas, x.in.rct, y.star.hat.in.rct, w.tilde) {
  f <- x.in.rct %*% as.matrix(eta, ncol=1)
  u <- (2*as.numeric(y.star.hat.in.rct > 0) - 1) * f
  cav <- sum(betas*u)
  cvx <- sum(phi_su(1, u) * abs(y.star.hat.in.rct) * w.tilde)
  return(cvx+cav)
}

cal_betas <- function(y.star.hat.in.rct, u, w.tilde) {
  betas <- - phi_0_grad(u) * abs(y.star.hat.in.rct) * w.tilde
  return(betas)
}

get_eta <- function(x.in.rct, w.tilde, y.star.hat.in.rct, eps=10^(-6), maxit=10^3) {
  # Estimate the linear rule coefficients using smooth ramp loss
  betas <- 2*abs(y.star.hat.in.rct)
  obj.t <- 0
  n <- dim(x.in.rct)[1]-1
  p <- dim(x.in.rct)[2]-1
  eta.hat <- rep(0, p+1)
  for (t in 1:maxit) {
    opt <- optim(eta.hat, smooth.ramp.loss, method = 'L-BFGS-B', control = list(maxit=10000),
                 betas=betas, x.in.rct=x.in.rct, y.star.hat.in.rct=y.star.hat.in.rct, w.tilde=w.tilde)
    eta.hat <- opt$par
    u <- (2*as.numeric(y.star.hat.in.rct > 0) - 1) * (x.in.rct %*% as.matrix(eta.hat, ncol=1))
    betas <- cal_betas(y.star.hat.in.rct, u, w.tilde)
    obj.tplus1 <- mean((phi_su(1, u) - phi_su(0, u)) * abs(y.star.hat.in.rct) * w.tilde) 
    # Uses objective function as criteria for convergence
    if (abs(obj.t - obj.tplus1) <= eps) {
      return(eta.hat)
    } else {
      obj.t <- obj.tplus1
    }
  }
  if (t==maxit) stop("Not converge within max iterations")
}

## B. Estimate contrast function value (tau)

est.contrast <- function(a.in.rct, x.in.rct, y.in.rct, meth=1, Qm=1, w.tilde=NULL) {
  # Estimate contrast function. 
  # meth: {1,2,3}, different methods of contrast value estimates, corresponding to IPW, 
  #       direct regresion and AIPW, respectively.
  if (meth==1) {
    # IPW_i
    phat <- mean(a.in.rct)
    y.star.hat.in.rct <- a.in.rct*y.in.rct/phat - (1-a.in.rct)*y.in.rct/(1-phat) # outcome adjusted
  } else if (meth==2) {
    # direct regression: Q(X_i,1) - Q(X_i,0)
    x <- x.in.rct[,-1]
    Q.reg <- Q.est(x, y.in.rct, a.in.rct, Qm, w.tilde=w.tilde)
    reg.mu0 <- Q.reg$reg.mu0
    reg.mu1 <- Q.reg$reg.mu1
    y.star.hat.in.rct <- reg.mu1 - reg.mu0
  } else if (meth==3) {
    # AIPW_i
    x <- x.in.rct[,-1]
    n <- length(a.in.rct)
    aipw_i <- v.aipw(x, y.in.rct, a.in.rct, Qm, w.tilde=w.tilde)$aipw_i
    y.star.hat.in.rct <- aipw_i
  } 
  return(y.star.hat.in.rct)
}

v.aipw <- function(x, y, A, Qm, w.tilde) {
  #phat <- mean(A) # equal propensity score for everyone, experimental data
  df <- as.data.frame(cbind(x,A))
  glm_model <- glm(formula = A ~ ., data = df, family = binomial)
  phat <- glm_model$fitted.values # different propensity score for each based on x, observational study
  n <- length(y)
  Q.reg <- Q.est(x, y, A, Qm, w.tilde)
  reg.mu0 <- Q.reg$reg.mu0
  reg.mu1 <- Q.reg$reg.mu1
  psi.aug <- y*A/phat-y*(1-A)/(1-phat)+(reg.mu1*(1-A/phat)-reg.mu0*(1-(1-A)/(1-phat)) )
  aipw <- mean(psi.aug)
  return(list(aipw=aipw, aipw_i=psi.aug))
}

library(ranger)

Q.est <- function(x, y, A, Qm, w.tilde) {
  if (Qm==1) {
    # Use random forest to est Q 
    loc.a1 <- which(A==1)
    loc.a0 <- which(A==0)
    dat <- data.frame(cbind(y, x))
    fit1 <- ranger(y~., data=dat[loc.a1,], case.weights = w.tilde[loc.a1]) 
    fit0 <- ranger(y~., data=dat[loc.a0,], case.weights = w.tilde[loc.a0])
    reg.mu1 <- predict(fit1, data = dat)$predictions
    reg.mu0 <- predict(fit0, data = dat)$predictions
  } else if (Qm==2) {
    n <- length(y)
    dat2 <- data.frame(cbind(y, x, A,  A*x))
    names(dat2)[1] <- "y"
    # Use linear model to est Q 
    lin <- glm(y ~ ., data=dat2, weights=w.tilde)
    newdf <- data.frame(x, rep(0, n), x*0)
    names(newdf) <- names(dat2)[-1]
    reg.mu0 <- predict(lin, newdata = newdf)
    newdf <-  data.frame(x, rep(1, n), x)
    names(newdf) <- names(dat2)[-1]
    reg.mu1 <- predict(lin, newdata = newdf)
  } else {
    stop("No such choice of Q function estimation methods")
  }
  return(list(reg.mu0=reg.mu0, reg.mu1=reg.mu1))
}


