################################ 
##  IMOR at exceptional laws  ## 
################################ 

rm(list=ls())
library(MASS)
source("gestCov_txRand_ZIPI.R")

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

regret <- function(b,stat,tx,tx0=0,k) {
   if(k==2) return(abs(b[3]+b[4]*stat+b[5]*tx0+b[6]*stat*tx0)*(tx - expit(30*(b[3]+b[4]*L2+b[5]*A1+b[6]*L2*A1)))^2)
   else return(abs(b[1]+b[2]*stat)*(tx - expit(30*(b[1]+b[2]*stat)))^2)
}

my.apply <- function(x,dim,fcn) {
   if(is.matrix(x)) return(apply(x,dim,fcn))
   else return(apply(matrix(x,ncol=1),dim,fcn))  
}  

myRHS.fcn <- function(Beta,Cval=0) {
   b.hat <- beta.hat
   B <- Beta
   if(length(Beta)==7) Cval <- Beta[7]
#   p.A1.coef <- coef(glm(A1~L1,family=binomial))
#   p.A1.given.L <- expit(cbind(1,L1)%*%p.A1.coef) 
#   p.A2.coef <- coef(glm(A2~L2+A1,family=binomial))
#   p.A2.given.L <- expit(cbind(1,L2,A1)%*%p.A2.coef)
   p.A1.given.L <- mean(A1) 
   p.A2.given.L <- mean(A2)
  
   mus <- cbind(regret(b.hat,L1,A1,k=1),regret(b.hat,L2,A2,A1,k=2)) 
   tv <- Y*0 
   tm <- my.apply(mus,1,sum) 

   tv <- tv + (Y + Cval + tm - mus[,1] + regret(B,L1,A1,k=1) - 
(regret(B,L1,0,k=1)*(1-p.A1.given.L) + regret(B,L1,1,k=1)*p.A1.given.L))^2 
   tv <- tv + (Y + Cval + tm - mus[,2] + regret(B,L2,A2,A1,k=2) - 
(regret(B,L2,0,A1,k=2)*(1-p.A2.given.L) + 
regret (B,L2,1,A1,k=2)*p.A2.given.L))^2
   sum(tv) 
}


#				
## Simulated data to compare ZIPI (80%) to IMOR
#
psi10 <- 1.6	
psi11 <- 0	
psi20 <- -3	
psi21 <- 2
psi22 <- 1	
psi23 <- 0	
psi<-c(psi10,psi11,psi20,psi21,psi22,psi23)
results <- psi
alpha<-c(.5,.5)
tol <- 0.00005  

n.sim <- 1000
n.sample <- 500#c(250,500,1000)#seq(250,1000,50)
n.boot <- 500
#bias.mat <- NULL

for(n in n.sample) {
   psiZIPI80 <- psiSM <- NULL
   varZIPI80 <- varSM <- NULL
   for(j in 1:n.sim) {
#if(j %in% seq(50,n.sim,50)) 
print(j)
      L1 <- rnorm(n,0,3)
      A1 <- rbinom(n,1,alpha[1])  
      L2 <- round(rnorm(n,.75,.5))
      A2 <- rbinom(n,1,alpha[2])  
      L <- cbind(L1,L2)
      A <- cbind(A1,A2)
      Y.opt <- rnorm(n,1,1) 
      Y <- Y.opt - abs(psi10 + psi11*L1)*(A1 - as.numeric(psi10 + 
             psi11*L1>0))^2 - abs(psi20 + psi21*L2 + psi22*A1 + psi23*L2*A1)*
             (A2 - as.numeric(psi20 + psi21*L2 + psi22*A1 + psi23*L2*A1>0))^2      

### ZIPI
       temp <- ZIPI_rand.fcn(Y,L,A,return.U=F,z.val=qnorm(0.9)) # 80% CI zeroing
       psiZIPI80 <- cbind(psiZIPI80,c(temp$psi.hat[2,1:2],temp$psi.hat[1,]))   
       varZIPI80 <- cbind(varZIPI80,c(temp$psi.var[2,1:2],temp$psi.var[1,])) 

### IMOR  
      B.c <- c(temp$psi.hat[2,1:2],temp$psi.hat[1,],0)  
      beta.hat <- B.c[1:6]
      minRHS <- nlm(myRHS.fcn,p=B.c,iterlim=1000)
      B.c <- minRHS$estimate
      while(max(abs(B.c[1:6] - beta.hat)) > tol) {
         beta.hat <- B.c[1:6]
         B.c[7] <- rnorm(1,B.c[7],1)
         minRHS <- nlm(myRHS.fcn,p=B.c,iterlim=1000)
         B.c <- minRHS$estimate
      }
      psiSM <- cbind(psiSM,B.c)

### Bootstrap variance
      bootZIPI80 <- bootSM <- NULL
      Y.keep <- Y
      L.keep <- L
      A.keep <- A
      for(i in 1:n.boot) {
       boot.samp <- sample(1:n,replace=T)
       Y <- Y.keep[boot.samp]
       L <- L.keep[boot.samp,]
       A <- A.keep[boot.samp,]

       B.c <- c(temp$psi.hat[2,1:2],temp$psi.hat[1,],0)  
       beta.hat <- B.c[1:6]
       minRHS <- nlm(myRHS.fcn,p=B.c,iterlim=1000)
       B.c <- minRHS$estimate
       while(max(abs(B.c[1:6] - beta.hat)) > tol) {
         beta.hat <- B.c[1:6]
         B.c[7] <- rnorm(1,B.c[7],1)
         minRHS <- nlm(myRHS.fcn,p=B.c,iterlim=1000)
         B.c <- minRHS$estimate
       }
       bootSM <- cbind(bootSM,B.c)
      }
      varSM <- cbind(varSM,apply(bootSM,1,var)) 
   }

   t1g <- apply(psiZIPI80,1,mean)
   t2g <- apply(sqrt(varZIPI80),1,mean) # calculated SE
   t3g <- apply(sqrt(varZIPI80 + (psiZIPI80 - psi)^2),1,mean)  # rMSE
   t4g <- apply((psiZIPI80 - 1.96*sqrt(varZIPI80) < psi) & (psiZIPI80 + 1.96*sqrt(varZIPI80) > psi),1,sum)/n.sim*100 # Wald CI coverage

   t1a <- apply(psiSM[1:6,],1,mean)
   t2a <- apply(sqrt(varSM[1:6,]),1,mean) # calculated SE
   t3a <- apply(sqrt(varSM[1:6,] + (psiSM[1:6,] - psi)^2),1,mean)  # rMSE
   t4a <- apply((psiSM[1:6,] - 1.96*sqrt(varSM[1:6,]) < psi) & (psiSM[1:6,] + 1.96*sqrt(varSM[1:6,]) > psi),1,sum)/n.sim*100 # Wald CI coverage
   print("psi (SE)   psi (SE)")
   print(round(cbind(psi,t1g,t2g,t3g,t4g,t1a,t2a,t3a,t4a),3)) 
}