R/imputeCellwise.R

In VIM: Visualization and Imputation of Missing Values

#' Cellwise-robust iterative regression imputation for mixed data
#'
#' Extends IRMI (Templ, Kowarik, and Filzmoser, 2011) with cellwise
#' contamination handling.  Each conditional regression uses a cell-weighted
#' IRWLS engine where per-cell weights in the design matrix downweight
#' contaminated cells without discarding entire observations.
#'
#' @details
#' The algorithm works iteratively: in each outer iteration, every variable
#' with missing values is used as response in a conditional regression on
#' all remaining variables.  For continuous responses, the custom
#' \code{cellIRWLS()} engine fits a weighted regression where each cell in
#' the design matrix receives its own weight reflecting potential cellwise
#' contamination.  For categorical responses, a weighted multinomial model
#' is used.  After each regression, cell weights for the response variable
#' are updated from the residuals.
#'
#' The algorithm proceeds as follows:
#' \enumerate{
#'   \item Missing values are initialised using \code{\link{initialise}}.
#'   \item Initial cell weights are computed with \code{cellWeights()} on
#'     all continuous variables in the initialised data.
#'   \item \strong{Outer loop} (up to \code{maxit} iterations):
#'     \itemize{
#'       \item For each variable \eqn{j} with missing values:
#'         \itemize{
#'           \item Form predictor matrix \eqn{X} (all other variables) and
#'             response \eqn{y} (variable \eqn{j}).
#'           \item If \eqn{j} is continuous: fit \code{cellIRWLS(X, y,
#'             w_cell, w_response)} and impute missing values in \eqn{j}
#'             using the fitted model plus uncertainty.
#'           \item If \eqn{j} is categorical: fit \code{nnet::multinom()}
#'             with row weights derived from the cell weight matrix and
#'             impute by sampling from predicted probabilities.
#'           \item Update cell weights for \eqn{j} from residuals via
#'             \code{cellWeightsFromResiduals()}.
#'         }
#'       \item Check convergence: relative change in imputed values
#'         falls below \code{eps}.
#'     }
#' }
#'
#' @param data a \code{data.frame} with missing values (mixed continuous
#'   and categorical variables are supported).
#' @param method weight function: \code{"tukey"} (default, Tukey bisquare)
#'   or \code{"huber"} (Huber).
#' @param alpha tuning constant.  \code{NULL} (default) uses 1.345 for
#'   Huber and 4.685 for Tukey, giving 95% efficiency at the normal model.
#' @param maxit maximum number of outer IRMI iterations (default: 100).
#' @param maxit_irwls maximum number of inner IRWLS iterations per
#'   regression (default: 50).
#' @param eps convergence tolerance for the outer loop (default: 5e-3).
#'   Convergence is declared when the relative change in imputed values
#'   falls below this threshold.
#' @param eps_irwls convergence tolerance for the inner IRWLS
#'   (default: 1e-6).
#' @param uncert imputation uncertainty method: \code{"pmm"} (predictive
#'   mean matching, default), \code{"normalerror"} (add normal noise), or
#'   \code{"resid"} (bootstrap residual).
#' @param weight_update strategy for updating cell weights between outer
#'   iterations: \code{"multivariate"} (default) uses an MCD-based
#'   multivariate update for weight coherence across variables, or
#'   \code{"univariate"} updates each variable independently from its
#'   residuals.
#' @param init_weights method for initialising cell weights, one of
#'   \code{"ddc"} (default; DetectDeviatingCells, requires the \pkg{cellWise}
#'   package and falls back to univariate weights when it is unavailable),
#'   \code{"univariate"} (per-column median/MAD standardisation), or
#'   \code{"mcd"} (minimum covariance determinant on the continuous block).
#'   The default is \code{"ddc"} because \code{"mcd"} downweights
#'   high-leverage points that carry the regression signal, which can make
#'   imputation worse than unconditional median imputation.
#' @param hard_threshold numeric in \eqn{[0, 1]}.  After convergence,
#'   cells with weight below this value are flagged as contaminated
#'   (default: 0.5).
#' @param trace logical; if \code{TRUE}, print progress information.
#'
#' @return A list with components:
#'   \item{data_imputed}{the imputed \code{data.frame}.}
#'   \item{cellweights}{\eqn{n \times p} matrix of final cell weights
#'     (1 = clean, 0 = fully downweighted).  Categorical columns always
#'     have weight 1.}
#'   \item{converged}{logical indicating whether the outer loop converged.}
#'   \item{iterations}{number of outer iterations used.}
#'
#' @author Matthias Templ
#' @references
#' Templ, M., Kowarik, A. and Filzmoser, P. (2011).
#' Iterative stepwise regression imputation using standard and robust
#' methods.  \emph{Computational Statistics & Data Analysis},
#' \strong{55}(10), 2793--2806.
#'
#' @family imputation methods
#' @seealso \code{\link{imputeCellM}}, \code{\link{imputeCellEM}},
#'   \code{\link{initialise}}, \code{\link{irmi}}
#'
#' @examples
#' \donttest{
#' data(sleep, package = "VIM")
#' result <- imputeCellIRMI(sleep)
#' head(result$data_imputed)
#' image(result$cellweights, main = "Cell weights")
#'
#' # With Huber weights (less aggressive downweighting)
#' result2 <- imputeCellIRMI(sleep, method = "huber", trace = TRUE)
#'
#' # Mixed data example
#' data(testdata)
#' result3 <- imputeCellIRMI(testdata$wna)
#' }
#'
#' @export
#' @importFrom stats model.matrix median mad rnorm predict sd as.formula
imputeCellIRMI <- function(data, method = "tukey", alpha = NULL,
                            maxit = 100, maxit_irwls = 50,
                            eps = 5e-3, eps_irwls = 1e-6,
                            uncert = "pmm",
                            weight_update = "multivariate",
                            init_weights = "ddc",
                            hard_threshold = 0.5,
                            trace = FALSE) {

  ## ---- input validation ----
  check_data(data)
  if (!is.data.frame(data)) {
    if (is.matrix(data))
      data <- as.data.frame(data)
    else
      stop("data must be a data.frame or matrix")
  }
  method <- match.arg(method, c("huber", "tukey"))
  uncert <- match.arg(uncert, c("pmm", "normalerror", "resid", "none"))
  weight_update <- match.arg(weight_update, c("multivariate", "per-variable"))
  init_weights <- match.arg(init_weights, c("mcd", "univariate", "ddc"))

  if (is.null(alpha)) {
    alpha <- if (method == "huber") 1.345 else 4.685
  }
  if (ncol(data) < 2) stop("Need at least 2 variables.")
  if (!any(is.na(data))) {
    message("No missing values in data. Nothing to impute.")
    n <- nrow(data)
    p <- ncol(data)
    W <- matrix(1, nrow = n, ncol = p,
                dimnames = list(rownames(data), colnames(data)))
    return(list(data_imputed = data, cellweights = W,
                converged = TRUE, iterations = 0L))
  }
  if (any(rowSums(!is.na(data)) == 0))
    stop("Unit non-responses (entire row missing) detected. Remove them first.")

  ## ---- detect variable types ----
  rn <- rownames(data)
  n <- nrow(data)
  p <- ncol(data)

  class1 <- function(x) class(x)[1]
  types <- vapply(data, class1, character(1), USE.NAMES = FALSE)

  # convert character to factor
  if (any(types == "character")) {
    chr_ind <- which(types == "character")
    warning("At least one character variable is converted into a factor")
    for (ind in chr_ind) {
      data[, ind] <- as.factor(data[, ind])
      types[ind] <- "factor"
    }
  }

  # refine factor types
  ind_fac <- which(types == "factor")
  for (ind in ind_fac) {
    fac_nlevels <- nlevels(data[[ind]])
    if (fac_nlevels < 2)
      stop(sprintf("Factor with less than 2 levels detected: '%s'",
                    names(data)[ind]))
    types[ind] <- ifelse(fac_nlevels == 2, "binary", "nominal")
  }
  ind_ord <- which(types == "ordered")
  for (ind in ind_ord) {
    fac_nlevels <- nlevels(data[[ind]])
    if (fac_nlevels == 2) types[ind] <- "binary"
  }

  is_continuous <- types %in% c("numeric", "integer")
  is_categorical <- !is_continuous

  ## ---- record missingness pattern ----
  M <- is.na(data)
  vars_miss <- which(colMeans(M) > 0)
  if (length(vars_miss) == 0) {
    W <- matrix(1, nrow = n, ncol = p,
                dimnames = list(rn, colnames(data)))
    return(list(data_imputed = data, cellweights = W,
                converged = TRUE, iterations = 0L))
  }

  ## work under safe positional names: pasted model formulas break with
  ## duplicated or non-syntactic column names (the response can resolve
  ## to the wrong column); original names are restored on exit
  cn <- colnames(data)
  colnames(data) <- .cw_safe_names(p)

  ## ---- step 1: initialise missing values ----
  data <- initialise(data, mixed = NULL, method = "median")

  ## ---- step 2: compute initial cell weights ----
  W <- matrix(1, nrow = n, ncol = p,
              dimnames = list(rn, cn))

  if (any(is_continuous)) {
    X_cont_init <- as.matrix(data[, is_continuous, drop = FALSE])
    if (init_weights == "mcd") {
      W[, is_continuous] <- cellWeightsMCD(X_cont_init,
                                            method = method, alpha = alpha)
    } else if (init_weights == "ddc" &&
               requireNamespace("cellWise", quietly = TRUE)) {
      # DDC-based initialization: detect cells, give flagged cells weight 0
      X_ddc <- X_cont_init
      colnames(X_ddc) <- .cw_safe_names(ncol(X_ddc))
      for (jj in seq_len(ncol(X_ddc))) {
        na_jj <- is.na(X_ddc[, jj])
        if (any(na_jj)) X_ddc[na_jj, jj] <- median(X_ddc[, jj], na.rm = TRUE)
      }
      ddc_res <- tryCatch({
        # checkDataSet() inside DDC prints via cat() even when silent = TRUE
        utils::capture.output(
          res <- cellWise::DDC(X_ddc,
                               DDCpars = list(fastDDC = TRUE, silent = TRUE))
        )
        res
      }, error = function(e) NULL)
      if (!is.null(ddc_res)) {
        # checkDataSet() inside DDC may refuse rows and columns (constant,
        # too discrete, too many NAs, ...), so stdResid covers only the
        # retained submatrix. Map its entries back via colInAnalysis /
        # rowInAnalysis; refused rows/columns keep univariate weights.
        W_cont <- cellWeights(X_cont_init, method = method, alpha = alpha)
        keep_col <- ddc_res$colInAnalysis
        if (is.null(keep_col)) keep_col <- seq_len(ncol(X_ddc))
        keep_row <- ddc_res$rowInAnalysis
        if (is.null(keep_row)) keep_row <- seq_len(nrow(X_ddc))
        for (jj in seq_along(keep_col)) {
          u <- abs(ddc_res$stdResid[, jj])
          w_jj <- .apply_weight_fun(u, method = method, alpha = alpha)
          w_jj[!is.finite(w_jj)] <- 1
          W_cont[keep_row, keep_col[jj]] <- w_jj
        }
        # cells missing in the original data are neutral
        W_cont[is.na(X_cont_init)] <- 1
        if (trace && length(keep_col) < ncol(X_ddc)) {
          message(sprintf(
            "  DDC init: %d of %d continuous columns analysed; univariate weights for the rest",
            length(keep_col), ncol(X_ddc)))
        }
        W[, is_continuous] <- W_cont
      } else {
        W[, is_continuous] <- cellWeightsMCD(X_cont_init,
                                              method = method, alpha = alpha)
      }
    } else {
      # Univariate fallback
      W[, is_continuous] <- cellWeights(X_cont_init,
                                         method = method, alpha = alpha)
    }
  }

  ## ---- step 2b: hard thresholding (detect-once) ----
  ## Set observed cells with low initial weight to NA and re-initialise.
  ## This removes cellwise outliers before the iterative imputation,
  ## analogous to the DDC detect-then-impute strategy.
  if (!is.null(hard_threshold) && hard_threshold > 0 && any(is_continuous)) {
    n_removed <- 0L
    for (jj in which(is_continuous)) {
      flagged_j <- !M[, jj] & (W[, jj] < hard_threshold)
      if (any(flagged_j)) {
        data[[jj]][flagged_j] <- NA
        M[flagged_j, jj] <- TRUE
        W[flagged_j, jj] <- 1  # treated as missing now
        n_removed <- n_removed + sum(flagged_j)
      }
    }
    if (n_removed > 0) {
      # Re-initialise the newly-missing cells
      data <- initialise(data, mixed = NULL, method = "median")
      # Update vars_miss
      vars_miss <- which(colMeans(M) > 0)
      if (trace) message(paste("  hard thresholding removed", n_removed,
                                "outlying cells (set to NA)"))
    }
  }

  ## ---- step 3: outer loop ----
  converged <- FALSE
  iterations <- 0L
  d <- Inf

  while (d > eps && iterations < maxit) {
    iterations <- iterations + 1L
    if (trace) {
      message("--------------------------------------")
      message(paste("cellIRMI: start of iteration", iterations))
    }
    data_previous <- data

    ## ---- inner loop: iterate over variables with missings ----
    for (j in vars_miss) {
      if (trace) {
        message(paste("  imputing variable:", j,
                      "(", cn[j], ") -",
                      ifelse(is_continuous[j], "continuous",
                             "categorical")))
      }

      miss_j <- M[, j]  # logical: which rows are missing for variable j
      n_miss <- sum(miss_j)
      if (n_miss == 0) next

      # predictor columns
      pred_cols <- setdiff(seq_len(p), j)

      if (is_continuous[j]) {
        ## ---- continuous response: cellIRWLS ----
        y <- data[[j]]

        # build numeric design matrix
        X <- .build_design_matrix(data, pred_cols)

        # cell weights for predictors (map to design matrix columns)
        w_cell_raw <- W[, pred_cols, drop = FALSE]
        w_cell <- .expand_cell_weights(data, pred_cols, w_cell_raw)

        # cell weights for response
        w_response <- W[, j]

        # fit cell-weighted IRWLS
        fit <- cellIRWLS(X, y, w_cell = w_cell, w_response = w_response,
                         method = method, alpha = alpha,
                         maxit = maxit_irwls, eps = eps_irwls)

        # compute predictions for all rows
        pred_all <- as.numeric(cbind(1, X) %*% fit$coefficients)
        sigma_hat <- fit$sigma

        # impute: use deterministic predictions during iteration
        # (uncertainty is added only to the final output)
        data[miss_j, j] <- pred_all[miss_j]

        # update cell weights for column j from residuals
        # (skip if multivariate update will overwrite for continuous vars)
        if (weight_update == "per-variable" || !is_continuous[j]) {
          resid_j <- y - pred_all
          w_new <- cellWeightsFromResiduals(
            resid_j, sigma = sigma_hat,
            method = method, alpha = alpha
          )
          # adaptive damping: lambda = 0.3 + 0.7 * t/t_max
          lambda <- 0.3 + 0.7 * min(iterations / maxit, 1)
          w_damped <- (1 - lambda) * W[, j] + lambda * w_new
          W[!miss_j, j] <- w_damped[!miss_j]
        }
        W[miss_j, j] <- 1

      } else {
        ## ---- categorical response: weighted multinomial ----
        y <- data[[j]]

        # row weights = geometric mean of cell weights across predictors
        w_row <- apply(W[, pred_cols, drop = FALSE], 1, function(ww) {
          exp(mean(log(pmax(ww, 1e-10))))
        })

        # build formula
        form_j <- as.formula(
          paste0(colnames(data)[j], " ~ ",
                 paste(colnames(data)[pred_cols], collapse = " + "))
        )

        # fit weighted multinomial
        multimod <- tryCatch({
          suppressMessages(
            nnet::multinom(form_j, data = data, weights = w_row,
                           maxit = 50, trace = FALSE, MaxNWts = 50000)
          )
        }, error = function(e) {
          warning(paste("Multinomial model failed for variable",
                        cn[j], ":", e$message,
                        "- using unweighted model"))
          suppressMessages(
            nnet::multinom(form_j, data = data,
                           maxit = 50, trace = FALSE, MaxNWts = 50000)
          )
        })

        # predict probabilities for missing rows
        prob_pred <- predict(multimod,
                             newdata = data[miss_j, , drop = FALSE],
                             type = "probs")

        # handle edge case: single missing row or binary factor
        if (is.null(dim(prob_pred))) {
          if (nlevels(y) == 2) {
            prob_pred <- matrix(c(1 - prob_pred, prob_pred),
                                nrow = n_miss)
            colnames(prob_pred) <- levels(y)
          } else {
            prob_pred <- matrix(prob_pred, nrow = 1)
            colnames(prob_pred) <- levels(y)
          }
        }

        # use mode (most probable category) during iteration
        # stochastic sampling is added in the final uncertainty step
        imputed_cats <- apply(prob_pred, 1, function(pp) {
          pp <- pmax(pp, 0)
          colnames(prob_pred)[which.max(pp)]
        })

        if (is.factor(data[[j]])) {
          data[miss_j, j] <- factor(imputed_cats,
                                     levels = levels(data[[j]]))
        } else {
          data[miss_j, j] <- imputed_cats
        }
        # categorical columns keep W[,j] = 1
      }
    }  # end inner loop

    ## ---- multivariate weight update (Option B) ----
    if (weight_update == "multivariate" && any(is_continuous)) {
      X_cont_now <- as.matrix(data[, is_continuous, drop = FALSE])
      W_mv <- cellWeightsMCD(X_cont_now, method = method, alpha = alpha)
      # Adaptive damping on the multivariate weights too
      lambda <- 0.3 + 0.7 * min(iterations / maxit, 1)
      W_old_cont <- W[, is_continuous, drop = FALSE]
      W_blended <- (1 - lambda) * W_old_cont + lambda * W_mv
      # Only update observed cells; imputed cells keep weight 1
      for (jj in which(is_continuous)) {
        local_j <- match(jj, which(is_continuous))
        obs_j <- !M[, jj]
        W[obs_j, jj] <- W_blended[obs_j, local_j]
        W[M[, jj], jj] <- 1
      }
      if (trace) {
        n_flagged <- sum(W[, is_continuous] < hard_threshold, na.rm = TRUE)
        message(paste("  multivariate weight update: flagged cells =", n_flagged))
      }
    }

    ## ---- check convergence (only on imputed cells, normalized) ----
    d <- 0
    n_imputed <- 0
    if (any(is_continuous)) {
      cont_cols <- which(is_continuous)
      for (jj in cont_cols) {
        miss_jj <- M[, jj]
        if (any(miss_jj)) {
          prev_vals <- data_previous[[jj]][miss_jj]
          curr_vals <- data[[jj]][miss_jj]
          denom <- sum(prev_vals^2) + 1e-10
          d <- d + sum((prev_vals - curr_vals)^2) / denom
          n_imputed <- n_imputed + 1L
        }
      }
    }
    if (any(is_categorical)) {
      cat_cols <- which(is_categorical)
      for (jj in cat_cols) {
        miss_jj <- M[, jj]
        if (any(miss_jj)) {
          d <- d + sum(data_previous[[jj]][miss_jj] !=
                         data[[jj]][miss_jj]) / sum(miss_jj)
          n_imputed <- n_imputed + 1L
        }
      }
    }
    if (n_imputed > 0) d <- d / n_imputed

    if (trace) {
      message(paste("  convergence criterion:", round(d, 6)))
    }

    if (d <= eps) {
      converged <- TRUE
      if (trace) message("cellIRMI converged.")
    }
  }  # end outer loop

  if (!converged && trace) {
    message(paste("cellIRMI did not converge after", maxit, "iterations.",
                  "Final criterion:", round(d, 6)))
  }

  ## ---- add imputation uncertainty to final output ----
  if (uncert != "none") {
    for (j in vars_miss) {
      miss_j <- M[, j]
      if (!any(miss_j)) next
      if (is_continuous[j]) {
        pred_cols <- setdiff(seq_len(p), j)
        X <- .build_design_matrix(data, pred_cols)
        w_cell_raw <- W[, pred_cols, drop = FALSE]
        w_cell <- .expand_cell_weights(data, pred_cols, w_cell_raw)
        w_response <- W[, j]
        fit <- cellIRWLS(X, data[[j]], w_cell = w_cell,
                          w_response = w_response,
                          method = method, alpha = alpha,
                          maxit = maxit_irwls, eps = eps_irwls)
        pred_all <- as.numeric(cbind(1, X) %*% fit$coefficients)
        data[miss_j, j] <- .add_uncertainty(
          pred = pred_all[miss_j],
          y_obs = data[[j]][!miss_j],
          pred_obs = pred_all[!miss_j],
          sigma = fit$sigma,
          uncert = uncert
        )
      }
      # categorical variables already have stochastic sampling
    }
  }

  rownames(data) <- rn
  colnames(data) <- cn
  list(
    data_imputed = data,
    cellweights  = W,
    converged    = converged,
    iterations   = iterations
  )
}


