require(GJRM)
require(mvtnorm)
mice.impute.hecknorm <- function(y, ry, x, ...) {
  
  # 1. Estimate the Heckman's model parameters
  res <- copulaSampleSel(list(ry ~ X1 + X2 + X3, 
                              y ~ X1 + X2),data=cbind(ry,y,x))
  
  # 2. Draw q.star
  q.star<-rmvnorm(1, res$coefficients, res$Vb, method = "chol")
  
  # 3. Impute y.star
  # 3.a. Calculate linear predictors for the selection equation
  var_sel<-res$X1[!ry,]
  c1<-dim(var_sel)[2]
  betaxstar_sel<-apply(var_sel,MARGIN=1,FUN=
                         function(var_sel){sum(q.star[1:c1]*var_sel)})
  # 3.b. Calculate linear predictors for the outcome equation
  var_out<-cbind(1,x[!ry,names(as.data.frame(res$X2))[2:dim(res$X2)[2]]])
  c2<-dim(var_out)[2]
  betaxstar_out<-apply(var_out,MARGIN=1,FUN=
                         function(var_out){sum(q.star[(c1+1):(c1+c2)]*var_out)})
  
  # 3.c. In copulaParSampleSel rho and sigma are transformed, recalculate it
  transtheta<-function(x){(exp(2*x)-1)/(1+exp(2*x))}
  q.star[length(q.star)]<-pmin(100,q.star[length(q.star)]) 
  # transtheta(100)==1
  q.star[length(q.star)]<-pmax(-100,q.star[length(q.star)])
  # transtheta(-100)==-1
  rho.star<-transtheta(q.star[length(q.star)])
  sigma.star<-exp(q.star[,(length(q.star)-1)])
  
  y.star<-betaxstar_out+sigma.star*rho.star*
    (-dnorm(betaxstar_sel)/(1-pnorm(betaxstar_sel)))
  + rnorm(sum(!ry),0, sigma.star)
  return(y.star)
  }