library(np)
library(quantreg)

#########################################################################
# BALANCE TEST BASED ON QUANTILE REGRESSION AND RESAMPLING (HUBER 2011) #
#########################################################################

###########################
# procedure "balancetest" #
###########################

# input for procedure "balancetest"
# x: covariate (vector) the balance of which is of interest 
# p: propensity score
# d: binary treatment indicator
# m: sample size used in bootstrap replications, the default is the original sample size
# grid: ranks at which the conditional quantiles/test statistics are to be evaluated, between 0 and 1; default is {0.10, 0.15,..., 0.85, 0.90}}
# eval: p-scores at which the conditional quantiles/test statistics are to be evaluated, between 0 and 1; default are the empirical values  
# order: defines the polynomial order of the p-score on which x is regressed; higher orders make the regression more flexible, but increase variance; default is 3 
# boot: number of bootstrap replications; default is 499

# output of procedure "balancetest"
# stat.ks/stat.cvm: Kolmogorov-Smirnov/Cramer-von-Mises test statistics (non-pivotal!)
# pvalue.ks/pvalue.cvm: p-value of Kolmogorov-Smirnov/Cramer-von-Mises test

###########################################################################################################################
# procedure "gcv"  (to find the optimal order of the polynomial function on the p-score via generalized cross-valdiation) #
###########################################################################################################################

# input for procedure "gcv"
# x: covariate (vector) the balance of which is of interest 
# d: binary treatment indicator
# p: propensity score
# tau: rank at which the conditional quantiles are to be evaluated; default is 0.5
# maxorder: maximum order to be considered in the cross-validation; default is 6

# output of procedure "gcv"
# gcv1/gcv0: generalized cross-validation criterions (for orders 1 to maxorder) under treatment/non-treatment
# optorder1/optorder0: Optimal order of the polynomial under treatment/nontreatment

##############
# PROCEDURES # 
##############

balancetest<-function(x,p,d,m=length(d),grid=seq(0.1,0.9,0.05),eval=p, order=3,boot=499){
n=length(d)
n1=length(d[d==1])
n0=length(d[d==0])
q1<-c()
q0<-c()
pp<-c()
evalp<-c()
for (j in 1:order){
	pp<-cbind(pp,p^j)
	evalp<-cbind(evalp,eval^j)
}
for (i in 1:length(grid)){
	if (order==1) {
	temp1<-rq(x[d==1]~pp[d==1], tau=grid[i])$coefficients
	}
	if (order>1) {
	temp1<-rq(x[d==1]~pp[d==1,], tau=grid[i])$coefficients
	}
fit1<-cbind(1,evalp)%*%temp1
q1<-c(q1,fit1)
	if (order==1) {
	temp0<-rq(x[d==0]~pp[d==0], tau=grid[i])$coefficients
	}
	if (order>1) {
	temp0<-rq(x[d==0]~pp[d==0,], tau=grid[i])$coefficients
	}
fit0<-cbind(1,evalp)%*%temp0
}
q0<-c(q0,fit0)
diff<-q1-q0
rboot1<-matrix(,0,length(eval)*length(grid))
rboot0<-matrix(,0,length(eval)*length(grid))
while(dim(rboot1)[1]<boot){
sboot<-sample(1:n,m,TRUE)
db<-d[sboot]
xb<-x[sboot]
if (order==1) {
pb<-pp[sboot]
}
if (order>1) {
pb<-pp[sboot,]
}
q1b<-c()
q0b<-c()
for (i in 1:length(grid)){
if (order==1) {
temp1<-rq(xb[db==1]~pb[db==1], tau=grid[i])$coefficients
}
if (order>1) {
temp1<-rq(xb[db==1]~pb[db==1,], tau=grid[i])$coefficients
}
fit1<-cbind(1,evalp)%*%temp1
q1b<-c(q1b,fit1)
if (order==1) {
temp0<-rq(xb[db==0]~pb[db==0], tau=grid[i])$coefficients
}
if (order>1) {
temp0<-rq(xb[db==0]~pb[db==0,], tau=grid[i])$coefficients
}
fit0<-cbind(1,evalp)%*%temp0
q0b<-c(q0b,fit0)
}
rboot1<-rbind(rboot1,q1b)
rboot0<-rbind(rboot0,q0b)
}
rboot<-rboot1-rboot0
varb<-sd(rboot[,1:dim(rboot)[2]])^2*m/n
tvalue<-n1*n0/(n1+n0)*mean((q1-q0)^2/varb)
test<-c()
for(i in 1:boot) test<-c(test,n1*n0/(n1+n0)*mean((rboot[i,]-(q1-q0))^2/varb))
cvm=mean(tvalue<test)
stat.cvm<-tvalue
tvalue<-(n1*n0/(n1+n0))^0.5*max(abs(q1-q0)/varb^0.5)
test<-c()
for(i in 1:boot) test<-c(test,(n1*n0/(n1+n0))^0.5*max(abs(rboot[i,]-(q1-q0))/varb^0.5))
ks=mean(tvalue<test)
stat.ks<-tvalue
list(stat.ks=stat.ks,stat.cvm=stat.cvm,pvalue.ks=ks,pvalue.cvm=cvm)}

#cross-validation to find optimal order
optimalordergcv<-function(x,d, p, tau=0.5 ,maxorder=6){
pp<-p
for(i in 2:maxorder) pp<-cbind(pp,p^i)
gcv1<-c()
gcv0<-c()
count<-c(1:maxorder)
for(i in 1:maxorder){
temp1<-rq(x[d==1]~pp[d==1,1:i], tau=tau)$coefficients
temp0<-rq(x[d==0]~pp[d==0,1:i], tau=tau)$coefficients
gcv1<-c(gcv1, gcvq(x[d==1], cbind(1,pp[d==1,1:i]), temp1, tau))
gcv0<-c(gcv0, gcvq(x[d==0], cbind(1,pp[d==0,1:i]), temp1, tau))
	}
optorder1=count[gcv1==min(gcv1)]
optorder0=count[gcv0==min(gcv0)]
list (gcv1=gcv1, gcv0=gcv0, optorder1=optorder1, optorder0=optorder0)
}

#check function
check<-function(lambda,q) (q-(lambda<0))*lambda

#generalized cross-validation for quantiles
gcvq<-function(dep,reg,beta,quantile){
pred<-reg%*%beta
H=reg%*%solve(t(reg)%*%reg)%*%t(reg)
mean(check(dep-pred,quantile))/(1-sum(diag(H))/length(dep))^2}


##############
# SIMULATION #
##############

n=500  # sample size
rep=1000 # number of Monte Carlo simulations

testks<-c()
testcvm<-c()

for (i in 1:rep){
set.seed(i)
z<-runif(n,0,3)
x<- runif(n,0,3)
d<-((-3+0.3*x^3+0.5*z+5*rnorm(n))>0)
p.est<-glm(d~x+z,family=binomial(probit))$fitted
upper<-uocquantile(p.est, 0.75)
lower<-uocquantile(p.est, 0.25)
temp<-balancetest(x=x,p=p.est,d=d,m=length(p.est),grid=seq(0.25,0.75,0.05),eval=seq(lower,upper,(upper-lower)/9), order=3,boot=299)
testks<-c(testks,temp$pvalue.ks)
testcvm<-c(testcvm,temp$pvalue.cvm)
}

mean(testcvm)  #average p-values
mean(testks)
mean(testcvm<0.05) #probability to reject at 5% level of significance
mean(testks<0.05)
mean(testcvm<0.1)
mean(testks<0.1)