############################################################
#
#    Reproducibility Material for
#
#    Joint count transformation models with covariate-dependent correlations
#    -- Appendix A: Simulation Study
#
#    by Lukas Graz, Luisa Barbanti, Roland Brandl and Torsten Hothorn
#
#############################################################

# ==============================================================================
# 0. SETUP & LIBRARIES
# ==============================================================================
testrun <- FALSE # light version for development

# Control whether caching is used for computing simulations
if (exists("use_sim_cache")) {
  dir.create("R/cache", showWarnings = FALSE)
} else {
  message(
    "Warning: Simulation took a few days on hundret cores. Consider setting use_sim_cache <- TRUE to load results from cache (if available) or to enable caching for future runs."
  )
  message(
    "Caching is disabled for simulations. Set use_sim_cache <- TRUE to enable caching."
  )
  use_sim_cache <- FALSE
}

# set ncores
options(mc.cores = min(parallel::detectCores() - 2, 250))

check_min_version <- function(package_name, minimum_version) {
  if (!requireNamespace(package_name, quietly = TRUE)) {
    stop(sprintf(
      "Cannot find required package %s. Please install %s >= %s.",
      package_name,
      package_name,
      minimum_version
    ))
  }

  installed_version <- packageVersion(package_name)
  if (installed_version < minimum_version) {
    stop(sprintf(
      "Please update %s package to version %s or higher.",
      package_name,
      minimum_version
    ))
  } else {
    suppressPackageStartupMessages(
      library(package_name, character.only = TRUE)
    )
  }
}

# Core data & plotting libraries
check_min_version("dplyr", "1.2.1")
check_min_version("tidyr", "1.3.2")
check_min_version("ggplot2", "4.0.3")
check_min_version("patchwork", "1.3.2")
check_min_version("latex2exp", "0.9.8")
check_min_version("scales", "1.4.0")
check_min_version("cowplot", "1.2.0")

# Modeling & underlying dependencies
check_min_version("qrng", "0.0-11")
check_min_version("cotram", "0.6-0")
check_min_version("tram", "1.4-2")
check_min_version("mvtnorm", "1.4-0")
check_min_version("mlt", "1.8-0")
check_min_version("basefun", "1.2-6")
check_min_version("variables", "1.1-2")

# Data & styling namespaces
check_min_version("TH.data", "1.1-5")
check_min_version("viridis", "0.6.5")

# ==============================================================================
# 1. HELPER FUNCTIONS (formerly ctm-helpfun.R)
# ==============================================================================

# turns out that the marginal coefs can be extracted via:
#   M$models$parm(coef(M, fixed = TRUE))
get_Mcoefs <- function(M) {
  Mmodel_coefs <- list()
  Lambda_coefs <- c()
  Model_coefs_ind <- list()
  for (mod in M$models$models) {
    response <- mod$model$response
    Model_coefs_ind[[response]] <- match(
      paste0(response, ".", names(mod$theta)),
      names(M$theta)
    )
    tmp <- M$theta[Model_coefs_ind[[response]]]
    names(tmp) <- gsub(paste0("^", response, "\\."), "", names(tmp))
    Mmodel_coefs[[response]] <- tmp
  }
  Lambda_coefs <- M$theta[-unlist(Model_coefs_ind)]
  return(list(
    Mmodels = Mmodel_coefs,
    Lambdas = Lambda_coefs
  ))
}

get_Mmodels <- function(M) {
  CoefMMs <- get_Mcoefs(M)$Mmodels
  Mmodels <- list()
  for (i in seq_along(M$models$models)) {
    mod <- M$models$models[[i]]
    call. <- mod$call
    call.$theta <- CoefMMs[[i]]
    call.$dofit <- FALSE # keep marginally fitted coefs
    Mmodels[[mod$model$response]] <- eval(call.)
  }
  Mmodels
}

O2L <- function(O) {
  L <- O %*% diag(1 / diag(O))
  dimnames(L) <- dimnames(O)
  L
}

# returns L s.t. t(L) %*% L = A
mychol <- function(A) {
  as.array(
    chol(mvtnorm::syMatrices(
      A[lower.tri(A, diag = TRUE)],
      diag = TRUE
    ))
  )[,, 1]
}

