# 
## G-estimation for general K-interval problem with linear blip - variance from (3.14)
## 		Written by E Moodie, 14 March, 2005 		 
#  

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

source("1int_Cov_txRand.R")
source("covJR_txRand2.R")
source("Kinterval_xi.R")
source("temp.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))
}

jr.fcn <- function(Y,status,action,est.psi=NULL,return.U=F,true.psi,true.a) {
   n <- length(Y)
   K <- ifelse(is.vector(action),1,dim(action)[2])

   if(K==1) return(g.cov.1int(Y,status,action))
   if(K==2) return(jr_2ints.fcn(Y,status,action,est.psi,return.U,true.psi,true.a))
   if(K>3) {
      print("Can't do this")
      return(NULL)
   }

   if(is.null(est.psi)) {
      temp <- g.xi.Kint(Y,status,action)
      psi.hat <- c(temp[3,1:2],temp[2,],temp[1,])
   } else { psi.hat<- est.psi }

   Y.dash <- Y
   L <- status
   A <- action

   W1 <- cbind(1,L[,1],0,0,0,0,0,0,0,0)
   W2 <- cbind(0,0,1,L[,2],A[,1],L[,2]*A[,1],0,0,0,0)
   W3 <- cbind(0,0,0,0,0,0,1,L[,3],A[,2],L[,3]*A[,2])

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

   S.part <- cbind(A[,1]/p.A.given.L[,1] - (1-A[,1])/(1-p.A.given.L[,1]), 
      A[,2]/p.A.given.L[,2] - (1-A[,2])/(1-p.A.given.L[,2]),
      A[,3]/p.A.given.L[,3] - (1-A[,3])/(1-p.A.given.L[,3]))
   d.alpha.S <- -diag(c(mean(A[,1]/p.A.given.L[,1]^2 + (1-A[,1])/(1-p.A.given.L[,1])^2),
      mean(A[,2]/p.A.given.L[,2]^2 + (1-A[,2])/(1-p.A.given.L[,2])^2),
      mean(A[,3]/p.A.given.L[,3]^2 + (1-A[,3])/(1-p.A.given.L[,3])^2))) 

   W1.sq.inv <- W2.sq.inv <- W3.sq.inv <- W.sq.inv <- matrix(0,ncol=10,nrow=10)
   W1.sq.inv[1:2,1:2] <- solve(t(cbind(1,L[,1])) %*% cbind(1,L[,1])/n)
   W2.sq.inv[3:6,3:6] <- solve(t(cbind(1,L[,2],A[,1],L[,2]*A[,1]))%*%cbind(1,L[,2],A[,1],L[,2]*A[,1])/n)
   W3.sq.inv[7:10,7:10] <- solve(t(cbind(1,L[,3],A[,2],L[,3]*A[,2]))%*%cbind(1,L[,3],A[,2],L[,3]*A[,2])/n)
   W.sq.inv <- W1.sq.inv + W2.sq.inv + W3.sq.inv

   Hj3 <- as.vector(Y.dash - A[,3]*(W3 %*% psi.hat))
   Hj.ctr3 <- as.vector(Hj3 - W3 %*% (W3.sq.inv %*% apply(as.vector(Hj3) * W3,2,mean)))

   d.opt.minus.A3 <- d.opt.k(psi.hat[7:10],L[,3],A[,2]) - A[,3]
   Hj2 <- as.vector(Y.dash - A[,2]*(W2 %*% psi.hat) + d.opt.minus.A3*(W3 %*% psi.hat))
   Hj.ctr2 <- as.vector(Hj2 - W2 %*% (W2.sq.inv %*% apply(as.vector(Hj2) * W2,2,mean))) 

   d.opt.minus.A2 <- d.opt.k(psi.hat[3:6],L[,2],A[,1]) - A[,2]
   Hj1 <- as.vector(Y.dash - A[,1]*(W1 %*% psi.hat) + d.opt.minus.A2*(W2 %*% psi.hat) + d.opt.minus.A3*(W3 %*% psi.hat))
   Hj.ctr1 <- as.vector(Hj1 - W1 %*% (W1.sq.inv %*% apply(as.vector(Hj1) * W1,2,mean))) 

   U <- Hj.ctr1*(as.vector(aa[,1])*W1) + Hj.ctr2*(as.vector(aa[,2])*W2) + Hj.ctr3*(as.vector(aa[,3])*W3)

   d.alpha.U <- -cbind(t(W1)%*%Hj.ctr1/n, t(W2)%*%Hj.ctr2/n, t(W3)%*%Hj.ctr3/n) 
   d.varsigma.U <- -t(W1)%*%(W1*aa[,1])/n -t(W2)%*%(W2*aa[,2])/n -t(W3)%*%(W3*aa[,3])/n

   U.adj <- t((Hj.ctr1*as.vector(aa[,1]))*W1 + (Hj.ctr2*as.vector(aa[,2]))*W2 + (Hj.ctr3*as.vector(aa[,3]))*W3) - 
       d.alpha.U %*% solve(d.alpha.S) %*% t(S.part) + 
       d.varsigma.U %*% W.sq.inv %*% t(Hj.ctr1 * W1 + Hj.ctr2 * W2 + Hj.ctr3 * W3)

   d.psi.U <- -t(W1)%*%(W1*aa[,1]*A[,1])/n + t(as.vector(aa[,1])*W1)%*%(d.opt.minus.A2*W2)/n + 
       t(as.vector(aa[,1])*W1)%*%(d.opt.minus.A3*W3)/n -
       t(W2)%*%(W2*aa[,2]*A[,2])/n + t(as.vector(aa[,2])*W2)%*%(d.opt.minus.A3*W3)/n -
       t(W3)%*%(W3*aa[,3]*A[,3])/n -
       d.varsigma.U %*% W.sq.inv %*% (t(W1)%*%(A[,1]*W1) + t(W2)%*%(A[,2]*W2) + t(W3)%*%(A[,3]*W3))/n  # Chain rule!

   d.psi.U.adj <- d.psi.U - (t(W1)%*%(W1*A[,1]/d.alpha.S[1,1]*mean(S.part[,1]))/n + 
       t(W2)%*%(W2*A[,2]/d.alpha.S[2,2]*mean(S.part[,2]))/n + 
       t(W3)%*%(W3*A[,3]/d.alpha.S[3,3]*mean(S.part[,3]))/n) + 
       d.varsigma.U %*% W.sq.inv  %*% (t(W1)%*%(W1*A[,1])/n + t(W2)%*%(W2*A[,2])/n + t(W3)%*%(W3*A[,3])/n)

   marg.cov <- (solve(d.psi.U.adj) %*% U.adj %*% t(U.adj) %*% t(solve(d.psi.U.adj)))/(n^2)
   var.psi <- c(diag(marg.cov))

   ## Testing U with TRUE psi's  
   if(return.U) {
      Hj3.true <- as.vector(Y.dash - A[,3]*(W3 %*% true.psi))
      Hj.ctr3 <- as.vector(Hj3.true - W3 %*% (W3.sq.inv %*% apply(as.vector(Hj3) * W3,2,mean)))

      d.opt.minus.A3 <- d.opt.k(true.psi[7:10],L[,3],A[,2]) - A[,3]
      Hj2.true <- as.vector(Y.dash - A[,2]*(W2 %*% true.psi) + d.opt.minus.A3*(W3 %*% true.psi))
      Hj.ctr2 <- as.vector(Hj2.true - W2 %*% (W2.sq.inv %*% apply(as.vector(Hj2) * W2,2,mean))) 

      d.opt.minus.A2 <- d.opt.k(true.psi[3:6],L[,2],A[,1]) - A[,2]
      Hj1.true <- as.vector(Y.dash - A[,1]*(W1 %*% true.psi) + d.opt.minus.A2*(W2 %*% true.psi) + d.opt.minus.A3*(W3 %*% true.psi))
      Hj.ctr1 <- as.vector(Hj1.true - W1 %*% (W1.sq.inv %*% apply(as.vector(Hj1) * W1,2,mean))) 

      U.dagger <- Hj.ctr1*(as.vector(aa[,1])*W1) + Hj.ctr2*(as.vector(aa[,2])*W2) + Hj.ctr3*(as.vector(aa[,3])*W3)
      U.hat <- apply(U.dagger,2,mean)
      
      d.a.u <- -cbind(t(W1)%*%Hj.ctr1/n, t(W2)%*%Hj.ctr2/n, t(W3)%*%Hj.ctr3/n) 
      uu <- t(U.dagger) - d.a.u %*% solve(d.alpha.S) %*% t(S.part) + 
        d.varsigma.U %*% W.sq.inv %*% t(Hj.ctr1 * W1 + Hj.ctr2 * W2 + Hj.ctr3 * W3)
      var.U.hat <- diag(uu %*% t(uu)/(n^2))

      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))
}