## Bootstrap source detection demo
library(MASS)
library(glmnet)
library(DoubleML)
library(mlr3)
library(mlr3learners)
library(PartialTL)

n_t <- 60
n_s1 <- n_s2 <- n_s3 <- n_s4 <- 500

p <- 200  # nuisance dimension including zeros
nonzero.Beta_t <- c(0.3, 0.9, 0.8, 0.7, seq(0.1, 0.9, length.out = 16))
nonzero.Gamma_t <- seq(0.1, 0.9, length.out = 10)

# source 1--4: nuisance coefficients
nonzero.Beta_s1 <- c(0.3, 0.9, 0.8, 0.7, seq(0.1, 0.9, length.out = 16))
nonzero.Gamma_s1 <- seq(0.6, 1.4, length.out = 10)
nonzero.Beta_s2 <- c(1.3, 1.9, 1.8, 1.7, seq(1.1, 1.9, length.out = 16))
nonzero.Gamma_s2 <- seq(0.6, 1.4, length.out = 10)
nonzero.Beta_s3 <- c(5.3, 5.9, 5.8, 5.7, seq(0.1, 0.9, length.out = 16))
nonzero.Gamma_s3 <- seq(0.6, 1.4, length.out = 10)
nonzero.Beta_s4 <- c(10.3, 10.9, 10.8, 10.7, seq(0.1, 0.9, length.out = 16))
nonzero.Gamma_s4 <- seq(0.6, 1.4, length.out = 10)

num_source <- 4
p.nonzero.Beta <- length(nonzero.Beta_t)
p.nonzero.Gamma <- length(nonzero.Gamma_t)

causal_t <- c(-0.8)
q <- length(causal_t)
causal_s1 <- c(0.8)
causal_s2 <- c(0.4)
causal_s3 <- c(1.8)
causal_s4 <- c(1.8)

# DML 'learner$task_type' must be 'regr'
lgr::get_logger("mlr3")$set_threshold("warn")
learner_la = lrn("regr.cv_glmnet")
ml_f = learner_la$clone()
ml_g = learner_la$clone()

sim <- 5
B <- 10
detected.source <- matrix(NA, nrow = sim, ncol = num_source)
source_sizes <- c(n_s1, n_s2, n_s3, n_s4)

for (s in 1:sim) {
  ################################ target data #############################
  Beta_t <- matrix(c(nonzero.Beta_t, rep(0, p - p.nonzero.Beta)))
  Gamma_t <- matrix(c(nonzero.Gamma_t, rep(0, p - p.nonzero.Gamma)))
  d_t <- DGP(n_t, q, p, p.nonzero.Beta, causal_t, Beta_t, Gamma_t,
             mu = rep(10, p), sigma = 0.5, f_func = f_0, g_func = g_0)
  D_t <- d_t$D
  X_t <- d_t$X
  Y_t <- d_t$Y
  ################################ source 1--4 data ########################
  Beta_s1 <- matrix(c(nonzero.Beta_s1, rep(0, p - p.nonzero.Beta)))
  Gamma_s1 <- matrix(c(nonzero.Gamma_s1, rep(0, p - p.nonzero.Gamma)))
  d_s1 <- DGP(n_s1, q, p, p.nonzero.Beta, causal_s1, Beta_s1, Gamma_s1,
              mu = rep(10, p), sigma = 0.5, f_func = f_k, g_func = g_k)
  Beta_s2 <- matrix(c(nonzero.Beta_s2, rep(0, p - p.nonzero.Beta)))
  Gamma_s2 <- matrix(c(nonzero.Gamma_s2, rep(0, p - p.nonzero.Gamma)))
  d_s2 <- DGP(n_s2, q, p, p.nonzero.Beta, causal_s2, Beta_s2, Gamma_s2,
              mu = rep(-10, p), sigma = 0.2, f_func = f_k, g_func = g_k)
  Beta_s3 <- matrix(c(nonzero.Beta_s3, rep(0, p - p.nonzero.Beta)))
  Gamma_s3 <- matrix(c(nonzero.Gamma_s3, rep(0, p - p.nonzero.Gamma)))
  d_s3 <- DGP(n_s3, q, p, p.nonzero.Beta, causal_s3, Beta_s3, Gamma_s3,
              mu = rep(10, p), sigma = 0.8, f_func = f_k, g_func = g_k)
  Beta_s4 <- matrix(c(nonzero.Beta_s4, rep(0, p - p.nonzero.Beta)))
  Gamma_s4 <- matrix(c(nonzero.Gamma_s4, rep(0, p - p.nonzero.Gamma)))
  d_s4 <- DGP(n_s4, q, p, p.nonzero.Beta, causal_s4, Beta_s4, Gamma_s4,
              mu = rep(-10, p), sigma = 0.8, f_func = f_k, g_func = g_k)
  D_s_all <- rbind(d_s1$D, d_s2$D, d_s3$D, d_s4$D)
  X_s_all <- rbind(d_s1$X, d_s2$X, d_s3$X, d_s4$X)
  Y_s_all <- rbind(d_s1$Y, d_s2$Y, d_s3$Y, d_s4$Y)
  ############################### detection #################################
  res <- boot_detection(D_t, X_t, Y_t, D_s_all, X_s_all, Y_s_all, source_sizes, B, ml_f, ml_g)
  detected.source[s, ] <- sapply(res$detected.source, function(x) mean(x) >= 0.5)

  cat("\nSim ", s, " - Detection: ", paste(res$Source, "=", detected.source[s, ], collapse = ", "), "\n")
}