#' @responses: character vector of responses, e.g. c("y1", "y2", "y3")
#' @marginalformulas: List of formuals for marginal models
#' @data: data.frame
#' @marginalFun: character vector of marginal moels, e.g. c("Polr", "Lm", "Colr")
#' @marginalArgs: list of lists of additional arguments for marginal models,
#' @startValues: character,
#'
#' @...: additional arguments for Mmlt
Mmlt_with_marginals <- function(
  responses,
  marginalformulas,
  data,
  marginalFun,
  marginalArgs = NULL,
  startValues = c(
    "none",
    "countApprox",
    "auglag",
    "spg",
    "nloptr",
    "constrOptim",
    "nlminb",
    "optim"
  ),
  # args for xfun::cache_rds():
  rerun = FALSE,
  cache = TRUE,
  dir = "R/cache/",
  file = "mmltpar.rds",
  ...
) {
  all_args <- c(as.list(environment()), list(...))
  mmltargs <- list(data = data, ...)
  info <- list() # attached later to M

  ## create hash for Mmlt fit
  all_args_hash <- all_args
  if (!is.null(all_args_hash$args)) {
    all_args_hash$args <- rlang::hash(all_args_hash$args) # list with mltoptim-args
  }
  all_args_hash$data <- rlang::hash(data)
  all_args_hash$THIS.OBJ.Classes <- sapply(all_args_hash, class)
  hash <- as.list(as.character(all_args_hash))
  names(hash) <- names(all_args_hash)

  ## Marginal models
  mm <- marginal_models(
    responses,
    marginalformulas,
    data,
    marginalFun,
    marginalArgs
  )

  ## startValues
  startValues <- match.arg(startValues)
  if (startValues != "none") {
    stopifnot(is.null(mmltargs$theta)) # othewise overwritten
    info$startValues <- startValues

    if (length(startValues) == 1 && startValues %in% names(mltoptim())) {
      ## different optimizers
      all_args$optim <- mltoptim()[startValues]
      all_args$startValues <- "none"
      m.start <- do.call(Mmlt_with_marginals, all_args)
      info$startValues.time <- m.start$info$mmlttime
      info$startValues.par <- m.start$par
      mmltargs$theta <- m.start$par
    } else if (startValues == "countApprox") {
      ## countApprox: get start values by approximating counts as numeric and fitting Mmlt with data = d.numeric
      ## get numeric data:
      integer_responses <- responses[sapply(data[responses], is.integer)]
      if (length(integer_responses) == 0) {
        warning("No integer responses found for count approximation")
      }
      d.numeric <- data
      d.numeric[integer_responses] <- lapply(data[integer_responses], \(x) {
        as.numeric(x) + 0.5
      })

      # set parameters by running this very function but with data = d.numeric
      all_args$data <- d.numeric
      all_args$startValues <- "none"
      m.start <- do.call(Mmlt_with_marginals, all_args)
      info$startValues.time <- m.start$info$mmlttime
      info$startValues.par <- m.start$par
      info$cts_approx_ll <- logLik(m.start)
      mmltargs$theta <- m.start$par
    } else {
      stop("Unknown startValues method")
    }
  }

  ## Mmlt
  M <- do.call(
    cached_Mmlt,
    c(
      mm,
      mmltargs,
      list(
        hash = hash,
        rerun = rerun,
        dir = dir,
        file = file,
        cache = cache
      )
    )
  )
  info <- c(info, M$info)

  ## if model was loaded from cache, adjust mmlttime to include original mmlttime
  ## add LogLik fun
  info$m_nofit_high_M <- local({
    M <- 1e+6
    mmltargs$args <- list(
      M = M,
      w = t(ghalton(n = M, d = length(responses) - 1L))
    )
    mmltargs$dofit <- FALSE
    do.call("Mmlt", c(mm, mmltargs))
  })

  M$info <- info
  M
}

marginal_models <- function(
  responses,
  marginalformulas,
  data,
  marginalFun,
  marginalArgs = NULL,
  isLatent = stop()
) {
  ## if marginal formulas is a single RHS-formula create list of formulas with same RHS but different LHS
  if (
    is.character(marginalformulas) &&
      length(marginalformulas) == 1 &&
      grepl("^~", marginalformulas)
  ) {
    marginalformulas <- lapply(responses, function(r) {
      as.formula(paste0(r, marginalformulas))
    })
  }

  ## checks
  stopifnot(length(responses) == length(marginalformulas))
  stopifnot(is.list(marginalformulas))

  ## marginal models
  mm <- list()
  for (i in seq_along(responses)) {
    ff <- as.formula(marginalformulas[[i]])
    mm[[i]] <- do.call(
      marginalFun[i],
      c(list(formula = ff, data = data), marginalArgs[[i]])
    )
  }
  mm
}

