R/multiple_model.R

In trajeR: Group Based Modeling Trajectory

#' @title Fitting Multivariate Longitudinal Mixture Models
#'
#' @description Fits a multivariate (joint) group-based trajectory model for multiple longitudinal outcomes.
#'
#' @param formulas A list of formulas, one for each longitudinal outcome (e.g., \code{list(Y1 ~ Time, Y2 ~ Time + I(Time^2))}).
#' @param ng A list of integers specifying the number of latent groups for each outcome.
#' @param models A character vector of distribution models for each outcome (e.g., c("CNORM", "LOGIT")). Supported models are "CNORM", "LOGIT", "ZIP", "BETA", "POIS".
#' @param Method The estimation method. "L" for Likelihood or "EM" for Expectation-Maximization.
#' @param Risk An optional one-sided formula for global risk factors that influence group membership (e.g., \code{~ X1 + X2}).
#' @param TCOV Formula or Matrix. Optional time-dependent covariates. 
#' @param data data.frame. The dataset containing all variables.
#' @param degre.nu A list or vector specifying the polynomial degree for the zero-inflation part of "ZIP" models.
#' @param degre.phi A list or vector specifying the polynomial degree for the precision parameter (phi) of "BETA" models.
#' @param ssigma Logical. If TRUE, assumes a common standard deviation across groups for "CNORM" models. Can be a vector of logicals, one for each CNORM model. Default is TRUE.
#' @param ymax A list or vector of maximum values for censoring in "CNORM" models.
#' @param ymin A list or vector of minimum values for censoring in "CNORM" models.
#' @param paraminit An optional list of numeric vectors for user-defined initial parameters. The list should have one element per model, plus an optional last element for the psi parameters.
#' @param hessian Logical. If TRUE, computes the Hessian matrix to estimate standard errors. Default is FALSE.
#' @param control A list of control parameters for the optimization algorithm (e.g., \code{list(maxit=100)}).
#'
#' @return An object of class \code{trajectory.mult} containing the fitted model parameters for each outcome, the interaction parameters (psi), likelihood, and convergence status.
#'
#' @export
trajeR_mult <- function(
  formulas,
  ng,
  models,
  Method = "L",
  Risk = NULL,
  TCOV = NULL,
  data = NULL,
  degre.nu = NULL,
  degre.phi = NULL,
  ssigma = TRUE,
  ymax = NULL,
  ymin = NULL,
  paraminit = NULL,
  control = list(maxit = 100),
  hessian = FALSE
) {
  cli::cli_rule("Fitting the trajeR multiple model")
  step_id <- cli::cli_progress_step("In progress...")

  if (!Method %in% c("L", "EM")) {
    cli::cli_progress_done(id = step_id, result = "failed")
    stop(
      " "
    )
  }

  if (is.function(data)) {
    data <- NULL
  }
  env_eval <- if (is.null(data)) parent.frame() else data
  n_models <- length(formulas)
  n_ind <- NULL
  lY <- list()
  lA <- list()
  lnbeta <- list()
  lnnu <- list()
  lymin <- list()
  lymax <- list()
  lnw <- list()
  lX <- list()
  lparam <- list()
  lvp <- integer(n_models)
  lextra <- integer(n_models)

  for (m in 1:n_models) {
    form_m <- formulas[[m]]
    if (!is.list(form_m)) {
      form_groups <- rep(list(form_m), ng[[m]])
    } else {
      form_groups <- form_m
    }
    current_form_1 <- as.formula(form_groups[[1]])
    Y_mat <- as.matrix(eval(current_form_1[[2]], envir = env_eval))
    A_mat <- NULL

    for (g_idx in 1:ng[[m]]) {
      current_form_g <- as.formula(form_groups[[g_idx]])
      labels_g <- attr(terms(current_form_g), "term.labels")
      if (length(labels_g) > 0) {
        A_mat <- tryCatch(
          {
            as.matrix(eval(parse(text = labels_g[1])[[1]], envir = env_eval))
          },
          error = function(e) NULL
        )
        if (!is.null(A_mat)) break
      }
    }

    if (is.null(A_mat)) {
      A_mat <- matrix(0, nrow = nrow(Y_mat), ncol = ncol(Y_mat))
    }
    if (nrow(Y_mat) == 0 || ncol(Y_mat) == 0) {
      stop(paste("matrix Y of the model", m, "is empty."))
    }
    if (is.null(n_ind)) {
      n_ind <- nrow(Y_mat)
    }

    A_mat <- A_mat[, 1:ncol(Y_mat), drop = FALSE]
    lY[[m]] <- Y_mat
    lA[[m]] <- A_mat

    min_y <- min(as.numeric(Y_mat), na.rm = TRUE)
    max_y <- max(as.numeric(Y_mat), na.rm = TRUE)
    val_ymin <- if (is.null(ymin)) {
      min_y - 1
    } else if (is.list(ymin)) {
      ymin[[m]]
    } else if (length(ymin) >= m) {
      ymin[m]
    } else {
      ymin[1]
    }
    val_ymax <- if (is.null(ymax)) {
      max_y + 1
    } else if (is.list(ymax)) {
      ymax[[m]]
    } else if (length(ymax) >= m) {
      ymax[m]
    } else {
      ymax[1]
    }
    lymin[[m]] <- as.numeric(val_ymin)
    lymax[[m]] <- as.numeric(val_ymax)
    lnw[[m]] <- 0L
    ssigma_m <- if (length(ssigma) > 1) ssigma[[m]] else ssigma

    degre_val <- integer(ng[[m]])
    for (k in 1:ng[[m]]) {
      degre_val[k] <- length(attr(
        terms(as.formula(form_groups[[k]])),
        "term.labels"
      ))
    }
    lnbeta[[m]] <- as.integer(degre_val + 1L)

    d_nu_m <- if (is.list(degre.nu)) {
      degre.nu[[m]]
    } else if (!is.null(degre.nu)) {
      degre.nu
    } else {
      0
    }
    d_phi_m <- if (is.list(degre.phi)) {
      degre.phi[[m]]
    } else if (!is.null(degre.phi)) {
      degre.phi
    } else {
      0
    }
    lnnu[[m]] <- if (models[m] == "ZIP") {
      as.integer(d_nu_m + 1L)
    } else {
      rep(0L, ng[[m]])
    }

    if (models[m] == "CNORM") {
      lextra[m] <- as.integer(ng[[m]])
    } else if (models[m] == "ZIP") {
      lextra[m] <- sum(lnnu[[m]])
    } else if (models[m] %in% c("LOGIT", "POIS")) {
      lextra[m] <- 0L
    } else if (models[m] == "BETA") {
      lextra[m] <- sum(d_phi_m + 1L)
    }

    if (is.null(paraminit)) {
      rep_univ <- tryCatch(
        {
          utils::capture.output(
            res <- trajeR(
              Y = lY[[m]],
              A = lA[[m]],
              ng = ng[[m]],
              degre = as.vector(degre_val),
              Model = models[m],
              Method = "L",
              ssigma = ssigma_m,
              ymax = lymax[[m]],
              ymin = lymin[[m]],
              hessian = FALSE
            )
          )
          res
        },
        error = function(e) NULL
      )

      if (!is.null(rep_univ$theta)) {
        theta_init <- if (length(rep_univ$theta) == ng[[m]]) {
          rep_univ$theta[-1]
        } else {
          rep_univ$theta
        }
      } else if (!is.null(rep_univ$pi)) {
        theta_init <- log(rep_univ$pi[-1] / rep_univ$pi[1])
      } else {
        theta_init <- rep(0, ng[[m]] - 1)
      }

      if (!is.null(rep_univ$beta)) {
        old_beta <- rep_univ$beta
        if (is.matrix(old_beta)) {
          beta_init <- c()
          for (k in 1:ng[[m]]) {
            beta_init <- c(beta_init, old_beta[k, 1:lnbeta[[m]][k]])
          }
        } else if (is.list(old_beta)) {
          beta_init <- c()
          for (k in 1:ng[[m]]) {
            beta_init <- c(beta_init, old_beta[[k]][1:lnbeta[[m]][k]])
          }
        } else {
          beta_init <- as.numeric(old_beta)
        }
      } else {
        beta_init <- rep(0, sum(lnbeta[[m]]))
      }
      beta_init[is.na(beta_init)] <- 0

      if (models[m] == "CNORM") {
        alpha_init <- if (!is.null(rep_univ$sigma)) {
          log(rep_univ$sigma)
        } else {
          rep(0, ng[[m]])
        }
        if (length(alpha_init) == 1) {
          alpha_init <- rep(alpha_init, ng[[m]])
        }
        lparam[[m]] <- as.numeric(c(theta_init, beta_init, alpha_init))
      } else if (models[m] == "ZIP") {
        nu_init <- if (!is.null(rep_univ$nu)) {
          unlist(rep_univ$nu)
        } else {
          rep(0, sum(lnnu[[m]]))
        }
        lparam[[m]] <- as.numeric(c(theta_init, beta_init, nu_init))
      } else if (models[m] %in% c("LOGIT", "POIS")) {
        lparam[[m]] <- as.numeric(c(theta_init, beta_init))
      } else if (models[m] == "BETA") {
        phi_init <- if (!is.null(rep_univ$phi)) {
          unlist(rep_univ$phi)
        } else {
          rep(0, lextra[m])
        }
        lparam[[m]] <- as.numeric(c(theta_init, beta_init, phi_init))
      }
    }
  }

  X_global <- if (is.null(Risk)) {
    matrix(1, nrow = n_ind, ncol = 1)
  } else {
    as.matrix(stats::model.matrix(
      Risk,
      stats::model.frame(Risk, data = env_eval, na.action = stats::na.pass)
    ))
  }
  nx_global <- ncol(X_global)

  if (!is.null(paraminit)) {
    for (m in 1:n_models) {
      lparam[[m]] <- paraminit[[m]]
    }
  }

  for (m in 1:n_models) {
    lX[[m]] <- X_global
    lvp[m] <- ((ng[[m]] - 1) * nx_global) + sum(lnbeta[[m]]) + lextra[m]
  }

  mk <- as.matrix(expand.grid(lapply(ng, function(g) 0:(g - 1))))
  n_psi <- 0
  for (j in 1:(n_models - 1)) {
    for (h in (j + 1):n_models) {
      n_psi <- n_psi + (ng[[j]] - 1) * (ng[[h]] - 1)
    }
  }

  if (n_psi > 0) {
    lparam[[n_models + 1]] <- if (is.null(paraminit)) {
      rep(0, n_psi)
    } else {
      paraminit[[n_models + 1]]
    }
    lvp <- c(lvp, n_psi)
  }

  vparam <- as.numeric(unlist(lparam))
  vp_c <- as.integer(c(0, lvp))
  ln_safe <- lapply(1:n_models, function(x) as.integer(n_ind))
  lnx_safe <- lapply(1:n_models, function(x) as.integer(nx_global))
  lTCOV_safe <- lapply(1:n_models, function(x) {
    matrix(0, nrow = n_ind, ncol = 0)
  })
  lssigma <- if (length(ssigma) == 1) rep(ssigma, n_models) else ssigma

  if (!is.null(paraminit)) {
    expected_length <- sum(lvp)
    provided_length <- length(vparam)
    if (expected_length != provided_length) {
      cli::cli_progress_done(id = step_id, result = "failed")
      stop(sprintf(
        "Error: 'paraminit' contains %d parameters, but your settings require %d.",
        provided_length,
        expected_length
      ))
    }
  }

  optim_control <- list(fnscale = -1)
  if (!is.null(control$maxit)) {
    optim_control$maxit <- control$maxit
  }
  if (!is.null(control$reltol)) {
    optim_control$reltol <- control$reltol
  }
  if (!is.null(control$trace)) {
    optim_control$trace <- control$trace
  }

  if (Method == "L") {
    opt_res <- stats::optim(
      par = vparam,
      fn = likelihoodMult_cpp,
      gr = difLalphaMult_cpp,
      method = "BFGS",
      hessian = FALSE,
      control = optim_control,
      lng = lapply(ng, as.integer),
      lnx = lnx_safe,
      lnbeta = lnbeta,
      ln = ln_safe,
      lA = lA,
      lY = lY,
      lX = lX,
      lnnu = lnnu,
      model = as.character(models),
      lTCOVinit = lTCOV_safe,
      lnw = lnw,
      vp = vp_c,
      mk = mk,
      lymin = lymin,
      lymax = lymax,
      ssigma = lssigma
    )

    sol_vec <- opt_res$par
    final_likelihood <- opt_res$value
    final_convergence <- opt_res$convergence
  } else if (Method == "EM") {
    itermax <- if (!is.null(control$maxit)) control$maxit else 100
    EMIRLS <- if (!is.null(control$EMIRLS)) control$EMIRLS else FALSE

    em_res <- EMMult_cpp(
      lparam = lparam,
      lng = lapply(ng, as.integer),
      lnx = lnx_safe,
      lnbeta = lnbeta,
      lnnu = lnnu,
      ln = ln_safe,
      lA = lA,
      lY = lY,
      lX = lX,
      lymin = lymin,
      lymax = lymax,
      lTCOVinit = lTCOV_safe,
      lnw = lnw,
      mk = mk,
      vp = as.numeric(vp_c),
      model = as.character(models),
      itermax = itermax,
      EMIRLS = EMIRLS
    )

    sol_vec <- as.numeric(unlist(em_res))

    final_likelihood <- likelihoodMult_cpp(
      vparam = sol_vec,
      lng = lapply(ng, as.integer),
      lnx = lnx_safe,
      lnbeta = lnbeta,
      ln = ln_safe,
      lA = lA,
      lY = lY,
      lX = lX,
      lymin = lymin,
      lymax = lymax,
      lTCOVinit = lTCOV_safe,
      lnw = lnw,
      vp = vp_c,
      mk = mk,
      lnnu = lnnu,
      model = as.character(models),
      ssigma = lssigma
    )
    final_convergence <- 0
  }

  se_vec <- rep(NA, length(sol_vec))
  varcov_mat <- NULL

  if (hessian) {
    hess_mat <- tryCatch(
      {
        stats::optimHess(
          par = sol_vec,
          fn = likelihoodMult_cpp,
          gr = difLalphaMult_cpp,
          lng = lapply(ng, as.integer),
          lnx = lnx_safe,
          lnbeta = lnbeta,
          ln = ln_safe,
          lA = lA,
          lY = lY,
          lX = lX,
          lnnu = lnnu,
          model = as.character(models),
          lTCOVinit = lTCOV_safe,
          lnw = lnw,
          vp = vp_c,
          mk = mk,
          lymin = lymin,
          lymax = lymax,
          ssigma = lssigma
        )
      },
      error = function(e) NULL
    )

    if (!is.null(hess_mat)) {
      invH <- tryCatch(solve(-hess_mat), error = function(e) {
        MASS::ginv(-hess_mat)
      })
      varcov_mat <- invH

      var_diag <- diag(invH)
      var_diag[var_diag < 0] <- NA
      raw_SE <- sqrt(var_diag)

      vpcum <- cumsum(vp_c)
      for (m in 1:n_models) {
        th_len <- (ng[[m]] - 1) * nx_global
        b_len <- sum(lnbeta[[m]])
        offset <- vpcum[m] + th_len + b_len

        idx_theta_beta <- (vpcum[m] + 1):offset
        se_vec[idx_theta_beta] <- raw_SE[idx_theta_beta]

        if (models[m] == "CNORM") {
          idx_sigma <- (offset + 1):(offset + lextra[m])
          se_vec[idx_sigma] <- exp(sol_vec[idx_sigma]) * raw_SE[idx_sigma]
        } else if (models[m] %in% c("ZIP", "BETA")) {
          idx_extra <- (offset + 1):(offset + lextra[m])
          se_vec[idx_extra] <- raw_SE[idx_extra]
        }
      }
      if (n_psi > 0) {
        idx_psi <- (vpcum[n_models + 1] + 1):vpcum[n_models + 2]
        se_vec[idx_psi] <- raw_SE[idx_psi]
      }
    }
  }

  vpcum <- cumsum(vp_c)
  results <- list(
    models = list(),
    psi = NULL,
    likelihood = final_likelihood,
    convergence = final_convergence
  )

  for (m in 1:n_models) {
    th_len <- (ng[[m]] - 1) * nx_global
    b_len <- sum(lnbeta[[m]])
    idx_start <- vpcum[m] + 1
    m_parts <- sol_vec[idx_start:vpcum[m + 1]]
    m_se <- if (hessian) {
      se_vec[idx_start:vpcum[m + 1]]
    } else {
      rep(NA, length(m_parts))
    }

    results$models[[paste0("Model_", m)]] <- list(
      theta = m_parts[1:th_len],
      beta = m_parts[(th_len + 1):(th_len + b_len)]
    )

    if (hessian) {
      results$models[[paste0("Model_", m)]]$se_theta <- m_se[1:th_len]
      results$models[[paste0("Model_", m)]]$se_beta <- m_se[
        (th_len + 1):(th_len + b_len)
      ]
    }

    G <- ng[[m]]
    if (G == 1) {
      results$models[[paste0("Model_", m)]]$pi <- 1.0
      if (hessian) results$models[[paste0("Model_", m)]]$se_pi <- 0.0
    } else {
      idx_theta_int <- seq(1, th_len, by = nx_global)
      theta_int <- m_parts[idx_theta_int]
      exp_theta <- exp(c(0, theta_int))
      pi_est <- exp_theta / sum(exp_theta)
      results$models[[paste0("Model_", m)]]$pi <- pi_est

      if (hessian && !is.null(varcov_mat)) {
        global_idx_theta_int <- (vpcum[m] + 0) + idx_theta_int
        cov_theta <- varcov_mat[
          global_idx_theta_int,
          global_idx_theta_int,
          drop = FALSE
        ]
        J <- matrix(0, nrow = G, ncol = G - 1)
        for (i in 1:G) {
          for (j in 1:(G - 1)) {
            k <- j + 1
            if (i == 1) {
              J[i, j] <- -pi_est[1] * pi_est[k]
            } else if (i == k) {
              J[i, j] <- pi_est[i] * (1 - pi_est[i])
            } else {
              J[i, j] <- -pi_est[i] * pi_est[k]
            }
          }
        }
        cov_pi <- J %*% cov_theta %*% t(J)
        var_pi_diag <- diag(cov_pi)
        var_pi_diag[var_pi_diag < 0] <- 0
        results$models[[paste0("Model_", m)]]$se_pi <- sqrt(var_pi_diag)
      }
    }

    offset <- th_len + b_len
    if (models[m] == "CNORM") {
      ssigma_m <- if (length(ssigma) > 1) ssigma[[m]] else ssigma
      if (ssigma_m) {
        results$models[[paste0("Model_", m)]]$sigma <- exp(m_parts[offset + 1])
        if (hessian) {
          results$models[[paste0("Model_", m)]]$se_sigma <- m_se[offset + 1]
        }
      } else {
        results$models[[paste0("Model_", m)]]$sigma <- exp(m_parts[
          (offset + 1):(offset + lextra[m])
        ])
        if (hessian) {
          results$models[[paste0("Model_", m)]]$se_sigma <- m_se[
            (offset + 1):(offset + lextra[m])
          ]
        }
      }
    } else if (models[m] == "ZIP") {
      results$models[[paste0("Model_", m)]]$nu <- m_parts[
        (offset + 1):(offset + sum(lnnu[[m]]))
      ]
      if (hessian) {
        results$models[[paste0("Model_", m)]]$se_nu <- m_se[
          (offset + 1):(offset + sum(lnnu[[m]]))
        ]
      }
    } else if (models[m] == "BETA") {
      results$models[[paste0("Model_", m)]]$phi <- m_parts[
        (offset + 1):(offset + lextra[m])
      ]
      if (hessian) {
        results$models[[paste0("Model_", m)]]$se_phi <- m_se[
          (offset + 1):(offset + lextra[m])
        ]
      }
    }
  }

  if (n_psi > 0) {
    results$psi <- sol_vec[(vpcum[n_models + 1] + 1):vpcum[n_models + 2]]
    if (hessian) {
      results$se_psi <- se_vec[(vpcum[n_models + 1] + 1):vpcum[n_models + 2]]
    }
  }

  results$ng <- ng
  results$model_types <- models
  results$lnbeta <- lnbeta
  results$Method <- Method
  if (hessian && !is.null(varcov_mat)) {
    results$varcov <- varcov_mat
  }

  cli::cli_progress_done(id = step_id)
  cli::cli_rule("Result")
  class(results) <- "trajectory.mult"
  return(results)
}


