library(mvtnorm)

#####################################################################################
# IV test of Huber and Mellace based on Bennett (2009) and Chen and Szroeter (2012) #
#####################################################################################

# INPUT: 
# y: outcome
# d: binary treatment
# z: binary instrument
# n_boot1: number of bootstrap replications for the first bootstrap in Bennett (2009) procedure
# n_boot2: number of bootstrap replications for the second bootstrap in Bennett (2009) procedure
 

# OUTPUT:
# fullr.pval: p-value under full recentering using Bennett (2009)
# partialr.pval: p-value under partial recentering using Bennett (2009)
# ChSz.pval: p-value using Chen and Szroeter (2012)
# stdif.treated: Standardized difference between nearest bound and point estimate for treated (positive values indicate violation of the null)
# stdif.nontreated: Standardized difference between nearest bound and point estimate for nontreated (positive values indicate violation of the null)


# IV test based on Bennett (2009) and Chen and Szroeter (2012)
ivtest<-function(y,d,z,n_boot1=999, n_boot2=999){
n=length(y)
p1_1<-mean(d[z==1])
p1_0<-mean(d[z==0])
p0_1<-1-mean(d[z==1])
p0_0<-1-mean(d[z==0])
z1<-z[d==1]
z0<-z[d==0]
y1<-y[d==1]
y0<-y[d==0]
y11=sort(y1[z1==1])
y01=sort(y1[z1==0])
y10=sort(y0[z0==1])
y00=sort(y0[z0==0])

y11temp<-sort(runif(length(y11)))
y00temp<-sort(runif(length(y00)))

y11min<-(y11[y11temp<=quantile(y11temp,p1_0/p1_1)]) 
y11max<-(y11[y11temp>=quantile(y11temp, 1-p1_0/p1_1)])
y00min<-(y00[y00temp<=quantile(y00temp, p0_1/p0_0)])
y00max<-(y00[y00temp>=quantile(y00temp,1-p0_1/p0_0)])

#compute the normalized differences
normdif.treated<-((mean(y11)>mean(y01))*(mean(y11min)-mean(y01))-(mean(y11)<mean(y01))*(mean(y11max)-mean(y01)))/sd(y)   
normdif.nontreated<-((mean(y00)>mean(y10))*(mean(y00min)-mean(y10))-(mean(y00)<mean(y10))*(mean(y00max)-mean(y10)))/sd(y)

#first bootstrap 
Tn1b = c()
Tn2b = c()
Tn3b = c()
Tn4b = c()

while(length(Tn1b)<n_boot1){
sboot<-sample(1:length(y),length(y),TRUE)
zb<-z[sboot]
db<-d[sboot]
yb<-y[sboot]

obsb=length(yb)
obs1b<-sum(db)
obs0b<-sum(1-db)

p1_1b<-mean(db[zb==1])
p1_0b<-mean(db[zb==0])
p0_1b<-1-mean(db[zb==1])
p0_0b<-1-mean(db[zb==0])

if (p1_0b/p1_1b<=1){
z1b<-zb[db==1]
z0b<-zb[db==0]
y1b<-yb[db==1]
y0b<-yb[db==0]

y11b=sort(y1b[z1b==1])
y01b=sort(y1b[z1b==0])
y10b=sort(y0b[z0b==1])
y00b=sort(y0b[z0b==0])

y11tempb<-sort(runif(length(y11b)))
y00tempb<-sort(runif(length(y00b)))

y11minb<-mean(y11b[y11tempb<=quantile(y11tempb,p1_0b/p1_1b)]) 
y11maxb<-mean(y11b[y11tempb>=quantile(y11tempb, 1-p1_0b/p1_1b)])
y00minb<-mean(y00b[y00tempb<=quantile(y00tempb, p0_1b/p0_0b)])
y00maxb<-mean(y00b[y00tempb>=quantile(y00tempb,1-p0_1b/p0_0b)])

y01b=mean(y01b)
y10b=mean(y10b)

Tn1b = c(Tn1b,(y11minb-y01b)) 
Tn2b = c(Tn2b,(y01b-y11maxb)) 
Tn3b = c(Tn3b,(y00minb-y10b)) 
Tn4b = c(Tn4b,(y10b-y00maxb)) 
}
}

# variances
varTn1b=var(Tn1b)
varTn2b=var(Tn2b)
varTn3b=var(Tn3b)
varTn4b=var(Tn4b)

# create objects for partial and full recentering
FCbsample1 <-matrix(0,nrow=1,ncol=n_boot2)
FCbsample2 <-matrix(0,nrow=1,ncol=n_boot2)
FCbsample3 <-matrix(0,nrow=1,ncol=n_boot2)
FCbsample4 <-matrix(0,nrow=1,ncol=n_boot2)

PCbsample1 <-matrix(0,nrow=1,ncol=n_boot2)
PCbsample2 <-matrix(0,nrow=1,ncol=n_boot2)
PCbsample3 <-matrix(0,nrow=1,ncol=n_boot2)
PCbsample4 <-matrix(0,nrow=1,ncol=n_boot2)

FCp_vals1 <- matrix(0, nrow = n_boot2, ncol = 1)
FCp_vals2 <- matrix(0, nrow = n_boot2, ncol = 1)
FCp_vals3 <- matrix(0, nrow = n_boot2, ncol = 1)
FCp_vals4 <- matrix(0, nrow = n_boot2, ncol = 1)

PCp_vals1 <- matrix(0, nrow = n_boot2, ncol = 1)
PCp_vals2 <- matrix(0, nrow = n_boot2, ncol = 1)
PCp_vals3 <- matrix(0, nrow = n_boot2, ncol = 1)
PCp_vals4 <- matrix(0, nrow = n_boot2, ncol = 1)

#Fix value of sequence delta_n
delta_n = sqrt(2*log(log(n))/n)

# compute test statistics
Tn1 = (mean(y11min)-mean(y01)) 
Tn2 = (mean(y01)-mean(y11max)) 
Tn3 = (mean(y00min)-mean(y10)) 
Tn4 = (mean(y10)-mean(y00max))
Tn=c(Tn1,Tn2, Tn3, Tn4) 

# compute recentering terms
maxVec1=( Tn1 >= (-delta_n*sqrt(varTn1b))  )*Tn1    + ( Tn1 <  (-delta_n*sqrt(varTn1b))  )* (-delta_n*sqrt(varTn1b) )
maxVec2=( Tn2 >= (-delta_n*sqrt(varTn2b))  )*Tn2   +( Tn2 < (-delta_n*sqrt(varTn2b))  )* (-delta_n*sqrt(varTn2b) )
maxVec3=( Tn3 >= (-delta_n*sqrt(varTn3b))  )*Tn3    + ( Tn3 <  (-delta_n*sqrt(varTn3b))  )* (-delta_n*sqrt(varTn3b) )
maxVec4=( Tn4 >= (-delta_n*sqrt(varTn4b))  )*Tn4   +( Tn4 < (-delta_n*sqrt(varTn4b))  )* (-delta_n*sqrt(varTn4b) )

Tn1 = sqrt(n)*Tn1 
Tn2 = sqrt(n)*Tn2
Tn3 = sqrt(n)*Tn3 
Tn4 = sqrt(n)*Tn4

Tn1b=sqrt(n)*Tn1b
Tn2b=sqrt(n)*Tn2b
Tn3b=sqrt(n)*Tn3b
Tn4b=sqrt(n)*Tn4b

# Generate bootstrap samples (using first bootstrap)- full recentering
FCBootstrap1<-matrix(Tn1b-rep(Tn1,n_boot1),1,n_boot1)
FCBootstrap2<-matrix(Tn2b-rep(Tn2,n_boot1),1,n_boot1)
FCBootstrap3<-matrix(Tn3b-rep(Tn3,n_boot1),1,n_boot1)
FCBootstrap4<-matrix(Tn4b-rep(Tn4,n_boot1),1,n_boot1)

# Generate bootstrap samples (using first bootstrap)- partial recentering
PCBootstrap1<-matrix(Tn1b-rep((sqrt(n)* maxVec1),n_boot1),1,n_boot1)
PCBootstrap2<-matrix(Tn2b-rep((sqrt(n)* maxVec2),n_boot1),1,n_boot1)
PCBootstrap3<-matrix(Tn3b-rep((sqrt(n)* maxVec3),n_boot1),1,n_boot1)
PCBootstrap4<-matrix(Tn4b-rep((sqrt(n)* maxVec4),n_boot1),1,n_boot1)

# Generate empirical distribution functions - full recentering
FC_ECDF1 <- apply(FCBootstrap1, 1, ecdf)
FC_ECDF2 <- apply(FCBootstrap2, 1, ecdf)
FC_ECDF3 <- apply(FCBootstrap3, 1, ecdf)
FC_ECDF4 <- apply(FCBootstrap4, 1, ecdf)

# Generate empirical distribution functions - partial recentering
PC_ECDF1 <- apply(PCBootstrap1, 1, ecdf)
PC_ECDF2 <- apply(PCBootstrap2, 1, ecdf)
PC_ECDF3 <- apply(PCBootstrap3, 1, ecdf)
PC_ECDF4 <- apply(PCBootstrap4, 1, ecdf)

# evaluate the test statistics at the empirical distributions to compute the p-values
p_val_1 =  1- FC_ECDF1[[1]](Tn1)
p_val_2 =  1- FC_ECDF2[[1]](Tn2)
p_val_3 =  1- FC_ECDF3[[1]](Tn3)
p_val_4 =  1- FC_ECDF4[[1]](Tn4) 

# compute the minimum p-value based on the original sample
p_val = min(p_val_1,p_val_2,p_val_3,p_val_4)

# second bootstrap - draw observations from first bootstrap
bsample2 <- sample(n_boot1, n_boot2, replace=TRUE)

# draw samples of recentered test statistics out of first bootstrap - full recentering
FCbsample1 = matrix(FCBootstrap1[1,bsample2],1,n_boot2)
FCbsample2 = matrix(FCBootstrap2[1,bsample2],1,n_boot2)
FCbsample3 = matrix(FCBootstrap3[1,bsample2],1,n_boot2)
FCbsample4 = matrix(FCBootstrap4[1,bsample2],1,n_boot2)

# draw samples of recentered test statistics out of first bootstrap - partial recentering
PCbsample1 = matrix(PCBootstrap1[1,bsample2],1,n_boot2)
PCbsample2 = matrix(PCBootstrap2[1,bsample2],1,n_boot2)
PCbsample3 = matrix(PCBootstrap3[1,bsample2],1,n_boot2)
PCbsample4 = matrix(PCBootstrap4[1,bsample2],1,n_boot2)

# empirical distributions of fully recentered p-values
FCp_vals1[1:n_boot2,1] = 1-FC_ECDF1[[1]](FCbsample1[1,1:n_boot2])	
FCp_vals2[1:n_boot2,1] = 1-FC_ECDF2[[1]](FCbsample2[1,1:n_boot2])
FCp_vals3[1:n_boot2,1] = 1-FC_ECDF3[[1]](FCbsample3[1,1:n_boot2])
FCp_vals4[1:n_boot2,1] = 1-FC_ECDF4[[1]](FCbsample4[1,1:n_boot2])
FCp_vals=cbind(FCp_vals1,FCp_vals2,FCp_vals3,FCp_vals4)

# empirical distributions of partially recentered p-values
PCp_vals1[1:n_boot2,1] = 1-FC_ECDF1[[1]](PCbsample1[1,1:n_boot2])
PCp_vals2[1:n_boot2,1] = 1-FC_ECDF2[[1]](PCbsample2[1,1:n_boot2])
PCp_vals3[1:n_boot2,1] = 1-FC_ECDF3[[1]](PCbsample3[1,1:n_boot2])
PCp_vals4[1:n_boot2,1] = 1-FC_ECDF4[[1]](PCbsample4[1,1:n_boot2])
PCp_vals=cbind(PCp_vals1,PCp_vals2,PCp_vals3,PCp_vals4)

# compute adjusted p-values - full recentering 
min_p_vec <- apply(FCp_vals,1,min) # minimum p-value vector
edfF_n <- ecdf(min_p_vec)    # empirical distribution function
FCpval = edfF_n(p_val)  

# compute adjusted p-values - partial recentering 
min_p_vec <- apply(PCp_vals,1,min)
edfF_n <- ecdf(min_p_vec)
PCpval = edfF_n(p_val) 

# CHEN AND SZROETER TEST
J=c(n*varTn1b,n*varTn2b,n*varTn3b, n*varTn4b)
J[J==0]=0.000000000001
theta=1/sqrt(J)
#Tuning parameter for the sample size
K=1/delta_n
# Smooth the indicator function
psi=pnorm(K*theta*Tn);
Lambda=(dnorm(K*theta*Tn)*K)/sqrt(n)
# Create a diagonal matrix with the variances of the elements of Tn
delta=diag(theta);
# Compute the P-value
Q1=sqrt(n)*t(psi)%*%(delta%*%Tn)-matrix(1,1,length(Tn))%*%(Lambda);
Q2=sqrt(t(psi)%*%delta%*%t(t(J)%*%delta*psi))
Pc=min(1-pnorm(Q1/max(Q2,0.00000001)),1)

list(ChSz.pval=Pc, fullr.pval=FCpval, partialr.pval=PCpval, normdif.treated=normdif.treated, normdif.nontreated=normdif.nontreated)
} 