#' Mmlt(...)
cached_Mmlt <- function(
  ...,
  # xfun::cache_rds arguments:
  rerun = FALSE,
  dir = stop("Cache directory is required"),
  file = "mmltpar.rds",
  hash = stop("Hash is required"),
  cache = TRUE
) {
  tic <- Sys.time()
  args <- list(...)
  cache_used <- FALSE

  # cached Mmlt(...)
  if (cache) {
    cache_used <- TRUE
    mmltpar <- xfun::cache_rds(
      {
        cache_used <- FALSE
        tmp <- do.call("Mmlt", args)
        c(
          list(mmlttime = Sys.time() - tic),
          tmp[c("par", "optim_hessian", "convergence")]
        )
      },
      file = file,
      dir = dir,
      hash = hash,
      rerun = rerun,
      clean = FALSE
    )
    if (cache_used) {
      message("mmltpar loaded from cache")
    }
    args$theta <- mmltpar[["par"]]
    args$dofit <- FALSE
  }
  M <- do.call("Mmlt", c(args))

  if (cache) {
    M$optim_hessian <- mmltpar[["optim_hessian"]]
    M$convergence <- mmltpar[["convergence"]]
  }

  M$info <- list(
    cache_used = cache_used,
    mmlttime = if (cache) {
      mmltpar[["mmlttime"]]
    } else {
      Sys.time() - tic
    },
    hash = hash
  )
  M
}

print.mmlt <- function(M) {
  cf <- c(M$par, M$fixed)[M$parnames]
  prm <- M$parm(cf)

  ## marginal coefs
  for (j in seq_along(M$names)) {
    nm <- M$names[j]
    names(prm[[j]]) <- sub(paste0(nm, "."), "", names(prm[[j]]), fixed = TRUE)
  }
  names(prm)[seq_along(M$names)] <- M$names

  ## Lambda coefs
  Lpar <- prm[[length(prm)]]
  Lpar <- ltMatrices(t(Lpar), byrow = TRUE, diag = FALSE, names = M$names)
  diagonals(Lpar) <- 0
  prm[[length(prm)]] <- Lpar
  names(prm)[length(prm)] <- paste(
    c("Lambda()", as.character(M$formula)),
    collapse = " "
  )

  ## print
  cat("Mmlt coefficients:\n")
  print(prm)
  invisible(prm)
}


# ==============================================================================
# 2. SIMULATION SETUP & EXECUTION (formerly sim.r)
# ==============================================================================
nsim <- if (testrun) 1 else 15
Nmax <- 10000
Jmax <- 20 # changing this changes the whole data == invalidated cache
P <- Jmax * (Jmax - 1) / 2
nm <- paste0("sp", seq_len(Jmax))
names(nm) <- nm

FitSetup <- rbind(
  ## check Genz x J
  expand.grid(
    startValues = "countApprox",
    randQuadrature = FALSE,
    nsim = seq_len(nsim),
    optim = "auglag",
    N = 500,
    Genz = if (testrun) 125 else c(125, 250, 500, 1000, 2000, 5000), # number of quadrature points for Mmlt
    J = if (testrun) 3:4 else c(3, 5, 7, 10, 12, 15),
    counts = c("low", "high")
  ),
  ## consistency --> increase N
  expand.grid(
    startValues = "countApprox",
    randQuadrature = FALSE,
    nsim = seq_len(nsim),
    optim = "auglag",
    N = if (testrun) c(125, 250, 1000) else c(125, 250, 1000, 2500, 5000, 1e+4),
    Genz = if (testrun) 125 else 500, # number of quadrature points for Mmlt
    J = if (testrun) 4 else 10,
    counts = c("low", "high")
  )
)
FitSetup$startValues <- as.character(FitSetup$startValues)
FitSetup$optim <- as.character(FitSetup$optim)

set.seed(1)
df_species_low <- sample(3:9, Jmax, replace = TRUE)
names(df_species_low) <- nm

