################################################################################
## Evaluator Functions
################################################################################

#' Compute glmmTMB covariance matrix
#'
#' @description Computes the heterogeneous compound symmetry covariance matrix
#'   from the glmmTMB wrapper function.
#'
#' @param fit A glmmTMB fit object.
#'
#' @return The compound symmetry covariance matrix estimate obtained by glmmTMB.
compute_glmmtmb_covar_mat <- function(fit) {
  ## extract components of heterogeneous compound symmetry covariance matrix
  ## estimate
  corr_fit_attributes <- attributes(glmmTMB::VarCorr(fit)[["cond"]]$participant)

  ## get the correlation matrix: off-diagonals are the correlation estimate,
  ## diagonals are equal to one
  corr_mat <- corr_fit_attributes$correlation

  ## compute the unweighted covariance matrix
  unweighted_covar_mat <- outer(
    corr_fit_attributes$stddev,
    corr_fit_attributes$stddev
  )

  ## haddamard product of correlation coefficient with unweighted covariance
  ## matrix
  covar_mat <- corr_mat * unweighted_covar_mat

  return(covar_mat)
}

#' Compute PROC MIXED covariance matrix from CovParams
#'
#' @description Computes the heterogeneous compound symmetry covariance matrix
#'   from the CovParams output by PROC MIXED by way of sasr.
#'
#' @param fit A data.frame returned by the proc_mixed_fun() PROC MIXED wrapper
#'   function. Assuming p repeated measures, the first p rows of this dataframe
#'   contain the variance estimates of these repeated emasures. The last row
#'   contains the heterogeneous compound symmetry correlation estimate.
#'
#' @return The compound symmetry covariance matrix estimate returned by PROC
#'   MIXED.
compute_sasr_covar_mat <- function(fit) {
  ## number of repeated measures
  num_rep_meas <- nrow(fit) - 1

  ## compute the covariance matrix estimate, not accounting for correlation
  var_ests <- fit$Estimate[1:num_rep_meas]
  std_devs <- sqrt(var_ests)
  unweighted_covar_mat <- outer(std_devs, std_devs)

  ## weight covariance matrix by the estimated correlation
  cor_est <- fit$Estimate[num_rep_meas + 1]
  covar_mat <- unweighted_covar_mat * cor_est
  diag(covar_mat) <- var_ests # variances shouldn't be weighted

  return(covar_mat)
}

#' Extract covariance matrix from model fits
#'
#' @description Extracts the repeated measures' covariance matrix estimated by
#'   the various fitting procedures.
#'
#' @param fit An object returned by either the mrmm_wrapper_fun(),
#'   glmmTMB_wrapper_fun() or nlme_wrapper_fun().
#'
#' @return The fit argument's estimated repeated measures covariance matrix.
get_covar_mat <- function(fit) {
  ## extract the covariance matrix from the fit object
  if (class(fit)[1] == "gls") {
    covar_mat <- nlme::getVarCov(fit)
  } else if (class(fit)[1] == "glmmTMB") {
    covar_mat <- compute_glmmtmb_covar_mat(fit)
  } else if (class(fit)[1] == "mmrm") {
    covar_mat <- mmrm::VarCorr(fit)
  } else if (class(fit)[1] == "data.frame") {
    # SAS PROC MIXED
    covar_mat <- compute_sasr_covar_mat(fit)
  }

  ## standardize row and column names
  rownames(covar_mat) <- NULL
  colnames(covar_mat) <- NULL

  ## standardize matrix type
  covar_mat <- as.matrix(covar_mat)

  return(covar_mat)
}

#' Compute the Frobenius norm losses
#'
#' @description This function computes the Frobenius norm loss of each
#'   covariance matrix estimate. This loss is the sum of element-wise squared
#'   errors.
#'
#' @param fit_results A tibble of simulation replicates generated by simChef. It
#'   contains the fitted mixed models for all model implementations.
#'  @param true_covar_mat_ls A named list containing the true repeated measures
#'   covariance matrices associated with each data-generating process. The
#'   matrices' names must match those of the data-generating processes.
#'
#' @return A tibble of simulation replicates containing the Frobenius loss of
#'   every model fit.
frobenius_loss_fun <- function(fit_results, true_covar_mat_ls) {
  fit_results %>%
    dplyr::mutate(
      frobenius_loss = purrr::map2_dbl(
        fit,
        .dgp_name,
        function(f, dgp_name) {
          est_covar_mat <- get_covar_mat(f)
          true_covar_mat <- true_covar_mat_ls[[dgp_name]]
          norm(est_covar_mat - true_covar_mat, type = "F")
        }
      )
    ) %>%
    dplyr::select(-fit)
}

#' Compute the Spectral norm losses
#'
#' @description This function computes the spectral norm loss of each covariance
#'   matrix estimate. This loss is the absolute largest eigenvalue of the
#'   difference between the estimated and true repeated measures covariance
#'   matrix.
#'
#' @param fit_results A tibble of simulation replicates generated by simChef. It
#'   contains the fitted mixed models for all model implementations.
#'  @param true_covar_mat_ls A named list containing the true repeated measures
#'   covariance matrices associated with each data-generating process. The
#'   matrices' names must match those of the data-generating processes.
#'
#' @return A tibble of simulation replicates containing the Frobenius loss of
#'   every model fit.
spectral_loss_fun <- function(fit_results, true_covar_mat_ls) {
  fit_results %>%
    dplyr::mutate(
      spectral_loss = purrr::map2_dbl(
        fit,
        .dgp_name,
        function(f, dgp_name) {
          est_covar_mat <- get_covar_mat(f)
          true_covar_mat <- true_covar_mat_ls[[dgp_name]]
          norm(est_covar_mat - true_covar_mat, type = "2")
        }
      )
    ) %>%
    dplyr::select(-fit)
}


