### Code for Table 5: simulation study

################################################################################

library(tmvtnorm)
source('MUSEC_trial_estimates.R')


N = 10^5   # Number of bootstrap replicates

delta = 0.14  # Assumed true effect size: set equal to 0.10, 0.14 and 0.18

mu = delta*sqrt(c(Ihat_1,Ihat_2))

set.seed(7)

Sys.time()

### Unconditional perspective

Z = rmvnorm(N, mean = mu, sigma = cov_matrix)
MLE1sim = Z[,1]/sqrt(Ihat_1)
MLEsim = MUEsim = ubcMLEsim = UMVUEsim = rep(NA, N)
stop1 = rep(0, N)

for(i in 1:N){
  
  if(Z[i,1] > e1 | Z[i,1] < f1){
    
    MLEsim[i] = UMVUEsim[i] = MUEsim[i] = ubcMLEsim[i] = Z[i,1]/sqrt(Ihat_1)
    
    stop1[i] = 1
    
  } else {
    
    MLEsim[i] = Z[i,2]/sqrt(Ihat_2)
    
    MLE = MLEsim[i]
    Z2 = Z[i,2]

    MUEsim[i] = optimize(MUEobj,interval = c(-1,1))$minimum

    ubcMLEsim[i] = optimize(ubcMLE_obj, interval = c(-1,1))$minimum

    UMVUEsim[i] = MLEsim[i] - sqrt(Ihat_2star/(Ihat_1*Ihat_2))*
      (dnorm((e1-sqrt(Ihat_1/Ihat_2)*Z2)/sqrt(Ihat_2star/Ihat_2)) -
       dnorm((f1-sqrt(Ihat_1/Ihat_2)*Z2)/sqrt(Ihat_2star/Ihat_2)))/
      (pnorm((e1-sqrt(Ihat_1/Ihat_2)*Z2)/sqrt(Ihat_2star/Ihat_2)) -
       pnorm((f1-sqrt(Ihat_1/Ihat_2)*Z2)/sqrt(Ihat_2star/Ihat_2)))
  }
}

sum(stop1)/N

uncond_results = cbind(c(mean(MLEsim), mean(MLE1sim),
                         mean(MUEsim), mean(ubcMLEsim), mean(UMVUEsim)),
                       c(sd(MLEsim), sd(MLE1sim),
                         sd(MUEsim), sd(ubcMLEsim), sd(UMVUEsim)))

rownames(uncond_results) = c('MLE', 'MLE1', 'MUE', 'Bias-corrected MLE',
                             'UMVUE')

print(uncond_results)


# Means for trials that stop early

uncond_stop_early = c(mean(MLE1sim[stop1==1]), sd(MLE1sim[stop1==1]))

print(uncond_stop_early)


# Means for trials that continue

uncond_continue = cbind(c(mean(MLEsim[stop1==0]), mean(MLE1sim[stop1==0]),
                            mean(MUEsim[stop1==0]), mean(ubcMLEsim[stop1==0]),
                            mean(UMVUEsim[stop1==0])),
                          c(sd(MLEsim[stop1==0]), sd(MLE1sim[stop1==0]),
                            sd(MUEsim[stop1==0]), sd(ubcMLEsim[stop1==0]),
                            sd(UMVUEsim[stop1==0])))

rownames(uncond_continue) = c('MLE', 'MLE1', 'MUE', 'Bias-corrected MLE',
                                'UMVUE')

print(uncond_continue)


### Conditional perspective

Z = rtmvnorm(N, mean = mu, sigma = cov_matrix,
             lower = c(f1, -Inf), upper = c(e1, Inf))

cMLE1sim = Z[,1]/sqrt(Ihat_1)
cMLEsim = Z[,2]/sqrt(Ihat_2)

cMLE2sim = (Z[,2]*sqrt(Ihat_2)-Z[,1]*sqrt(Ihat_1))/Ihat_2star

cbcMLEsim = cMUEsim = rep(NA, N)

