# 
## G-estimation for general K-interval problem with linear blip
## 	     **VARIABLE EFFECTIVE SAMPLE SIZE ALGORITHM**
## 	 Uses my recursive method of calculating the variance
## 	          (Completely randomized treatment) 
## 	     Written by Erica E M Moodie, November, 2005 		 
#  

source("gestCov1int_txRand.R")

d.opt.k <- function(psi.hat.k,stat1,tx0=0) {
   if(length(as.vector(psi.hat.k))==2) return(as.numeric(cbind(1,stat1) %*% psi.hat.k > 0))
   return(as.numeric(cbind(1,stat1,tx0,stat1*tx0) %*% psi.hat.k > 0))
}

VESS_rand.fcn <- function(Y,status,action,no.stat.vbls=1,plot.H=F,return.U=F,true.psi,z.val = 1.96) {
   Y.all <- Y
   status.all <- status
   action.all <- action
   n <- length(Y.all)
   use.who <- 1:n
   K <- ifelse(is.vector(action),1,dim(action)[2])
   if(length(no.stat.vbls)==1) n.stat.vars <- rep(no.stat.vbls,K)
   if(plot.H) par(mfrow=c(2,round(K/2)))

   if(K==1) {
      ret.u <- return.U
      p <- true.psi
      return(g_estCov_1int.fcn(Y,status,action,return.U=ret.u,true.psi=p))
   }
   psi.hat <- NULL
   var.psi <- NULL
   U.hat <- NULL
   var.U <- var.U.hat <- NULL

   for(k in 1:(K-1)) {
      n.j <- length(use.who)
      Y.dash <- Y.all[use.who]
      n.stat <- n.stat.vars[K-k+1]
      n.stat.sum <- sum(n.stat.vars[(K-k+1):K])
      L <- status[use.who,(K-n.stat.sum+1):(K-n.stat.sum+n.stat)]
      if(!is.matrix(L)) L <- matrix(L,ncol=1)
      A <- action[use.who,(K-k):(K-k+1)]

      p.A.given.L <- mean(A[,2])
      aa <- A[,2] - p.A.given.L

      B <- mean(Y.dash*aa) 
   	for(i in 1:n.stat) B <- rbind(B,mean(Y.dash*L[,i]*aa))
	B <- rbind(B,mean(Y.dash*A[,1]*aa))
   	for(i in 1:n.stat) B <- rbind(B,mean(Y.dash*A[,1]*L[,i]*aa))

   	X <- matrix(nrow=2*(n.stat+1),ncol=2*(n.stat+1))
   	X[1,1] <- mean(A[,2]*aa)
   	X[n.stat+2,1] <- mean(A[,2]*A[,1]*aa)
   	X[n.stat+2,n.stat+2] <- mean(A[,2]*A[,1]^2*aa)
   	for(i in 1:n.stat) {
   	   X[i+1,1] <- X[1,i+1] <- mean(A[,2]*L[,i]*aa)
   	   X[n.stat+2+i,1] <- X[n.stat+2,i+1] <- mean(A[,2]*A[,1]*L[,i]*aa)
   	   X[n.stat+2+i,n.stat+2] <- X[n.stat+2,n.stat+2+i] <- mean(A[,2]*A[,1]^2*L[,i]*aa)
   	   for(j in i:n.stat) {
		X[i+1,j+1] <- X[j+1,i+1] <- mean(A[,2]*L[,i]*L[,j]*aa)
		X[n.stat+2+i,j+1] <- X[n.stat+2+j,i+1] <- mean(A[,2]*A[,1]*L[,i]*L[,j]*aa)
		X[n.stat+2+i,n.stat+j+2] <- X[n.stat+2+j,n.stat+2+i] <- mean(A[,2]*A[,1]^2*L[,i]*L[,j]*aa)
	   }
   	}
	X[1:(n.stat+1),(n.stat+2):(2*(n.stat+1))] <- X[(n.stat+2):(2*(n.stat+1)),1:(n.stat+1)]

      W <- cbind(1,L,A[,1],L*A[,1])
      W2.inv <- solve(t(W)%*%W/n.j)

      C1.temp <- apply(matrix(rep(cbind(mean(Y),mean(Y*L),mean(Y*A[,1]),mean(Y*L*A[,1])) 
         %*% W2.inv,each=n.j),ncol=4) * cbind(1,L,A[,1],L*A[,1]),1,sum)
      C1 <- rbind(mean(C1.temp*aa),mean(C1.temp*L*aa),mean(C1.temp*A[,1]*aa),mean(C1.temp*L*A[,1]*aa))

      C2.temp <- W2.inv %*% t(A[,2]*W) %*% W/n.j
      C2 <- matrix(rep(0,16),ncol=4)

     if(n.j < 500) {
         for(i in 1:n.j) C2 <- C2 + rbind(aa[i],(L*aa)[i],(A[,1]*aa)[i],(L*A[,1]*aa)[i]) %*% 
            cbind(1,L[i],A[i,1],L[i]*A[i,1]) %*% C2.temp
         C2 <- C2/n.j
      } else {
         Mat1<-rbind(aa,L[,1]*aa,A[,1]*aa,L[,1]*A[,1]*aa)
   	   Mat2<-cbind(L[,1]*0+1,L[,1],A[,1],L[,1]*A[,1])
	   FMat<-matrix(0,4,4)
	   nblocks<-n.j/500
	   for(i in 1:nblocks){
	      i1<-(i-1)*500+1
	      i2<-i*500
	      sub1<-Mat1[,i1:i2]
	      sub2<-Mat2[i1:i2,]
	      FMat<-FMat+sub1 %*% (sub2 %*% C2.temp)
	   }
         C2 <- FMat/n.j
      }

      psi.temp <- solve(X-C2)%*%(B-C1)
      psi.hat <- rbind(psi.hat,as.vector(psi.temp))
      
      S.part <- A[,2]/p.A.given.L - (1-A[,2])/(1-p.A.given.L)  
      d.alpha.S <- -mean(A[,2]/p.A.given.L^2 + (1-A[,2])/(1-p.A.given.L)^2) 

      Hj <- Y.dash - A[,2]*(W %*% psi.temp)
      Hj.ctr <- Hj - W %*% (W2.inv %*% rbind(mean(Hj),mean(Hj*L),mean(Hj*A[,1]),mean(Hj*L*A[,1]))) 

      d.alpha.U <- -t(Hj.ctr)%*%W/n.j
      d.varsigma.U <- -t(W)%*%(W*aa)/n.j 

      if(k==1) {
         adj.U.psi.OLD <- list(matrix(0,ncol=4,nrow=n.j))
         U.OLD <- list(matrix(0,ncol=4,nrow=n.j))
         SCORE.adj.U.psi.OLD <- list(matrix(0,ncol=4,nrow=n.j))
         SCORE.U.OLD <- list(matrix(0,ncol=4,nrow=n.j))
      }
      adj.temp <- matrix(0,ncol=n.j,nrow=4)
      for(i in 1:length(adj.U.psi.OLD)) {
         Uj.stuff <- matrix(unlist(adj.U.psi.OLD[i]),ncol=4)
         Ujminus1.stuff <- matrix(unlist(U.OLD[i]),nrow=4)
         adj.temp <- adj.temp + (t(aa*W) %*% Uj.stuff) %*% Ujminus1.stuff/n.j
      }

      U.adj <- t(cbind(Hj.ctr*aa,Hj.ctr*aa*L,Hj.ctr*aa*A[,1],Hj.ctr*aa*L*A[,1])) - 
          t(d.alpha.U/d.alpha.S) %*% t(as.vector(S.part)) + 
          d.varsigma.U %*% W2.inv %*% t(cbind(Hj.ctr,Hj.ctr*L,Hj.ctr*A[,1],Hj.ctr*L*A[,1])) - 
          adj.temp 

      d.psi.U <- -t(W)%*%(W*aa*A[,2])/n.j - 
          d.varsigma.U %*% W2.inv %*% t(cbind(1,L,A[,1],L*A[,1]))%*%(cbind(1,L,A[,1],L*A[,1])*A[,2])/n.j  #Chain rule
      d.psi.U.adj <- d.psi.U - 
          t(cbind(1,L,A[,1],L*A[,1]))%*%(cbind(1,L,A[,1],L*A[,1])*A[,2])/n * (1/d.alpha.S) * mean(S.part) + 
          d.varsigma.U %*% W2.inv %*% t(cbind(1,L,A[,1],L*A[,1])) %*%(cbind(1,L,A[,1],L*A[,1])*A[,2])/n.j

      var.mat <- (solve(d.psi.U.adj) %*% U.adj %*% t(U.adj) %*% t(solve(d.psi.U.adj)))/(n.j^2)
      var.psi <- rbind(var.psi,diag(var.mat))

      if(plot.H & n.stat==1) plot(L,Hj,ylab="H_j(psi)",xlab="L_j",main=paste("j = ",K-k+1,sep=""))

      ## Testing U with TRUE psi's  
      if(return.U) {
         if(k==1) Y.dash.true <- Y.dash
         m <- length(true.psi)
         Hj.true <- Y.dash.true - A[,2]*(W %*% matrix(true.psi[(m-k*4+1):(m-(k-1)*4)],nrow=4))
         Hj.ctr <- Hj.true - W %*% (W2.inv %*% rbind(mean(Hj),mean(Hj*L),mean(Hj*A[,1]),mean(Hj*L*A[,1]))) 
         U.dagger <- t(cbind(Hj.ctr*aa,Hj.ctr*aa*L,Hj.ctr*aa*A[,1],Hj.ctr*aa*L*A[,1])) 
         U.hat <- rbind(U.hat,apply(U.dagger,1,mean))

         d.a.u <- -t(Hj.ctr)%*%W/n.j
         uu <- U.dagger - t(d.a.u/d.alpha.S) %*% t(as.vector(S.part)) + 
           d.varsigma.U %*% W2.inv %*% t(cbind(Hj.ctr,Hj.ctr*L,Hj.ctr*A[,1],Hj.ctr*L*A[,1])) 
         var.U.hat <- rbind(var.U.hat,diag(uu %*% t(uu)/(n.j^2)))

         true.opt.minus.A <- d.opt.k(as.vector(true.psi[(m-k*4+1):(m-(k-1)*4)]),L,A[,1]) - A[,2]
         Y.dash.true <- Y.dash.true + true.opt.minus.A*(W %*% true.psi[(m-k*4+1):(m-(k-1)*4)])
         d.p.U <- -t(W)%*%(W*aa*A[,2])/n - d.varsigma.U %*% W2.inv %*% t(W)%*%(W*A[,2])/n  #Chain rule
         d.p.U.adj <- d.p.U - t(W)%*%(W*A[,2])/n * (1/d.alpha.S) * mean(S.part) + 
            d.varsigma.U %*% W2.inv %*% t(W) %*%(W*A[,2])/n.j
         SCORE.adj.U.psi.OLD[k] <- list(true.opt.minus.A*W) 
         SCORE.U.OLD[k] <- list(solve(d.p.U.adj) %*% U.dagger) 
      }

      L <- status[,(K-n.stat.sum+1):(K-n.stat.sum+n.stat)]
      if(!is.matrix(L)) L <- matrix(L,ncol=1)
      A <- action[,(K-k):(K-k+1)]
      W <- cbind(1,L,A[,1],L*A[,1])
      d.opt.minus.A <- d.opt.k(as.vector(psi.temp),L,A[,1]) - A[,2]
      Y.all <- Y.all + d.opt.minus.A*(W %*% psi.temp)

      adj.U.psi.OLD[k] <- list(d.opt.minus.A*W) 
      U.OLD[k] <- list(solve(d.psi.U.adj) %*% U.adj) 

      ##############################################
      ##  Check if CI for each person contains 0  ##
      ##############################################
      individ.se <- as.vector(sqrt(diag(var.mat) %*% t(cbind(1,L^2,A[,1]^2,L^2*A[,1]^2)) + 
           2*(var.mat[1,2:4] %*% t(cbind(L,A[,1],L*A[,1])) + 
              var.mat[2,3:4] %*% t(cbind(L*A[,1],L^2*A[,1])) + 
              var.mat[3,4] * as.vector(L*A[,1]^2))))
      individ.cov0 <- (W %*% psi.temp - z.val*individ.se < 0) & (W %*% psi.temp + z.val*individ.se > 0)
      use.who <- (1:n)[!individ.cov0]
#      print(sum(individ.cov0)) 
   }
   
   n.j <- length(use.who)
   Y.dash <- Y.all[use.who]
   n.stat <- n.stat.vars[1]
   L <- status[use.who,1:n.stat]
   if(!is.matrix(L)) L <- matrix(L,ncol=1)
   A <- action[use.who,1]

   p.A.given.L <- mean(A)
   aa <- A - p.A.given.L

   B <- mean(Y.dash*aa)
   B <- rbind(B,mean(Y.dash*L*aa))
   X <- matrix(nrow=2,ncol=2)   
   X[1,1] <- mean(A*aa)
   X[2,1] <- X[1,2] <- mean(A*L*aa)
   X[2,2] <- mean(A*L^2*aa)

   W <- cbind(1,L)
   W2.inv <- solve(matrix(c(1,mean(L),mean(L),mean(L^2)),ncol=2))

   C1.temp <- apply(matrix(rep(cbind(mean(Y.dash),mean(Y.dash*L)) %*% 
      solve(matrix(c(1,mean(L),mean(L),mean(L^2)),ncol=2)),each=n.j),ncol=2) * 
      cbind(1,L),1,sum)
   C1 <- rbind(mean(C1.temp*aa),mean(C1.temp*L*aa))

   C2.temp <- W2.inv %*% matrix(c(mean(A),mean(A*L),mean(A*L),mean(A*L^2)),ncol=2) 
   C2 <- matrix(rep(0,4),ncol=2)
   if(n.j<500) {
      for(i in 1:n.j) C2 <- C2 + rbind(aa[i],(L*aa)[i]) %*% cbind(1,L[i]) %*% C2.temp
      C2 <- C2/n.j
   } else {
      Mat1<-rbind(aa,L[,1]*aa)
      Mat2<-cbind(0*L[,1]+1,L[,1])
      FMat<-matrix(0,2,2)
      nblocks<-n.j/500
      for(i in 1:nblocks){
         i1<-(i-1)*500+1
         i2<-i*500
         sub1<-Mat1[,i1:i2]
         sub2<-Mat2[i1:i2,]
         FMat<-FMat+sub1 %*% (sub2 %*% C2.temp)
      }
      C2<-FMat/n.j
   }

   psi.temp <- solve(X-C2)%*%(B-C1)

   psi.hat <- rbind(psi.hat,c(as.vector(psi.temp),NA,NA))
   dimnames(psi.hat) <- list(1:K,c("psi0","psi1","psi2","psi3"))

   S.part <- A/p.A.given.L - (1-A)/(1-p.A.given.L)
   d.alpha.S <- -mean(A/p.A.given.L^2 + (1-A)/(1-p.A.given.L)^2) 

   Hj <- Y.dash - (A*W %*% psi.temp)
   Hj.ctr <- Hj - W %*% (W2.inv %*% rbind(mean(Hj),mean(Hj*L))) 

   d.alpha.U <- -t(Hj.ctr)%*%W/n.j
   d.varsigma.U <- -t(W)%*%(W*aa)/n.j

   adj.temp <- matrix(0,ncol=n.j,nrow=2)
   for(i in 1:length(adj.U.psi.OLD)) {
      Uj.stuff <- matrix(unlist(adj.U.psi.OLD[i]),ncol=4)
      Ujminus1.stuff <- matrix(unlist(U.OLD[i]),nrow=4)
      adj.temp <- adj.temp + (t(aa*W) %*% Uj.stuff[use.who,]) %*% Ujminus1.stuff[,use.who]/n.j
   }

   U.adj <- t(cbind(Hj.ctr*aa,Hj.ctr*aa*L)) - t(d.alpha.U/d.alpha.S) %*% t(as.vector(S.part)) + 
       d.varsigma.U %*% W2.inv %*% t(cbind(Hj.ctr,Hj.ctr*L)) - adj.temp 

   d.psi.U <- -t(W)%*%(W*aa*A)/n.j - 
       d.varsigma.U %*% W2.inv %*% t(W)%*%(W*A)/n.j  # Chain rule
   d.psi.U.adj <- d.psi.U - t(W)%*%(W*A)/n * (1/d.alpha.S) * mean(S.part) + 
       d.varsigma.U %*% W2.inv %*% t(W)%*%(W*A)/n.j

   var.psi <- rbind(var.psi,c(diag((solve(d.psi.U.adj) %*% U.adj %*% t(U.adj) %*% t(solve(d.psi.U.adj)))/(n.j^2)),NA,NA))

   if(plot.H & n.stat==1) plot(L,Y.dash-A*cbind(1,L)%*% psi.temp[1:2],ylab="H_j(psi)",xlab="L_j",main="j = 1")

   ## Testing U with TRUE psi's  
   if(return.U) {
      m <- length(true.psi)
      Hj.true <- Y.dash.true - A*(cbind(1,L) %*% true.psi[1:2])
      Hj.ctr <- Hj.true - W %*% (W2.inv %*% apply(as.vector(Hj) * W,2,mean))
      U.dagger <- t(cbind(Hj.ctr*aa,Hj.ctr*aa*L)) 
      U.hat <- rbind(U.hat,c(apply(U.dagger,1,mean),NA,NA))

      d.a.u <- -t(Hj.ctr)%*%cbind(1,L)/n
      uu <- U.dagger - t(d.a.u/d.alpha.S) %*% t(as.vector(S.part)) + 
        d.varsigma.U %*% W2.inv %*% t(cbind(Hj.ctr,Hj.ctr*L)) 
      var.U.hat <- rbind(var.U.hat,c(diag(uu %*% t(uu)/(n.j^2)),NA,NA))

      return(list(psi.hat=psi.hat,psi.var=var.psi,U.hat=U.hat,U.var.hat=var.U.hat))
   }
   return(list(psi.hat=psi.hat,psi.var=var.psi))

}