#' Squared error loss of heterogeneous compound symmetry matrix parameters
#'
#' @description This function computes the squared error loss of the individual
#'   compound symmetry matrix parameters. These parameters are the variance
#'   terms at each time point, as well as the correlation coefficient.
#'
#' @param fit_results A tibble of simulation replicates generated by simChef. It
#'   contains the fitted mixed models for all model implementations.
#'  @param true_covar_mat_ls A named list containing the true repeated measures
#'   covariance matrices associated with each data-generating process. The
#'   matrices' names must match those of the data-generating processes.
#'
#' @return A tibble of simulation replicates containing the squared error loss
#'   for each parameter of the intra-observation covariance matrix.
csh_param_sq_err_loss_fun <- function(fit_results, true_covar_mat_ls) {
  ## extract the true variance parameter values
  hom_vars <- diag(true_covar_mat_ls[[1]])
  het_vars <- diag(true_covar_mat_ls[[2]])

  ## extract the true correlation parameter value
  hom_corr <- true_covar_mat_ls[[1]][1, 2] / sqrt(hom_vars[1] * hom_vars[2])
  het_corr <- true_covar_mat_ls[[2]][1, 2] / sqrt(het_vars[1] * het_vars[2])

  ## compute the squared error losses
  fit_results %>%
    dplyr::mutate(
      sq_err_loss = purrr::map2_dfr(
        fit,
        .dgp_name,
        function(f, dgp_name) {
          ## extract the estimated variances and correlation
          est_covar_mat <- get_covar_mat(f)
          vars_est <- diag(est_covar_mat)
          corr_est <- est_covar_mat[1, 2] / sqrt(vars_est[1] * vars_est[2])

          ## assemble the true parameter vector
          if (dgp_name == "hom_rct") {
            true_params <- c(hom_vars, hom_corr)
          } else {
            true_params <- c(het_vars, het_corr)
          }

          ## compute the squared error loss
          loss <- (c(vars_est, corr_est) - true_params)^2
          names(loss) <- c(paste0("var_t0", seq_len(9)), "var_t10", "corr")
          return(loss)
        }
      )
    ) %>%
    dplyr::select(-fit) %>%
    tidyr::unnest(cols = sq_err_loss)
}

#' Bias of heterogeneous compound symmetry matrix parameters
#'
#' @description This function computes the bias of the individual
#'   compound symmetry matrix parameters. These parameters are the variance
#'   terms at each time point, as well as the correlation coefficient.
#'
#' @param fit_results A tibble of simulation replicates generated by simChef. It
#'   contains the fitted mixed models for all model implementations.
#'  @param true_covar_mat_ls A named list containing the true repeated measures
#'   covariance matrices associated with each data-generating process. The
#'   matrices' names must match those of the data-generating processes.
#'
#' @return A tibble containing the empirical bias for each parameter of the
#'   intra-observation covariance matrix, stratified by data-generating process,
#'   method and sample size.
csh_param_bias_fun <- function(fit_results, true_covar_mat_ls) {
  ## extract the true variance parameter values
  hom_vars <- diag(true_covar_mat_ls[[1]])
  het_vars <- diag(true_covar_mat_ls[[2]])

  ## extract the true correlation parameter value
  hom_corr <- true_covar_mat_ls[[1]][1, 2] / sqrt(hom_vars[1] * hom_vars[2])
  het_corr <- true_covar_mat_ls[[2]][1, 2] / sqrt(het_vars[1] * het_vars[2])

  ## compute the biases
  strata_vars <- c(".dgp_name", ".method_name", "num_part")
  fit_results %>%
    dplyr::mutate(
      accuracy = purrr::map2_dfr(
        fit,
        .dgp_name,
        function(f, dgp_name) {
          ## extract the estimated covariance matrix from the model fit
          est_covar_mat <- get_covar_mat(f)

          ## extract the estimated variances
          vars_est <- diag(est_covar_mat)

          ## compute the estimated correlation coefficient
          ## NOTE: This correlation coefficient is identical across all time
          ## points. We use the covariance of the first and second time points
          ## here, but any abitrary pair of time points could be used instead.
          corr_est <- est_covar_mat[1, 2] / sqrt(vars_est[1] * vars_est[2])

          ## assemble the true parameter vector
          if (dgp_name == "hom_rct") {
            true_params <- c(hom_vars, hom_corr)
          } else {
            true_params <- c(het_vars, het_corr)
          }

          ## compute the squared error loss
          accuracy <- c(vars_est, corr_est) - true_params
          names(accuracy) <- c(paste0("var_t0", seq_len(9)), "var_t10", "corr")
          return(accuracy)
        }
      )
    ) %>%
    dplyr::select(-fit) %>%
    tidyr::unnest(cols = accuracy) %>%
    dplyr::group_by(dplyr::across({{ strata_vars }})) %>%
    dplyr::summarize(
      var_t01 = mean(var_t01),
      var_t02 = mean(var_t02),
      var_t03 = mean(var_t03),
      var_t04 = mean(var_t04),
      var_t05 = mean(var_t05),
      var_t06 = mean(var_t06),
      var_t07 = mean(var_t07),
      var_t08 = mean(var_t08),
      var_t09 = mean(var_t09),
      var_t10 = mean(var_t10),
      corr = mean(corr),
      .groups = "drop"
    )
}