#' Cellwise M-estimation imputation
#'
#' Impute missing values using a cell-weighted M-estimation approach. Each
#' cell in the predictor matrix receives its own weight reflecting potential
#' cellwise contamination, so that contaminated predictor cells are
#' downweighted without discarding entire observations.
#'
#' The function has two interfaces: with a model formula, a single response
#' variable is imputed from the specified predictors; with a
#' \code{data.frame} (or matrix) as first argument, all variables with
#' missing values are imputed by chained equations, i.e. each such variable
#' is regressed on all remaining variables and the sweeps are iterated
#' until the imputed values stabilise.
#'
#' @param formula a model formula (e.g., \code{y ~ x1 + x2}) describing a
#'   single response to impute, or a \code{data.frame}/matrix with missing
#'   values; in the latter case all variables with missing values are
#'   imputed by chained equations and \code{data} must not be supplied.
#' @param data data.frame containing the data (formula interface only).
#' @param method weight function: \code{"tukey"} (default) or \code{"huber"}.
#'   Tukey bisquare is recommended because the consistency proof requires
#'   redescending weights.
#' @param alpha tuning constant. \code{NULL} (default) uses 4.685 for Tukey
#'   and 1.345 for Huber.
#' @param maxit_irwls maximum IRWLS iterations (default: 50).
#' @param eps_irwls convergence tolerance for IRWLS (default: 1e-6).
#' @param uncert imputation uncertainty method: \code{"pmm"} (default),
#'   \code{"normalerror"}, \code{"resid"}, or \code{"none"} (deterministic
#'   predictions; categorical variables are imputed by the most probable
#'   category).
#' @param value_back \code{"all"} (default) returns the complete dataset,
#'   or \code{"ymiss"} returns only the imputed values (formula interface
#'   only; ignored with a data.frame first argument).
#' @param maxit maximum number of chained-equation sweeps (data.frame
#'   interface only; default: 10).
#' @param eps convergence tolerance for the chained sweeps (data.frame
#'   interface only; default: 5e-3). Convergence is declared when the
#'   relative change in imputed values falls below this threshold.
#' @param trace logical; if \code{TRUE}, print progress of the chained
#'   sweeps (data.frame interface only).
#'
#' @return If \code{value_back = "ymiss"}, a named vector of imputed values
#'   (for rows that were originally missing) is returned. Otherwise, a list
#'   with components:
#'   \describe{
#'     \item{data_imputed}{the imputed data.frame (same structure as input)}
#'     \item{cellweights}{n x p matrix of final cell weights (1 = clean,
#'       0 = fully downweighted). Categorical columns always have weight 1.}
#'     \item{converged}{logical; always \code{TRUE} for single-formula
#'       imputation, convergence of the sweeps for the data.frame interface}
#'     \item{iterations}{integer; always \code{1L} for single-formula
#'       imputation, the number of sweeps for the data.frame interface}
#'   }
#'
#' @details
#' The formula interface is a lightweight single-response alternative to
#' \code{\link{imputeCellIRMI}}. It fits one cell-weighted IRWLS regression
#' using \code{cellIRWLS()} and imputes the missing values in the response
#' variable. This is appropriate when only one variable needs imputation
#' and a specific model formula is desired.
#'
#' The data.frame interface runs the same per-variable machinery as a
#' chained-equations algorithm: missing values are initialised
#' (median/mode), then each variable with missing values in turn is used as
#' response in a formula containing all remaining variables. Sweeps use
#' deterministic predictions and are iterated until the relative change of
#' the imputed values falls below \code{eps} (or \code{maxit} is reached);
#' the requested \code{uncert} step is applied once after convergence.
#'
#' For categorical response variables, a weighted multinomial model via
#' \code{\link[nnet]{multinom}} is fitted instead. Categorical predictors
#' are not subject to the cellwise contamination model (their cell weights
#' are always 1).
#'
#' @note Model uncertainty via bootstrap (Rubin's combining rules for
#'   multiple imputation) is not yet implemented. The current version
#'   provides single imputation with stochastic uncertainty (PMM or
#'   residual draw). For valid multiple imputation, call the function
#'   repeatedly with different seeds and combine using Rubin's rules.
#'
#' @author Matthias Templ
#' @references
#' M. Templ, A. Kowarik, P. Filzmoser (2011) Iterative stepwise regression
#' imputation using standard and robust methods. \emph{Computational
#' Statistics & Data Analysis}, Vol. 55, pp. 2793-2806.
#'
#' @family imputation methods
#' @seealso \code{\link{imputeCellIRMI}}, \code{\link{imputeRobust}}
#'
#' @examples
#' \donttest{
#' data(sleep, package = "VIM")
#' # Impute Dream using BodyWgt and BrainWgt as predictors
#' result <- imputeCellM(Dream ~ BodyWgt + BrainWgt, data = sleep)
#' head(result)
#'
#' # Return only imputed values
#' impvals <- imputeCellM(Dream ~ BodyWgt + BrainWgt, data = sleep,
#'                        value_back = "ymiss")
#'
#' # Huber weights (less aggressive downweighting)
#' result2 <- imputeCellM(Dream ~ BodyWgt + BrainWgt, data = sleep,
#'                        method = "huber")
#'
#' # Chained-equations interface: impute all variables with missings
#' result3 <- imputeCellM(sleep)
#' head(result3$data_imputed)
#' }
#'
#' @export
#' @importFrom stats model.frame model.extract model.matrix as.formula
imputeCellM <- function(formula, data, method = "tukey", alpha = NULL,
                        maxit_irwls = 50, eps_irwls = 1e-6,
                        uncert = "pmm", value_back = "all",
                        maxit = 10, eps = 5e-3, trace = FALSE) {

  ## ---- data.frame interface: chained equations over all variables ----
  if (!missing(formula) && !inherits(formula, "formula")) {
    if (!missing(data))
      stop("when the first argument is a data.frame/matrix, ",
           "'data' must not be supplied")
    if (!identical(value_back, "all"))
      warning("'value_back' applies to the formula interface only ",
              "and is ignored")
    return(.imputeCellM_chain(formula, method = method, alpha = alpha,
                              maxit_irwls = maxit_irwls,
                              eps_irwls = eps_irwls, uncert = uncert,
                              maxit = maxit, eps = eps, trace = trace))
  }

  ## ---- input validation ----
  check_data(data)
  if (!is.data.frame(data)) {
    if (is.matrix(data))
      data <- as.data.frame(data)
    else
      stop("data must be a data.frame or matrix")
  }
  method <- match.arg(method, c("huber", "tukey"))
  uncert <- match.arg(uncert, c("pmm", "normalerror", "resid", "none"))
  value_back <- match.arg(value_back, c("all", "ymiss"))

  if (is.null(alpha)) {
    alpha <- if (method == "huber") 1.345 else 4.685
  }
  if (sum(vapply(data, is.numeric, logical(1))) > nrow(data)) {
    warning("more continuous variables than observations; results may be unstable. Consider regularization.")
  }

  rn <- rownames(data)

  ## ---- parse formula ----
  y_var <- all.vars(formula)[1]
  x_vars <- all.vars(formula)[-1]

  if (!y_var %in% colnames(data))
    stop(paste("Response variable", y_var, "not found in data."))
  if (!all(x_vars %in% colnames(data)))
    stop("Not all predictor variables found in data.")

  missindex <- is.na(data[[y_var]])
  n_miss <- sum(missindex)

  if (n_miss == 0) {
    message(paste("No missing values in", y_var,
                  "- nothing to impute."))
    if (value_back == "ymiss") return(numeric(0))
    n <- nrow(data)
    p_d <- ncol(data)
    W <- matrix(1, nrow = n, ncol = p_d,
                dimnames = list(rn, colnames(data)))
    return(list(data_imputed = data, cellweights = W,
                converged = TRUE, iterations = 0L))
  }

  ## ---- initialise missing values in predictors ----
  pred_cols <- match(x_vars, colnames(data))
  data_work <- data
  for (col_idx in pred_cols) {
    v <- data_work[[col_idx]]
    if (any(is.na(v))) {
      if (is.numeric(v)) {
        data_work[[col_idx]][is.na(v)] <- median(v, na.rm = TRUE)
      } else if (is.factor(v) || is.character(v)) {
        mode_val <- names(which.max(table(v, useNA = "no")))
        data_work[[col_idx]][is.na(v)] <- mode_val
      }
    }
  }
  # also initialise response for complete-data operations
  y_orig <- data[[y_var]]
  if (is.numeric(y_orig)) {
    data_work[[y_var]][missindex] <- median(y_orig, na.rm = TRUE)
  } else if (is.factor(y_orig) || is.character(y_orig)) {
    mode_val <- names(which.max(table(y_orig, useNA = "no")))
    if (is.factor(y_orig)) {
      data_work[[y_var]][missindex] <- factor(mode_val,
                                               levels = levels(y_orig))
    } else {
      data_work[[y_var]][missindex] <- mode_val
    }
  }

  ## ---- determine response type ----
  y_is_continuous <- is.numeric(data[[y_var]])
  n <- nrow(data_work)

  if (y_is_continuous) {
    ## ---- continuous response: cellIRWLS ----
    y <- data_work[[y_var]]

    # build design matrix from predictors
    X <- .build_design_matrix(data_work, pred_cols)

    # compute cell weights for predictors
    w_cell_raw <- .compute_predictor_cellweights(data_work, pred_cols,
                                                  method = method,
                                                  alpha = alpha)
    w_cell <- .expand_cell_weights(data_work, pred_cols, w_cell_raw)

    # initial response weights (all 1)
    w_response <- rep(1, n)

    # fit cellIRWLS
    fit <- cellIRWLS(X, y, w_cell = w_cell, w_response = w_response,
                     method = method, alpha = alpha,
                     maxit = maxit_irwls, eps = eps_irwls,
                     init = "s-estimator")

    # predictions
    pred_all <- as.numeric(cbind(1, X) %*% fit$coefficients)
    sigma_hat <- fit$sigma

    # impute with uncertainty
    if (uncert == "none") {
      ymiss <- pred_all[missindex]
    } else {
      ymiss <- .add_uncertainty(
        pred = pred_all[missindex],
        y_obs = y[!missindex],
        pred_obs = pred_all[!missindex],
        sigma = sigma_hat,
        uncert = uncert
      )
    }

  } else {
    ## ---- categorical response: weighted multinomial ----
    y <- data_work[[y_var]]

    # compute row weights from cell weights of continuous predictors
    cont_pred <- pred_cols[vapply(data_work[pred_cols], is.numeric,
                                  logical(1))]
    if (length(cont_pred) > 0) {
      w_cell_raw <- .compute_predictor_cellweights(
        data_work, cont_pred, method = method, alpha = alpha
      )
      # row weights = product of predictor cell weights (Eq. 11)
      w_row <- apply(w_cell_raw, 1, function(ww) {
        exp(mean(log(pmax(ww, 1e-10))))
      })
    } else {
      w_row <- rep(1, n)
    }

    # fit weighted multinomial; the weights must live in `data`, because
    # model.frame() evaluates the weights argument in `data` and then in
    # environment(formula) -- the caller's environment, where w_row does
    # not exist
    data_fit <- data_work
    data_fit[[".cw_row_weights"]] <- w_row
    multimod <- tryCatch({
      suppressMessages(
        nnet::multinom(formula, data = data_fit, weights = .cw_row_weights,
                       maxit = 50, trace = FALSE, MaxNWts = 50000)
      )
    }, error = function(e) {
      warning(paste("Multinomial model failed:", e$message,
                    "- using unweighted model"))
      suppressMessages(
        nnet::multinom(formula, data = data_work,
                       maxit = 50, trace = FALSE, MaxNWts = 50000)
      )
    })

    prob_pred <- predict(multimod,
                         newdata = data_work[missindex, , drop = FALSE],
                         type = "probs")

    lvls <- levels(data[[y_var]])
    if (is.null(lvls)) lvls <- sort(unique(na.omit(data[[y_var]])))

    # handle edge cases
    if (is.null(dim(prob_pred))) {
      if (length(lvls) == 2) {
        prob_pred <- matrix(c(1 - prob_pred, prob_pred),
                            nrow = n_miss)
        colnames(prob_pred) <- lvls
      } else {
        prob_pred <- matrix(prob_pred, nrow = 1)
        colnames(prob_pred) <- lvls
      }
    }

    ymiss <- apply(prob_pred, 1, function(pp) {
      pp <- pmax(pp, 0)
      if (uncert == "none") {
        # deterministic: most probable category
        colnames(prob_pred)[which.max(pp)]
      } else {
        pp <- pp / sum(pp)
        sample(colnames(prob_pred), size = 1, prob = pp)
      }
    })

    if (is.factor(data[[y_var]])) {
      ymiss <- factor(ymiss, levels = levels(data[[y_var]]))
    }
  }

  ## ---- return ----
  if (value_back == "ymiss") {
    return(ymiss)
  } else {
    data[missindex, y_var] <- ymiss
    rownames(data) <- rn

    # Compute cell weights for the full data
    n <- nrow(data)
    p <- ncol(data)
    W <- matrix(1, nrow = n, ncol = p,
                dimnames = list(rn, colnames(data)))
    num_cols <- which(vapply(data, is.numeric, logical(1)))
    if (length(num_cols) > 0) {
      W[, num_cols] <- cellWeights(
        as.matrix(data[, num_cols, drop = FALSE]),
        method = method, alpha = alpha
      )
    }

    return(list(
      data_imputed = data,
      cellweights  = W,
      converged    = TRUE,
      iterations   = 1L
    ))
  }
}


