## Gibbs Sampler for Gaussian Selection Model
## Li, Poirier and Tobias (2004) JAE

sample.selection <- function(Data, Prior, Mcmc)
{
## Set up Data

xmat <- as.matrix(Data$X)
zmat <- as.matrix(Data$Z)
y <- as.matrix(Data$Y)
DD <- as.matrix(Data$DD)
nobs<-length(y)

## Set up counters based on data
nvarx <- ncol(xmat)
nvarz <- ncol(zmat)

kbar <- 2*nvarx+nvarz
kz <- nvarz
kx <- nvarx

## Assign Priors

rho <- Prior$rho
Sigma <- Prior$Sigma
beta1<-Prior$beta1p	#priors added for parameters
beta0<-Prior$beta0p
theta<-Prior$thetap


## Set up MCMC parameters
iter <- Mcmc$iter
burn <- Mcmc$burn

## Setting initials

mub1 <- mub0 <- rep(0,nvarx)
mut <- rep(0,nvarz)
Vb1 <- 100*diag(nvarx)
Vb0 <- 100*diag(nvarx)
Vt <- 100*diag(nvarz)
R <- diag(3)
covm <- list(Vt,Vb1,Vb0)
covarian <- blockdiag(covm)
inv_parms <- solve(covarian)
big_mean <- matrix(c(mut,mub1,mub0),,1)

ate.draw<-NULL
att.draw<-NULL
atut.draw<-NULL

## Matrices to be used later

ZZ <- t(zmat)%*%zmat
ZX <- t(zmat)%*%xmat
XX <- t(xmat)%*%xmat
XZ <- t(xmat)%*%zmat

## Set Vector Lenghts and Initial conditions

b1_final <- b0_final <- matrix(0,nvarx,iter-burn)
t_final <- matrix(0,nvarz, iter-burn)
sig1_final <- sig0_final <- corr1fin <- corr0fin <- corr10fin <- matrix(0,iter-burn,1)

Dstar <- ifelse(DD==1,rtrun(0,1,0,100),rtrun(0,1,-100,0)) ## check truncnorm implementation

Sigma <- Sigma
Siginv <- solve(Sigma)

sig1 <- sqrt(Sigma[2,2])
sig0 <- sqrt(Sigma[3,3])
sig1D <- Sigma[1,2]
sig0D <- Sigma[1,3]
sig10 <- Sigma[2,3]

## Begin Gibbs Sampler
itime = proc.time()[3]
for (j in 2:iter)
{
## Draw the missing outcome

  denom0 <- (sig0^2)-(sig0D^2)
  denom1 <- (sig1^2)-(sig1D^2)
  middle <- sig10*sig0D*sig1D
  
  mu_1 <- (xmat%*%beta1) + (Dstar-(zmat%*%theta))*(((sig0^2)*sig1D-sig10*sig0D)/denom0)+
          (y-(xmat%*%beta0))*((sig10-sig0D*sig1D)/denom0)

  mu_0 <- (xmat%*%beta0) + (Dstar-(zmat%*%theta))*(((sig1^2)*sig0D-sig10*sig1D)/denom1)+
          (y-(xmat%*%beta1))*((sig10-sig0D*sig1D)/denom1)
          
  omega_1 <- (sig1^2) - (((sig1D^2)*(sig0^2)-2*middle +(sig10^2))/denom0)
  omega_0 <- (sig0^2) - (((sig0D^2)*(sig1^2)-2*middle +(sig10^2))/denom1)
  
  mu_miss <- (1-DD)*mu_1 + (DD)*mu_0
  var_miss <- (1-DD)*omega_1 + (DD)*omega_0
  
  y_miss <- mu_miss+sqrt(var_miss)*rnorm(nobs)
  
  ## Draw the Latent Data
  
  denomd <- (sig1^2)*(sig0^2) - (sig10^2)
  mu_d <- (zmat%*%theta)+(DD*y+(1-DD)*(y_miss)-(xmat%*%beta1))*(((sig0^2)*sig1D-sig10*sig0D)/denomd)+
          (DD*y_miss + (1-DD)*y -(xmat%*%beta0))*(((sig1^2)*sig0D-sig10*sig1D)/denomd)
  omega_d <- 1-(((sig1D^2)*(sig0^2)-2*sig10*sig0D*sig1D + (sig1^2)*(sig0D^2))/denomd)
  Dstar <- ifelse(DD==1,rtrun(mu_d,omega_d,0,Inf),rtrun(mu_d,omega_d,-Inf,0))
    
  ## Draw Coefficients
  
  s11 <- Siginv[1,1]
  s21 <- Siginv[2,1]
  s31 <- Siginv[3,1]
  s12 <- Siginv[1,2]
  s22 <- Siginv[2,2]
  s32 <- Siginv[3,2]
  s13 <- Siginv[1,3]
  s23 <- Siginv[2,3]
  s33 <- Siginv[3,3]
  
  yu1 <- t(Dstar)
  yu2 <- t(DD*y+(1-DD)*y_miss)
  yu3 <- t((1-DD)*y+(DD)*y_miss)
  yuse <- t(cbind(yu1,yu2,yu3))
  
  zm1 <- t(zmat)*s11
  zm2 <- t(zmat)*s12
  zm3 <- t(zmat)*s13
  zmata <- cbind(zm1,zm2,zm3)
  
  x1m1 <- t(xmat)*s21
  x1m2 <- t(xmat)*s22
  x1m3 <- t(xmat)*s23
  x1mata <-cbind(x1m1,x1m2,x1m3) 

  x0m1 <- t(xmat)*s31
  x0m2 <- t(xmat)*s32
  x0m3 <- t(xmat)*s33
  x0mata <- cbind(x0m1,x0m2,x0m3)
  
  m1 <- ZZ*s11
  m2 <- ZX*s12
  m3 <- ZX*s13
  m4 <- XZ*s21
  m5 <- XX*s22
  m6 <- XX*s23
  m7 <- XZ*s31
  m8 <- XX*s32
  m9 <- XX*s33
  
  Xmatpart <- rbind(cbind(m1,m2,m3), cbind(m4,m5,m6), cbind(m7,m8,m9))
  
  ymp1 <- zmata%*%yuse
  ymp2 <- x1mata%*%yuse
  ymp3 <- x0mata%*%yuse
  
  ymatpart <- rbind(ymp1,ymp2,ymp3)
  covmatpart <- chol2inv(chol(Xmatpart+inv_parms))
  meanpart <- covmatpart%*%(ymatpart+(inv_parms%*%big_mean))
  
  betas <- meanpart+t(chol(covmatpart))%*%rnorm(kbar)
  
  theta <- betas[1:kz]
  beta1 <- betas[(kz+1):(kz+kx)]
  beta0 <- betas[(kz+kx+1):(kbar)]
  
  ### Draw the Covariance Matrix
  
  y_one <- DD*y + (1-DD)*y_miss
  y_zero <- DD*y_miss + (1-DD)*y
  
  e_D <- Dstar - (zmat%*%theta)
  e_1 <- y_one - (xmat%*%beta1)
  e_0 <- y_zero - (xmat%*%beta0)
  
  e00 <- t(e_0)%*%e_0
  e01 <- t(e_0)%*%e_1
  e02 <- t(e_0)%*%e_D
  e10 <- t(e_1)%*%e_0
  e11 <- t(e_1)%*%e_1
  e12 <- t(e_1)%*%e_D
  e20 <- t(e_D)%*%e_0
  e21 <- t(e_D)%*%e_1
  e22 <- t(e_D)%*%e_D
  
  
  term3 <- matrix(c(e00,e10,e20,e01,e11,e21,e02,e12,e22),3,3)

  Sig_temp <- nobile.wishart((nobs+rho),(term3+(rho*R)))
    
  st33 <- Sig_temp[3,3] 
  st32 <- Sig_temp[3,2]  
  st31 <- Sig_temp[3,1]  
  st23 <- Sig_temp[2,3]
  st22 <- Sig_temp[2,2]
  st21 <- Sig_temp[2,1]
  st13 <- Sig_temp[1,3]
  st12 <- Sig_temp[1,2]
  st11 <- Sig_temp[1,1]
  
  Sig <- rbind(cbind(st33,st32,st31),cbind(st23,st22,st21),cbind(st13,st12,st11))
  
  Siginv <- chol2inv(chol(Sig))

  
  sig1 <- sqrt(Sig[2,2])
  sig0 <- sqrt(Sig[3,3])
  sig10 <- Sig[2,3]
  sig1D <- Sig[1,2]
  sig0D <- Sig[1,3]
  
  ## Save Draws
  
  if (j>burn)
  {
    b1_final[,j-burn] <- beta1
    b0_final[,j-burn] <- beta0
    t_final[,j-burn] <- theta
    sig1_final[j-burn,1] <- sig1
    sig0_final[j-burn,1] <- sig0
    corr1fin[j-burn,1] <- sig1D/sig1
    corr0fin[j-burn,1] <- sig0D/sig0
    corr10fin[j-burn,1] <- sig10/(sig1*sig0)
    
    ate.draw[j-burn]<-mean(y_one-y_zero)
    att.draw[j-burn]<-sum(DD*(y_one-y_zero))/sum(DD)
    atut.draw[j-burn]<-sum((1-DD)*(y_one-y_zero))/sum(1-DD)
   }
reps <- j

if(reps%%100==0)
{
  ctime <- proc.time()[3]
  timetoend <-((ctime-itime)/reps)*(iter-reps)
  cat(" ",reps,"(",round(timetoend/60,1),")",fill=TRUE)
  ctime=proc.time()[3]
  cat("  Total Time Elapsed: ", round((ctime-itime)/60,2),"\n")

 }
}
list(b1fin=matrix(b1_final,,nvarx,byrow=T),b0fin=matrix(b0_final,,nvarx,byrow=T),
tfin=matrix(t_final,,nvarz,byrow=T),y1draw=y_one, y0draw=y_zero,ate.draw=ate.draw,att.draw=att.draw,atut.draw=atut.draw,c1fin=corr1fin,c0fin=corr0fin,c10fin=corr10fin)
}