### get mmlt `M_true` object for simulations
## 1. get marginal parameters h_xy
x <- runif(Nmax, 0, 1)
beta_x <- runif(Jmax, -1, 1)
names(beta_x) <- nm
h_xy <- lapply(nm, function(n) {
  tmp <- c(1, x = unname(beta_x[n]), "(Intercept)" = 0)
  names(tmp)[1] <- n
  tmp
})
df_species_high <- sample(30:50, Jmax, replace = TRUE)
names(df_species_high) <- nm

## 2. get Lambda parameters via one Wishart sample
set.seed(1)
W.x_equal_0 <- rWishart(n = 1, df = Jmax, Sigma = diag(Jmax))[,, 1] |> cov2cor()
W.x_equal_1 <- rWishart(n = 1, df = Jmax, Sigma = diag(Jmax))[,, 1] |> cov2cor()
dimnames(W.x_equal_0) <- list(nm, nm)
dimnames(W.x_equal_1) <- list(nm, nm)

L.x_equal_0 <- O2L(solve(mychol(W.x_equal_0)))
L.x_equal_1 <- O2L(solve(mychol(W.x_equal_1)))
dimnames(L.x_equal_0) <- list(nm, nm)
dimnames(L.x_equal_1) <- list(nm, nm)

L0 <- ltMatrices(L.x_equal_0[lower.tri(L.x_equal_0, diag = FALSE)], names = nm)
L1 <- ltMatrices(L.x_equal_1[lower.tri(L.x_equal_1, diag = FALSE)], names = nm)
L0 <- ltMatrices(L0, byrow = TRUE) ### compare with coef(<mmlt>, type = "Lambda", newdata = (data.frame(x = 0))
L1 <- ltMatrices(L1, byrow = TRUE) ### compare with coef(<mmlt>, type = "Lambda", newdata = (data.frame(x = 1))

## Lambda(x)
## = L0par       + x * L1par
## = L.x_equal_0 + x * (L.x_equal_1 - L.x_equal_0)
Lx <- ltMatrices(
  c(unclass(L0)) + c(unclass(L1 - L0)) %o% x,
  names = nm,
  byrow = TRUE
)
Ox <- standardize(invchol = Lx)

gen_data <- function(df_qchisq = df_species_low) {
  Z <- matrix(rnorm(Jmax * Nmax), ncol = Nmax) ### N(0, I)
  Z <- solve(Ox, Z) ### N(0, Sigma(x))
  Zx <- Z + beta_x %o% x
  Y <- lapply(nm, function(n) {
    as.integer(ceiling(qchisq(
      pnorm(Zx[n, ], log.p = TRUE),
      df = df_qchisq[n],
      log.p = TRUE
    )))
  })
  Y$x <- x
  d <- as.data.frame(Y)
  d
}

# simulate new data
set.seed(1)
datasets_low <- list()
datasets_high <- list()
cat("Simulating", nsim, "datasets:\n")
for (i in 1:nsim) {
  cat(".")
  datasets_low[[i]] <- gen_data()
  datasets_high[[i]] <- gen_data(df_qchisq = df_species_high)
}
cat("\n")
datasets <- list(
  low = datasets_low,
  high = datasets_high
)

# END SimSetup -----------------------------------------------------------------

fit_model_from_FitSetup_row <- function(
  FitSetup_i,
  cache = FALSE,
  rerun = FALSE
) {
  J <- FitSetup_i$J
  Genz <- FitSetup_i$Genz
  N <- FitSetup_i$N
  d <- datasets[[FitSetup_i$counts]][[FitSetup_i$nsim]][1:N, c(nm[1:J], "x")]

  QuadraturePoints <- if (FitSetup_i$randQuadrature) {
    t(ghalton(n = Genz, d = Jmax - 1L))
  } else {
    set.seed(1)
    matrix(runif((J - 1L) * Genz * N), nrow = J - 1L)
  }

  M <- Mmlt_with_marginals(
    responses = nm[1:J],
    marginalformulas = "~ x",
    data = d,
    marginalFun = rep("BoxCox", J),
    marginalArgs = NULL,
    startValues = FitSetup_i$startValues,
    marginalOnly = FALSE,
    cache = cache,
    dir = "R/cache/",
    file = "mmltpar.rds",
    rerun = rerun,
    formula = ~ 1 + x,
    args = list(M = Genz, w = QuadraturePoints),
    optim = mltoptim()[FitSetup_i$optim] # turn of computation of Hessian
  )
  M
}

