library(MASS)

set.seed(4338)

#Expit transform
expit<-function(x){return(1/(1+exp(-x)))}

n<-2000
p<-10

x.index<-1:p
al.select<-sample(0:1,size=p,rep=T)
be.select<-sample(0:1,size=p,rep=T)

#Confounders
confound.list<-x.index[al.select == 1 & be.select==1]

#Instruments
intrument.list<-x.index[al.select ==1 & be.select==0]

#Predictors of outcome only
outcomeonly.list<-x.index[al.select ==0  & be.select==1]

#Spurious variables
spurious.list<-x.index[al.select ==0  & be.select==0]

#Predictors of outcome
outcome.list<-sort(c(confound.list,outcomeonly.list))

#Predictors of exposure
treatment.list<-x.index[al.select ==1]

################################################

set.seed(43973)
#Generate alpha parameters (for treatment allocation)
al<-rnorm(p+1,0,2.0)

#Generate beta parameters (for outcome regression)
be<-rnorm(p+1,0,2.0)

theta<-2.0
sig.Y<-1

#Covariance matrix of predictors
Sig<-rWishart(1,p+1,diag(p)/(p+1))[,,1]


#Run replicates

nrep<-1000
ests.mat.ipw<-matrix(0,nrow=nrep,ncol=6)

for(irep in 1:nrep){
	#Predictors
	X<-mvrnorm(n,mu=rep(0,p),Sigma=Sig)

	#Data simulation
	Xmat<-cbind(rep(1,n),X)
	eta.Z<-Xmat %*% (c(1,al.select)*al)
	p.Z<-expit(eta.Z)
	Z<-rbinom(n,1,p.Z)
	eta.Y<-Z*theta+Xmat %*% (c(1,be.select)*be)
	
	Y<-eta.Y+rnorm(n)*sig.Y
	#boxplot(Y~Z)

	##########################################################

	#Single robust
	#IPW using: all predictors

	p1<-glm(Z~X,family=binomial)
	ps.vals<-fitted(p1)

	ests.mat.ipw[irep,1]<-mean(Y*Z/ps.vals)-mean(Y*(1-Z)/(1-ps.vals))
	ests.mat.ipw[irep,2]<-sum(Y*Z/ps.vals)/sum(Z/ps.vals)-
					sum(Y*(1-Z)/(1-ps.vals))/sum((1-Z)/(1-ps.vals))
	
	#IPW using: predictors of treatment only

	p1<-glm(Z~X[,treatment.list],family=binomial)
	ps.vals<-fitted(p1)

	ests.mat.ipw[irep,3]<-mean(Y*Z/ps.vals)-mean(Y*(1-Z)/(1-ps.vals))
	ests.mat.ipw[irep,4]<-sum(Y*Z/ps.vals)/sum(Z/ps.vals)-
					sum(Y*(1-Z)/(1-ps.vals))/sum((1-Z)/(1-ps.vals))

	
	#IPW using: confounders only

	p1<-glm(Z~X[,confound.list],family=binomial)
	ps.vals<-fitted(p1)

	ests.mat.ipw[irep,5]<-mean(Y*Z/ps.vals)-mean(Y*(1-Z)/(1-ps.vals))
	ests.mat.ipw[irep,6]<-sum(Y*Z/ps.vals)/sum(Z/ps.vals)-
					sum(Y*(1-Z)/(1-ps.vals))/sum((1-Z)/(1-ps.vals))
	

}

mse.vals<-n*apply((ests.mat.ipw-theta)^2,2,mean)
name.list<-c('IPW-All0','IPW-All','IPW-Tmt0','IPW-Tmt','IPW-Conf0','IPW-Conf')
boxplot(ests.mat.ipw,names=name.list,ylim=range(-2,6),main='IPW-Estimates')
abline(h=theta,col='red')
text(1:ncol(ests.mat.ipw),rep(6,ncol(ests.mat.ipw)),format(mse.vals,digits=4))
dev.print(device = pdf, file='PS-IPW.pdf',width=11)

boxplot(Y~Z,names=c('Untreated','Treated')
mean(Y[Z==1])-mean(Y[Z==0])