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

# ==============================================================================
# 1. SETUP & LIBRARIES
# ==============================================================================
# check if variable run_hmsc exists, if not set to FALSE
if (!exists("run_hmsc")) {
  run_hmsc <- FALSE
  message(
    "Variable `run_hmsc` not found. Set `run_hmsc <- TRUE` to run HMSC models. This will take several days and is turned off by default."
  )
}

if (!exists("Genz_M_Monte_Carlo")) {
  Genz_M_Monte_Carlo <- 50
  message(
    "Variable `Genz_M_Monte_Carlo` not found. Default to 50 for faster computation (~2min). Set to 250 to match the manuscript (~ 10min)."
  )
}

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
    ))
  }
}

check_min_version("cotram", "0.6-1")
check_min_version("mvtnorm", "1.4-0")
check_min_version("tram", "1.4-2")
check_min_version("mlt", "1.8-0")
check_min_version("TH.data", "1.1-5")
check_min_version("viridis", "0.6.2")
library(tram)

# ==============================================================================
# 2. DATA LOADING & PREPROCESSING
# ==============================================================================
# load aquabirds
if (requireNamespace("TH.data", quietly = TRUE)) {
  load(system.file("rda/aquabirds.rda", package = "TH.data"))
} else {
  stop(
    "Cannot find aquabirds data!
    Install `TH.data` package via install.packages('TH.data')"
  )
}

# add day of Year
aquabirds$dayY <- as.numeric(format(aquabirds$Date, "%j"))
# add factor Year
aquabirds$Year <- as.factor(format(aquabirds$Date, "%Y"))
# Add cyclic basis for day of Year
cb <- cyclic_basis(numeric_var("dayY"), order = 4, frequency = 365)
XTag <- as.data.frame(model.matrix(cb, data = aquabirds))
colnames(XTag) <- paste0("t", 1:ncol(XTag))
tvars <- paste0("t", 1:ncol(XTag))
toy <- paste0(tvars, collapse = "+") #time of Year
# final Data D -- one can remove below
aquabirds <- cbind(aquabirds, XTag)
aquabirds$GreatCrestedGrebe1p <- 1 + aquabirds$GreatCrestedGrebe
aquabirds$GreatCormorant1p <- 1 + aquabirds$GreatCormorant
aquabirds$Goosander1p <- 1 + aquabirds$Goosander
rm(XTag, cb)

cor.cols <- c("#1E88E5", "#DCA707", "#DE1760")

## newdata
cb <- cyclic_basis(numeric_var("dayY"), order = 4, frequency = 365)
nd <- data.frame(dayY = 1:365)
X <- model.matrix(cb, data = nd)
colnames(X) <- paste0("t", 1:ncol(X))
nd <- cbind(nd, as.data.frame(X))

# Dates for x-axis
ats <- c(
  min(aquabirds$Date),
  "2003-01-01",
  "2004-01-01",
  "2005-01-01",
  "2006-01-01",
  "2007-01-01",
  "2008-01-01",
  "2009-01-01",
  "2010-01-01",
  "2011-01-01",
  "2012-01-01",
  "2013-01-01",
  "2014-01-01",
  "2015-01-01",
  "2016-01-01",
  max(aquabirds$Date)
)
labs <- c(
  "05.2002",
  " ",
  "2004",
  " ",
  "2006",
  " ",
  "2008",
  " ",
  "2010",
  " ",
  "2012",
  " ",
  "2014",
  " ",
  "2016",
  "11.2016"
)
xtick <- cumsum(c(1, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31))
xlabs <- c(
  "Jan",
  "Feb",
  "Mar",
  "Apr",
  "May",
  "Jun",
  "Jul",
  "Aug",
  "Sep",
  "Oct",
  "Nov",
  "Dec",
  "Jan"
)