resultsM <- function(M) {
  J <- length(M$models$models)

  # mean absolute deviation of lower triangular enries for x=0 and x=1
  S0 <- as.array(coef(M, type = "Cor", newdata = data.frame(x = 0)))[,, 1]
  S1 <- as.array(coef(M, type = "Cor", newdata = data.frame(x = 1)))[,, 1]
  sigma_mae <- mean(abs(c(
    (W.x_equal_0[1:J, 1:J] - S0)[lower.tri(S0)],
    (W.x_equal_1[1:J, 1:J] - S1)[lower.tri(S1)]
  )))
  beta_mae <- mean(abs(M$par[paste0(M$names[1:J], ".x")] - beta_x[1:J]))

  list(
    sigma_mae = sigma_mae,
    beta_mae = beta_mae,
    converged = M$converged,
    ll = logLik(M),
    high_M_ll = logLik(M$info$m_nofit_high_M, parm = M$par),
    high_M_ll_approx = if (!is.null(M$info$startValues.par)) {
      logLik(M$info$m_nofit_high_M, parm = M$info$startValues.par)
    } else {
      NA
    },
    cts_approx_ll = if (!is.null(M$info$cts_approx_ll)) {
      M$info$cts_approx_ll
    } else {
      NA
    },
    mean_abs_gradient = mean(abs(Gradient(M)[!grepl("Bs", names(M$par))])),
    mean_abs_gradient_high_M = mean(abs(
      Gradient(M$info$m_nofit_high_M, parm = M$par)[!grepl("Bs", names(M$par))]
    )),
    mean_abs_gradient_high_M_approx = if (!is.null(M$info$startValues.par)) {
      mean(abs(
        Gradient(M$info$m_nofit_high_M, parm = M$info$startValues.par)[
          !grepl("Bs", names(M$par))
        ]
      ))
    } else {
      NA
    },
    mmlttime = as.double(M$info$mmlttime, units = "secs"),
    startValues.time = if (!is.null(M$info$startValues.time)) {
      as.double(M$info$startValues.time, units = "secs")
    } else {
      NA
    },
    par = M$par,
    startValues.par = if (!is.null(M$info$startValues.par)) {
      M$info$startValues.par
    } else {
      NA
    }
  )
}

set.seed(1)
permutation <- sample(nrow(FitSetup)) # debias parallelization by randomizing order of computations

# Pass use_sim_cache to control caching dynamically
results <- parallel::mclapply(
  permutation,
  function(x) {
    M <- fit_model_from_FitSetup_row(
      FitSetup[x, ],
      cache = use_sim_cache,
      rerun = FALSE
    )
    resultsM(M)
  },
  mc.preschedule = FALSE
)[match(1:nrow(FitSetup), permutation)]

# --- Optional RDS caching lines (Commented out) ---
# if (!testrun) {
#   saveRDS(results, "R/sim-results.rds")
#   saveRDS(FitSetup, "R/sim-FitSetup.rds")
# }
#
# FitSetup <- readRDS("R/sim-FitSetup.rds")
# results <- readRDS("R/sim-results.rds")
# ---------------------------------------------------

## analyze results =============================================================
options(digits = 3)
options(width = 160)
options(max.print = 10000)
results.na <- sapply(results, inherits, "error")

results <- lapply(results, function(r) {
  if (inherits(r, "error")) {
    list(
      sigma_mae = -Inf,
      beta_mae = -Inf,
      ll = -Inf,
      high_M_ll = -Inf,
      ll_approx = -Inf,
      high_M_ll_approx = -Inf,
      cts_approx_ll = -Inf,
      mmlttime = Inf,
      startValues.time = Inf,
      par = NA,
      startValues.par = NA,
      mean_abs_gradient = Inf,
      mean_abs_gradient_high_M = Inf,
      mean_abs_gradient_high_M_approx = Inf
    )
  } else {
    r
  }
})

results_vecs <- lapply(results, \(r) {
  c(unlist(r[c(
    "sigma_mae",
    "beta_mae",
    "ll",
    "high_M_ll",
    "ll_approx",
    "high_M_ll_approx",
    "cts_approx_ll",
    "mmlttime",
    "startValues.time",
    "mean_abs_gradient",
    "mean_abs_gradient_high_M",
    "mean_abs_gradient_high_M_approx"
  )]))
})