##################
# LATE estimator #
##################

late<-function(y,d,z){
late<-(mean(y[z==1])-mean(y[z==0]))/(mean(d[z==1])-mean(d[z==0]))
list (late=late)
}


##############
# Simulation #
##############

mctestpartial<-c( )
mctestfull<-c( )
mcChenSzroeter<-c( )
mcdiff1<-c( )
mcdiff2<-c( )
mclate<-c()
mcnaive<-c()
mccomplier<-c()

n<-1000
sigma <- matrix(c(1,0.5,0.5,1),2,2)
repetitions<-1000
i=1
while(length(mctestpartial)<repetitions){
set.seed(i)
e<-(rmvnorm(n,rep(0,2),sigma))
z<-rnorm(n)>0
d<-(0.5*z+e[,1])>0
y<-d+z+e[,2]

p1_1<-mean(d[z==1])
p1_0<-mean(d[z==0])
p0_1<-1-mean(d[z==1])
p0_0<-1-mean(d[z==0])

z1<-z[d==1]
z0<-z[d==0]
y1<-y[d==1]
y0<-y[d==0]
y11=y1[z1==1]
y01=y1[z1==0]
y10=y0[z0==1]
y00=y0[z0==0]
mclate=c(mclate, late(y,d,z)$late)
mcnaive<-c(mcnaive, mean(y[d==1])-mean(y[d==0]))
mccomplier<-c(mccomplier, mean(d[z==1])-mean(d[z==0]))
temp<-ivtest(y,d,z, n_boot1=499, n_boot2=499)
mctestpartial=c(mctestpartial, temp$partialr.pval)
mctestfull=c(mctestfull, temp$fullr.pval)
mcChenSzroeter<-c(mcChenSzroeter,  temp$ChSz.pval)
mcdiff1<-c(mcdiff1, temp$normdif.treated) 
mcdiff2<-c(mcdiff2, temp$normdif.nontreated)
i=i+1
}

mean(mclate)
mean(mcnaive)
mean(mcdiff1)
mean(mcdiff2)
mean(mctestpartial<0.05)
mean(mctestfull<0.05)
mean(mcChenSzroeter<0.05)
i