#' Compute Posterior Probabilities and Group Assignments for Multivariate Trajectory Models
#'
#' @description
#' Calculates the posterior probabilities of group membership for each individual
#' based on a fitted multivariate trajectory model. It assigns each individual to
#' the most likely joint trajectory group and provides the specific sub-group
#' assignments across all dimensions.
#'
#' @param object An object of class \code{trajectory.mult} returned by the \code{trajeR_mult} function.
#' @param lY List of matrices representing the longitudinal responses (one matrix per model). Must match the data used for fitting.
#' @param lA List of matrices representing the time or age of observation (one matrix per model). Must match the data used for fitting.
#' @param Risk An optional one-sided formula for global risk factors, matching the one used for fitting.
#' @param lymin An optional list of minimum boundaries for censoring. If \code{NULL}, it is calculated automatically.
#' @param lymax An optional list of maximum boundaries for censoring. If \code{NULL}, it is calculated automatically.
#' @param lTCOV **(Not implemented)** List of time-dependent covariates.
#' @param ldelta **(Not implemented)** List of matrices indicating the presence of time-dependent covariates.
#' @param lnw **(Not implemented)** List of integers representing weights.
#'
#' @return A list containing the following components:
#' \item{\code{post_prob}}{A matrix of posterior probabilities for each individual and each joint group.}
#' \item{\code{joint_group}}{A numeric vector indicating the most likely joint group for each individual.}
#' \item{\code{group_assignments}}{A matrix showing the sub-group assignment for each joint group.}
#'
#' @export
get_group_probabilities <- function(
  object,
  lY,
  lA,
  Risk = NULL,
  lymin = NULL,
  lymax = NULL,
  lTCOV = NULL,
  ldelta = NULL,
  lnw = NULL
) {
  n_models <- length(object$models)
  n <- nrow(lY[[1]])
  lng <- object$ng
  lnbeta <- object$lnbeta
  lpsi <- object$psi

  ltheta2 <- list()
  lbeta <- list()
  lsigma <- list()

  for (m in 1:n_models) {
    mod_data <- object$models[[m]]
    ltheta2[[m]] <- c(0, mod_data$theta)
    lbeta[[m]] <- mod_data$beta

    if (!is.null(mod_data$sigma)) {
      if (length(mod_data$sigma) == 1) {
        lsigma[[m]] <- rep(mod_data$sigma, lng[[m]])
      } else {
        lsigma[[m]] <- mod_data$sigma
      }
    } else {
      lsigma[[m]] <- numeric(0)
    }
  }

  ltmp_1 <- lapply(lng, function(g) 1:g)
  mk_1 <- as.matrix(expand.grid(ltmp_1))
  mode(mk_1) <- "numeric"

  ltmp_0 <- lapply(lng, function(g) 0:(g - 1))
  mk_0 <- as.matrix(expand.grid(ltmp_0))
  mode(mk_0) <- "numeric"

  if (is.null(lymin)) {
    lymin <- lapply(lY, function(Y) min(Y, na.rm = TRUE) - 1)
  }
  if (is.null(lymax)) {
    lymax <- lapply(lY, function(Y) max(Y, na.rm = TRUE) + 1)
  }
  if (is.null(lnw)) {
    lnw <- rep(list(0L), n_models)
  }
  if (is.null(lTCOV)) {
    lTCOV <- rep(list(NULL), n_models)
  }
  if (is.null(ldelta)) {
    ldelta <- rep(list(NULL), n_models)
  }

  for (m in 1:n_models) {
    lA[[m]] <- as.matrix(lA[[m]])
    mode(lA[[m]]) <- "numeric"
    lY[[m]] <- as.matrix(lY[[m]])
    mode(lY[[m]]) <- "numeric"
  }

  if (is.null(Risk)) {
    X_global <- matrix(1, nrow = n, ncol = 1)
  } else {
    mf_risk <- stats::model.frame(Risk, na.action = stats::na.pass)
    X_global <- as.matrix(stats::model.matrix(Risk, mf_risk))
  }
  mode(X_global) <- "numeric"
  lX <- rep(list(X_global), n_models)

  if (is.null(lpsi) || length(lpsi) == 0) {
    mpsi <- matrix(0, nrow = 1, ncol = 1)
    mode(mpsi) <- "numeric"
  } else {
    mpsi <- mPsi_cpp(as.numeric(lpsi), lng)
  }

  groups_posterior <- matrix(0, nrow = n, ncol = nrow(mk_1))
  groups_prior <- matrix(0, nrow = n, ncol = nrow(mk_1))

  for (i in 1:n) {
    tmp_vec <- numeric(nrow(mk_1))
    for (l in 1:nrow(mk_1)) {
      prod_val <- 1
      for (m in 1:n_models) {
        beta_list <- unname(split(
          lbeta[[m]],
          rep(seq_along(lnbeta[[m]]), lnbeta[[m]])
        ))

        prod_val <- prod_val *
          gkCNORM_cpp(
            beta_list,
            lsigma[[m]],
            i,
            mk_1[l, m],
            lnbeta[[m]],
            lA[[m]],
            lY[[m]],
            lymin[[m]],
            lymax[[m]],
            lTCOV[[m]],
            ldelta[[m]],
            lnw[[m]]
          )
      }

      prior_val <- piikMult_cpp(
        ltheta2,
        mpsi,
        i - 1,
        as.numeric(mk_0[l, ]),
        lng,
        lX,
        mk_0
      )

      groups_prior[i, l] <- prior_val
      tmp_vec[l] <- prod_val * prior_val
    }

    sum_tmp <- sum(tmp_vec)
    if (is.na(sum_tmp) || sum_tmp == 0) {
      sum_tmp <- 1e-300
    }
    groups_posterior[i, ] <- tmp_vec / sum_tmp
  }

  grmul <- apply(groups_posterior, 1, which.max)
  grmul2 <- mk_1[grmul, , drop = FALSE]

  ng_vec <- unlist(lng)
  prob_vec_prior <- colMeans(groups_prior)
  prob_vec_post <- colMeans(groups_posterior)

  joint_probs <- data.frame(mk_1)
  colnames(joint_probs) <- paste0("Model_", 1:n_models)
  joint_probs$Prior_Pct <- round(prob_vec_prior * 100, 5)
  joint_probs$Posterior_Pct <- round(prob_vec_post * 100, 5)

  marginals_prior <- list()
  marginals_post <- list()
  for (m in 1:n_models) {
    marg_pr <- sapply(1:ng_vec[m], function(g) {
      sum(prob_vec_prior[mk_1[, m] == g])
    })
    marg_po <- sapply(1:ng_vec[m], function(g) {
      sum(prob_vec_post[mk_1[, m] == g])
    })

    marginals_prior[[paste0("Model_", m)]] <- round(marg_pr * 100, 5)
    marginals_post[[paste0("Model_", m)]] <- round(marg_po * 100, 5)
  }

  conditionals_prior <- list()
  conditionals_post <- list()

  if (n_models >= 2) {
    for (m_target in 1:n_models) {
      for (m_given in 1:n_models) {
        if (m_target != m_given) {
          cond_mat_prior <- matrix(
            0,
            nrow = ng_vec[m_target],
            ncol = ng_vec[m_given]
          )
          cond_mat_post <- matrix(
            0,
            nrow = ng_vec[m_target],
            ncol = ng_vec[m_given]
          )

          rownames(cond_mat_prior) <- rownames(cond_mat_post) <- paste0(
            "M",
            m_target,
            "_G",
            1:ng_vec[m_target]
          )
          colnames(cond_mat_prior) <- colnames(cond_mat_post) <- paste0(
            "M",
            m_given,
            "_G",
            1:ng_vec[m_given]
          )

          for (g_given in 1:ng_vec[m_given]) {
            p_given_prior <- marginals_prior[[paste0("Model_", m_given)]][
              g_given
            ] /
              100
            p_given_post <- marginals_post[[paste0("Model_", m_given)]][
              g_given
            ] /
              100

            for (g_target in 1:ng_vec[m_target]) {
              idx_matching <- (mk_1[, m_target] == g_target &
                mk_1[, m_given] == g_given)

              p_joint_prior <- sum(prob_vec_prior[idx_matching])
              p_joint_post <- sum(prob_vec_post[idx_matching])

              cond_mat_prior[g_target, g_given] <- if (p_given_prior > 0) {
                round((p_joint_prior / p_given_prior) * 100, 5)
              } else {
                0
              }
              cond_mat_post[g_target, g_given] <- if (p_given_post > 0) {
                round((p_joint_post / p_given_post) * 100, 5)
              } else {
                0
              }
            }
          }

          label <- paste0("Model_", m_target, "_given_Model_", m_given)
          conditionals_prior[[label]] <- cond_mat_prior
          conditionals_post[[label]] <- cond_mat_post
        }
      }
    }
  }

  return(list(
    post_prob = groups_posterior,
    prior_prob = groups_prior,
    joint_group = grmul,
    group_assignments = grmul2,
    Joint_Table = joint_probs,
    Marginal_Prior = marginals_prior,
    Marginal_Posterior = marginals_post,
    Conditional_Posterior = conditionals_post
  ))
}