# colors with alpha:
qs <- c(seq(5, 95, by = 5)) / 100
qs <- c(0.01, qs, 0.99)
col1 <- viridis::plasma((length(qs) + 1) / 2)
cols_quantiles <- c(col1, rev(col1[-length(col1)]))
alphas <- c(qs * I(qs <= 0.5) + (1 - qs) * I(qs > 0.5))
alphas <- alphas + 0.08


# ==============================================================================
# 3. HELPER FUNCTIONS
# ==============================================================================

# 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
}

s_rho <- function(x) {
  return(6 * asin(x / 2) / pi)
}

plot_corr.df <- function(
  df,
  newplot = TRUE,
  main = NULL,
  legend = TRUE,
  lwd = NULL,
  ylab = expression(paste("Spearman\'s  ", "", rho)),
  ...
) {
  if (is.null(lwd)) {
    lwd <- 2.5
  }

  df <- as.data.frame(df)
  if (!"dayY" %in% colnames(df) && nrow(df) == 365) {
    df$dayY <- 1:nrow(df)
  }
  stopifnot(
    c(
      "dayY",
      "GreatCormorant.GreatCrestedGrebe",
      "Goosander.GreatCrestedGrebe",
      "Goosander.GreatCormorant"
    ) %in%
      colnames(df)
  )
  if (newplot) {
    plot(
      df$dayY,
      df$GreatCormorant.GreatCrestedGrebe,
      ylim = c(-1, 1),
      type = "l",
      xlab = "",
      ylab = ylab,
      las = 1,
      col = cor.cols[1],
      main = main,
      xaxt = "n",
      lwd = lwd,
      ...
    )
    axis(side = 1, at = xtick, labels = xlabs)
  } else {
    lines(
      df$dayY,
      df$GreatCormorant.GreatCrestedGrebe,
      col = cor.cols[1],
      lwd = lwd,
      ...
    )
  }
  lines(
    df$dayY,
    df$Goosander.GreatCrestedGrebe,
    col = cor.cols[2],
    lwd = lwd,
    ...
  )
  lines(df$dayY, df$Goosander.GreatCormorant, col = cor.cols[3], lwd = lwd, ...)
  abline(h = 0, lty = 2, col = "grey")
  if (newplot && legend) {
    legend(
      "bottomleft",
      col = cor.cols,
      lty = 1,
      lwd = lwd,
      legend = c(
        "Great Crested Grebe - Great Cormorant",
        "Great Crested Grebe - Goosander",
        "Great Cormorant - Goosander"
      ),
      bty = "n"
    )
  }
}

plot_corr.Mmlt <- function(Mmlt_obj, main = NULL, ...) {
  stopifnot(inherits(Mmlt_obj, "Mmlt"))
  tmp <- cbind(nd, t(unclass(coef(Mmlt_obj, newdata = nd, type = "Spearman"))))

  if (any(grepl("^R\\.", colnames(tmp)))) {
    colnames(tmp) <- gsub("..", ".", colnames(tmp), fixed = TRUE)
    colnames(tmp) <- gsub("R.", "", colnames(tmp), fixed = TRUE)
    colnames(tmp) <- gsub("\\.$", "", colnames(tmp))
  }

  colnames(tmp) <- gsub("GreatCormorant1p", "GreatCormorant", colnames(tmp))
  colnames(tmp) <- gsub(
    "GreatCrestedGrebe1p",
    "GreatCrestedGrebe",
    colnames(tmp)
  )
  colnames(tmp) <- gsub("Goosander1p", "Goosander", colnames(tmp))

  plot_corr.df(tmp, main = main, ...)
}

# add bootstrap lines to plot_corr.Mmlt
# one entry of df_list is one model
plot_corr.boot.Mmlt <- function(df_list, lwd = 0.7, alpha = 0.7, ...) {
  cor.cols <- adjustcolor(cor.cols, alpha = 0.3)
  lapply(df_list, function(df) {
    lines(
      df$dayY,
      df$GreatCormorant.GreatCrestedGrebe,
      col = cor.cols[1],
      lwd = lwd,
      ...
    )
    lines(
      df$dayY,
      df$Goosander.GreatCrestedGrebe,
      col = cor.cols[2],
      lwd = lwd,
      ...
    )
    lines(
      df$dayY,
      df$Goosander.GreatCormorant,
      col = cor.cols[3],
      lwd = lwd,
      ...
    )
  })
  invisible(NULL)
}

