R/helperfunctions_results.R

In alexvolkmann/multifammPaper: Convenience Functions for the Analysis of the multiFAMM Paper

################################################################################
################################################################################
##                                                                            ##
##                   Helperfunctions for the Presented Results                ##
##                                                                            ##
################################################################################
################################################################################


# Aim:
# Functions to evaluate or prepare the resulting objects for evaluation. Also
# contains plot functions.


#------------------------------------------------------------------------------#
# Root (Relative) Mean Squared Error for Scalar Estimates
#------------------------------------------------------------------------------#
#' Root (Relative) Mean Squared Error for Scalar Estimates
#'
#' This function calculates the root (relative) mean squared
#' error for scalar quantities.
#'
#' @param theta_true True component (scalar).
#' @param theta_estim Estimated component (scalar).
#' @param relative FALSE if the absolute MSE is to be computed. Defaults to the
#'   relative MSE (TRUE).
#' @export
rrMSE <- function (theta_true, theta_estim, relative = TRUE) {

  if (relative) {
    sqrt(mean((theta_true - theta_estim)^2) / mean(theta_true^2))
  } else {
    sqrt(mean((theta_true - theta_estim)^2))
  }

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Multivariate Root Relative Mean Squared Error for Functions
#------------------------------------------------------------------------------#
#' Multivariate Root (Relative) Mean Squared Error for Functions
#'
#' This function calculates the multivariate root (relative) mean
#' squared error for functions. For eigenfunctions, which are only defined up to
#' a sign change, there is the option to flip the functions first if the flipped
#' functions seems to be closer to the true function
#'
#' @param fun_true MultiFunData object containing the true functions.
#' @param fun_estim MultiFunData object containing the estimated functions.
#' @param flip Are the estimated functions to be flipped before calculating the
#'  mrrMSE? Defaults to FALSE.
#' @inheritParams rrMSE
#' @export
mrrMSE <- function (fun_true, fun_estim, flip = FALSE, relative = TRUE) {

  if (flip == TRUE) {
    fun_estim <- funData::flipFuns(refObject = fun_true, newObject = fun_estim)
  }

  if (relative) {
    sqrt(mean(funData::norm(fun_true - fun_estim)) /
           mean(funData::norm(fun_true)))
  } else {
    sqrt(mean(funData::norm(fun_true - fun_estim)))
  }

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Univariate Root (Relative) Mean Squared Error for Functions
#------------------------------------------------------------------------------#
#' Univariate Root (Relative) Mean Squared Error for Functions
#'
#' This function calculates the univariate root (relative) mean
#' squared error for functions. For eigenfunctions, which are only defined up to
#' a sign change, there is the option to flip the functions first if the flipped
#' functions seems to be closer to the true function
#'
#' @inheritParams mrrMSE
#' @export
urrMSE <- function (fun_true, fun_estim, flip = FALSE, relative = TRUE) {

  if (flip == TRUE) {
    fun_estim <- flipFuns(refObject = fun_true, newObject = fun_estim)
  }

  if (relative) {
    lapply(seq_along(fun_true@.Data), function (x) {
      sqrt(mean(funData::norm(fun_true@.Data[[x]] - fun_estim@.Data[[x]])) /
             mean(funData::norm(fun_true@.Data[[x]])))
    })

  } else {
    lapply(seq_along(fun_true@.Data), function (x) {
      sqrt(mean(funData::norm(fun_true@.Data[[x]] - fun_estim@.Data[[x]])))
    })

  }

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Flip Estimated Eigenscores According to the Flipping of Eigenfunctions
#------------------------------------------------------------------------------#
#' Flip Estimated Eigenscores According to the Flipping of Eigenfunctions
#' @param fun_true True eigenfunctions as a funData object.
#' @param fun_estim Estimated eigenfunctions as a funData object.
#' @param score_estim Estimated scores.
#' @export
flip_scores <- function (fun_true, fun_estim, score_estim) {

  # Flip the estimated eigenfunctions
  flipped <- flipFuns(refObject = fun_true, newObject = fun_estim)

  # Information of which eigenfunctions were flipped on each dimension
  flips <- lapply(seq_along(flipped@.Data), function (x) {
    flipped@.Data[[x]]@X[, 1] != fun_estim@.Data[[x]]@X[, 1]
  })

  # If the same eigenfunctions were flipped on each dimension, flip the scores
  if(length(unique(flips)) == 1) {
    flips <- flips[[1]]
    score_estim[, flips] <- score_estim[, flips] * (-1)
  } else {
    warning("No flipping because different flips on the dimensions.")
  }

  # Output
  score_estim

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Flip Estimated Eigenscores According to the Flipping of Univariate
# Eigenfunctions
#------------------------------------------------------------------------------#
#' Flip Estimated Eigenscores According to the Flipping of Univariate
#' Eigenfunctions
#' @inheritParams flip_scores
#' @export
flip_scores_uni <- function (fun_true, fun_estim, score_estim) {

  # Arguments
  # fun_true    : True eigenfunctions
  # fun_estim   : Estimated eigenfunctions
  # score_estim : Estimated scores

  # Flip the estimated eigenfunctions
  flipped <- flipFuns(refObject = fun_true, newObject = fun_estim)

  # Information of which eigenfunctions were flipped on each dimension
  flips <- flipped@X[, 1] != fun_estim@X[, 1]

  score_estim[, flips] <- score_estim[, flips] * (-1)

  # Output
  score_estim

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Extract the Estimated Eigenfunctions for One Specific Variance Component
#------------------------------------------------------------------------------#
#' Extract the Estimated Eigenfunctions for One Specific Variance Component
#' @param number Number of eigenfunction to be plotted.
#' @param component Name of variance component.
#' @param m_true_comp True model component.
#' @param eigenfcts Simulation results of eigenfcts.
#' @param flip Are the eigenfunctions to be flipped?
#' @export
extract_Eigenfct_sim <- function (number, component, m_true_comp, eigenfcts,
                                  flip = TRUE) {

  # Flip the eigenfunctions if necessary
  eig <- lapply(eigenfcts$mul, function (x) {
    if (flip) {
      flipFuns(refObject = m_true_comp$eigenfcts[[component]],
               newObject = x[[component]])
    } else {
      x[[component]]
    }
  })

  # Extract the observations and create the X-matrix
  estims <- lapply(eig, function (x) {
    lapply(x@.Data, function (y){
      y@X[number, ]
    })
  })
  estims <- do.call(rbind, estims)
  estims <- lapply(1:ncol(estims), function (x) do.call(rbind, estims[, x]))

  # Concatenate the true eigenfunction
  estims <- lapply(seq_along(estims), function (x) {
    rbind(estims[[x]],
          m_true_comp$eigenfcts[[component]]@.Data[[x]]@X[number, ])
  })

  # Construct the multiFunData Object
  multiFunData(lapply(estims, function (x) {
    funData(argvals = argvals(eigenfcts$mul[[1]][[component]][[1]]),
            X = x)
  }))

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Extract the Estimated Eigenfunctions for One Specific Variance Component for
# a different number of estimated fcts
#------------------------------------------------------------------------------#
#' Extract the Estimated Eigenfunctions for One Specific Variance Component for
#' a different number of estimated fcts
#' @inheritParams extract_Eigenfct_sim
#' @export
extract_Eigenfct_sim_dim <- function (number, component, m_true_comp, eigenfcts,
                                      flip = TRUE) {

  # Flip the eigenfunctions if necessary
  eig <- lapply(eigenfcts$filled, function (x) {
    if (flip) {
      flipFuns(refObject = m_true_comp$eigenfcts[[component]],
               newObject = x[[component]])
    } else {
      x[[component]]
    }
  })

  # Extract the observations and create the X-matrix
  estims <- lapply(eig, function (x) {
    lapply(x@.Data, function (y){
      y@X[number, ]
    })
  })
  estims <- do.call(rbind, estims)
  estims <- lapply(1:ncol(estims), function (x) do.call(rbind, estims[, x]))

  # Concatenate the true eigenfunction
  estims <- lapply(seq_along(estims), function (x) {
    rbind(estims[[x]],
          m_true_comp$eigenfcts[[component]]@.Data[[x]]@X[number, ])
  })

  # Construct the multiFunData Object
  multiFunData(lapply(estims, function (x) {
    funData(argvals = argvals(eigenfcts$mul[[1]][[component]][[1]]),
            X = x)
  }))

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Extract the Estimated Eigenfunctions for One Specific Variance Component of
# the Univariate Estimation
#------------------------------------------------------------------------------#
#' Extract the Estimated Eigenfunctions for One Specific Variance Component of
#' the Univariate Estimation
#' @param eigenfcts Simulation results of eigenfcts that have been filled up.
#' @inheritParams extract_Eigenfct_sim
#' @export
extract_Eigenfct_sim_uni <- function (number, component, eigenfcts) {

  # Arguments
  # number      : Number of Eigenfunction to be plotted
  # component   : Name of Variance Component
  # m_true_comp : True model components
  # eigenfcts   : Simulation results of eigenfcts that have been filled up

  # Flip the eigenfunctions if necessary
  eig <- lapply(eigenfcts$filled_uni, function (x) {
    out <- lapply(seq_along(x[[component]]), function (y) {
      flipFuns(refObject = x[[component]][[y]]$tru,
               newObject = x[[component]][[y]]$est)
    })
    names(out) <- names(x[[component]])
    out
  })

  # Extract the observations and create the X-matrix
  estims <- lapply(eig, function (x) {
    lapply(x, function (y){
      y@X[number, ]
    })
  })
  estims <- do.call(rbind, estims)
  estims <- lapply(1:ncol(estims), function (x) do.call(rbind, estims[, x]))

  # Concatenate the true eigenfunction
  estims <- lapply(seq_along(estims), function (x) {
    rbind(estims[[x]],
          eigenfcts$filled_uni[[1]][[component]][[x]]$tru@X[number, ])
  })

  # Construct the multiFunData Object
  multiFunData(lapply(estims, function (x) {
    funData(argvals =
              argvals(eigenfcts$filled_uni[[1]][[component]][[1]]$tru),
            X = x)
  }))

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Extract the Estimated Covariate Effect for One Specific Model Term
#------------------------------------------------------------------------------#
#' Extract the Estimated Covariate Effect for One Specific Model Term
#'
#' This function extracts one single effect
#' function from the simulated effects and adds the true observation as the last
#' observation.
#'
#' @param term Name of the covariate effect function to extract.
#' @param m_true_comp True model component from which to extract the true
#'  observation.
#' @param cov_preds List of simulation results from which to extract (as given
#'  by e.g. load_sim_results()).
#' @param se TRUE if the standard error is to be extracted instead of the
#'  estimated function. Defaults to FALSE.
#' @export
extract_Covfct_sim <- function (term, m_true_comp, cov_preds, se = FALSE) {

  se <- if (se) "se.fit" else "fit"

  covs <- lapply(cov_preds$mul, function (x) x[[se]][[term]])

  # Extract the observations and create the X-matrix
  estims <- lapply(covs, function (x) {
    lapply(x@.Data, function (y){
      y@X
    })
  })
  estims <- do.call(rbind, estims)
  estims <- lapply(1:ncol(estims), function (x) do.call(rbind, estims[, x]))

  # Concatenate the true eigenfunction
  estims <- lapply(seq_along(estims), function (x) {
    rbind(estims[[x]],
          m_true_comp$cov_preds[[se]][[term]]@.Data[[x]]@X)
  })

  # Construct the multiFunData Object
  multiFunData(lapply(estims, function (x) {
    funData(argvals = argvals(cov_preds$mul[[1]][[se]][[term]][[1]]),
            X = x)
  }))

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Extract the Estimated Covariate Effect for One Specific Model Term of the
# Univariate Estimation
#------------------------------------------------------------------------------#
#' Extract the Estimated Covariate Effect for One Specific Model Term of the
#' Univariate Estimation
#' @param term Name of covariate effect term.
#' @param term_uni Name of covariate in univariate estimation ("int",
#'   "famm_cb_mean", "famm_cb_covariate.1", "famm_cb_inter_1_2").
#' @param m_true_comp True model components.
#' @param cov_preds Simulation results of covariate effects.
#' @export
extract_Covfct_sim_uni <- function (term, term_uni, m_true_comp, cov_preds) {

  # Arguments
  # term        :
  # term_uni    : Name of covariate in univariate estimation ("int",
  #                 "famm_cb_mean", "famm_cb_covariate.1", "famm_cb_inter_1_2")
  # m_true_comp : True model components
  # eigenfcts   :

  covs <- lapply(cov_preds$uni, function (x) {
    lapply(x, function (y) y$fit[[term_uni]])
  })

  # Extract the observations and create the X-matrix
  estims <- lapply(covs, function (x) {
    lapply(x, function (y){
      y@X
    })
  })
  estims <- do.call(rbind, estims)
  estims <- lapply(1:ncol(estims), function (x) do.call(rbind, estims[, x]))

  # Concatenate the true eigenfunction
  estims <- lapply(seq_along(estims), function (x) {
    rbind(estims[[x]],
          m_true_comp$cov_preds$fit[[term]]@.Data[[x]]@X)
  })

  # Construct the multiFunData Object
  multiFunData(lapply(estims, function (x) {
    funData(argvals = argvals(cov_preds$mul[[1]]$fit[[term]][[1]]),
            X = x)
  }))

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Compute the True Values to Compare with Fitted Values
#------------------------------------------------------------------------------#
#' Compute the True Values to Compare with Fitted Values
#'
#' This function takes the fitted curves object of the
#' simulation and recalculates the true curves.
#'
#' @param fitted_cu Object saved from the simulation.
#' @inheritParams sim_eval_components
#' @export
compute_fitted_sim <- function (fitted_cu, I = 10, J = 16, reps = 5,
                                nested = FALSE) {

  reps_B <- rep(reps*J, times = I)
  reps_C <- rep(reps, times = J)

  # For Random Intercept of Subject
  if ("B" %in% names(fitted_cu$tru[[1]]$re)) {
    re_B_true <- lapply(seq_along(fitted_cu$tru), function (x) {
      multiFunData(lapply(fitted_cu$tru[[x]]$re$B, function (y) {
        funData(argvals = argvals(y),
                X = y@X[rep(1:nrow(y@X), times = reps_B), ])
      }))
    })
  } else {
    # Zero object
    re_B_true <- lapply(seq_along(fitted_cu$tru), function (x) {
      argvals <- argvals(fitted_cu$tru[[x]]$mu[[1]])
      multiFunData(lapply(seq_along(fitted_cu$tru[[x]]$mu), function (y) {
        funData(argvals = argvals,
                X = matrix(0, nrow = nObs(fitted_cu$tru[[x]]$mu),
                           ncol = length(unlist(argvals))))
      }))
    })
  }

  # Has not been tested for C in names()
  if ("C" %in% names(fitted_cu$tru[[1]]$re)) {
    if(!nested) {
      re_C_true <- lapply(seq_along(fitted_cu$tru), function (x) {
        multiFunData(lapply(fitted_cu$tru[[x]]$re$C, function (y) {
          funData(argvals = argvals(y),
                  X = y@X[rep(rep(1:nrow(y@X), times = reps_C), times = I), ])
        }))
      })
    } else {
      re_C_true <- lapply(seq_along(fitted_cu$tru), function (x) {
        multiFunData(lapply(fitted_cu$tru[[x]]$re$C, function (y) {
          funData(argvals = argvals(y),
                  X = y@X[rep(1:nrow(y@X), each = reps), ])
        }))
      })
    }
  } else {
    # Zero object
    re_C_true <- lapply(seq_along(fitted_cu$tru), function (x) {
      argvals <- argvals(fitted_cu$tru[[x]]$mu[[1]])
      multiFunData(lapply(seq_along(fitted_cu$tru[[x]]$mu), function (y) {
        funData(argvals = argvals,
                X = matrix(0, nrow = nObs(fitted_cu$tru[[x]]$mu),
                           ncol = length(unlist(argvals))))
      }))
    })
  }

  mapply(function (x, y, z) {x$mu + x$re$E + y + z},
         fitted_cu$tru, re_B_true, re_C_true, SIMPLIFY = FALSE)
}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Transform a funData object to an Data.Frame
#------------------------------------------------------------------------------#
#' Transform a funData object to an Data.Frame
#'
#' This function transform a funData object to a data.frame
#' so that it is easy to plot it using ggplot. It internally checks if the
#' object is a multiFunData object.
#'
#' @param fundata FunData or MultiFundata object to transform.
#' @return A data.frame with four variables
#'   \itemize{
#'     \item \code{t} (num): the functional argument.
#'     \item \code{y} (num): the functional value.
#'     \item \code{dim} (int): the number of the dimension (position in the
#'     multiFunData object).
#'     \item \code{obs} (int): the number of the observation in the funData
#'     object.}
#' @export
funData2DataFrame <- function(fundata) {

  # Automatic checking if the funData object belongs to class multiFunData
  multifun <- "multiFunData" %in% class(fundata)

  data_list <-if (multifun == TRUE) {
    lapply(fundata@.Data, function (x) {
      list(argvals = unlist(x@argvals), x = x@X)
    })
  } else {
    # maybe here list(list(...)) otherwise doesn't work for univariate funData?
    list(argvals = unlist(fundata@argvals), x = fundata@X)
  }

  dat <- data.frame(
    t = do.call(c, lapply(data_list, function (x) {
      rep(x$argvals, times = nrow(x$x))
    })),
    y = do.call(c, lapply(data_list, function (x) {
      as.vector(t(x$x))
    })),
    dim = rep(seq_along(data_list), times = sapply(data_list, function (x) {
      length(x$x)
    })),
    obs = do.call(c, lapply(data_list, function (x) {
      rep(seq_len(nrow(x$x)), each = length(x$argvals))
    })))

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Evaluate the model components in the simulation
#------------------------------------------------------------------------------#
#' Evaluate the model components in the simulation
#'
#' This function takes the folder containing the
#' results of a simulation scenario and returns a data frame with relative
#' MSE values for the separate model components.
#'
#' @param folder Folder with saved objects from the simulation.
#' @param m_true_comp True model components used for the simulation.
#' @param label_cov Vector of labels for the covariates.
#' @param relative Is the relative error measure to be used or the absolute?
#'   Defaults to TRUE.
#' @param I Number of levels for component B. Defaults to 9.
#' @param J Number of levels for component C. Defaults to 16.
#' @param reps Number of repetitions of the levels of B and C. Defaults to 5.
#' @param nested TRUE if the model component C is nested. Defaults to FALSE.
#' @param fixed_fpc FALSE if the number of FPCs in the model is selected by the
#'   model not by the user. This means that eigenvalues, scores and
#'   eigenfunctions are not evaluated. Defaults to TRUE.
#' @param sigma2 Vector of true values can be supplied if they differ from the
#'   values in m_true_comp. Defaults to NULL.
#' @export
sim_eval_components <- function (folder, m_true_comp, label_cov,
                                 relative = TRUE, I = 9, J = 16, reps = 5,
                                 nested = FALSE, fixed_fpc = TRUE,
                                 sigma2 = NULL) {

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  error_var <- load_sim_results(folder = folder, component = "error_var")

  if(is.null(sigma2)) {
    if (length(m_true_comp$error_var$modelweights) == 1) {
      true_sig <- m_true_comp$error_var$modelsig2
    } else {
      true_sig <- m_true_comp$error_var$modelsig2 /
        m_true_comp$error_var$modelweights
    }
  } else {
    true_sig <- sigma2
  }

  dat_err <- do.call(rbind, lapply(seq_along(error_var$mul),
                                   function (x, true) {

    sigma_hat <- error_var$mul[[x]]$modelsig2 / error_var$mul[[x]]$modelweights
    if (length(sigma_hat) != length(true)) {
      max_w <- max(sigma_hat)
      zero_var <- which(error_var$mul[[x]]$uni_vars == 0)
      min_var <- which.min(error_var$mul[[x]]$uni_vars[-zero_var])
      rest <- sigma_hat[-which.max(sigma_hat)]
      sigma_hat <- vector(mode = "numeric", length = length(true))
      sigma_hat[c(zero_var, min_var)] <- max_w
      ind <- 1
      for (i in seq_along(sigma_hat)) {
        if (sigma_hat[i] == 0) {
          sigma_hat[i] <- rest[ind]
          ind <- ind + 1
        }
      }
    }
    data.frame(it = rep(x, times = length(true_sig)),
               hat = sigma_hat,
               no = factor(seq_along(true_sig),
                           labels = if(length(true_sig) == 1) {
                             "sigma^2"
                           } else {
                             c(paste0("sigma[dim", seq_along(true_sig),
                                      "]^2"))}),
               true = true,
               y = mapply(function (x, y) {
                 rrMSE(theta_true = x, theta_estim = y, relative = relative)
               }, true, sigma_hat),
               comp = factor("sigma^2"))

  }, true = true_sig))

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  eigenvals <- load_sim_results(folder = folder, component = "eigenvals")

  if (fixed_fpc) {
    dat_val <- do.call(rbind, lapply(names(m_true_comp$eigenvals),
                                     function (x) {
      do.call(rbind, lapply(seq_along(eigenvals$mul), function (y, true) {
        vals <- eigenvals$mul[[y]][[x]]
        data.frame(it = rep(y, times = length(vals)),
                   hat = vals,
                   no = factor(1:length(vals),
                               labels = paste0("upsilon[", seq_along(vals),
                                               "]^", x)),
                   true = true,
                   y = mapply(function (u, v) {
                     rrMSE(theta_true = u, theta_estim = v, relative = relative)
                   }, true, vals),
                   comp = "Eigenvalues")
      }, true = m_true_comp$eigenvals[[x]]))
    }))
  }

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  eigscores <- load_sim_results(folder = folder, component = "eigscores")
  eigenfcts <- load_sim_results(folder = folder, component = "eigenfcts")

  if (fixed_fpc) {
    eigscores$mul_flip <- lapply(seq_along(eigscores$mul), function (y) {
      scores <- lapply(names(eigscores$mul[[1]]), function (x) {
        flip_scores(fun_true = m_true_comp$eigenfcts[[x]],
                    fun_estim = eigenfcts$mul[[y]][[x]],
                    score_estim = eigscores$mul[[y]][[x]])
      })
      names(scores) <- names(eigscores$mul[[1]])
      scores
    })

    dat_sco <- do.call(rbind, lapply(names(eigscores$mul_flip[[1]]),
                                     function (x) {
     do.call(rbind, lapply(seq_along(eigscores$mul_flip), function (y) {
       scores <- eigscores$mul_flip[[y]][[x]]
       true <- eigscores$tru[[y]][[x]]
       data.frame(it = rep(y, times = ncol(scores)),
                  no = factor(1:ncol(scores),
                              labels = paste0("rho[",seq_len(ncol(scores)),
                                              "]^", x)),
                  y = sapply(1:ncol(scores), function (z) {
                    rrMSE(theta_true = true[, z], theta_estim = scores[, z],
                          relative = relative)
                  }),
                  comp = "Scores")
     }))
   }))
  }

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  # Eigenfunctions
  if(fixed_fpc) {
    dat_fun <- do.call(rbind, lapply(names(eigenfcts$mul[[1]]), function (x) {
      do.call(rbind, lapply(seq_along(eigenfcts$mul), function (y) {
        fcts <- eigenfcts$mul[[y]][[x]]
        data.frame(it = rep(y, times = funData::nObs(fcts)),
                   no = factor(1:funData::nObs(fcts),
                               labels = paste0("psi[",
                                               seq_len(funData::nObs(fcts)),
                                               "]^", x)),
                   y = sapply(1:funData::nObs(fcts), function (z) {
                     mrrMSE(
                       fun_true = funData::extractObs(
                         m_true_comp$eigenfcts[[x]], obs = z),
                       fun_estim = funData::extractObs(fcts, obs = z),
                       flip = TRUE, relative = relative)
                   }),
                   comp = "Eigenfunctions")
      }))
    }))
  }

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  # Covariance operator
  auto_cross <- c(lapply(1:length(m_true_comp$fitted_curves),
                         function (x) c(x, x)),
                  combn(1:length(m_true_comp$fitted_curves),
                        m = 2, simplify = FALSE))
  true_covs <- lapply(names(m_true_comp$eigenvals), function (x) {
    lapply(auto_cross, function (y) {
      funData::funData(argvals = list(seq(0,1,length.out = 100),
                                      seq(0,1,length.out = 100)),
              X = array(covSurv(vals = m_true_comp$eigenvals[[x]],
                          fcts = m_true_comp$eigenfcts[[x]],
                          dim1 = y[1], dim2 = y[2]), dim = c(1,100,100)))
    })
  })
  names(true_covs) <- names(m_true_comp$eigenvals)

  dat_cop <- do.call(rbind, lapply(names(true_covs), function (x) {
    do.call(rbind, lapply(seq_along(auto_cross), function (y) {
      do.call(rbind, lapply(seq_along(eigenvals$mul), function (z) {
        estim_cov <- funData::funData(argvals = list(seq(0,1,length.out = 100),
                                            seq(0,1,length.out = 100)),
                                      X = array(covSurv(
                                        vals = eigenvals$mul[[z]][[x]],
                                        fcts = eigenfcts$mul[[z]][[x]],
                                        dim1 = auto_cross[[y]][1],
                                        dim2 = auto_cross[[y]][2]),
                                        dim = c(1,100,100)))
        data.frame(it = z,
                   no = factor(paste0("K[",
                                      paste(auto_cross[[y]], collapse = ""),
                                      "]^", x)),
                   y = mrrMSE(fun_true = true_covs[[x]][[y]],
                              fun_estim = estim_cov,
                              flip = FALSE, relative = relative),
                   comp = "Covariance")
      }))
    }))
  }))

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  # Random effects
  ran_preds <- load_sim_results(folder = folder, component = "ran_preds")

  dat_ran <- do.call(rbind, lapply(names(ran_preds$mul[[1]]), function (x) {
    do.call(rbind, lapply(seq_along(ran_preds$mul), function (y) {
      randef <- ran_preds$mul[[y]][[x]]
      rantru <- ran_preds$tru[[y]][[x]]
      data.frame(it = y,
                 no = factor(1, labels = x),
                 y = mrrMSE(fun_true = rantru, fun_estim = randef,
                            flip = FALSE, relative = relative),
                 comp = "Fit")
    }))
  }))

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  # Fitted values
  fitted_cu <- load_sim_results(folder = folder, component = "fitted_cu")

  fit_true <- compute_fitted_sim(fitted_cu = fitted_cu, I = I, J = J,
                                 reps = reps, nested = nested)

  dat_fit <- do.call(rbind, lapply(seq_along(fitted_cu$mul), function (x) {
    data.frame(it = x,
               no = factor(1, labels = "Y"),
               y = mrrMSE(fun_true = fit_true[[x]],
                          fun_estim = fitted_cu$mul[[x]],
                          flip = FALSE, relative = relative),
               comp = "Fit")
  }))


  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  # Mean
  mu_estim <- compute_mu_sim(fitted_cu = fitted_cu, ran_preds = ran_preds,
                             I = I, J = J, reps = reps, nested = nested)

  dat_mu <- do.call(rbind, lapply(seq_along(fitted_cu$tru), function (x) {
    data.frame(it = x,
               no = factor(1, labels = "mu"),
               y = mrrMSE(fun_true = fitted_cu$tru[[x]]$mu,
                          fun_estim = mu_estim[[x]],
                          flip = FALSE, relative = relative),
               comp = "Fit")
  }))

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  # Covariate effects
  cov_preds <- load_sim_results(folder = folder, component = "cov_preds")
  cov_preds$mul <- lapply(cov_preds$mul, function (it) {
    lapply(it, function (se_fit) {
      se_fit[[2]] <- se_fit[[1]] + se_fit[[2]]
      se_fit[[1]] <- NULL
      se_fit
    })
  })

  m_true_comp$cov_preds$fit[[2]] <- m_true_comp$cov_preds$fit[[1]] +
    m_true_comp$cov_preds$fit[[2]]
  m_true_comp$cov_preds$fit[[1]] <- NULL

  names <- label_cov

  dat_cov <- do.call(rbind, lapply(seq_along(cov_preds$mul), function (x) {
    do.call(rbind, lapply(seq_along(cov_preds$mul[[x]]$fit), function (y) {
      data.frame(it = x,
                 no = factor(names[y]),
                 y = mrrMSE(fun_true = m_true_comp$cov_preds$fit[[y]],
                            fun_estim = cov_preds$mul[[x]]$fit[[y]],
                            relative = relative),
                 comp = "Effectfunctions")
    }))
  }))

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#

  comb <- c("it", "no", "y", "comp")
  dat <- rbind(dat_fit, dat_ran, dat_mu, dat_err[, comb], dat_cop,
               if (fixed_fpc) dat_val[,comb],
               if (fixed_fpc) dat_fun, if (fixed_fpc) dat_sco, dat_cov)

  dat

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Evaluate the model components per dimension in the simulation
#------------------------------------------------------------------------------#
#' Evaluate the model components per dimension in the simulation
#'
#' This function takes the folder containing the
#' results of a simulation scenario and returns a data frame with relative
#' MSE values for the separate model components and the separate dimensions.
#'
#' @inheritParams sim_eval_components
#' @param uni_compare TRUE if the simulation scenario includes univariate model
#'   estimates that should be also evaluated. Defaults to FALSE.
#' @export
sim_eval_dimensions <- function (folder, m_true_comp, label_cov,
                                 relative = TRUE, uni_compare = FALSE,
                                 I = 9, J = 16, reps = 5, nested = FALSE) {

  dim_names <- paste0("dim", seq_along(m_true_comp$fitted_curves))

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  error_var <- load_sim_results(folder = folder, component = "error_var",
                                uni = uni_compare)

  true_sig <- m_true_comp$error_var$modelsig2 /
    m_true_comp$error_var$modelweights

  dat_err_u <- if (uni_compare) {
    do.call(rbind, lapply(seq_along(error_var$mul),
       function (x, true) {
         sigma_hat <- unlist(error_var$uni[[x]])
         data.frame(it = rep(x, times = length(sigma_hat)),
                    hat = sigma_hat,
                    no = factor(1, labels = c("sigma^2")),
                    true = true,
                    y = mapply(function (x, y) {
                      rrMSE(theta_true = x, theta_estim = y,
                            relative = relative)
                    }, true, sigma_hat),
                    comp = factor("sigma^2"),
                    dim = if (length(true_sig) == 1) "cross" else dim_names,
                    method = factor("uni"))
       }, true = true_sig))
  } else NULL

  dat_err_m <- do.call(rbind, lapply(seq_along(error_var$mul),
                                     function (x, true) {
     sigma_hat <- error_var$mul[[x]]$modelsig2 / error_var$mul[[x]]$modelweights
     if (length(sigma_hat) != length(true)) {
       max_w <- max(sigma_hat)
       zero_var <- which(error_var$mul[[x]]$uni_vars == 0)
       min_var <- which.min(error_var$mul[[x]]$uni_vars[-zero_var])
       rest <- sigma_hat[-which.max(sigma_hat)]
       sigma_hat <- vector(mode = "numeric", length = length(true))
       sigma_hat[c(zero_var, min_var)] <- max_w
       ind <- 1
       for (i in seq_along(sigma_hat)) {
         if (sigma_hat[i] == 0) {
           sigma_hat[i] <- rest[ind]
           ind <- ind + 1
         }
       }
     }
     data.frame(it = rep(x, times = length(sigma_hat)),
                hat = sigma_hat,
                no = factor(1, labels = c("sigma^2")),
                true = true,
                y = mapply(function (x, y) {
                  rrMSE(theta_true = x, theta_estim = y, relative = relative)
                }, true, sigma_hat),
                comp = factor("sigma^2"),
                dim = if (length(true_sig) == 1) "cross" else dim_names,
                method = factor("mul"))
   }, true = true_sig))

  dat_err <- rbind(dat_err_m, dat_err_u)

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  ran_preds <- load_sim_results(folder = folder, component = "ran_preds",
                                uni = uni_compare)

  dat_ran_m <- do.call(rbind, lapply(names(ran_preds$mul[[1]]), function (x) {
    do.call(rbind, lapply(seq_along(ran_preds$mul), function (y) {
      randef <- ran_preds$mul[[y]][[x]]
      rantru <- ran_preds$tru[[y]][[x]]
      data.frame(it = rep(y, times = length(randef)),
                 dim = dim_names,
                 y = unlist(urrMSE(fun_true = rantru, fun_estim = randef,
                                   flip = FALSE, relative = relative)),
                 no = factor(x),
                 comp = factor("Fit"),
                 method = factor("mul"))
    }))
  }))

  dat_ran_u <- if (uni_compare) {
    do.call(rbind, lapply(names(ran_preds$mul[[1]]), function (x) {
      do.call(rbind, lapply(seq_along(ran_preds$mul), function (y) {
        randef <- lapply(ran_preds$uni[[y]], function (z) z[[x]])
        rantru <- ran_preds$tru[[y]][[x]]
        do.call(rbind, lapply(seq_along(dim_names), function (z) {
          data.frame(it = y,
                     dim = dim_names[z],
                     no = factor(x),
                     y = mrrMSE(fun_true = rantru[[z]],
                                fun_estim = randef[[z]],
                                flip = FALSE, relative = relative),
                     comp = factor("Fit"),
                     method = factor("uni"))
        }))
      }))
    }))
  } else NULL

  dat_ran <- rbind(dat_ran_m, dat_ran_u)

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  # Fitted values
  fitted_cu <- load_sim_results(folder = folder, component = "fitted_cu",
                                uni = uni_compare)

  fit_true <- compute_fitted_sim(fitted_cu = fitted_cu, I = I, J = J,
                                 reps = reps, nested = nested)

  dat_fit_m <- do.call(rbind, lapply(seq_along(fitted_cu$mul), function (x) {
      data.frame(it = x,
                 dim = dim_names,
                 no = factor("Y"),
                 y = unlist(urrMSE(fun_true = fit_true[[x]],
                                   fun_estim = fitted_cu$mul[[x]],
                                   flip = FALSE, relative = relative)),
                 comp = "Fit",
                 method = factor("mul"))
  }))

  dat_fit_u <- if (uni_compare) {
    do.call(rbind, lapply(seq_along(fitted_cu$uni), function (x) {
      do.call(rbind, lapply(seq_along(fitted_cu$uni[[1]]), function (y) {
        data.frame(it = x,
                   dim = dim_names[y],
                   no = factor("Y"),
                   y = mrrMSE(fun_true = fit_true[[x]][[y]],
                              fun_estim = fitted_cu$uni[[x]][[y]],
                              flip = FALSE, relative = relative),
                   comp = factor("Fit"),
                   method = factor("uni"))
      }))
    }))
  } else NULL


  dat_fit <- rbind(dat_fit_m, dat_fit_u)

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  # Mean
  mu_mul <- compute_mu_sim(fitted_cu = fitted_cu, ran_preds = ran_preds,
                           I = I, J = J, reps = reps, nested = nested,
                           mul = TRUE)
  if (uni_compare) {
    mu_uni <- compute_mu_sim(fitted_cu = fitted_cu, ran_preds = ran_preds,
                             I = I, J = J, reps = reps, nested = nested,
                             mul = FALSE)
  }

  dat_mu_m <- do.call(rbind, lapply(seq_along(fitted_cu$tru), function (x) {
    data.frame(it = x,
               dim = dim_names,
               no = factor("mu"),
               y = unlist(urrMSE(fun_true = fitted_cu$tru[[x]]$mu,
                          fun_estim = mu_mul[[x]],
                          flip = FALSE, relative = relative)),
               comp = "Fit",
               method = factor("mul"))
  }))

  dat_mu_u <- if (uni_compare) {
    do.call(rbind, lapply(seq_along(mu_uni), function (x) {
      do.call(rbind, lapply(seq_along(mu_uni[[1]]), function (y) {
        data.frame(it = x,
                   dim = dim_names[y],
                   no = factor("mu"),
                   y = mrrMSE(fun_true = fitted_cu$tru[[x]]$mu[[y]],
                              fun_estim = mu_uni[[x]][[y]],
                              flip = FALSE, relative = relative),
                   comp = factor("Fit"),
                   method = factor("uni"))
      }))
    }))
  } else NULL

  dat_mu <- rbind(dat_mu_m, dat_mu_u)

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  # Covariance operator
  eigenvals <- load_sim_results(folder = folder, component = "eigenvals",
                                uni = uni_compare)
  eigenfcts <- load_sim_results(folder = folder, component = "eigenfcts",
                                uni = uni_compare)


  auto_cross <- c(lapply(1:length(m_true_comp$fitted_curves),
                         function (x) c(x, x)),
                  combn(1:length(m_true_comp$fitted_curves),
                        m = 2, simplify = FALSE))
  true_covs <- lapply(names(m_true_comp$eigenvals), function (x) {
    lapply(auto_cross, function (y) {
      funData::funData(argvals = list(seq(0,1,length.out = 100),
                                      seq(0,1,length.out = 100)),
                       X = array(covSurv(vals = m_true_comp$eigenvals[[x]],
                                         fcts = m_true_comp$eigenfcts[[x]],
                                         dim1 = y[1], dim2 = y[2]),
                                 dim = c(1,100,100)))
    })
  })
  names(true_covs) <- names(m_true_comp$eigenvals)

  dat_cop_m <- do.call(rbind, lapply(names(true_covs), function (x) {
    do.call(rbind, lapply(seq_along(auto_cross), function (y) {
      do.call(rbind, lapply(seq_along(eigenvals$mul), function (z) {
        estim_cov <- funData::funData(argvals = list(seq(0,1,length.out = 100),
                                                     seq(0,1,length.out = 100)),
                             X = array(covSurv(
                               vals = eigenvals$mul[[z]][[x]],
                               fcts = eigenfcts$mul[[z]][[x]],
                               dim1 = auto_cross[[y]][1],
                               dim2 = auto_cross[[y]][2]),
                               dim = c(1,100,100)))
        data.frame(it = z,
                   dim = c(dim_names,
                           rep("cross", choose(length(dim_names), 2)))[y],
                   no = factor(paste0("K[",
                                      paste(auto_cross[[y]], collapse = ""),
                                      "]^", x)),
                   y = mrrMSE(fun_true = true_covs[[x]][[y]],
                              fun_estim = estim_cov,
                              flip = FALSE, relative = relative),
                   comp = "Covariance",
                   method = "mul")
      }))
    }))
  }))

  dat_cop_u <- if (uni_compare) {
    do.call(rbind, lapply(names(true_covs), function (x) {
      do.call(rbind, lapply(seq_along(auto_cross), function (y) {
        do.call(rbind, lapply(seq_along(eigenvals$uni), function (z) {

          vals <- lapply(eigenvals$uni[[z]], "[[", x)
          fcts <- lapply(eigenfcts$uni[[z]], "[[", x)

          estim_cov <- funData::funData(argvals =
                                          list(seq(0,1,length.out = 100),
                                               seq(0,1,length.out = 100)),
                                        X = array(covSurv(vals = vals,
                                              fcts = fcts,
                                              dim1 = auto_cross[[y]][1],
                                              dim2 = auto_cross[[y]][2],
                                              multi = FALSE),
                                              dim = c(1,100,100)))
          data.frame(it = z,
                     dim = c(dim_names,
                             rep("cross", choose(length(dim_names), 2)))[y],
                     no = factor(paste0("K[",
                                        paste(auto_cross[[y]], collapse = ""),
                                        "]^", x)),
                     y = mrrMSE(fun_true = true_covs[[x]][[y]],
                                fun_estim = estim_cov,
                                flip = FALSE, relative = relative),
                     comp = "Covariance",
                     method = "uni")
        }))
      }))
    }))
  } else NULL

  dat_cop <- rbind(dat_cop_m, dat_cop_u)

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  # Covariate effects
  cov_preds <- load_sim_results(folder = folder, component = "cov_preds",
                                uni = uni_compare)
  cov_preds$mul <- lapply(cov_preds$mul, function (it) {
    lapply(it, function (se_fit) {
      se_fit[[2]] <- se_fit[[1]] + se_fit[[2]]
      se_fit[[1]] <- NULL
      se_fit
    })
  })
  m_true_comp$cov_preds$fit[[2]] <- m_true_comp$cov_preds$fit[[1]] +
    m_true_comp$cov_preds$fit[[2]]
  m_true_comp$cov_preds$fit[[1]] <- NULL

  if (uni_compare) {
    cov_preds$uni <- lapply(cov_preds$uni, function (it) {
      lapply(it, function (dim) {
        lapply(dim, function (se_fit) {
          se_fit[[2]] <- se_fit[[1]] + se_fit[[2]]
          se_fit[[1]] <- NULL
          se_fit
        })
      })
    })
  }
  names <- label_cov

  dat_cov_m <- do.call(rbind, lapply(seq_along(cov_preds$mul), function (x) {
    do.call(rbind, lapply(seq_along(cov_preds$mul[[x]]$fit), function (y) {
      data.frame(it = x,
                 no = factor(names[y]),
                 dim = dim_names,
                 y = unlist(urrMSE(fun_true = m_true_comp$cov_preds$fit[[y]],
                                   fun_estim = cov_preds$mul[[x]]$fit[[y]],
                                   flip = FALSE, relative = relative)),
                 comp = factor("Effectfunctions"),
                 method = factor("mul"))
    }))
  }))

  dat_cov_u <- if (uni_compare) {
    do.call(rbind, lapply(seq_along(cov_preds$mul), function (x) {
      do.call(rbind, lapply(seq_along(cov_preds$uni[[x]]), function (y) {
        do.call(rbind, lapply(seq_along(cov_preds$uni[[x]][[y]]$fit),
          function (z) {
            aha <- names(cov_preds$uni[[x]][[y]]$fit)
            aha <- sub("famm_cb", "s(t):", aha)
            aha <- gsub("_", ".", aha)
            names(cov_preds$uni[[x]][[y]]$fit) <- sub(".mean|.covariate", "",
                                                      aha)

            u <- names(cov_preds$uni[[x]][[y]]$fit)[[z]]
            v <- which(names(m_true_comp$cov_preds$fit) == u)

            data.frame(it = x,
                       no = factor(names[v]),
                       dim = dim_names[y],
                       y = mrrMSE(fun_true =
                                    m_true_comp$cov_preds$fit[[u]][[y]],
                                  fun_estim = cov_preds$uni[[x]][[y]]$fit[[z]],
                                  relative = relative),
                       comp = factor("Effectfunctions"),
                       method = factor("uni"))
          }))
      }))
    }))
  } else NULL

  dat_cov <- rbind(dat_cov_m, dat_cov_u)

  #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
  comb <- c("it", "no", "y", "comp", "dim", "method")
  dat <- rbind(dat_fit, dat_ran, dat_mu, dat_err[, comb], dat_cop, dat_cov)

  dat

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Estimate Covariance Surface
#------------------------------------------------------------------------------#
#' Estimate Covariance Surface
#' @param vals Eigenvalues (Vector).
#' @param fcts Eigenfunctions (MultiFunData).
#' @param dim1 Choose first dimension (Scalar).
#' @param dim2 Choose second dimension (Scalar).
#' @param multi Is the input multivariate (if not then vals and fcts are lists).
#' @export
covSurv <- function (vals, fcts, dim1, dim2, multi = TRUE) {

  if (multi == TRUE) {

    # Extract the matrix of Eigenfunctions
    mat <- do.call(rbind, lapply(fcts, function (x) t(x@X)))

    # Construct a diagonal matric containing the Eigenvalues
    dia <- diag(vals, nrow = nObs(fcts))

  } else {

    # Create the matrix of Eigenfunctions
    mat <- rbind(cbind(t(fcts[[1]]@X),
                       matrix(0, ncol = nObs(fcts[[2]]), nrow = 100)),
                 cbind(matrix(0, ncol = nObs(fcts[[1]]), nrow = 100),
                       t(fcts[[2]]@X)))

    # Construct a diagonal matric containing the Eigenvalues
    dia <- diag(unlist(vals), nrow = length(unlist(vals)))

  }

  # Compute the Auto- and Cross-covariance
  cov <- mat %*% dia %*% t(mat)

  d1 <- 1:100 + (dim1 - 1)*100
  d2 <- 1:100 + (dim2 - 1)*100

  # Extract the wanted covariance surface
  cov[d1, d2]


}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Plot Covariate Effects Together
#------------------------------------------------------------------------------#
#' Plot Covariate Effects Together
#' @param m_true_comp True model components.
#' @param effects Number of effect function to be plotted.
#' @param m_fac Multiplication factor.
#' @param limits Y-limits for the plot.
#' @param dimlabels Labels for the dimensions in the plot.
#' @export
covariate_plot <- function (m_true_comp, effects = c(4, 5, 6), m_fac = 2,
                            limits = limits, dimlabels = c("ACO", "EPG")) {

  # Extract covariate information
  covars <- lapply(effects, function (x) {
    funData2DataFrame(fundata = m_true_comp$cov_preds$fit[[x]])
  })

  # Extract information for standard deviation
  covars.se <- lapply(effects, function (x) {
    funData2DataFrame(fundata = m_true_comp$cov_preds$se.fit[[x]])
  })

  # Create the corresponding labels of the plot
  labels <- sapply(effects, function (x) {
    paste0("beta[", x-2, "](t)")
  })

  # Construct the data.frame
  dat <- data.frame(
    t = do.call(c, lapply(covars, "[[", "t")),
    y = do.call(c, lapply(covars, "[[", "y")),
    y_plus = do.call(c, lapply(seq_along(covars), function (x) {
      covars[[x]]$y + m_fac*covars.se[[x]]$y
    })),
    y_minu = do.call(c, lapply(seq_along(covars), function (x) {
      covars[[x]]$y - m_fac*covars.se[[x]]$y
    })),
    effect = factor(rep(effects, times = sapply(covars, nrow)),
                    labels = labels),
    dim = factor(do.call(c, lapply(covars, "[[", "dim")),
                 labels = dimlabels)
  )

  # Plot the data
  ggplot2::ggplot(data = dat, aes(x = t)) +
    geom_line(aes(y = y), size = 0.25) +
    geom_line(aes(y = y_plus), linetype = "dashed", size = 0.25) +
    geom_line(aes(y = y_minu), linetype = "dashed", size = 0.25) +
    geom_hline(yintercept = 0, linetype = "dotted", size = 0.3) +
    facet_grid(dim ~ effect, labeller = label_parsed) +
    xlab("Normalized Time (t)") +
    ylab(expression("Effect Function (" ~ beta^(d)~"(t) )")) +
    theme_grey(base_size = 8) +
    scale_y_continuous(limits = limits)

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Predict The Covariate Effects of a Univariate sparseFLMM Model
#------------------------------------------------------------------------------#
#' Predict The Covariate Effects of a Univariate sparseFLMM Model
#'
#' This function helps to compare univariate sparseFLMM
#' models to multiFAMM models. This functions predicts the covariate effects of
#' a model that has the structure of a regression model fitted by sparseFLMM().
#'
#' This functions' name in the thesis was predict_model().
#'
#' @param model sparseFLMM model with covariates.
#' @param type Gam terms to be predicted.
#' @param unconditional Std conditional on lambda.
#' @param grid Grid of evaluation points. Defaults to observations on [0,1].
#' @export
predict_sparseFLMM_covar <- function (model, type = "terms",
                                      unconditional = FALSE,
                                      grid = seq(0, 1, length.out = 100)) {

  # Use first row so that there are reasonable values for all the
  # variables
  dat <- model[[grep("^fpc_famm_hat", names(model))]]$
    famm_estim$model[1, ]

  # Set all covariable and interaction values to 1
  name <- c("^covariate", "^inter_")
  name <- grepl(paste(name, collapse = "|"), names(dat))
  dat[, name] <- 1

  # Blow up data set
  dat <- dat[rep(1, times = length(grid)), ]
  dat$yindex.vec <- grid

  # Predict data set
  out <- mgcv::predict.bam(model[[grep("^fpc_famm_hat",
                                       names(model))]]$famm_estim,
                           newdata = dat, type = type, se.fit = TRUE,
                           unconditional = unconditional)
  out$fit <- cbind(out$fit,
        intercept = rep(model[[grep("^fpc_famm_hat",
                                    names(model))]]$famm_estim$coefficients[1],
                         times = length(grid)))
  out$se.fit <- cbind(out$se.fit,
        intercept = rep(sqrt(model[[grep("^fpc_famm_hat",
                                         names(model))]]$famm_estim$Vp[1,1])))
  out

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Transform Output of Covariate Predictions of Univariate Models to Data Frame
#------------------------------------------------------------------------------#
#' Transform Output of Covariate Predictions of Univariate Models to Data Frame
#'
#' This function helps to compare univariate sparseFLMM
#' models to multiFAMM models. This functions converts the predictions of the
#' covariate effects of one or two models as given by the function
#' predict_sparseFLMM_covar() to a data.frame.
#'
#' This functions' name in the thesis was predict2DataFrame().
#'
#' @param aco_pr Output of the first model (dimension aco).
#' @param epg_pr Output of the second model (dimension epg). Can be NULL.
#' @param effect Which effect to extract. If intercept is specified (1), then
#'   the scalar intercept is added.
#' @param grid Grid of evaluation points. Defaults to observations on [0,1].
#' @export
predictedUnivar2DataFrame <- function (aco_pr, epg_pr, effect,
                                       grid = seq(0, 1, length.out = 100)) {

  # If only one model is given, keep the rest of the code and produce NA output
  if(is.null(epg_pr)){
    epg_pr <- aco_pr
    epg_pr[, ] <- NA
    mostattributes(epg_pr) <- attributes(aco_pr)
  }

  # Handle intercept differently
  if (effect == 1) {
    y <- c(aco_pr[, effect] + aco_pr[, ncol(aco_pr)],
           epg_pr[, effect] + epg_pr[, ncol(epg_pr)])
  } else {
    y <- c(aco_pr[, effect],
           epg_pr[, effect])
  }

  # Output of data.frame
  data.frame(
    t = rep(grid, times = 2),
    y = y,
    dim = rep(c(1, 2), each = length(grid))
  )

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Plot the Comparison of Univariate and Multivariate Covariate Effects
#------------------------------------------------------------------------------#
#' Plot the Comparison of Univariate and Multivariate Covariate Effects
#' @param m_true_comp Model components of the multivariate model.
#' @param aco_pr Predictions for aco model.
#' @param epg_pr Predictions for epg model.
#' @param effects Index number of the effects to be plotted from multivariate
#'   model.
#' @param effects_uni Index number of the effects to be plotted from the
#'   univariate
#' @param m_fac Multiplication factor.
#' @param limits Y-limits of the Plot.
#' @export
cova_comp_plot <- function (m_true_comp, aco_pr, epg_pr,
                            effects = c(4, 5, 6), effects_uni = c(5, 6, 7),
                            m_fac = 2, limits = limits) {

  # List of data.frames of covariate effects for multivariate model
  covars <- lapply(effects, function (x) {
    funData2DataFrame(fundata = m_true_comp$cov_preds$fit[[x]])
  })

  # List of data.frames of covariate effects for univariate models
  covars_u <- lapply(effects_uni, function (x) {
    predictedUnivar2DataFrame(aco_pr = aco_pr$fit, epg_pr = epg_pr$fit,
                              effect = x)
  })

  # List of covariate credibility intervals for multivariate model
  covars.se <- lapply(effects, function (x) {
    funData2DataFrame(fundata = m_true_comp$cov_preds$se.fit[[x]])
  })

  # List of covariate credibility intervals forunivariate models
  covars_u.se <- lapply(effects_uni, function (x) {
    predictedUnivar2DataFrame(aco_pr = aco_pr$se.fit, epg_pr = epg_pr$se.fit,
                      effect = x)
  })

  # Create the labels for the plot
  labels <- sapply(effects, function (x) {
    paste0("beta[", x-2, "](t)")
  })

  # Which approach yields "better" results
  difs <- mapply(function (x, y) {
    (x$y - y$y < 0)
  }, covars.se, covars_u.se, SIMPLIFY = FALSE)

  # Construct data.frame
  dat <- data.frame(
    t = c(do.call(c, lapply(covars, "[[", "t")),
          do.call(c, lapply(covars_u, "[[", "t"))),
    y = c(do.call(c, lapply(covars, "[[", "y")),
          do.call(c, lapply(covars_u, "[[", "y"))),
    y_plus = c(do.call(c, lapply(seq_along(covars), function (x) {
      covars[[x]]$y + m_fac*covars.se[[x]]$y
    })), do.call(c, lapply(seq_along(covars_u), function (x) {
      covars_u[[x]]$y + m_fac*covars_u.se[[x]]$y
    }))),
    y_minu = c(do.call(c, lapply(seq_along(covars), function (x) {
      covars[[x]]$y - m_fac*covars.se[[x]]$y
    })), do.call(c, lapply(seq_along(covars_u), function (x) {
      covars_u[[x]]$y - m_fac*covars_u.se[[x]]$y
    }))),
    effect = factor(c(rep(effects, times = sapply(covars, nrow)),
                      rep(effects, times = sapply(covars_u, nrow))),
                    labels = labels),
    dim = factor(c(do.call(c, lapply(covars, "[[", "dim")),
                   do.call(c, lapply(covars_u, "[[", "dim"))),
                 labels = c("ACO", "EPG")),
    method = factor(c(rep(1, times = do.call(sum, lapply(covars, nrow))),
                      rep(rep(c(2, 3), each = 100), times = length(covars_u))),
                    labels = c("multi", "aco", "epg")),
    dif = factor(rep(do.call(c, difs), times = 2),
                 labels = c("worse", "better")),
    x_start = rep(rep(c(min(grid), box_grid), times = 2*length(effects)),
                  times = 2),
    x_end = rep(rep(c(box_grid, max(grid)), times = 2*length(effects)),
                times = 2),
    min_val = limits[1],
    max_val = limits[2]
  )

  # Plot the data
  ggplot2::ggplot(data = dat, aes(x = t)) +
    geom_line(aes(y = y, col = method), size = 0.25) +
    geom_line(aes(y = y_plus, col = method), linetype = "dashed", size = 0.25) +
    geom_line(aes(y = y_minu, col = method), linetype = "dashed", size = 0.25) +
    geom_hline(yintercept = 0, linetype = "dotted", size = 0.3) +
    geom_rect(aes(ymin = min_val, ymax = max_val, xmin = x_start, xmax = x_end,
                  fill = dif), alpha = 0.15) +
    facet_grid(dim ~ effect, labeller = label_parsed) +
    xlab("Normalized Time (t)") +
    ylab(expression("Effect Function (" ~ beta^(d)~"(t) )")) +
    theme_grey(base_size = 8)  +
    theme(legend.position = "none") +
    scale_color_manual(values = c("black", "tomato2", "steelblue3")) +
    scale_fill_manual(values = c("grey80", NA)) +
    scale_y_continuous(expand = c(0, 0))

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Convert Univariate Estimates of Eigenfunctions to Data.frame
#------------------------------------------------------------------------------#
#' Convert Univariate Estimates of Eigenfunctions to Data.frame
#' @param phi_aco Univariate eigenfunction for model aco.
#' @param phi_epg Univariate eigenfunction for model epg.
#' @param grid Grid of evaluation points.
#' @export
fpc2DataFrame <- function(phi_aco, phi_epg, grid = seq(0,1, length.out = 100)) {

  # Arguments
  # phi_aco : Univariate eigenfunction for model aco
  # phi_epg : Univariate eigenfunction for model epg
  # grid    : Grid of evaluation points

  # Construct the data.frame
  dat <- data.frame(
    t = c(rep(grid, times = ncol(phi_aco)), rep(grid, times = ncol(phi_epg))),
    y = c(as.vector(phi_aco), as.vector(phi_epg)),
    dim = c(rep("ACO", length(phi_aco)), rep("EPG", length(phi_epg))),
    obs = c(rep(seq_len(ncol(phi_aco)), each = length(grid)),
            rep(seq_len(ncol(phi_epg)), each = length(grid)))
  )

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Compute the Surface of a Covariance Compoent and Convert it to Data.Frame
#------------------------------------------------------------------------------#
#' Compute the Surface of a Covariance Compoent and Convert it to Data.Frame
#' @param m_aco Univariate model on dimension aco.
#' @param m_epg  Univariate model on dimension epg.
#' @param component Covariance component to extract
cov2DataFrame <- function(m_aco, m_epg, component) {

  dimnames <- c("ACO", "EPG")

  # Restructure model for easier coding if from univariate source
  model <- list(model_indep = list(m_aco, m_epg))
  names(model$model_indep) <- c("ACO", "EPG")

  # Extract the variance component wanted
  fpc <- lapply(model$model_indep, function (x) {
    y <- list()
    y[[1]] <- x[[grep("^fpc_hat", names(x))]][[paste0("phi_", component,
                                                      "_hat_grid")]]
    y[[2]] <- x[[grep("^fpc_hat", names(x))]][[paste0("nu_", component,
                                                      "_hat")]]
    y
  })

  # Compute the Auto-covariance on each dimension separately
  cov <- lapply(fpc, function (x) {
    x[[1]] %*% diag(x[[2]], nrow = length(x[[2]])) %*% t(x[[1]])
  })

  # Fill in covariance matrix with missing values for the Cross-covariance
  cov <- rbind(cbind(cov[[1]], matrix(NA, ncol = ncol(cov[[1]]),
                                      nrow = nrow(cov[[1]]))),
               cbind(matrix(NA, ncol = ncol(cov[[2]]), nrow = nrow(cov[[2]])),
                     cov[[2]]))

  # Extract the grid on which the fPCs are computed
  grid <- model$model_indep[[1]]$my_grid

  # Name the rows and columns for matrix transformation to data.frame
  rownames(cov) <- colnames(cov) <- rep(grid, times = 2)

  # Create data.frame for plotting
  dat <- data.frame(row_dim = rep(rep(dimnames, each = length(grid)),
                                  times = 2*length(grid)),
                    col_dim = rep(dimnames,
                                  each = 2 * length(grid) * length(grid)),
                    row = as.numeric(rownames(cov)[row(cov)]),
                    col = as.numeric(colnames(cov)[col(cov)]),
                    value = c(cov))

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Load All Simulation Results of Model Components of a Folder
#------------------------------------------------------------------------------#
#' Load All Simulation Results of Model Components of a Folder
#'
#' This function takes the a folder containing simulation
#' results and what part of the model is to be loaded. It then loads all
#' corresponding results and concatenates them.
#'
#' @param folder Folder with saved objects from the simulation.
#' @param component String of the model components to be loaded.
#' @param uni Extract also the univariate part of the components. Defaults to
#'   FALSE.
#' @export
load_sim_results <- function (folder, component, uni = FALSE) {

  # load all elements of the folder that have the component in its name
  files <- list.files(path = folder)
  out <- lapply(files[grepl(component, files)], function (x) {
    load(file.path(folder, x))
    get(ls()[-grep("^x$", ls())])
  })

  # Combine the different simulation runs to one object
  out <- switch(component,
    "cov_preds" = ,
    "eigenfcts" = ,
    "eigenvals" = ,
    "error_var" = {
      list("mul" = do.call(c, lapply(out, "[[", "mul")),
           "uni" = if (!uni) NULL else do.call(c, lapply(out, "[[", "uni")))
      },
    "eigscores" = ,
    "fitted_cu" = ,
    "ran_preds" = {
      list("mul" = do.call(c, lapply(out, "[[", "mul")),
           "uni" = if (!uni) NULL else do.call(c, lapply(out, "[[", "uni")),
           "tru" = do.call(c, lapply(out, "[[", "tru")))
    })
  out[lengths(out) != 0]

}
#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#
# Return Number of Selected FPCS per Component and Iteration
#------------------------------------------------------------------------------#
#' Return Number of Selected FPCS per Component and Iteration
#'
#' This function takes the a folder containing simulation
#' results and returns a data set which contains the number of selected FPCS
#' per component.
#'
#' @inheritParams load_sim_results
#' @export
return_number_fpcs <- function (folder, uni = FALSE) {

  # Use the eigenvalues to count the number of selected components
  eigenvals <- load_sim_results(folder = folder, component = "eigenvals",
                                uni = uni)

  # Extract the number for the multivariate model
  dat_m <- do.call(rbind, lapply(seq_along(eigenvals$mul), function (it) {
    data.frame(it = it,
               no = factor(names(eigenvals$mul[[it]])),
               method = factor("mul"),
               y = sapply(eigenvals$mul[[it]], length))
  }))

  # It the simulation scenario contains univariate models, also extract their
  # number
  if (uni) {
    dat_u <-  do.call(rbind, lapply(seq_along(eigenvals$uni), function (it) {
      do.call(rbind, lapply(seq_along(eigenvals$uni[[it]]), function (dim) {
        data.frame(it = it,
                   no = factor(names(eigenvals$uni[[it]][[dim]])),
                   method = factor("uni"),
                   y = sapply(eigenvals$uni[[it]][[dim]], length),
                   dim = factor(names(eigenvals$uni[[it]])[dim]))
      }))
    }))
    dat_m$dim <- factor("mul")
  } else {
    dat_u <- NULL
  }

  rbind(dat_m, dat_u)

}

#------------------------------------------------------------------------------#
# Compute the Fitted Mean Trajectories
#------------------------------------------------------------------------------#
#' Compute the Fitted Mean Trajectories
#'
#' This function takes the fitted curves object of the
#' simulation and the random effects object and calculates the estimated mu
#' functions.
#'
#' @param fitted_cu Object of fitted curves saved from the simulation.
#' @param ran_preds Object of random effects saved from the simulation.
#' @param mul FALSE if the univariate estimate is extracted from fitted_cu.
#'   Defaults to TRUE which is the multivariate estimate.
#'
#' @inheritParams sim_eval_components
#' @export
compute_mu_sim <- function (fitted_cu, ran_preds, I = 10, J = 16, reps = 5,
                            nested = FALSE, mul = TRUE) {

  list_element <- if (mul) "mul" else "uni"
  reps_B <- rep(reps*J, times = I)
  reps_C <- rep(reps, times = J)

  if (!mul) {
    ran_preds$uni <- lapply(seq_along(ran_preds$uni), function (it) {
      obj <- lapply(names(ran_preds$uni[[it]][[1]]), function (comp) {
        multiFunData(lapply(ran_preds$uni[[it]], "[[", comp))
      })
      names(obj) <- names(ran_preds$uni[[it]][[1]])
      obj
    })
  }

  # For Random Intercept of Subject
  if ("B" %in% names(ran_preds$mul[[1]])) {
    re_B <- lapply(seq_along(ran_preds$mul), function (x) {
      multiFunData(lapply(ran_preds[[list_element]][[x]]$B, function (y) {
        funData(argvals = argvals(y),
                X = y@X[rep(1:nrow(y@X), times = reps_B), ])
      }))
    })
  } else {
    # Zero object
    re_B <- lapply(seq_along(ran_preds$mul), function (x) {
      argvals <- argvals(ran_preds[[list_element]][[x]]$mu[[1]])
      multiFunData(lapply(seq_along(ran_preds$mul[[x]]$mu), function (y) {
        funData(argvals = argvals,
                X = matrix(0, nrow = nObs(ran_preds[[list_element]][[x]]$mu),
                           ncol = length(unlist(argvals))))
      }))
    })
  }

  # Has not been tested for C in names()
  if ("C" %in% names(ran_preds$mul[[1]])) {
    if(!nested) {
      re_C <- lapply(seq_along(ran_preds$mul), function (x) {
        multiFunData(lapply(ran_preds[[list_element]][[x]]$C, function (y) {
          funData(argvals = argvals(y),
                  X = y@X[rep(rep(1:nrow(y@X), times = reps_C), times = I), ])
        }))
      })
    } else {
      re_C <- lapply(seq_along(ran_preds$mul), function (x) {
        multiFunData(lapply(ran_preds[[list_element]][[x]]$C, function (y) {
          funData(argvals = argvals(y),
                  X = y@X[rep(1:nrow(y@X), each = reps), ])
        }))
      })
    }
  } else {
    # Zero object
    re_C <- lapply(seq_along(ran_preds$mul), function (x) {
      argvals <- argvals(ran_preds$mul[[1]]$E[[1]])
      multiFunData(lapply(seq_along(ran_preds$mul[[x]]$E), function (y) {
        funData(argvals = argvals,
                X = matrix(0, nrow = nObs(ran_preds[[list_element]][[x]]$E),
                           ncol = length(unlist(argvals))))
      }))
    })
  }

  if (mul) {
    mapply(function (fit, x, y, z) {
      fit - x$E - y - z
    }, fitted_cu$mul, ran_preds$mul, re_B, re_C, SIMPLIFY = FALSE)
  } else {
    mapply(function (fit, x, y, z) {
      funData::multiFunData(fit) - x$E - y - z
    }, fitted_cu$uni, ran_preds$uni, re_B, re_C, SIMPLIFY = FALSE)
  }

}
#------------------------------------------------------------------------------#