#' Chained-equations engine behind the imputeCellM data.frame interface
#'
#' Runs IRMI-style sweeps: every variable with missing values in turn is
#' the response of a single-formula \code{imputeCellM()} fit on all
#' remaining variables. Sweeps are deterministic (\code{uncert = "none"})
#' and iterated until the imputed values stabilise; the requested
#' uncertainty step is applied once after convergence.
#'
#' @param data data.frame or matrix with missing values
#' @inheritParams imputeCellM
#' @return list with data_imputed, cellweights, converged, iterations
#' @keywords internal
#' @noRd
.imputeCellM_chain <- function(data, method = "tukey", alpha = NULL,
                               maxit_irwls = 50, eps_irwls = 1e-6,
                               uncert = "pmm", maxit = 10, eps = 5e-3,
                               trace = FALSE) {

  ## ---- input validation ----
  check_data(data)
  if (!is.data.frame(data)) {
    if (is.matrix(data))
      data <- as.data.frame(data)
    else
      stop("data must be a data.frame or matrix")
  }
  method <- match.arg(method, c("huber", "tukey"))
  uncert <- match.arg(uncert, c("pmm", "normalerror", "resid", "none"))
  if (is.null(alpha)) {
    alpha <- if (method == "huber") 1.345 else 4.685
  }
  if (ncol(data) < 2) stop("Need at least 2 variables.")

  rn <- rownames(data)
  cn <- colnames(data)
  n <- nrow(data)
  p <- ncol(data)

  if (!any(is.na(data))) {
    message("No missing values in data. Nothing to impute.")
    W <- matrix(1, nrow = n, ncol = p, dimnames = list(rn, cn))
    return(list(data_imputed = data, cellweights = W,
                converged = TRUE, iterations = 0L))
  }
  if (any(rowSums(!is.na(data)) == 0))
    stop("Unit non-responses (entire row missing) detected. Remove them first.")

  ## ---- variable types (as in imputeCellIRMI) ----
  chr_ind <- which(vapply(data, is.character, logical(1)))
  if (length(chr_ind) > 0) {
    warning("At least one character variable is converted into a factor")
    for (ind in chr_ind) data[[ind]] <- as.factor(data[[ind]])
  }
  for (j in seq_len(p)) {
    if (is.factor(data[[j]]) && nlevels(data[[j]]) < 2)
      stop(sprintf("Factor with less than 2 levels detected: '%s'", cn[j]))
  }

  M <- is.na(data)
  vars_miss <- which(colMeans(M) > 0)

  ## work under safe positional names: pasted formulas break with
  ## duplicated or non-syntactic column names
  cn_int <- .cw_safe_names(p)
  colnames(data) <- cn_int

  ## ---- initialise and sweep ----
  data <- initialise(data, mixed = NULL, method = "median")

  impute_var <- function(data, j, uncert_j) {
    work <- data
    work[[j]][M[, j]] <- NA
    form_j <- stats::as.formula(
      paste0(cn_int[j], " ~ ", paste(cn_int[-j], collapse = " + "))
    )
    imputeCellM(form_j, work, method = method, alpha = alpha,
                maxit_irwls = maxit_irwls, eps_irwls = eps_irwls,
                uncert = uncert_j, value_back = "ymiss")
  }

  converged <- FALSE
  iterations <- 0L
  d <- Inf

  while (d > eps && iterations < maxit) {
    iterations <- iterations + 1L
    data_previous <- data

    for (j in vars_miss) {
      # deterministic sweeps; uncertainty is added after convergence
      data[M[, j], j] <- impute_var(data, j, uncert_j = "none")
    }

    ## convergence on imputed cells (criterion as in imputeCellIRMI)
    d <- 0
    n_imputed <- 0L
    for (j in vars_miss) {
      miss_j <- M[, j]
      prev_vals <- data_previous[[j]][miss_j]
      curr_vals <- data[[j]][miss_j]
      if (is.numeric(curr_vals)) {
        d <- d + sum((prev_vals - curr_vals)^2) / (sum(prev_vals^2) + 1e-10)
      } else {
        d <- d + sum(prev_vals != curr_vals) / sum(miss_j)
      }
      n_imputed <- n_imputed + 1L
    }
    if (n_imputed > 0) d <- d / n_imputed

    if (trace) {
      message(sprintf("cellM chain: iteration %d, criterion %g",
                      iterations, d))
    }
    if (d <= eps) converged <- TRUE
  }

  if (!converged && trace) {
    message(sprintf("cellM chain did not converge after %d iterations.",
                    maxit))
  }

  ## ---- final pass: add the requested imputation uncertainty ----
  if (uncert != "none") {
    for (j in vars_miss) {
      data[M[, j], j] <- impute_var(data, j, uncert_j = uncert)
    }
  }

  ## ---- restore names, compute final cell weights ----
  colnames(data) <- cn
  rownames(data) <- rn
  W <- matrix(1, nrow = n, ncol = p, dimnames = list(rn, cn))
  num_cols <- which(vapply(data, is.numeric, logical(1)))
  if (length(num_cols) > 0) {
    W[, num_cols] <- cellWeights(
      as.matrix(data[, num_cols, drop = FALSE]),
      method = method, alpha = alpha
    )
  }

  list(
    data_imputed = data,
    cellweights  = W,
    converged    = converged,
    iterations   = iterations
  )
}