#' Plot aquabirds density predictions from tram model
#' @param gather_years Logical, gather data from all years? reasonable if tram
#'                      model has no year effect
plot_aquabirds_denisty.tram <- function(model, gather_years = FALSE, ...) {
  stopifnot(inherits(model, "tram"))

  pred <- predict(as.mlt(model), type = "quantile", smooth = TRUE, prob = qs)
  if (inherits(pred, "response")) {
    pred <- matrix(
      pred$approx,
      nrow = length(qs),
      dimnames = list(prob = as.character(qs), NULL)
    )
  }
  if (!is.null(model$log_first) && model$log_first) {
    pred <- pred - 1
  }

  x_axis <- if (gather_years) aquabirds$dayY else aquabirds$Date

  # discrete count quantiles
  pred <- floor(pred)

  # data points
  plot(
    x_axis,
    model$response$approxy + 1,
    log = "y",
    ylab = "counts",
    xaxt = "n",
    yaxt = "n",
    xlab = " ",
    col = rgb(.1, .1, .1, .35),
    las = 1,
    ...
  )

  desired_labels <- c(0, 1, 2, 5, 10, 20, 50, 100, 200)
  axis(side = 2, at = desired_labels + 1, labels = desired_labels, las = 1)

  if (gather_years) {
    axis(side = 1, at = xtick, labels = xlabs)
  } else {
    axis(side = 1, at = ats, labels = labs)
  }

  # quantile lines
  ind <- if (gather_years) {
    aquabirds$Year == "2003"
  } else {
    rep(TRUE, nrow(aquabirds))
  }
  for (i in seq_along(qs)) {
    lines(
      x_axis[ind],
      pred[i, ind] + 1,
      col = adjustcolor(cols_quantiles[i], alpha = alphas[i] * 2.5),
      lwd = 1.3
    )
  }
}

plot_triple <- function(h, k, s, gather_years = FALSE) {
  # get marginal models first, if Mmlt provided
  if (inherits(h, "mmlt")) {
    m <- h
    MMs <- get_Mmodels(h)
    return(do.call(plot_triple, c(unname(MMs), gather_years = gather_years)))
  }

  # Set up layout with 4 rows, where the legend (4th row) is smaller
  layout_matrix <- matrix(c(1, 2, 3, 4), nrow = 4, ncol = 1)
  layout(layout_matrix, heights = c(1, 1, 1, 0.36)) # Legend gets 0.36 rel-height

  # Set margins for the main plots
  mymai <- c(1, 1, 1, 1)
  par(mai = c(1.02, 0.82, 0.82, 0.42))
  par(mai = c(.3, .6, .3, .3))

  # Plot the three main plots
  plot_aquabirds_denisty.tram(
    h,
    main = "Great Crested Grebe",
    gather_years = gather_years
  )
  plot_aquabirds_denisty.tram(
    k,
    main = "Great Cormorant",
    gather_years = gather_years
  )
  plot_aquabirds_denisty.tram(
    s,
    main = "Goosander",
    gather_years = gather_years
  )

  # Add the legend with smaller margins
  par(mai = c(0, 0.6, 0, 0.6)) # Smaller top/bottom margins for legend
  # smaller spaces between box and title/x-labels

  add_leg()

  # Reset layout
  layout(1)
}

