# 
## G-estimation for general K-interval problem with linear blip
## 	     PROBABILITY OF A UNIQUE RULE WEIGHTING SCHEME
## 	 Uses my recursive method of calculating the variance
## 	           (Completely randomized treatment) 
## 	    Written by Erica E M Moodie, 20 December, 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))
}

PUR_rand.fcn <- function(Y,status,action,no.stat.vbls=1,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)
   wts <- rep(1,n)

   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
   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.A.given.L <- mean(A[,2])
      aa <- A[,2] - p.A.given.L

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

   	X <- matrix(nrow=2*(n.stat+1),ncol=2*(n.stat+1))
   	X[1,1] <- (A[,2]*aa) %*% wts/sum(wts)
   	X[n.stat+2,1] <- (A[,2]*A[,1]*aa) %*% wts/sum(wts)
   	X[n.stat+2,n.stat+2] <- (A[,2]*A[,1]^2*aa) %*% wts/sum(wts)
   	for(i in 1:n.stat) {
   	   X[i+1,1] <- X[1,i+1] <- (A[,2]*L[,i]*aa) %*% wts/sum(wts)
   	   X[n.stat+2+i,1] <- X[n.stat+2,i+1] <- (A[,2]*A[,1]*L[,i]*aa) %*% wts/sum(wts)
   	   X[n.stat+2+i,n.stat+2] <- X[n.stat+2,n.stat+2+i] <- (A[,2]*A[,1]^2*L[,i]*aa) %*% wts/sum(wts)
   	   for(j in i:n.stat) {
		X[i+1,j+1] <- X[j+1,i+1] <- (A[,2]*L[,i]*L[,j]*aa) %*% wts/sum(wts)
		X[n.stat+2+i,j+1] <- X[n.stat+2+j,i+1] <- (A[,2]*A[,1]*L[,i]*L[,j]*aa) %*% wts/sum(wts)
		X[n.stat+2+i,n.stat+j+2] <- X[n.stat+2+j,n.stat+2+i] <- (A[,2]*A[,1]^2*L[,i]*L[,j]*aa) %*% wts/sum(wts)
	   }
   	}
	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 <- t((cbind(mean(Y.dash),mean(Y.dash*L),mean(Y.dash*A[,1]),mean(Y.dash*L*A[,1])) 
              %*% W2.inv) %*% t(cbind(wts,wts*L,wts*A[,1],wts*L*A[,1])) %*% 
              cbind(aa,L*aa,A[,1]*aa,L*A[,1]*aa)/sum(wts))

      C2 <- (t(cbind(aa,L*aa,A[,1]*aa,L*A[,1]*aa)) %*% cbind(wts,wts*L,wts*A[,1],wts*L*A[,1])/sum(wts)) %*%
            ((W2.inv %*% t(cbind(A[,2],A[,2]*L,A[,2]*A[,1],A[,2]*L*A[,1])) %*% cbind(1,L,A[,1],L*A[,1]))/n)  

      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
      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 
      U.adj <- U.adj*wts

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

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

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

      #############################################
      ##  Update weights by prob. unique regime  ##
      #############################################
      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.z <- (W %*% psi.temp)/individ.se
      wts <- as.vector(wts*(1-2*pnorm(-abs(individ.z))))
      #plot(individ.z,pnorm(abs(individ.z)))
      #plot(individ.z,wts)
   }
   
   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 <- t(Y.dash*aa) %*% wts/sum(wts)
   B <- rbind(B,t(Y.dash*L*aa) %*% wts/sum(wts))
   X <- matrix(nrow=2,ncol=2)   
   X[1,1] <- (A*aa) %*% wts/sum(wts)
   X[2,1] <- X[1,2] <- t(A*L*aa) %*% wts/sum(wts)
   X[2,2] <- t(A*L^2*aa) %*% wts/sum(wts)

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

   C1 <- t((cbind(mean(Y.dash),mean(Y.dash*L)) %*% W2.inv) %*% t(cbind(wts,wts*L)) %*% cbind(aa,L*aa)/sum(wts))
   C2 <- (t(cbind(aa,L*aa)) %*% cbind(wts,wts*L)/sum(wts)) %*%
            ((W2.inv %*% t(cbind(A,A*L)) %*% cbind(1,L))/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 <- 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 
   U.adj <- U.adj*wts

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

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

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

}