#' Update Test Set Representations, Cluster Assignments, and Cluster Centers
#'
#' This function updates the representations, cluster assignments, and cluster centers for a test dataset based on a trained model and a reference training dataset.
#'
#' @param args A list of arguments, typically containing configurations such as the number of hidden features (\code{n_hidden_fea}) and whether to use CUDA (\code{cuda}).
#' @param model A trained model, typically an autoencoder, that will be used to generate new representations for the test data.
#' @param data_test A list representing the test dataset. This list will be updated with new representations, cluster assignments, and cluster centers.
#' @param dataloader_test An object representing the test data loader, which provides batches of test data to the model.
#' @param data_train A list representing the training dataset, which contains reference cluster centers and other attributes needed for updating the test dataset.
#'
#' @return The function updates the following attributes within the \code{data_test} list:
#' \itemize{
#'   \item \code{rep}: The new representations (embeddings) for the test dataset.
#'   \item \code{n_cat}: The number of clusters, set to match the training dataset.
#'   \item \code{M}: The cluster centers, copied from the training dataset.
#'   \item \code{pred_C}: The predicted cluster assignments for each test sample, based on the updated representations and cluster centers.
#' }
#'
#' @details
#' The \code{update_testset_R_C_M_K} function performs the following steps:
#' \enumerate{
#'   \item Initializes a tensor, \code{final_embed}, to store the embeddings for the test dataset.
#'   \item Sets the model to evaluation mode.
#'   \item Iterates through the test data batches using \code{dataloader_test}, processing each batch through the model to obtain the embeddings.
#'   \item Updates the \code{rep} attribute of \code{data_test} with the new embeddings.
#'   \item Updates \code{data_test$n_cat} and \code{data_test$M} to match the training dataset.
#'   \item For each test sample, computes the closest cluster center (based on Euclidean distance) and assigns the corresponding cluster to \code{pred_C}.
#' }
#' @examples
#' \donttest{
#' ## Called internally by DICEr() to assign test-set cluster labels.
#' ## Requires a trained DICEr model and data objects; see ?DICEr.
#' }
#' @import torch
#' @export


# Define the update_testset_R_C_M_K function
update_testset_R_C_M_K <- function(args, model, data_test, dataloader_test, data_train) {
  message("-----------------")
  message("    update_testset_R_C_M_K")

  # Initialize final_embed tensor
  final_embed <- torch_randn(nrow(data_test$data_x), args$n_hidden_fea, dtype = torch_float())

  # Set the model to evaluation mode
  model$eval()

  #dataloader_test$reset()  # Reset the iterator
  # Initialize a counter to track the current index in final_embed
  current_index <- 1
  batch_iter <- dataloader_test$.iter()  # Create an iterator

  while (current_index <= final_embed$size()[1]) {
    # To get the batch data, use an iterator
    batch_idx  <- batch_iter$.next()  # Get the first batch

    # If batch_idx is NULL, break out of the loop
    if (is.null(batch_idx)) break

    data_x <- batch_idx[[1]]
    data_x <- torch_tensor(data_x, requires_grad = FALSE)
    data_x <- data_x$unsqueeze(1)  # Add a dimension to make the shape [1, 1, 4]

    # Perform a forward pass through the model in autoencoder mode
    output <- model$forward(data_x, "autoencoder")
    enc <- output[[1]]
    pred <- output[[2]]

    # Get the embeddings
    embed <- enc$cpu()[, 1, , drop = FALSE]

    #embed <- enc$data$cpu()[,1,,drop=FALSE]
    final_embed[current_index,] <- embed

    # Increment the index counter for the next batch
    current_index <- current_index + 1
  }

  # Update data_test attributes
  data_test$rep <- final_embed
  message("        update data_test R!")

  data_test$n_cat <- data_train$n_cat
  message("        update data_test n_cat")

  data_test$M <- data_train$M
  message("        update data_test M")

  # Update data_test.C
  representations <- data_test$rep

  pred_C <- torch_zeros(data_test$rep$size()[1], dtype = torch_int())

  for (i in seq_len(data_test$rep$size()[1])) {
    embed <- representations[i,]
    trans_embed <- embed$view(c(embed$size(), 1))
    # crasshes at xj <- torrch_norm
    # Ensure the tensors are correctly aligned with 1-based indexing
    trans_embed <- embed$view(c(embed$size(), 1))  # Reshape the embedding

    xj <- torch_norm(trans_embed - data_train$data_x[i,], dim = 1)
    new_cluster <- torch_argmin(xj)
    pred_C[i] <- new_cluster
  }

  data_test$pred_C <- pred_C
  message("        update pred data_test C")
  return(data_test)
}
