# 
## G-estimation for general K-interval problem with linear blip
## 	Uses my recursive method of calculating the variance
## 	      (Sequentially randomized treatment) 
## 	    Written by Erica E M Moodie, 16 June, 2005 		 
#  

#rm(list=ls())
#Y<-scan("y_3ints.out")
#action<-read.table("a_3ints.out")
#status<-read.table("l_3ints.out")

source("gestCov1int_txSeq.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))
}

expit <- function(x) { exp(x)/(1+exp(x)) }

g_estCov_SeqRand.fcn <- function(Y,status,action,no.stat.vbls=1,plot.H=F,return.U=F,true.psi) {
   n <- length(Y)
   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_SeqRand.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
   Y.dash <- Y

   for(k in 1:(K-1)) {
      n.stat <- n.stat.vars[K-k+1]
      n.stat.sum <- sum(n.stat.vars[(K-k+1):K])
      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)]

      p.coef <- coef(glm(A[,2]~L+A[,1],family=binomial)) 
      p.A.given.L <- expit(cbind(1,L,A[,1])%*%p.coef)  
      p.temp <- exp(cbind(1,L,A[,1])%*%p.coef)  
      aa <- as.vector(A[,2] - p.A.given.L)
#print(p.coef)

      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)

      C1.temp <- apply(matrix(rep(cbind(mean(Y),mean(Y*L),mean(Y*A[,1]),mean(Y*L*A[,1])) %*% W2.inv,each=n),ncol=4) * cbind(rep(1,n),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
      C2 <- matrix(rep(0,16),ncol=4)

     if(n < 500) {
         for(i in 1:n) 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
      } 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/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
      }

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

      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 <- diag(-c(mean(Hj.ctr*p.A.given.L/(1+p.temp)),mean(Hj.ctr*p.A.given.L/(1+p.temp)*L^2),mean(Hj.ctr*p.A.given.L/(1+p.temp)*A[,1]^2)))
      d.alpha.U[1,2] <- d.alpha.U[2,1] <- -mean(Hj.ctr*p.A.given.L/(1+p.temp)*L)
      d.alpha.U[1,3] <- d.alpha.U[3,1] <- -mean(Hj.ctr*p.A.given.L/(1+p.temp)*A[,1])
      d.alpha.U[3,2] <- d.alpha.U[2,3] <- -mean(Hj.ctr*p.A.given.L/(1+p.temp)*L*A[,1])
      d.alpha.U <- rbind(d.alpha.U,-c(mean(Hj.ctr*p.A.given.L/(1+p.temp)*L*A[,1]),mean(Hj.ctr*p.A.given.L/(1+p.temp)*L^2*A[,1]),mean(Hj.ctr*p.A.given.L/(1+p.temp)*L*A[,1]^2)))
      d.varsigma.U <- -t(W)%*%(W*aa)/n 

      temp.a <- as.vector(solve(d.alpha.S) %*% apply(S.part,2,mean))
      x.term.for.alpha <- diag(c(mean(p.A.given.L/(1+p.temp)*(temp.a[1]*A[,2] + temp.a[2]*A[,2]*L + temp.a[3]*A[,2]*A[,1])),
        mean(p.A.given.L/(1+p.temp)*(temp.a[1]*A[,2]*L^2 + temp.a[2]*A[,2]*L^3 + temp.a[3]*A[,2]*A[,1]*L^2)),
        mean(p.A.given.L/(1+p.temp)*(temp.a[1]*A[,2]*A[,1]^2 + temp.a[2]*A[,2]*L*A[,1]^2 + temp.a[3]*A[,2]*A[,1]^3)),
        mean(p.A.given.L/(1+p.temp)*(temp.a[1]*A[,2]*L^2*A[,1]^2 + temp.a[2]*A[,2]*L^3*A[,1]^2 + temp.a[3]*A[,2]*L^2*A[,1]^3))))
      x.term.for.alpha[1,2] <- x.term.for.alpha[2,1] <- mean(p.A.given.L/(1+p.temp)*(temp.a[1]*A[,2]*L + 
        temp.a[2]*A[,2]*L^2 + temp.a[3]*A[,2]*L*A[,1]))
      x.term.for.alpha[1,3] <- x.term.for.alpha[3,1] <- mean(p.A.given.L/(1+p.temp)*(temp.a[1]*A[,2]*A[,1] + 
        temp.a[2]*A[,2]*L*A[,1] + temp.a[3]*A[,2]*A[,1]^2))
      x.term.for.alpha[1,4] <- x.term.for.alpha[4,1] <- x.term.for.alpha[2,3] <- x.term.for.alpha[3,2] <- 
        mean(p.A.given.L/(1+p.temp)*(temp.a[1]*A[,2]*L*A[,1] + temp.a[2]*A[,2]*L^2*A[,1] + temp.a[3]*A[,2]*L*A[,1]^2))
      x.term.for.alpha[2,4] <- x.term.for.alpha[4,2] <- mean(p.A.given.L/(1+p.temp)*(temp.a[1]*A[,2]*L^2*A[,1] + 
        temp.a[2]*A[,2]*L^3*A[,1] + temp.a[3]*A[,2]*L^2*A[,1]^2))
      x.term.for.alpha[4,3] <- x.term.for.alpha[3,4] <- mean(p.A.given.L/(1+p.temp)*(temp.a[1]*A[,2]*L*A[,1]^2 + 
        temp.a[2]*A[,2]*L^2*A[,1]^2 + temp.a[3]*A[,2]*L*A[,1]^3))

      if(k==1) {
         adj.U.psi.OLD <- list(matrix(0,ncol=4,nrow=n))
         U.OLD <- list(matrix(0,ncol=4,nrow=n))
         SCORE.adj.U.psi.OLD <- list(matrix(0,ncol=4,nrow=n))
         SCORE.U.OLD <- list(matrix(0,ncol=4,nrow=n))
      }
      adj.temp <- matrix(0,ncol=n,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
      }

      U.adj <- t(cbind(Hj.ctr*aa,Hj.ctr*aa*L,Hj.ctr*aa*A[,1],Hj.ctr*aa*L*A[,1])) - 
          d.alpha.U %*% solve(d.alpha.S) %*% t(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 - 
          d.varsigma.U %*% W2.inv %*% t(W)%*%(W*A[,2])/n  #Chain rule
      d.psi.U.adj <- d.psi.U - x.term.for.alpha + 
          d.varsigma.U %*% W2.inv %*% t(W) %*%(W*A[,2])/n

      var.mat <- (solve(d.psi.U.adj) %*% U.adj %*% t(U.adj) %*% t(solve(d.psi.U.adj)))/(n^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 <- diag(-c(mean(Hj.ctr*p.A.given.L/(1+p.temp)),mean(Hj.ctr*p.A.given.L/(1+p.temp)*L^2),mean(Hj.ctr*p.A.given.L/(1+p.temp)*A[,1]^2)))
         d.a.u[1,2] <- d.a.u[2,1] <- -mean(Hj.ctr*p.A.given.L/(1+p.temp)*L)
         d.a.u[1,3] <- d.a.u[3,1] <- -mean(Hj.ctr*p.A.given.L/(1+p.temp)*A[,1])
         d.a.u[3,2] <- d.a.u[2,3] <- -mean(Hj.ctr*p.A.given.L/(1+p.temp)*L*A[,1])
         d.a.u <- rbind(d.a.u,-c(mean(Hj.ctr*p.A.given.L/(1+p.temp)*L*A[,1]),mean(Hj.ctr*p.A.given.L/(1+p.temp)*L^2*A[,1]),mean(Hj.ctr*p.A.given.L/(1+p.temp)*L*A[,1]^2)))
         uu <- U.dagger - d.a.u %*% solve(d.alpha.S) %*% t(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^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 - x.term.for.alpha + d.varsigma.U %*% W2.inv %*% t(W) %*%(W*A[,2])/n
         SCORE.adj.U.psi.OLD[k] <- list(true.opt.minus.A*W) 
         SCORE.U.OLD[k] <- list(solve(d.p.U.adj) %*% U.dagger) 
      }

      d.opt.minus.A <- d.opt.k(as.vector(psi.temp),L,A[,1]) - A[,2]
      Y.dash <- Y.dash + 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) 
   }
   
   n.stat <- n.stat.vars[1]
   L <- status[,1:n.stat]
   if(!is.matrix(L)) L <- matrix(L,ncol=1)
   A <- action[,1]

   p.coef <- coef(glm(A~L,family=binomial)) 
   p.A.given.L <- expit(cbind(1,L)%*%p.coef)  
   p.temp <- exp(cbind(1,L)%*%p.coef)  
   aa <- as.vector(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),ncol=2) * cbind(rep(1,n),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<500) {
      for(i in 1:n) C2 <- C2 + rbind(aa[i],(L*aa)[i]) %*% cbind(1,L[i]) %*% C2.temp
      C2 <- C2/n
   } else {
      Mat1<-rbind(aa,L[,1]*aa)
      Mat2<-cbind(0*L[,1]+1,L[,1])
      FMat<-matrix(0,2,2)
      nblocks<-n/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
   }

   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 <- cbind(aa,aa*L)
   d.alpha.S <- diag(c(mean(-p.A.given.L/(1+p.temp)),mean(-p.A.given.L/(1+p.temp)*L^2)))
   d.alpha.S[1,2] <- d.alpha.S[2,1] <- mean(-p.A.given.L/(1+p.temp)*L)

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

   d.alpha.U <- diag(-c(mean(Hj.ctr*p.A.given.L/(1+p.temp)),mean(Hj.ctr*p.A.given.L/(1+p.temp)*L^2)))
   d.alpha.U[1,2] <- d.alpha.U[2,1] <- -mean(Hj.ctr*p.A.given.L/(1+p.temp)*L)
   d.varsigma.U <- -t(W)%*%(W*aa)/n

   adj.temp <- matrix(0,ncol=n,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) %*% Ujminus1.stuff/n
   }

   U.adj <- t(cbind(Hj.ctr*aa,Hj.ctr*aa*L)) - t(d.alpha.U) %*% solve(d.alpha.S) %*% t(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 - 
       d.varsigma.U %*% W2.inv %*% t(W)%*%(W*A)/n  # Chain rule
   temp.a <- as.vector(solve(d.alpha.S) %*% apply(S.part,2,mean))
   x.term.for.alpha <- diag(c(mean(temp.a[1]*A + temp.a[2]*A*L),mean(temp.a[1]*A*L^2 + temp.a[2]*A*L^3)))
   x.term.for.alpha[1,2] <- x.term.for.alpha[2,1] <- mean(temp.a[1]*A*L + temp.a[2]*A*L^2)
   d.psi.U.adj <- d.psi.U - x.term.for.alpha + d.varsigma.U %*% W2.inv %*% t(W)%*%(W*A)/n

   var.psi <- rbind(var.psi,c(diag((solve(d.psi.U.adj) %*% U.adj %*% t(U.adj) %*% t(solve(d.psi.U.adj)))/(n^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*(W %*% 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 <- diag(-c(mean(Hj.ctr*p.A.given.L/(1+p.temp)),mean(Hj.ctr*p.A.given.L/(1+p.temp)*L^2)))
      d.a.u[1,2] <- d.a.u[2,1] <- -mean(Hj.ctr*p.A.given.L/(1+p.temp)*L)
      uu <- U.dagger - t(d.a.u) %*% solve(d.alpha.S) %*% t(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^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))

}