# quantile-legend
add_leg <- function(x = TRUE) {
  xl <- 1
  yb <- 0.1
  xr <- 1.2
  yt <- 1.9

  n. <- length(qs)

  plot(
    NA,
    type = "n",
    ann = FALSE,
    xlim = c(0, 2),
    ylim = c(0, 2),
    xaxt = "n",
    yaxt = "n",
    bty = "n"
  )
  rect(
    head(seq(yb, yt, (yt - yb) / n.), -1),
    xl,
    tail(seq(yb, yt, (yt - yb) / n.), -1),
    xr,
    col = diag(sapply(alphas, function(alpha) {
      adjustcolor(col = cols_quantiles, alpha = alpha * 2.5)
    }))
  )

  mtext(
    qs,
    side = 1,
    at = tail(seq(yb, yt, (yt - yb) / n.), -1) - 0.05,
    las = 1,
    cex = 0.5,
    line = -2.5
  )
  mtext("Quantiles", side = 3, line = -2, cex = .7)
}

plot_aquabirds_data <- function(...) {
  desired_labels <- c(0, 1, 2, 5, 10, 20, 50, 100, 200)

  op <- par(mfrow = c(3, 1), mai = c(0.35, 0.6, 0.3, 0.1))

  plot(
    I(GreatCrestedGrebe + 1) ~ Date,
    main = "Great Crested Grebe",
    xaxt = "n",
    yaxt = "n",
    log = "y",
    ylab = "counts",
    ...
  )
  axis(side = 2, at = desired_labels + 1, labels = desired_labels, las = 1)
  axis(side = 1, at = ats, labels = labs)
  abline(v = ats, lty = 2, col = "lightgray")

  plot(
    I(GreatCormorant + 1) ~ Date,
    main = "Great Cormorant",
    xaxt = "n",
    yaxt = "n",
    log = "y",
    ylab = "counts",
    ...
  )
  axis(side = 2, at = desired_labels + 1, labels = desired_labels, las = 1)
  axis(side = 1, at = ats, labels = labs)
  abline(v = ats, lty = 2, col = "lightgray")

  plot(
    I(Goosander + 1) ~ Date,
    main = "Goosander",
    xaxt = "n",
    yaxt = "n",
    log = "y",
    ylab = "counts",
    ...
  )
  axis(side = 2, at = desired_labels + 1, labels = desired_labels, las = 1)
  axis(side = 1, at = ats, labels = labs)
  abline(v = ats, lty = 2, col = "lightgray")

  par(op)
}


# function to add confidence intervals to existing correlation plot
add_cor_CI <- function(samples, alpha = 0.25, cols = cor.cols) {
  stopifnot(
    class(samples) == "list",
    length(samples) == 3,
    all(sapply(samples, is.matrix)),
    all(sapply(samples, nrow) == 365)
  )

  for (i in 1:3) {
    rho <- s_rho(samples[[i]])
    q_low <- apply(rho, MARGIN = 1, FUN = function(x) {
      quantile(x, prob = 0.025)
    })
    q_upp <- apply(rho, MARGIN = 1, FUN = function(x) {
      quantile(x, prob = 0.975)
    })
    polygon(
      c(1:365, 365:1),
      c(q_low, rev(q_upp)),
      col = adjustcolor(cor.cols[i], alpha = alpha),
      border = NA
    )
  }
}

# Wald/asymptotic parametric bootstrap confidence bands for correlations
### sampling nsamp values from the asymptotic distribution of the parameters
sample.parboot <- function(mod, nsamp = 200) {
  set.seed(1)
  H <- Hessian(mod)
  H <- (H + t(H)) / 2 # ensure symmetry

  V <- solve(H)
  P <- rmvnorm(nsamp, mean = coef(mod), sigma = (V + t(V)) / 2)
  m_tmp <- mod
  CR <- vector(mode = "list", length = nrow(P))

  # setup temp-model and retrieve corr via coef()
  for (i in 1:nsamp) {
    m_tmp$par[] <- P[i, ]
    CR[[i]] <- t(unclass(coef(m_tmp, newdata = nd, type = "Corr")))
  }

  # save results to matricies
  r21 <- r31 <- r32 <- matrix(NA, nrow = nrow(nd), ncol = nsamp)
  for (l in 1:nsamp) {
    r21[, l] <- CR[[l]][, 1]
    r31[, l] <- CR[[l]][, 2]
    r32[, l] <- CR[[l]][, 3]
  }

  return(list(r21, r31, r32))
}