# ============================================================================
# Internal helper functions
# ============================================================================

#' Build a numeric design matrix from selected columns
#'
#' Handles both numeric and factor predictors. Factors are expanded into
#' dummy variables (treatment coding, dropping the first level as
#' reference). Does NOT include an intercept column.
#'
#' @param data data.frame
#' @param pred_cols integer vector of column indices
#' @return numeric matrix (n x q) where q is the total number of design
#'   columns after dummy expansion
#' @keywords internal
#' @noRd
.build_design_matrix <- function(data, pred_cols) {
  parts <- list()
  for (col in pred_cols) {
    v <- data[[col]]
    if (is.numeric(v)) {
      parts[[length(parts) + 1L]] <- matrix(
        v, ncol = 1,
        dimnames = list(NULL, colnames(data)[col])
      )
    } else if (is.factor(v) || is.character(v)) {
      if (is.character(v)) v <- as.factor(v)
      # create dummy variables (omit first level as reference)
      lvls <- levels(v)
      if (length(lvls) > 1) {
        dmat <- matrix(0, nrow = nrow(data), ncol = length(lvls) - 1)
        colnames(dmat) <- paste0(colnames(data)[col], lvls[-1])
        for (k in seq_along(lvls[-1])) {
          dmat[, k] <- as.numeric(v == lvls[k + 1])
        }
        parts[[length(parts) + 1L]] <- dmat
      }
    }
  }
  do.call(cbind, parts)
}


