##############################
# File to generate simulation data set
##############################

library(pacman)

p_load(dplyr)

logit <- function(x) {
    1 / (1 + exp(-x))
}


d1 <- function(N, ...) {
    B1 <- rbinom(N, 1, 0.5) # %>% as.factor
    C1 <- sample(1:10, N, replace = TRUE) %>% as.factor
    C2 <- sample(1:10, N, replace = TRUE) %>% as.factor
    B2 <- as.integer(as.numeric(C2) <= 5) # %>% as.factor
    N1 <- rnorm(N)
    Sigma <- matrix(0.9, 3, 3)
    diag(Sigma) <- 1
    v47 <- MASS::mvrnorm(N, rep(0, 3), Sigma = Sigma)

    N2 <- v47[, 1]
    N3 <- v47[, 2]
    N4 <- v47[, 3]
    
    u1 <- runif(N)
    u2 <- runif(N)
    S1 <- pmin(u1, u2)
    S2 <- abs(u1 - u2)
    S3 <- 1 - pmax(u1, u2)
    
    x <- data.frame(B1, B2, C1, C2,
                    N1, N2, N3, N4,
                    S1, S2, S3)
    #colnames(x) <- paste0('x', 1:11)
    return(x)
}

###

###
# Outcome model


e0 <- function(x, sigma=0.2, ...) {
    N = NROW(x)
    mu = rep(0, N)
    y = mu + rnorm(N, 0, sigma)
}

e1 <- function(x, sigma=0.2,...) {
    N = NROW(x)
    mu = 0.2 * x$N2
    y = mu + rnorm(N, sigma)
}

e2 <- function(x, sigma=0.2,...) {
    N = NROW(x)
    mu = 0.1 * (x$N1 + x$N2)
    y = mu + rnorm(N, sigma)
}

e3 <- function(x, sigma=0.2,...) {
    N = NROW(x)
    mu = 0.2 * x$B1
    y = mu + rnorm(N, sigma)
}

e4 <- function(x, sigma=0.2,...) {
    N = NROW(x)
    mu = 0.2 * x$B2
    y = mu + rnorm(N, sigma)
}


e5 <- function(x, sigma=0.2,...) {
    N = NROW(x)
    mu = 0.5 * ( (x$B1 == 0) * (as.numeric(x$C1) <= 5) + (x$B1 == 1) * (as.numeric(x$C1) > 5))
    y = mu + rnorm(N, sigma)
}


dmodels = list("D1" = d1)
omodels = list("E0" = e0, "E1" = e1, "E2" = e2,
               "E3" = e3, "E4" = e4, "E5" = e5)

data_generate <- function(N = 400, sigma = 0.2,
                          dm = 'D1', om = 'R1', ...) {
    dmf = match.fun(dmodels[[dm]])
    omf = match.fun(omodels[[om]])

    x = dmf(N, ...)
    y = omf(x, sigma, ...)

    list(y = y, x = x)
}