runSimulation <- function(nSim=1000,
                          n=500,
                          sigmasq=1,
                          pi0=0.9,
                          pi1=0.9,
                          beta_0=0,
                          beta_a=1,
                          beta_l=1,
                          beta_l_sq=0,
                          beta_r=1,
                          beta_lr=0,
                          lambda0=0,
                          lambdaA=1,
                          lambdaR=1,
                          sigma_l=1) {
  
  # first perform asymptotic variance calculations
  # do calculations
  
  #calculate means of L under control and treatment
  #note these are also the means when intervening on R
  #since causally we assume L affects R
  EL_A1 <- lambda0 + lambdaA + lambdaR*(1-pi1)
  EL_A0 <- lambda0 + lambdaR*(1-pi0)
  EL_A1minusA0 <- EL_A1-EL_A0
  
  # P(R=0)
  PR0 <- 0.5*pi0 + 0.5*pi1
  
  # E(A|R=0)=P(A=1|R=0)=P(A=1,R=0)/P(R=0)=(P(R=0|A=1)*0.5)/P(R=0)
  EA_R0 <- (0.5*pi1)/PR0
  var_A_R0 <- EA_R0*(1-EA_R0)
  # E(A|R=1)=P(A=1|R=1)=P(A=1,R=1)/P(R=1)=(P(R=1|A=1)*0.5)/P(R=1)=((1-P(R=0|A=1))*0.5)/(1-P(R=1))
  EA_R1 <- (0.5*(1-pi1))/(1-PR0)
  var_A_R1 <- EA_R1*(1-EA_R1)
  
  varL_R0 <- sigma_l + (lambdaA^2)*var_A_R0
  varL_R1 <- sigma_l + (lambdaA^2)*var_A_R1
  covAL_R0 <- var_A_R0*sigma_l
  covAL_R1 <- var_A_R1*sigma_l
  
  varR_A0 <- pi0*(1-pi0)
  varR_A1 <- pi1*(1-pi1)
  varL_A0 <- lambdaR^2*varR_A0 + sigma_l
  varL_A1 <- lambdaR^2*varR_A1 + sigma_l
  
  ELsq_A1 <- varL_A1 + EL_A1^2
  ELsq_A0 <- varL_A0 + EL_A0^2
  
  # component 1 - pre-ICE data only
  varAL_R0 <- array(c(var_A_R0, covAL_R0, covAL_R0, varL_R0), dim=c(2,2))
  varAL_R0
  vcov_R0 <- sigmasq*solve(PR0*varAL_R0)
  preICEcomp1 <- vcov_R0[1,1] + 2*vcov_R0[1,2]*EL_A1minusA0 + vcov_R0[2,2]*EL_A1minusA0^2
  
  # component 1 - pre+post ICE data
  varAL_R1 <- array(c(var_A_R1, covAL_R1, covAL_R1, varL_R1), dim=c(2,2))
  varAL_R1
  vcov_R0R1 <- sigmasq*solve(PR0*varAL_R0+(1-PR0)*varAL_R1)
  prepostICEcomp1 <- vcov_R0R1[1,1] + 2*vcov_R0R1[1,2]*EL_A1minusA0 + vcov_R0R1[2,2]*EL_A1minusA0^2
  prepostICEcomp1
  
  # component 2
  comp2 <- 2*beta_l1*(varL_A0+varL_A1)
  comp2 
  
  # true value of hypothetical estimand
  # for quadratic contribution need difference in mean of L^2
  # but since Var(L) is same in both arms, this equals difference in squared means of L
  trueEffect <- beta_a + beta_l1*EL_A1minusA0 + beta_l_sq*(ELsq_A1 - ELsq_A0)
  
  # now simulations
  # arrays to save estimates
  # imputation/G-formula only using pre-ICE data
  pre_linear <- array(0, dim=nSim)
  pre_nonlinear <- array(0, dim=nSim)
  # G-formula/G-estimator using pre+post ICE data
  prepost_linear <- array(0, dim=nSim)
  prepost_nonlinear <- array(0, dim=nSim)
  # Loh's G-estimator
  loh_linear <- array(0, dim=nSim)
  loh_nonlinear <- array(0, dim=nSim)
  # save estimates of R effects
  prepost_r_eff <- array(0, dim=nSim)
  loh_r_eff <- array(0, dim=nSim)
  
  for (i in 1:nSim) {
  
    a <- 1*(runif(n)<0.5)
    r <- 1*(runif(n)<(a*(1-pi1)+(1-a)*(1-pi0)))
    l <- lambda0 + lambdaA*a + lambdaR*r + rnorm(n, mean=0, sd=sqrt(sigma_l))
            
    y <- beta_0 + beta_a*a+beta_l1*l+beta_r*r+
      beta_lr*l*r+beta_l_sq*l^2 + rnorm(n,sd=sqrt(sigmasq))
    
    #estimators assuming linear effect of l
    # model fitted only to pre-ICE data
    preICEmod <- lm(y[r==0]~a[r==0]+l[r==0])
    # pre-ICE data estimator
    pre_linear[i] <- coef(preICEmod)[2]+coef(preICEmod)[3]*(mean(l[a==1])-mean(l[a==0]))
  
    # pre+post ICE data estimator
    prepostICEmod <- lm(y~a+l+r)
    prepost_linear[i] <- coef(prepostICEmod)[2]+coef(prepostICEmod)[3]*(mean(l[a==1])-mean(l[a==0]))
    prepost_r_eff[i] <- coef(prepostICEmod)[4]
    
    # Loh estimator
    rMod <- glm(r~a+l, family="binomial")
    rPr <- rMod$fitted.values
    lohYMod <- lm(y~a+l+r+rPr)
    adjY <- y-coef(lohYMod)[4]*r
    loh_linear[i] <- mean(adjY[a==1])-mean(adjY[a==0])
    loh_r_eff[i] <- coef(lohYMod)[4]
    
    #estimators allowing for quadratic effects
    # model fitted only to pre-ICE data
    preICEmod <- lm(y[r==0]~a[r==0]+l[r==0]+I(l[r==0]^2))
    # pre-ICE data estimator
    pre_nonlinear[i] <- coef(preICEmod)[2]+coef(preICEmod)[3]*(mean(l[a==1])-mean(l[a==0]))+
      coef(preICEmod)[4]*(mean(l[a==1]^2)-mean(l[a==0]^2))
    
    # pre+post ICE data estimator
    prepostICEmod <- lm(y~a+l+r+I(l^2))
    prepost_nonlinear[i] <- coef(prepostICEmod)[2]+coef(prepostICEmod)[3]*(mean(l[a==1])-mean(l[a==0]))+
      coef(prepostICEmod)[5]*(mean(l[a==1]^2)-mean(l[a==0]^2))
    
    #Loh
    lohYMod <- lm(y~a+l+r+I(l^2)+rPr)
    adjY <- y-coef(lohYMod)[4]*r
    loh_nonlinear[i] <- mean(adjY[a==1])-mean(adjY[a==0])
  }
  
  # return results
  list(n=n,
       trueEffect=trueEffect,
       preICEcomp1=preICEcomp1,
       prepostICEcomp1=prepostICEcomp1,
       comp2=comp2,
       pre_linear=pre_linear,
       prepost_linear=prepost_linear,
       loh_linear=loh_linear,
       pre_nonlinear=pre_nonlinear,
       prepost_nonlinear=prepost_nonlinear,
       loh_nonlinear=loh_nonlinear,
       prepost_r_eff=prepost_r_eff,
       loh_r_eff=loh_r_eff)
}