#' Expand per-variable cell weights to match the design matrix
#'
#' When factors are dummy-expanded, each dummy column inherits the cell
#' weight of the parent factor variable (which is 1 for categorical
#' variables).
#'
#' @param data data.frame
#' @param pred_cols integer vector of column indices in \code{data}
#' @param w_cell_raw n x length(pred_cols) matrix of per-variable
#'   cell weights
#' @return n x q matrix matching the number of design matrix columns
#' @keywords internal
#' @noRd
.expand_cell_weights <- function(data, pred_cols, w_cell_raw) {
  parts <- list()
  k <- 0L
  for (i in seq_along(pred_cols)) {
    col <- pred_cols[i]
    v <- data[[col]]
    k <- k + 1L
    if (is.numeric(v)) {
      parts[[length(parts) + 1L]] <- w_cell_raw[, k, drop = FALSE]
    } else if (is.factor(v) || is.character(v)) {
      if (is.character(v)) v <- as.factor(v)
      n_dummy <- max(nlevels(v) - 1L, 0L)
      if (n_dummy > 0) {
        # replicate the cell weight for each dummy column
        parts[[length(parts) + 1L]] <- matrix(
          rep(w_cell_raw[, k], n_dummy),
          nrow = nrow(data), ncol = n_dummy
        )
      }
    }
  }
  do.call(cbind, parts)
}