results_mat <- do.call("rbind", results_vecs)
results_mat <- cbind(
  results_mat,
  abs_rel_ll_diff = abs(
    (results_mat[, "ll"] - results_mat[, "high_M_ll"]) /
      results_mat[, "high_M_ll"]
  ),
  abs_rel_ll_diff_approx = abs(
    (results_mat[, "cts_approx_ll"] - results_mat[, "high_M_ll_approx"]) /
      results_mat[, "high_M_ll_approx"]
  )
)

results_df <- cbind(results_mat, FitSetup)
results_agg <- aggregate(
  cbind(
    mmlttime,
    startValues.time,
    high_M_ll,
    high_M_ll_approx,
    cts_approx_ll,
    abs_rel_ll_diff,
    abs_rel_ll_diff_approx,
    sigma_mae,
    beta_mae,
    mean_abs_gradient,
    mean_abs_gradient_high_M,
    mean_abs_gradient_high_M_approx
  ) ~ J + Genz + N + startValues + randQuadrature + optim + counts,
  data = results_df,
  FUN = median,
  na.rm = TRUE,
  na.action = na.pass
)

# order by J and Genz
results_agg <- results_agg[
  order(
    results_agg$counts,
    results_agg$J,
    results_agg$Genz,
    results_agg$startValues,
    results_agg$randQuadrature,
    results_agg$N
  ),
]

results_agg$time.total <- results_agg$mmlttime +
  ifelse(is.na(results_agg$startValues.time), 0, results_agg$startValues.time)

results_agg_prefiltered <- results_agg[
  results_agg$optim == "auglag" &
    results_agg$randQuadrature == FALSE &
    results_agg$startValues == "countApprox",
]

# --- Optional RDS caching lines (Commented out) ---
# saveRDS(results_agg_prefiltered, "R/sim-results-agg-prefiltered.rds")
# ---------------------------------------------------

# ==============================================================================
# 3. PLOTTING
# ==============================================================================

# 1. INCREASE FONT SIZE BY ~20-25% (Base size changed from 16 to 20)
theme_set(theme_half_open(font_size = 11))

# Globally inject the grid lines and clean up facets
theme_update(
  panel.grid.major = element_line(color = "gray92", linewidth = 0.5), # Changed size to linewidth to avoid deprecation warnings
  panel.grid.minor = element_blank(),
  strip.background = element_blank()
)

# Define the shared x-axis breaks for Figs 1 & 2
x_breaks_custom <- c(125, 250, 500, 1000, 2500, 5000, 10000)
x_breaks_custom2 <- c(125, 250, 500, 1000, 2000, 5000)

# Define custom hollow shapes for consistency across 2 and 3-level mappings
hollow_shapes <- c(1, 2, 0)
pt_size_custom <- 2.5

# Pass the strictly computed results directly into the plotting framework
# --- Optional RDS caching lines (Commented out) ---
# RES <- readRDS("R/sim-results-agg-prefiltered.rds")
RES <- results_agg_prefiltered

colnames(RES)[match(
  c(
    "J",
    "Genz",
    "N",
    "mmlttime",
    "startValues.time",
    "sigma_mae",
    "beta_mae",
    "abs_rel_ll_diff",
    "abs_rel_ll_diff_approx",
    "mean_abs_gradient",
    "mean_abs_gradient_high_M",
    "mean_abs_gradient_high_M_approx"
  ),
  colnames(RES)
)] <- c(
  "J",
  "M",
  "N",
  "t_M",
  "t_cts",
  "MAE_Sigma",
  "MAE_beta",
  "delta_ell_M",
  "delta_ell_cts",
  "MAG_M",
  "MAG_10_6_M",
  "MAG_10_6_cts"
)

RES <- RES %>%
  mutate(counts = recode(counts, "low" = "low counts", "high" = "high counts"))

RES_J <- RES %>% filter(N == 500, M == 500)
RES_N <- RES %>% filter(J == 10, M == 500)
RES_M <- RES %>% filter(N == 500)

