#######################################################################
#  Regression estimates with propensity covariates
#######################################################################

samp  <- read.table("Diet0001.dat", row.names="ID",
   col.names=c("ID", "DISTR.1", "BLACK", "NBHISP", "GRADE", "SLFHLTH",
   "SLFWGHT", "WORKHARD", "GOODQUAL", "PHYSFIT", "PROUD",
   "LIKESLF", "ACCEPTED", "FEELLOVD", "DIET", "Y.0", "Y.1", "DISTR.2",
   "PI.TRUE"))


pihat <- glm( DIET ~ DISTR.1 + BLACK + NBHISP + GRADE + 
   SLFHLTH + SLFWGHT + WORKHARD + GOODQUAL + PHYSFIT + PROUD + 
   LIKESLF + ACCEPTED + FEELLOVD, family=binomial, data=samp)$fitted

# create nine strata and eight dummy variables
stratum <- cut(pihat, 
   breaks=quantile(pihat, c(0,.20,.40,.60,.70,.80,.85,.90,.95,1)),
   include.lowest=T, labels=1:9)
samp$D1 <- 1*(stratum==1)
samp$D2 <- 1*(stratum==2)
samp$D3 <- 1*(stratum==3)
samp$D4 <- 1*(stratum==4)
samp$D5 <- 1*(stratum==5)
samp$D6 <- 1*(stratum==6)
samp$D7 <- 1*(stratum==7)
samp$D8 <- 1*(stratum==8)

# regression estimation
t <- samp$DIET
y <- samp$DISTR.2
samp0 <- samp[ t==0, ]
samp1 <- samp[ t==1, ]

tmp0 <- lm( DISTR.2 ~ DISTR.1 + BLACK + NBHISP + GRADE + 
   SLFHLTH + SLFWGHT + WORKHARD + GOODQUAL + PHYSFIT + PROUD + 
   LIKESLF + ACCEPTED + FEELLOVD + D1 + D2 + D3 + D4 +
   D5 + D6 + D7 + D8, data=samp0)
yhat0 <- predict( tmp0, newdata=samp)
yimp0 <- y
yimp0[t==1] <- yhat0[t==1]

tmp1 <- lm( DISTR.2 ~ DISTR.1 + BLACK + NBHISP + GRADE + 
   SLFHLTH + SLFWGHT + WORKHARD + GOODQUAL + PHYSFIT + PROUD + 
   LIKESLF + ACCEPTED + FEELLOVD + D1 + D2 + D3 + D4 +
   D5 + D6 + D7 + D8, data=samp1)
yhat1 <- predict( tmp1, newdata=samp)
yimp1 <- y
yimp1[t==0] <- yhat1[t==0]

mu0 <- mean(yimp0)
mu1 <- mean(yimp1)

x <- cbind( const=1,
   samp$DISTR.1,
   samp$BLACK,
   samp$NBHISP,
   samp$GRADE,
   samp$SLFHLTH,
   samp$SLFWGHT,
   samp$WORKHARD,
   samp$GOODQUAL,
   samp$PHYSFIT,
   samp$PROUD,
   samp$LIKESLF,
   samp$ACCEPTED,
   samp$FEELLOVD,
   samp$D1,
   samp$D2,
   samp$D3,
   samp$D4,
   samp$D5,
   samp$D6,
   samp$D7,
   samp$D8) 

N <- nrow(x)
p <- ncol(x)
score <- matrix(NA, N, 2*p+2)
score[,1:p] <- (1-t) * (y - yhat0)*x
score[,(p+1):(2*p)] <- t * (y-yhat1)*x
score[,(2*p+1)] <- (1-t)*(y-mu0) + t*(yhat0-mu0)
score[,(2*p+2)] <- t*(y-mu1) + (1-t)*(yhat1-mu1)
a <- c( rep(0,2*p), -1, 1)
B <- t(score) %*% score
A <- matrix(0, 2*p+2, 2*p+2 )
A[1:p,1:p] <- t(x*(1-t)) %*% x
A[(p+1):(2*p),(p+1):(2*p)] <- t( t*x ) %*% x
A[(2*p+1),1:p] <- - apply( t*x, 2, sum)
A[(2*p+2),(p+1):(2*p)] <- - apply( (1-t)*x, 2, sum)
A[2*p+1,2*p+1] <- N
A[2*p+2,2*p+2] <- N
Ainv <- solve(A)
ACE <- mu1 - mu0
SE <- sqrt( t(a) %*% Ainv %*% B %*% t(Ainv) %*% a )

# > print(ACE)
# [1] -0.001383053
# > print(SE)
#            [,1]
# [1,] 0.01639706


# ACE for treated
mu1 <- sum(t*y)/sum(t)
mu0 <- sum( t*yhat0 ) / sum(t)
ACE.1 <- mu1 - mu0
score <- matrix(NA, N, p+2)
score[,1:p] <- (1-t) * (y - yhat0)*x
score[,(p+1)] <- t*(yhat0-mu0)
score[,(p+2)] <- t*(y-mu1)
B <- t(score) %*% score
A <- matrix(0, p+2, p+2 )
A[1:p,1:p] <- t(x*(1-t)) %*% x
A[(p+1),1:p] <- - apply( t*x, 2, sum)
A[p+1,p+1] <- sum(t)
A[p+2,p+2] <- sum(t)
Ainv <- solve(A)
a <- c( rep(0,p), -1, 1)
SE.1 <- sqrt( t(a) %*% Ainv %*% B %*% t(Ainv) %*% a )

# > print(ACE.1)
# [1] -0.01977586
# > print(SE.1)
#            [,1]
# [1,] 0.01428649