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

source("gestCov1int_txRand.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_2ints.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) {
      print("Can't do this")
      return(NULL)
   }

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

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

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

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

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

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

   Hj2 <- as.vector(Y.dash - A[,2]*(W2 %*% psi.hat))
   Hj.ctr2 <- as.vector(Hj2 - W2 %*% (W2.sq.inv %*% apply(as.vector(Hj2) * W2,2,mean)))

   d.opt.minus.A <- 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.A*(W2 %*% 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)

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

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

   d.psi.U <- -t(W1)%*%(W1*aa[,1]*A[,1])/n + 
       t(as.vector(aa[,1])*W1)%*%(d.opt.minus.A*W2)/n -
       t(W2)%*%(W2*aa[,2]*A[,2])/n - 
       d.varsigma.U %*% W.sq.inv %*% (t(W1) %*% (A[,1]*W1) + t(W2) %*% (A[,2]*W2))/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) + 
       d.varsigma.U %*% W.sq.inv  %*% (t(W1)%*%(W1*A[,1])/n + t(W2)%*%(W2*A[,2])/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) {
      Hj2.true <- as.vector(Y.dash - A[,2]*(W2 %*% true.psi))
      Hj.ctr2 <- as.vector(Hj2.true - W2 %*% (W2.sq.inv %*% apply(as.vector(Hj2) * W2,2,mean))) # varsigma.hat!

      d.opt.minus.A <- 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.A*(W2 %*% 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) # alpha.hat!
      U.hat <- apply(U.dagger,2,mean)
      
      d.a.u <- -cbind(t(W1)%*%Hj.ctr1/n,t(W2)%*%Hj.ctr2/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)
      var.U.hat <- diag(uu %*% t(uu)/(n^2))

      aa <- cbind(A[,1] - true.a[1],A[,2] - true.a[2])
      Hj.ctr2 <- as.vector(Hj2.true - W2 %*% (W2.sq.inv %*% apply(as.vector(Hj2.true) * W2,2,mean)))
      Hj.ctr1 <- as.vector(Hj1.true - W1 %*% (W1.sq.inv %*% apply(as.vector(Hj1.true) * W1,2,mean))) 

      U.dagger <- Hj.ctr1*(as.vector(aa[,1])*W1) + Hj.ctr2*(as.vector(aa[,2])*W2)
      d.alpha.U <- -cbind(t(W1)%*%Hj.ctr1/n,t(W2)%*%Hj.ctr2/n) 
      d.varsigma.U <- -t(W1)%*%(W1*aa[,1])/n -t(W2)%*%(W2*aa[,2])/n

      U.dagger <- t(Hj.ctr1*(as.vector(aa[,1])*W1) + Hj.ctr2*(as.vector(aa[,2])*W2)) - d.alpha.U %*% solve(d.alpha.S) %*% t(S.part) + d.varsigma.U %*% W.sq.inv %*% t(Hj.ctr1 * W1 + Hj.ctr2 * W2)
      var.U <- diag((U.dagger %*% t(U.dagger))/(n^2))
      return(list(psi.hat=psi.hat,psi.var=var.psi,U.hat=U.hat,U.var.hat=var.U.hat,U.var=var.U))
   }

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