# HELPER FUNCTION: RESHAPE TO LONG FORMAT
prep_long_data <- function(df) {
  df_time <- df %>%
    pivot_longer(
      cols = c(t_M, t_cts),
      names_to = "Method",
      values_to = "Time"
    ) %>%
    mutate(Method = if_else(Method == "t_M", "M", "cts"))

  df_mae <- df %>%
    pivot_longer(
      cols = c(MAE_Sigma, MAE_beta),
      names_to = "Parameter",
      values_to = "MAE"
    ) %>%
    mutate(Parameter = if_else(Parameter == "MAE_Sigma", "\\Sigma", "\\beta"))

  df_delta <- df %>%
    pivot_longer(
      cols = c(delta_ell_M, delta_ell_cts),
      names_to = "Method",
      values_to = "delta_ell"
    ) %>%
    mutate(Method = if_else(Method == "delta_ell_M", "M", "cts"))

  df_mag <- df %>%
    pivot_longer(
      cols = starts_with("MAG"),
      names_to = "Estimator",
      values_to = "MAG"
    ) %>%
    mutate(
      Estimator = case_when(
        Estimator == "MAG_M" ~ "M",
        Estimator == "MAG_10_6_M" ~ "\\hat{\\psi}_M / 10^6",
        Estimator == "MAG_10_6_cts" ~ "\\hat{\\psi}_{cts}"
      )
    )

  list(time = df_time, mae = df_mae, delta = df_delta, mag = df_mag)
}

data_N <- prep_long_data(RES_N)
data_M <- prep_long_data(RES_M)
data_J <- prep_long_data(RES_J)

# BUILD FIGURE 1 (Based on N)
f1_r1 <- ggplot(
  data_N$time,
  aes(x = N, y = Time, shape = Method, linetype = Method)
) +
  geom_point(size = pt_size_custom, alpha = 0.8) +
  geom_line(aes(group = Method), alpha = 0.5) +
  facet_wrap(~counts) +
  scale_shape_manual(
    values = hollow_shapes,
    labels = unname(TeX(c("$\\hat{\\psi}_{cts}$", "$\\hat{\\psi}_M$")))
  ) +
  scale_linetype_discrete(
    labels = unname(TeX(c("$\\hat{\\psi}_{cts}$", "$\\hat{\\psi}_M$")))
  ) +
  labs(
    x = TeX("$N$"),
    y = TeX("Computational time [s]"),
    shape = "Estimating",
    linetype = "Estimating"
  ) +
  scale_x_log10(breaks = x_breaks_custom) +
  scale_y_log10()

f1_r2 <- ggplot(
  data_N$mae,
  aes(x = N, y = MAE, shape = Parameter, linetype = Parameter)
) +
  geom_point(size = pt_size_custom, alpha = 0.8) +
  geom_line(aes(group = Parameter), alpha = 0.5) +
  facet_wrap(~counts) +
  scale_shape_manual(
    values = hollow_shapes,
    labels = unname(TeX(c("$MAE_\\Sigma$", "$MAE_\\beta$")))
  ) +
  scale_linetype_discrete(
    labels = unname(TeX(c("$MAE_\\Sigma$", "$MAE_\\beta$")))
  ) +
  labs(
    x = TeX("$N$"),
    y = TeX("${MAE}$"),
    shape = "Metric",
    linetype = "Metric"
  ) +
  scale_x_log10(breaks = x_breaks_custom) +
  scale_y_log10()

Fig_1 <- f1_r1 / f1_r2

# BUILD FIGURE 2 (Based on M, colored by J)
f2_r1 <- ggplot(
  data_M$time,
  aes(x = M, y = Time, color = as.factor(J), shape = Method, linetype = Method)
) +
  geom_point(size = pt_size_custom, alpha = 0.8) +
  geom_line(aes(group = interaction(J, Method)), alpha = 0.5) +
  facet_wrap(~counts) +
  scale_shape_manual(
    values = hollow_shapes,
    labels = unname(TeX(c("$\\hat{\\psi}_{cts}$", "$\\hat{\\psi}_M$")))
  ) +
  scale_linetype_discrete(
    labels = unname(TeX(c("$\\hat{\\psi}_{cts}$", "$\\hat{\\psi}_M$")))
  ) +
  labs(
    x = TeX("$M$"),
    y = TeX("Computational time [s]"),
    color = TeX("$J$"),
    shape = "Estimating",
    linetype = "Estimating"
  ) +
  scale_x_log10(breaks = x_breaks_custom2) +
  scale_y_log10()

