source('GenerateSetUp.R')



#Run replicates

nrep<-1000
ests.mat.dr<-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)

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

	#DR robust
	#DR using: all predictors in both models

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

	f0<-lm(Y~Z*X)
	new0<-data.frame(Y,Z=rep(0,n),X);mufit0<-predict(f0,new0)
	new1<-data.frame(Y,Z=rep(1,n),X);mufit1<-predict(f0,new1)


	ests.mat.dr[irep,1]<-mean((Y-mufit1)*Z/ps.vals)-mean((Y-mufit0)*(1-Z)/(1-ps.vals))+
					mean(mufit1-mufit0)
	ests.mat.dr[irep,2]<-sum((Y-mufit1)*Z/ps.vals)/sum(Z/ps.vals)-
					sum((Y-mufit0)*(1-Z)/(1-ps.vals))/sum((1-Z)/(1-ps.vals))+
					mean(mufit1-mufit0)

	
	#DR using: predictors of treatment only for PS, predictors of outcome only for OR

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

	V<-X[,outcome.list]
	f1<-lm(Y~Z*V)
	new0<-data.frame(Y,Z=rep(0,n),V);mufit0<-predict(f1,new0)
	new1<-data.frame(Y,Z=rep(1,n),V);mufit1<-predict(f1,new1)

	ests.mat.dr[irep,3]<-mean((Y-mufit1)*Z/ps.vals)-mean((Y-mufit0)*(1-Z)/(1-ps.vals))+
					mean(mufit1-mufit0)
	ests.mat.dr[irep,4]<-sum((Y-mufit1)*Z/ps.vals)/sum(Z/ps.vals)-
					sum((Y-mufit0)*(1-Z)/(1-ps.vals))/sum((1-Z)/(1-ps.vals))+
					mean(mufit1-mufit0)
	
	#DR using: confounders only in both models

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

	V<-X[,confound.list]
	f2<-lm(Y~Z*V)
	new0<-data.frame(Y,Z=rep(0,n),V);mufit0<-predict(f2,new0)
	new1<-data.frame(Y,Z=rep(1,n),V);mufit1<-predict(f2,new1)

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

}

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

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