for(i in 1:N){
  
  MLE = cMLEsim[i]
  cbcMLEsim[i] = optimize(cbcMLE_obj, interval = c(-1,1))$minimum
  
  Z2 = Z[i,2]
  MLE2 = cMLE2sim[i]
  # cMUEsim[i] = optimize(cMUE_obj,interval = c(-1,1))$minimum
  cMUEsim[i] = optim(MLE2, cmue_optim, k = 2, z = Z2,
        I1 = Ihat_1, I2 = Ihat_2, l = f1, u = e1,
        method = "Nelder-Mead")$par
}

w1 = (Ihat_2star^(-1))/sqrt(Ihat_1^(-1) + Ihat_2star^(-1))
w2 = sqrt(Ihat_1^(-1) + Ihat_2star^(-1))/(Ihat_1^(-1))

UMVCUEsim = cMLEsim - w1*(dnorm(w2*(cMLEsim - f1/sqrt(Ihat_1))) -
                          dnorm(w2*(cMLEsim - e1/sqrt(Ihat_1))))/
                         (pnorm(w2*(cMLEsim - f1/sqrt(Ihat_1))) - 
                          pnorm(w2*(cMLEsim - e1/sqrt(Ihat_1))))

cond_results = cbind(c(mean(cMLEsim), mean(cMLE1sim), mean(cMLE2sim),
                       mean(cMUEsim), mean(UMVCUEsim), mean(cbcMLEsim)),
                     c(sd(cMLEsim), sd(cMLE1sim), sd(cMLE2sim),
                       sd(cMUEsim), sd(UMVCUEsim), sd(cbcMLEsim))
)

row.names(cond_results) = c('cMLE', 'cMLE1', 'cMLE2', 'cMUE', 'UMVCUE',
                            'cBias-corrected MLE')

print(cond_results)

Sys.time()

#### Plot sampling distribution of estimators

library(ggplot2)
library(reshape2)
library(patchwork)


uncond_est = data.frame(MLE = MLEsim, MLE1 = MLE1sim,
                        UMVUE = UMVUEsim, UBCMLE = ubcMLEsim, MUE = MUEsim)

cond_est = data.frame(MLE2 = cMLE2sim, UMVCUE = UMVCUEsim, CBCMLE = cbcMLEsim,
                      CMUE = cMUEsim)

save(uncond_est, cond_est, file = 'est_theta014.Rdata')

#######

uncond_est = melt(uncond_est)

uncond_label = labeller(variable = c("MLE" = "MLE (overall)",
                                     "MLE1" = "MLE (stage 1)",
                                     "UMVUE" = "UMVUE",
                                     "UBCMLE" = "Bias-corrected MLE (UBC-MLE)",
                                     "MUE" = "MUE"))

uncond_plot = ggplot(data = uncond_est, aes(x = value)) + 
  geom_histogram(binwidth = 0.01) + xlim(c(-0.3, 0.6)) + ylim(c(0,10^4)) + 
  facet_wrap(~variable, nrow = 5, labeller = uncond_label) + 
  ggtitle("Unconditional estimators") + xlab('Estimator value') + ylab('Count')


cond_est = melt(cond_est)

cond_label = labeller(variable = c("MLE2" = "MLE (stage 2)",
                                   "UMVCUE" = "UMVCUE",
                                   "CBCMLE" = "Bias-corrected MLE (CBC-MLE)",
                                   "CMUE" = "CMUE"))

cond_plot = ggplot(data = cond_est, aes(x = value)) + 
  geom_histogram(binwidth = 0.01) + xlim(c(-0.3, 0.6)) + ylim(c(0,10^4)) + 
  facet_wrap(~variable, nrow = 4, labeller = cond_label) + 
  ggtitle("Conditional estimators") + xlab('Estimator value') + ylab('Count')


layout = c(area(1,1,200), area(29,2,200))

est_plot = uncond_plot + cond_plot + plot_layout(design = layout)

print(est_plot)