library(Matching)
library(tableone)
library(survey)

fstep <- read.csv("FirstSteps_data.csv")
summary(fstep)

vars <- c("age","race","parity","welfare","smoker","edu")
tabUnmatched <- CreateTableOne(vars = vars, strata = "firstep", data = fstep, test = FALSE)
print(tabUnmatched, smd = TRUE)
SMD.table <- ExtractSmd(tabUnmatched)

ests.mat<-matrix(0,nrow=2,ncol=4)
Z<-fstep$firstep
Y<-fstep$bwt
Y.bin<-(fstep$bwt<2500)
pi.hat<-fitted(glm(firstep~age+race+parity+welfare+smoker+edu+I(edu^2),data=fstep,family=binomial))

ps.lr.weight <- fstep$firstep + (1-fstep$firstep)*pi.hat/(1-pi.hat)
fs <- svydesign(ids=~0, data=fstep, weights=ps.lr.weight)
tabIPW <- svyCreateTableOne(vars = vars, strata = "firstep", data = fs, test = FALSE)
print(tabIPW, smd = TRUE)
round(cbind(SMD.table,ExtractSmd(tabIPW)),3)



w.hat<-pi.hat/(1-pi.hat)

## Matching
# SR
 sr.match <- Match(Y=Y,Tr=Z,X=pi.hat,M=1)
 ests.mat[1,1]<-sr.match$est
 ests.mat[2,1]<-sr.match$se
# DR
 dr.match <- Match(Y=Y,Tr=Z,Z=cbind(fstep$age,fstep$race,fstep$parity,fstep$welfare,fstep$smoker,fstep$edu+I(fstep$edu^2)),X=pi.hat,M=1)
 ests.mat[1,3]<-dr.match $est
 ests.mat[2,3]<-dr.match$se

## Weighting
# SR
 ests.mat[1,2]<-sum((Z-(1-Z)*w.hat)*Y)/sum(Z)              #IPW
# DR
 dat0 <- fstep; dat0$firstep<-0
 mu0 <- predict(lm(bwt~firstep+age+race+parity+welfare+smoker+edu+I(edu^2),data=fstep),dat0)
 ests.mat[1,4]<-sum((Z-(1-Z)*w.hat)*(Y-mu0))/sum(Z)  #AIPW
# Bootstrap
B <- 1000
boot.ests <- matrix(0,nrow=B,ncol=2)
n<-length(fstep[,1])
for(b in 1:B) {
if(b %in% seq(100,B,100)) print(b)
  id <- sample(1:n,replace=T)
  newdat <- fstep[id,]
  newpi.hat<-fitted(glm(firstep~age+race+parity+welfare+smoker+edu+I(edu^2),data=newdat,family=binomial))
  new.w<-newpi.hat/(1-newpi.hat)
  boot.ests[b,1]<-sum((newdat$firstep-(1-newdat$firstep)*new.w)*newdat$bwt)/sum(newdat$firstep)

  newdat0 <- newdat; newdat0$firstep<-0
  newmu0 <- predict(lm(bwt~firstep+age+race+parity+welfare+smoker+edu+I(edu^2),data=newdat),newdat0)
  boot.ests[b,2]<-sum((newdat$firstep-(1-newdat$firstep)*new.w)*(newdat$bwt-newmu0))/sum(newdat$firstep)
}
ests.mat[2,2]<-sd(boot.ests[,1])
ests.mat[2,4]<-sd(boot.ests[,2])

ests.mat
cbind(ests.mat[1,],ests.mat[1,]-1.96*ests.mat[2,],ests.mat[1,]+1.96*ests.mat[2,])


## Redo with smoking as exposure


vars <- c("age","race","parity","welfare","firstep","edu")
tabUnmatched <- CreateTableOne(vars = vars, strata = "smoker", data = fstep, test = FALSE)
print(tabUnmatched, smd = TRUE)
SMD.table <- ExtractSmd(tabUnmatched)

ests.mat<-matrix(0,nrow=2,ncol=4)
Z<-fstep$smoker
Y<-fstep$bwt
Y.bin<-(fstep$bwt<2500)
pi.hat<-fitted(glm((smoker=='Y')~firstep+age+race+parity+welfare+edu+I(edu^2),data=fstep,family=binomial))

fstep$smk <- as.numeric(fstep$smoker=="Y")
ps.lr.weight <- fstep$smk + (1-fstep$smk)*pi.hat/(1-pi.hat)
fs <- svydesign(ids=~0, data=fstep, weights=ps.lr.weight)
tabIPW <- svyCreateTableOne(vars = vars, strata = "smk", data = fs, test = FALSE)
print(tabIPW, smd = TRUE)
round(cbind(SMD.table,ExtractSmd(tabIPW)),3)


w.hat<-pi.hat/(1-pi.hat)

## Matching
# SR
 sr.match <- Match(Y=Y,Tr=fstep$smk,X=pi.hat,M=1)
 ests.mat[1,1]<-sr.match$est
 ests.mat[2,1]<-sr.match$se
# DR
 dr.match <- Match(Y=Y,Tr=fstep$smk,Z=cbind(fstep$firstep,fstep$age,fstep$race,fstep$parity,fstep$welfare,fstep$edu+I(fstep$edu^2)),X=pi.hat,M=1)
 ests.mat[1,3]<-dr.match$est
 ests.mat[2,3]<-dr.match$se

## Weighting
# SR
 Z<-fstep$smk
 ests.mat[1,2]<-sum((Z-(1-Z)*w.hat)*Y)/sum(Z)              #IPW
# DR
 dat0 <- fstep; dat0$firstep<-0
 mu0 <- predict(lm(bwt~firstep+age+race+parity+welfare+smoker+edu+I(edu^2),data=fstep),dat0)
 ests.mat[1,4]<-sum((Z-(1-Z)*w.hat)*(Y-mu0))/sum(Z)  #AIPW
# Bootstrap
B <- 1000
boot.ests <- matrix(0,nrow=B,ncol=2)
n<-length(fstep[,1])
for(b in 1:B) {
if(b %in% seq(100,B,100)) print(b)
  id <- sample(1:n,replace=T)
  newdat <- fstep[id,]
  newpi.hat<-fitted(glm(smk~age+race+parity+welfare+smoker+edu+I(edu^2),data=newdat,family=binomial))
  new.w<-newpi.hat/(1-newpi.hat)
  boot.ests[b,1]<-sum((newdat$firstep-(1-newdat$firstep)*new.w)*newdat$bwt)/sum(newdat$firstep)

  newdat0 <- newdat; newdat0$firstep<-0
  newmu0 <- predict(lm(bwt~firstep+age+race+parity+welfare+smoker+edu+I(edu^2),data=newdat),newdat0)
  boot.ests[b,2]<-sum((newdat$firstep-(1-newdat$firstep)*new.w)*(newdat$bwt-newmu0))/sum(newdat$firstep)
}
ests.mat[2,2]<-sd(boot.ests[,1])
ests.mat[2,4]<-sd(boot.ests[,2])

ests.mat
cbind(ests.mat[1,],ests.mat[1,]-1.96*ests.mat[2,],ests.mat[1,]+1.96*ests.mat[2,])

## Message is that there is no effect with the DR AIPTW, big negative effect with the others
## But too few smokers to really get this right.