f2_r2 <- ggplot(
  data_M$delta,
  aes(
    x = M,
    y = delta_ell,
    color = as.factor(J),
    shape = Method,
    linetype = Method
  )
) +
  geom_point(size = pt_size_custom, alpha = 0.8) +
  geom_line(aes(group = interaction(J, Method)), alpha = 0.5) +
  facet_wrap(~counts) +
  scale_shape_manual(
    values = hollow_shapes,
    labels = unname(TeX(c(
      "$\\Delta L_{{rel, M}}$",
      "$\\Delta L_{{rel, cts.}}$"
    )))
  ) +
  scale_linetype_discrete(
    labels = unname(TeX(c(
      "$\\Delta L_{{rel, M}}$",
      "$\\Delta L_{{rel, cts.}}$"
    )))
  ) +
  scale_x_log10(breaks = x_breaks_custom2) +
  scale_y_log10(labels = trans_format("log10", math_format(10^.x))) +
  labs(
    x = TeX("$M$"),
    y = TeX("$\\Delta L_{rel}$"),
    color = TeX("$J$"),
    shape = "Metric",
    linetype = "Metric"
  )

f2_r3 <- ggplot(
  data_M$mag,
  aes(
    x = M,
    y = MAG,
    color = as.factor(J),
    shape = Estimator,
    linetype = Estimator
  )
) +
  geom_point(size = pt_size_custom, alpha = 0.8) +
  geom_line(aes(group = interaction(J, Estimator)), alpha = 0.5) +
  facet_wrap(~counts) +
  scale_x_log10(breaks = x_breaks_custom2) +
  scale_y_log10(labels = trans_format("log10", math_format(10^.x))) +
  scale_shape_manual(
    values = hollow_shapes,
    labels = unname(TeX(c(
      "$MAG_{M, \\hat{\\psi}_M}$",
      "$MAG_{10^6, \\hat{\\psi}_M}$",
      "$MAG_{10^6, \\hat{\\psi}_{\\cts.}}$"
    )))
  ) +
  scale_linetype_discrete(
    labels = unname(TeX(c(
      "$MAG_{M, \\hat{\\psi}_M}$",
      "$MAG_{10^6, \\hat{\\psi}_M}$",
      "$MAG_{10^6, \\hat{\\psi}_{\\cts.}}$"
    )))
  ) +
  labs(
    x = TeX("$M$"),
    y = TeX("${MAG}$"),
    color = TeX("$J$"),
    shape = "Metric",
    linetype = "Metric"
  )

Fig_2 <- f2_r1 / f2_r2 / f2_r3

# BUILD FIGURE 3 (Based on J)
f3_r1 <- ggplot(
  data_J$time,
  aes(x = J, y = Time, shape = Method, linetype = Method)
) +
  geom_point(size = pt_size_custom, alpha = 0.8) +
  geom_line(aes(group = Method), alpha = 0.5) +
  scale_shape_manual(
    values = hollow_shapes,
    labels = unname(TeX(c("$\\hat{\\psi}_{cts}$", "$\\hat{\\psi}_M$")))
  ) +
  scale_linetype_discrete(
    labels = unname(TeX(c("$\\hat{\\psi}_{cts}$", "$\\hat{\\psi}_M$")))
  ) +
  scale_y_log10() +
  facet_wrap(~counts) +
  labs(
    x = TeX("$J$"),
    y = TeX("Computa time [s]"),
    shape = "Estimating",
    linetype = "Estimating"
  )

f3_r2 <- ggplot(
  data_J$mae,
  aes(x = J, y = MAE, shape = Parameter, linetype = Parameter)
) +
  geom_point(size = pt_size_custom, alpha = 0.8) +
  geom_line(aes(group = Parameter), alpha = 0.5) +
  facet_wrap(~counts) +
  scale_shape_manual(
    values = hollow_shapes,
    labels = unname(TeX(c("$MAE_{\\Sigma}$", "$MAE_{\\beta}$")))
  ) +
  scale_linetype_discrete(
    labels = unname(TeX(c("$MAE_{\\Sigma}$", "$MAE_{\\beta}$")))
  ) +
  scale_y_log10() +
  labs(x = TeX("$J$"), y = TeX("$MAE$"), shape = "Metric", linetype = "Metric")

Fig_3 <- f3_r1 / f3_r2

# Render Figures
Fig_1
Fig_2
Fig_3