#' Compute cell weights for predictor columns
#'
#' Applies \code{cellWeights()} to the continuous predictors and returns
#' an n x length(pred_cols) matrix where categorical columns get weight 1.
#'
#' @param data data.frame
#' @param pred_cols integer vector of column indices
#' @param method \code{"huber"} or \code{"tukey"}
#' @param alpha tuning constant
#' @return n x length(pred_cols) matrix
#' @keywords internal
#' @noRd
.compute_predictor_cellweights <- function(data, pred_cols,
                                            method, alpha) {
  n <- nrow(data)
  W_pred <- matrix(1, nrow = n, ncol = length(pred_cols))

  is_num <- vapply(data[pred_cols], is.numeric, logical(1))
  cont_idx <- which(is_num)

  if (length(cont_idx) > 0) {
    X_cont <- as.matrix(data[, pred_cols[cont_idx], drop = FALSE])
    W_pred[, cont_idx] <- cellWeights(X_cont, method = method,
                                       alpha = alpha)
  }

  W_pred
}


#' Add imputation uncertainty
#'
#' Depending on the \code{uncert} method, adds noise or uses predictive
#' mean matching to introduce appropriate uncertainty into imputed values.
#'
#' @param pred numeric vector of predictions for the missing rows
#' @param y_obs numeric vector of observed response values
#' @param pred_obs numeric vector of predictions for observed rows
#' @param sigma estimated residual standard deviation
#' @param uncert one of \code{"pmm"}, \code{"normalerror"}, \code{"resid"}
#' @return numeric vector of imputed values (same length as \code{pred})
#' @keywords internal
#' @noRd
.add_uncertainty <- function(pred, y_obs, pred_obs, sigma, uncert) {
  n_miss <- length(pred)

  if (uncert == "normalerror") {
    ymiss <- pred + rnorm(n_miss, mean = 0, sd = sigma)
  } else if (uncert == "resid") {
    obs_resid <- y_obs - pred_obs
    ymiss <- pred + sample(obs_resid, size = n_miss, replace = TRUE)
  } else {
    # PMM: predictive mean matching with 5 donors
    ymiss <- .pmm_impute(pred, y_obs, pred_obs, n_donors = 5)
  }

  ymiss
}