resTablePrint <- function(simSet,dp=3) {
  
  resTable <- array(0, dim=c(14,length(pi1Vals)))
  for (i in 1:length(pi1Vals)) {
    resTable[1,i] <- pi1Vals[i]-0.1
    resTable[2,i] <- pi1Vals[i]
    # pre ICE data estimator
    resTable[3,i] <- mean(simSet[[i]]$pre_linear-simSet[[i]]$trueEffect)
    resTable[4,i] <- sd(simSet[[i]]$pre_linear)
    # estimated SE based on asymptotic variance expression
    #resTable[i,5] <- sqrt((simSet[[i]]$preICEcomp1+simSet[[i]]$comp2)/simSet[[i]]$n)
    
    # pre+post
    resTable[5,i] <- mean(simSet[[i]]$prepost_linear-simSet[[i]]$trueEffect)
    resTable[6,i] <- sd(simSet[[i]]$prepost_linear)
    # estimated SE based on asymptotic variance expression
    #resTable[i,8] <- sqrt((simSet[[i]]$prepostICEcomp1+simSet[[i]]$comp2)/simSet[[i]]$n)
    
    # Loh estimator
    resTable[7,i] <- mean(simSet[[i]]$loh_linear-simSet[[i]]$trueEffect)
    resTable[8,i] <- sd(simSet[[i]]$loh_linear)
    
    #non-linear L estimators
    # pre ICE data estimator
    resTable[9,i] <- mean(simSet[[i]]$pre_nonlinear-simSet[[i]]$trueEffect)
    resTable[10,i] <- sd(simSet[[i]]$pre_nonlinear)
    
    # pre+post
    resTable[11,i] <- mean(simSet[[i]]$prepost_nonlinear-simSet[[i]]$trueEffect)
    resTable[12,i] <- sd(simSet[[i]]$prepost_nonlinear)
    
    # Loh estimator
    resTable[13,i] <- mean(simSet[[i]]$loh_nonlinear-simSet[[i]]$trueEffect)
    resTable[14,i] <- sd(simSet[[i]]$loh_nonlinear)
    
    row.names(resTable) <- c("pi0","pi1",
                             "pre_lin mean", "pre_lin sd",
                             "pre_post_lin mean", "pre_post_lin sd",
                             "loh_lin mean", "loh_lin sd",
                             "pre_nonlin mean", "pre_nonlin sd",
                             "prepost_nonlin_mean", "prepost_nonlin sd",
                             "loh_nonlin mean", "loh_nonlin sd")
  }
  
  resTable
}