#' refit models on new data
refit <- function(newdata, m_m, m_h, m_k, m_s) {
  # first 3 marginal models
  stopifnot(inherits(m_m, "mmlt"))
  cl_h <- m_h$call
  cl_h$data <- quote(newdata)
  cl_h$dofit <- TRUE
  cl_h$theta <- quote(m_h$theta)
  m_h_new <- eval(cl_h)

  cl_k <- m_k$call
  cl_k$data <- quote(newdata)
  cl_k$dofit <- TRUE
  cl_k$theta <- quote(m_k$theta)
  m_k_new <- eval(cl_k)

  cl_s <- m_s$call
  cl_s$data <- quote(newdata)
  cl_s$dofit <- TRUE
  cl_s$theta <- quote(m_s$theta)
  m_s_new <- eval(cl_s)

  # now joint model
  cl_m <- m_m$call
  cl_m[2:4] <- list(quote(m_h_new), quote(m_k_new), quote(m_s_new))
  cl_m$data <- quote(newdata)
  cl_m$theta <- quote(m_m$theta)
  # faster optim without hessian
  args <- list(M = 200, seed = 1, type = "ghalton")
  op <- mltoptim(
    trace = FALSE,
    nloptr = list(algorithm = "NLOPT_LD_MMA", xtol_rel = 1.0e-3, maxeval = 50L)
  )[3:1]
  cl_m$args <- quote(args)
  cl_m$optim <- quote(op)

  return(eval(cl_m))
}

get_kendall_df <- function(window_length) {
  d <- 1:max(aquabirds$dayY)
  wd <- data.frame(
    dayY = d,
    GreatCormorant.GreatCrestedGrebe = NA,
    Goosander.GreatCrestedGrebe = NA,
    Goosander.GreatCormorant = NA
  )

  for (i in d) {
    tmp <- subset(
      aquabirds,
      pmin((i - aquabirds$dayY) %% 365, (aquabirds$dayY - i) %% 365) <=
        window_length / 2
    )
    CR <- suppressWarnings(cor(
      tmp[, c("GreatCrestedGrebe", "GreatCormorant", "Goosander")],
      meth = "spearman",
      use = "pairwise.complete.obs"
    ))
    wd[
      i,
      c(
        "GreatCormorant.GreatCrestedGrebe",
        "Goosander.GreatCrestedGrebe",
        "Goosander.GreatCormorant"
      )
    ] <- CR[lower.tri(CR)]
  }
  wd
}


# ==============================================================================
# 4. MODEL FITS
# ==============================================================================
m_h <- BoxCox(
  R(GreatCrestedGrebe) ~ t1 +
    t2 +
    t3 +
    t4 +
    t5 +
    t6 +
    t7 +
    t8 |
    t1 + t2 + t3 + t4 + t5 + t6 + t7 + t8,
  data = aquabirds,
  LRtest = FALSE,
  na.action = na.pass
)
m_k <- BoxCox(
  R(GreatCormorant) ~ t1 +
    t2 +
    t3 +
    t4 +
    t5 +
    t6 +
    t7 +
    t8 |
    t1 + t2 + t3 + t4 + t5 + t6 + t7 + t8,
  data = aquabirds,
  LRtest = FALSE,
  na.action = na.pass
)
m_s <- BoxCox(
  R(Goosander) ~ t1 +
    t2 +
    t3 +
    t4 +
    t5 +
    t6 +
    t7 +
    t8 |
    t1 + t2 + t3 + t4 + t5 + t6 + t7 + t8,
  data = aquabirds,
  LRtest = FALSE,
  na.action = na.pass
)