#' Predictive mean matching
#'
#' For each predicted value for a missing row, find the \code{n_donors}
#' observed values whose predictions are closest, then randomly sample
#' one observed value from those donors.
#'
#' @param pred_miss predictions for missing rows
#' @param y_obs observed response values
#' @param pred_obs predictions for observed rows
#' @param n_donors number of candidate donors (default: 5)
#' @return numeric vector of imputed values
#' @keywords internal
#' @noRd
.pmm_impute <- function(pred_miss, y_obs, pred_obs, n_donors = 5) {
  n_miss <- length(pred_miss)
  n_obs <- length(y_obs)
  n_donors <- min(n_donors, n_obs)

  ymiss <- numeric(n_miss)
  for (i in seq_len(n_miss)) {
    dists <- abs(pred_miss[i] - pred_obs)
    donor_idx <- order(dists)[seq_len(n_donors)]
    donors <- y_obs[donor_idx]
    ymiss[i] <- donors[sample.int(length(donors), 1)]
  }

  ymiss
}

#' Unified cellwise-robust imputation dispatcher
#'
#' Convenience wrapper that dispatches to one of the three cellwise-robust
#' imputation methods: \code{\link{imputeCellIRMI}}, \code{\link{imputeCellM}},
#' or \code{\link{imputeCellEM}}.
#'
#' @param data data.frame with missing values (mixed continuous + categorical).
#' @param method imputation method: \code{"cellIRMI"} (default),
#'   \code{"cellM"}, or \code{"cellEM"}.
#' @param ... additional arguments passed to the chosen method.
#'
#' @return The return value of the dispatched function. See the documentation
#'   of the individual methods for details.
#'
#' @note Model uncertainty via bootstrap (Rubin's combining rules for
#'   multiple imputation) is not yet implemented. The current version
#'   provides single imputation with stochastic uncertainty (PMM or
#'   residual draw). For valid multiple imputation, call the function
#'   repeatedly with different seeds and combine using Rubin's rules.
#'
#' @author Matthias Templ
#' @family imputation methods
#' @seealso \code{\link{imputeCellIRMI}}, \code{\link{imputeCellM}},
#'   \code{\link{imputeCellEM}}
#'
#' @examples
#' \donttest{
#' data(sleep, package = "VIM")
#' result <- imputeCellwise(sleep, method = "cellIRMI")
#' head(result$data_imputed)
#' }
#'
#' @export
imputeCellwise <- function(data, method = c("cellIRMI", "cellM", "cellEM"), ...) {
  method <- match.arg(method)
  switch(method,
    cellIRMI = imputeCellIRMI(data, ...),
    cellM = imputeCellM(data, ...),
    cellEM = imputeCellEM(data, ...)
  )
}