# 
## G-estimation for general K-interval problem with linear blip
## 	Uses my recursive method of calculating the variance
## 	         (Completely 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")

library(MASS)
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))
}

g_estCov.fcn <- function(Y,status,action,no.stat.vbls=1,plot.H=F,return.U=F,true.psi=0,true.a,plim.vs) {
   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.fcn(Y,status,action,return.U=ret.u,true.psi=p))
   }
   psi.hat <- NULL
   var.psi <- NULL

   U.true <- U.hat <- NULL
   var.U.true <- var.U.hat <- var.U.hat.notrue <- NULL

   OLSvs.true <- NULL
   varOLSvs.true <- varOLSvs.hat <- NULL

   Salpha.true <- NULL
   varSalpha.true <- varSalpha.hat <- NULL

   Y.dash <- Y

#k<-1
   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.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 <- ginv(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 <- ginv(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]))) 

      varsigma.hat <- W2.inv %*% t(W) %*% Hj/n

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

      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])) - 
          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 - 
          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  #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

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

      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]))) 
         varsigma.plim <- W2.inv %*% t(W) %*% Hj.true/n
         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
         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^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
         SCORE.adj.U.psi.OLD[k] <- list(true.opt.minus.A*W) 
         SCORE.U.OLD[k] <- list(ginv(d.p.U.adj) %*% U.dagger) 

         #
         ## Everything "true" plugged in...
         OLSvs.true <- (t(W) %*% Hj.true - t(W) %*% W %*% plim.vs)/n ## apply(W*as.vector(Hj.true) - W*as.vector(W %*% plim.vs),2,mean)
         varOLSvs.true <- t(W*as.vector(Hj.true) - W*as.vector(W %*% plim.vs)) %*% (W*as.vector(Hj.true) - W*as.vector(W %*% plim.vs))/(n^2)
         #OLSvs.hat <- (t(W) %*% Hj - t(W) %*% W %*% varsigma.hat)/n  
         varOLSvs.hat <- t(W*as.vector(Hj) - W*as.vector(W %*% varsigma.hat)) %*% (W*as.vector(Hj) - W*as.vector(W %*% varsigma.hat))/(n^2)

         Salpha.true <- mean(A[,2]/true.a[2] - (1-A[,2])/(1-true.a[2]))
         varSalpha.hat <- (t(S.part) %*% S.part)/n
         varSalpha.true <- (t(A[,2]/true.a[2] - (1-A[,2])/(1-true.a[2])) %*% (A[,2]/true.a[2] - (1-A[,2])/(1-true.a[2])))/n

         p.A.given.L <- true.a[2]
         aa <- A[,2] - p.A.given.L
         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) 
         d.v.U <- -t(W)%*%(W*aa)/n 
         Hj.ctr <- Hj.true - W %*% plim.vs
         U.dagger <- t(cbind(Hj.true*aa,Hj.true*aa*L,Hj.true*aa*A[,1],Hj.true*aa*L*A[,1])) 
         U.true <- apply(U.dagger,1,mean)
         uu <- U.dagger - t(d.a.u/d.alpha.S) %*% t(as.vector(S.part)) + 
           d.v.U %*% W2.inv %*% t(cbind(Hj.ctr,Hj.ctr*L,Hj.ctr*A[,1],Hj.ctr*L*A[,1])) 
         var.U.true <- diag(uu %*% t(uu)/(n^2))
		### For now, just look at the last interval
      }

      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(ginv(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.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 <- ginv(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)) %*% ginv(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 <- ginv(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
   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/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 - 
       d.varsigma.U %*% W2.inv %*% t(W)%*%(W*A)/n  # 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

   var.psi <- rbind(var.psi,c(diag((ginv(d.psi.U.adj) %*% U.adj %*% t(U.adj) %*% t(ginv(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*(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^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,U.hat=U.hat,U.var.hat=var.U.hat,U.true=U.true,var.U.true=var.U.true,
#                       var.U.hat.notrue=var.U.hat.notrue))
#      return(list(psi.hat=psi.hat,psi.var=var.psi,U.hat=U.hat,U.var.hat=var.U.hat,U.true=U.true,
#                       var.U.true=var.U.true,var.U.hat.notrue=var.U.hat.notrue,vs.true=varsigma.plim,vs.est=varsigma.hat))
      return(list(psi.hat=psi.hat,psi.var=var.psi,U.true=U.true,U.hat=U.hat,U.var.hat=var.U.hat,
                 var.U.true=var.U.true,var.U.hat.notrue=var.U.hat.notrue,
                 OLSvs.true=OLSvs.true,varOLSvs.true=varOLSvs.true,varOLSvs.hat=varOLSvs.hat,
                 Salpha.true=Salpha.true,varSalpha.true=varSalpha.true,varSalpha.hat=varSalpha.hat))
   }
   return(list(psi.hat=psi.hat,psi.var=var.psi))

}