m_m <- Mmlt(
  m_h,
  m_k,
  m_s,
  formula = ~ t1 + t2 + t3 + t4 + t5 + t6 + t7 + t8,
  data = aquabirds,
  args = list(M = Genz_M_Monte_Carlo, seed = 1, type = "ghalton")
)

Mods <- list(m_h = m_h, m_k = m_k, m_s = m_s, m_m = m_m)


# ==============================================================================
# 5. MAIN PLOTTING
# ==============================================================================
plot_aquabirds_data(
  data = aquabirds,
  xlab = "",
  pch = 1,
  cex = .6,
  las = 1,
  col = rgb(0, 0, 0, .25)
)

plot_triple(Mods[["m_m"]], gather_years = TRUE)

op <- par(
  mfrow = c(3, 1),
  adj = 0,
  mai = c(1.02 / 3, 0.82 / 1.5, 0.82 / 3, 0.42 / 2)
)

df <- get_kendall_df(window_length = 21)
plot_corr.df(
  df,
  ylab = expression(paste("Spearman\'s  ", rho)),
  main = "(a) Empirical trajectories"
)

plot_corr.Mmlt(
  Mods[["m_m"]],
  main = "(b) JCTM based trajectories",
  legend = FALSE
)
add_cor_CI(sample.parboot(Mods[["m_m"]]), alpha = 0.15)
plot_corr.Mmlt(Mods[["m_m"]], newplot = FALSE)