pi1Vals <- seq(from=0.5, to=0.9, by=0.1)
simSet1 <- vector(mode = "list", length = length(pi1Vals))
simSet2 <- vector(mode = "list", length = length(pi1Vals))
simSet3 <- vector(mode = "list", length = length(pi1Vals))
simSet4 <- vector(mode = "list", length = length(pi1Vals))


nSim <- 10000
set.seed(7127332)
for (i in 1:length(pi1Vals)) {
  print(i)
  simSet1[[i]] <- runSimulation(nSim=nSim,pi0=(pi1Vals[i]-0.1), pi1=pi1Vals[i])
  simSet2[[i]] <- runSimulation(nSim=nSim,pi0=(pi1Vals[i]-0.1), pi1=pi1Vals[i], beta_l_sq = 1)
  simSet3[[i]] <- runSimulation(nSim=nSim,pi0=(pi1Vals[i]-0.1), pi1=pi1Vals[i], beta_lr = 1)
  simSet4[[i]] <- runSimulation(nSim=nSim,pi0=(pi1Vals[i]-0.1), pi1=pi1Vals[i], beta_l_sq = 1, beta_lr = 1)
}

#additional simulation to explore behaviour under quadratic L effect when
#ICE rates differ substantially between arms
set.seed(912751)
diffICERates <- runSimulation(nSim=10000,n=500,pi0=0.3, pi1=0.9, beta_l_sq = 1)
#estimation of R effect
mean(diffICERates$prepost_r_eff)
mean(diffICERates$loh_r_eff)
#bias of pre-post-linear
mean(diffICERates$prepost_linear)-diffICERates$trueEffect
#bias of Loh linear
mean(diffICERates$loh_linear)-diffICERates$trueEffect


#plotting results
library(dplyr)
library(tidyr)
library(purrr)
library(ggplot2)

plotSimSet <- function(setNumber) {
  
simSetName <- paste0("simSet", setNumber)
simSet <- get(simSetName, envir = parent.frame())

summary_df <- map2_dfr(simSet, pi1Vals, ~{
  tibble(
    pi1 = .y,
    trueEffect = .x$trueEffect,
    
    mean_pre_linear = mean(.x$pre_linear),
    sd_pre_linear   = sd(.x$pre_linear),
    
    mean_prepost_linear = mean(.x$prepost_linear),
    sd_prepost_linear   = sd(.x$prepost_linear),
    
    mean_loh_linear = mean(.x$loh_linear),
    sd_loh_linear   = sd(.x$loh_linear),
    
    mean_pre_nonlinear = mean(.x$pre_nonlinear),
    sd_pre_nonlinear   = sd(.x$pre_nonlinear),
    
    mean_prepost_nonlinear = mean(.x$prepost_nonlinear),
    sd_prepost_nonlinear   = sd(.x$prepost_nonlinear),
    
    mean_loh_nonlinear = mean(.x$loh_nonlinear),
    sd_loh_nonlinear   = sd(.x$loh_nonlinear)
  )
})

plot_df <- summary_df %>%
  pivot_longer(
    cols = -c(pi1, trueEffect),
    names_to = c("stat", "estimator"),
    names_pattern = "(mean|sd)_(.*)",
    values_to = "value"
  ) %>%
  pivot_wider(
    names_from = stat,
    values_from = value
  ) %>%
  mutate(bias = mean - trueEffect)

plot_df$estimator <- factor(
  plot_df$estimator,
  levels = c(
    "pre_linear",
    "prepost_linear",
    "loh_linear",
    "pre_nonlinear",
    "prepost_nonlinear",
    "loh_nonlinear"
  )
)

pd <- position_dodge(width = 0.05)

ggplot(plot_df, aes(x = pi1, y = bias, colour = estimator)) +
  geom_hline(yintercept = 0, linetype = "dashed", colour = "black") +
  geom_point(position = pd, size = 2) +
  geom_errorbar(
    aes(ymin = bias - sd, ymax = bias + sd),
    width = 0.015,
    position = pd
  ) +
  theme_bw() +
  labs(x = "Proportion ICE free (active treatment arm)", y = "Bias (mean - true)") +
  scale_colour_discrete(labels = c(
    pre_linear         = "Pre-linear",
    prepost_linear     = "Pre-post linear",
    loh_linear         = "Loh-linear",
    pre_nonlinear      = "Pre-nonlinear",
    prepost_nonlinear  = "Pre-post nonlinear",
    loh_nonlinear      = "Loh-nonlinear"
  )) +
  theme(legend.position = "bottom", legend.title = element_blank()) +
  guides(colour = guide_legend(nrow = 2, byrow = TRUE)) 
}

#generate figure images and save
for (i in 1:4) {
  p <- plotSimSet(i)
  
  ggplot2::ggsave(
    filename = paste0("simSet", i, ".eps"),
    plot = p,
    device = "eps",
    width = 6, height = 4
  )
}