##################################################################################
# Program: MI-lasso.R                                                            #
# NOTE: This R program includes:                                                 #
# 1) the R function "MI.lasso" to implement the MI-lasso variable selection      #
#    method.                                                                     #
# 2) An example on how to use the "MI.lasso" function.                           #
#                                                                                #
# Reference:                                                                     #
#    Chen, Q. & Wang, S. (2013). "Variable selection for multiply-imputed data   #
#    with application to dioxin exposure study," Statistics in Medicine, 32:     #
#    3646-59.                                                                    #
##################################################################################


MI.lasso = function(mydata=mydata, D, lambda, maxiter=200, eps=1e-6) {

## D is the number of imputations
## mydata is in the array format: mydata[[1]] is the first imputed dataset...
## for each mydata[[d]], the first p columns are covariates X, and the last one is the outcome Y

## number of observations
n = dim(mydata[[1]])[1]

## number of covariates
p = dim(mydata[[1]])[2]-1

## Standardize covariates X and center outcome Y 
x = NULL
y = NULL
meanx = NULL
normx = NULL
meany = 0
for (d in 1:D) {
  x[[d]] = mydata[[d]][,1:p]
  y[[d]] = mydata[[d]][,(p+1)]
  meanx[[d]] = apply(x[[d]], 2, mean)
  x[[d]] = scale(x[[d]], meanx[[d]], FALSE)
  normx[[d]] = sqrt(apply(x[[d]]^2, 2, sum))
  x[[d]] = scale(x[[d]], FALSE, normx[[d]])
  meany[d] = mean(y[[d]])
  y[[d]] = y[[d]] - meany[d]
}

## Ordinary least squares (OLS) estimates of beta coefficients
b.ols = matrix(0,p,D)
for (d in 1:D) {
  b.ols[,d]  = qr.solve(t(x[[d]])%*%x[[d]])%*%t(x[[d]])%*%y[[d]]
}

## Estimate beta_dj in Equation (4) 
iter = 0
dif = 1
c = rep(1,p)
b = matrix(0,p,D) 						
while (iter<=maxiter & dif>=eps) {
  iter = iter + 1
  b.old = b
  # update beta_d by ridge regression
  for (d in 1:D) {
    xtx = t(x[[d]])%*%x[[d]]
    diag(xtx) = diag(xtx) + lambda/c
    xty = t(x[[d]])%*%y[[d]]
    b[,d] = qr.solve(xtx)%*%xty
  } 		
  # update c in Equation (4)
  c = sqrt(apply(b^2,1,sum)) 			
  c[c<sqrt(D)*1e-10] = sqrt(D)*1e-10
  dif = max(abs(b-b.old))
}
b[apply((b^2),1,sum)<=5*1e-8,] = 0

## Calculate BIC 
sse = 0
df = 0
nzero = seq(p)[apply(b^2, 1, sum)>0]
for (d in 1:D) {
  sse[d] = mean((y[[d]]-x[[d]]%*%b[,d])^2)
}
for (j in nzero) {
  norm1 = sqrt(sum(b[j,]^2))
  norm2 = sqrt(sum(b.ols[j,]^2))
  df = df+1+(D-1)*norm1/norm2
}
bic = log(mean(sse)) + log(D*n)*df/(D*n)

## transform beta and intercept back to original scale
b.scaled = b
b0 = 0
coefficients = matrix(0,p+1,D)
for (d in 1:D) {
  b[,d] = b[,d]/normx[[d]]
  b0[d] = meany[d] - b[,d]%*%meanx[[d]]
  coefficients[,d] = c(b0[d], b[,d])
}

## output the selected variables
## TRUE means "selected", FALSE means "NOT selected"
varsel = abs(b[,1])>0

return(list(coefficients=coefficients, bic=bic, varsel=varsel))
}


######################################################
#  An example on the use of the "MI.lasso" function	#
######################################################


## specify tuning parameter
lamvec = (2^(seq(-1,4,by=0.05)))^2/2

## select the optimal tuning parameter by BIC
bicvec = 0
for (i in 1:length(lamvec)) {
  fit = MI.lasso(newdata, 5, lambda=lamvec[i])
  bicvec[i] = fit$bic
}

## select tuning parameter with the smallest BIC
opt = order(bicvec)[1]

## fit the model with the optimal lambda
fit.opt = MI.lasso(newdata, 5, lambda=lamvec[opt])
selection = which(fit.opt$varsel==TRUE)
print(fit.opt$coefficients)