# ==============================================================================
# 6. HMSC CONDITIONAL BLOCK
# ==============================================================================
if (run_hmsc) {
  fraction <- 2
  samples <- 100000 / fraction
  transient <- 200000 / fraction
  thin <- 1
  nChains <- 4

  library("Hmsc")

  Sys.setenv(OPENBLAS_NUM_THREADS = 1)
  Sys.setenv(OMP_NUM_THREADS = 1)
  set.seed(270161)

  X <- model.matrix(as.formula(paste("~ Year + ", toy)), data = aquabirds)
  Y <- as.matrix(
    aquabirds[, c("GreatCrestedGrebe", "GreatCormorant", "Goosander")]
  )

  set.seed(270161)
  studyDesign <- data.frame(
    sample = as.factor(1:nrow(X)),
    dayY = as.factor(aquabirds$dayY)
  )

  rL2 <- HmscRandomLevel(
    xData = data.frame(
      intercept = rep(1, nrow(aquabirds)),
      t1 = aquabirds$t1,
      t2 = aquabirds$t2,
      t3 = aquabirds$t3,
      t4 = aquabirds$t4,
      t5 = aquabirds$t5,
      t6 = aquabirds$t6,
      t7 = aquabirds$t7,
      t8 = aquabirds$t8
    )
  )

  ## computeAssociations for complex model
  ## --> HMSC book: equation 7.12
  get_hmsc_omega_cor_avg_and_samples_for_CIs <- function(m_var) {
    OmegaCor <- vector("list", m_var$nr)
    start <- 1
    postList <- poolMcmcChains(m_var$postList, start = start, thin = m_var$thin)
    # get lambda array from postlist
    Lambdas <- array(
      NA,
      dim = c(
        dim(postList[[1]]$Lambda[[1]])[1],
        dim(postList[[1]]$Lambda[[1]])[2],
        dim(postList[[1]]$Lambda[[1]])[3],
        length(postList)
      ), # = #MCMC-samples
      dimnames = list(
        NULL,
        NULL,
        covariate = names(rL2$x),
        sample = paste0("MCMCsample", 1:length(postList))
      )
    )
    # Fill the array
    for (i in seq_along(postList)) {
      Lambdas[,,, i] <- postList[[i]]$Lambda[[1]]
    }
    # get an array of dim [1:2, 1:3, 1:nrow(xData), 1:8]
    #   where we multiply each Lambda[i,j,,k] with its xData[,m]
    ind <- match(1:365, aquabirds$dayY)
    tmp <- as.matrix(rL2$x)[ind, ] # design matrix for dayY
    LamX <- apply(Lambdas, c(1, 2, 4), function(a) c(sample.nr = tmp %*% a)) |>
      aperm(c(2, 3, 1, 4))
    Omega <- apply(LamX, 3:4, crossprod) # first dim (length 9) is 3x3 matrix
    Omega_cor <- apply(Omega, 2:3, \(x) cov2cor(matrix(x, 3, 3)))[
      c(2, 3, 6),
      ,
    ] # only unique pairs
    dimnames(Omega_cor)[1:2] <- list(
      pair = c(
        "GreatCormorant.GreatCrestedGrebe",
        "Goosander.GreatCrestedGrebe",
        "Goosander.GreatCormorant"
      ),
      dayY = 1:365
    )

    # average over MCMC samples
    hmsc_omega_cor_avg <- apply(Omega_cor, 1:2, mean) |> t()
    stopifnot(
      c(
        "GreatCormorant.GreatCrestedGrebe",
        "Goosander.GreatCrestedGrebe",
        "Goosander.GreatCormorant"
      ) %in%
        colnames(hmsc_omega_cor_avg)
    )

    set.seed(1)
    subset_ind <- sample(
      1:dim(Omega_cor)[3],
      size = min(4000, dim(Omega_cor)[3])
    )
    # get sample replicates for CIs
    hmsc_samples_for_CIs <- list(
      HK = Omega_cor[1, , subset_ind], # GreatCormorant - GreatCrestedGrebe
      HS = Omega_cor[2, , subset_ind], # Goosander - GreatCrestedGrebe
      KS = Omega_cor[3, , subset_ind] # Goosander - GreatCormorant
    )
    list(
      hmsc_omega_cor_avg = hmsc_omega_cor_avg,
      hmsc_samples_for_CIs = hmsc_samples_for_CIs
    )
  }

  # function for a posterori analysis of convergence for Lambda (covariate-dependent Omega)
  get_mcmc_list_Lambda <- function(m_var) {
    LAMBDAS <- list()
    mpost <- list()
    nChains <- length(m_var$postList)

    postLambda <- vector("list", nChains)
    for (chain in 1:nChains) {
      tmp1 <- lapply(m_var$postList[[chain]], \(x) as.vector(x$Lambda[[1]]))
      tmp2 <- do.call(rbind, tmp1)
      postLambda[[chain]] = mcmc(tmp2, thin = 1, start = 1, end = nrow(tmp2))
    }
    mpost <- list()
    mcmc.list(postLambda)
  }

  m0 <- Hmsc(
    Y = Y,
    X = X,
    distr = "lognormal poisson",
    studyDesign = studyDesign,
    ranLevels = list(sample = rL2)
  )

  system.time(
    m_var <- sampleMcmc(
      m0,
      thin = thin,
      samples = samples,
      transient = transient,
      nChains = nChains,
      nParallel = nChains
    )
  )

  # Process and plot directly in memory (no RDS)
  obj <- get_hmsc_omega_cor_avg_and_samples_for_CIs(m_var)

  plot_corr.df(
    obj$hmsc_omega_cor_avg,
    ylab = "Latent residual correlation",
    main = "(c) HMSC based trajectories (log-normal)",
    legend = FALSE
  )
  add_cor_CI(obj$hmsc_samples_for_CIs)
} else {
  # dummy empty plot if HMSC not run
  plot(
    NA,
    NA,
    xlim = c(1, 365),
    ylim = c(-1, 1),
    xlab = "",
    ylab = expression(paste("Spearman\'s  ", rho)),
    main = "(c) HMSC based trajectories (log-normal)",
    xaxt = "n",
    yaxt = "n",
    bty = "n"
  )
}

par(op)


# pkg/cotram/inst/simulation/SIM-JCTM.R
message("For results of the simulation study, please run: \n`source(system.file('simulation', 'SIM-JCTM.R', package = 'cotram'))`\n! This will take a few/several days on a Workstation")