#' Plot Multivariate Trajectories
#'
#' @description
#' Plots a grid of longitudinal trajectories for a multivariate model.
#' The grid displays each outcome (rows) and each latent group (columns).
#' Individual trajectories are colored by their assigned group and faded,
#' with the specific group highlighted and its average polynomial curve overlaid.
#'
#' @param object An object of class \code{trajectory.mult}.
#' @param lY List of matrices representing the longitudinal responses.
#' @param lA List of matrices representing the time or age of observation.
#' @param assignments A matrix of group assignments (e.g., the \code{group_assignments} output from \code{get_group_probabilities}).
#' @param base_colors Optional vector of hex colors for the groups.
#'
#' @export

plot_trajectory_mult <- function(
  object,
  lY,
  lA,
  assignments,
  base_colors = NULL
) {
  n_models <- length(object$models)
  ng_max <- max(unlist(object$ng))

  if (is.null(base_colors)) {
    base_colors <- rep("#000000", length.out = ng_max)
  } else {
    if (ng_max > length(base_colors)) {
      base_colors <- grDevices::colorRampPalette(base_colors)(ng_max)
    }
  }

  eval_pol <- function(beta, t_vals) {
    res <- rep(0, length(t_vals))
    for (d in seq_along(beta)) {
      res <- res + beta[d] * (t_vals^(d - 1))
    }
    return(res)
  }

  old_par <- par(
    mfrow = c(n_models, ng_max),
    oma = c(5, 4, 2, 0) + 0.1,
    mar = c(2, 2, 1, 1) + 0.1
  )
  on.exit(par(old_par))

  for (j in 1:n_models) {
    Y_mat <- as.matrix(lY[[j]])
    A_mat <- as.matrix(lA[[j]])
    mod_data <- object$models[[j]]
    lnbeta_j <- object$lnbeta[[j]]
    ng_j <- object$ng[[j]]

    beta_list <- unname(split(
      mod_data$beta,
      rep(seq_along(lnbeta_j), lnbeta_j)
    ))

    y_min <- min(Y_mat, na.rm = TRUE)
    y_max <- max(Y_mat, na.rm = TRUE)
    t_min <- min(A_mat, na.rm = TRUE)
    t_max <- max(A_mat, na.rm = TRUE)

    for (k in 1:ng_max) {
      if (k > ng_j) {
        plot.new()
        next
      }

      cols_faded <- paste0(base_colors, "15")
      cols_faded[k] <- paste0(base_colors[k], "80")

      plot(
        A_mat[1, ],
        Y_mat[1, ],
        ylim = c(y_min, y_max),
        xlim = c(t_min, t_max),
        type = "n",
        axes = FALSE,
        xlab = "",
        ylab = ""
      )

      axis(1)
      if (k == 1) {
        axis(2)
      }
      box(bty = "l")
      grid()

      for (i in 1:nrow(Y_mat)) {
        grp_assigned <- assignments[i, j]
        lines(A_mat[i, ], Y_mat[i, ], col = cols_faded[grp_assigned])
      }

      t_seq <- seq(t_min, t_max, length.out = 100)
      y_seq <- eval_pol(beta_list[[k]], t_seq)
      lines(t_seq, y_seq, col = base_colors[k], lwd = 3)

      if (j == 1) {
        mtext(paste("Group", k), side = 3, line = 0.5, cex = 0.9, font = 2)
      }
      if (k == 1) {
        mtext(paste("Outcome", j), side = 2, line = 2.5, cex = 0.9, font = 2)
      }
    }
  }
}