#' @title validate_args
#' @description A utility function to validate arguments passed to a att_gt()
#' @param data data.table used in function
#' @param args list of arguments to validate
#'
#' @return nothing, but throws an error if any of the arguments are not valid
#' @noRd
validate_args <- function(args, data){

  data_names <- names(data)

  # ---------------------- Error Checking ----------------------
  args$control_group <- args$control_group[1]
  # Flag for control group types
  control_group_message <- "control_group must be either 'nevertreated' or 'notyettreated'"
  dreamerr::check_set_arg(args$control_group, "match", .choices = c("nevertreated", "notyettreated"), .message = control_group_message, .up = 1)

  # Flag for tname, gname, yname
  name_message <- "__ARG__ must be a character scalar and a name of a column from the dataset."
  dreamerr::check_set_arg(args$tname, args$gname, args$yname, "match", .choices = data_names, .message = name_message, .up = 1)

  # Strip idname from clustervars before validation: users may naturally pass
  # clustervars = c(idname, extra_var), and idname clustering is implicit/redundant.
  # This must happen before the dreamerr check which enforces scalar length.
  if (!is.null(args$clustervars) && !is.null(args$idname) && (args$idname %in% args$clustervars)) {
    args$clustervars <- setdiff(args$clustervars, args$idname)
    if (length(args$clustervars) == 0L) args$clustervars <- NULL
  }

  # At most one cluster variable beyond idname is supported (clustering at the
  # unit level via idname is implicit). Checked before the dreamerr scalar check
  # below so the error wording matches the slow path (pre_process_did) and
  # mboot(), instead of a dreamerr message exposing internal argument names.
  if (length(args$clustervars) > 1) {
    stop("At most one cluster variable (beyond 'idname') is supported. Please reduce to one.")
  }

  # Flag for clustervars and weightsname
  checkvar_message <- "__ARG__ must be NULL or a character scalar that is a name of a column from the dataset."
  dreamerr::check_set_arg(args$weightsname, args$clustervars, "NULL | match", .choices = data_names, .message = checkvar_message, .up = 1)

  # check if time periods are numeric
  if(!data[, is.numeric(get(args$tname))]){stop("The time variable '", args$tname, "' must be numeric. Please convert it.")}

  # Check if gname is numeric
  if(!data[, is.numeric(get(args$gname))]){stop("The group variable '", args$gname, "' must be numeric. Please convert it.")}

  # Check if yname is numeric (logical 0/1 outcomes are also allowed)
  if(!data[, is.numeric(get(args$yname)) || is.logical(get(args$yname))]){stop("The outcome variable '", args$yname, "' must be numeric. Please convert it.")}

  # Flag for idname
  if(!is.null(args$idname)){
    # Check if idname is in the data
    name_message <- "__ARG__ must be a character scalar and a name of a column from the dataset."
    dreamerr::check_set_arg(args$idname, "match", .choices = data_names, .message = name_message, .up = 1)

    #  check if idname is numeric
    if(!data[, is.numeric(get(args$idname))]){stop("The id variable '", args$idname, "' must be numeric. Please convert it.")}

    # Check if gname is unique by idname: irreversibility of the treatment
    # Use direct column access instead of get() for speed
    id_g_unique <- unique(data[, c(args$idname, args$gname), with = FALSE])
    check_treatment_uniqueness <- anyDuplicated(id_g_unique[[1]]) == 0L
    if (!check_treatment_uniqueness) {
      stop("The value of gname (treatment variable) must be the same across all periods for each particular unit. The treatment must be irreversible.")
    }

    # Check if any combination of idname and tname is duplicated
    n_id_year <- anyDuplicated(data, by = c(args$idname, args$tname))
    # If any combination is duplicated, stop execution and throw an error
    if (n_id_year > 0) {
      stop("The value of idname must be unique (by tname). Some units are observed more than once in a period.")
    }
  }

  # Flag for base period: not in c("universal", "varying"), stop
  args$base_period <- args$base_period[1]
  base_period_message <- "base_period must be either 'universal' or 'varying'."
  dreamerr::check_set_arg(args$base_period, "match", .choices = c("universal", "varying"), .message = base_period_message, .up = 1)

  # Flags for cluster variable
  # Note: idname was already stripped from clustervars and the at-most-one check
  # was enforced above (before the dreamerr check)
  if (!is.null(args$clustervars)) {
    # Check that cluster variables do not vary over time within each unit. For true repeated cross-sections
    # the user may omit idname (an internal ".rowid" is created later, one observation per row), so there is
    # no within-unit time variation to check and idname is not yet available here -- skip the check then.
    if (length(args$clustervars) > 0 && !is.null(args$idname)) {
      # Efficiently check for time-varying cluster variables
      clust_tv <- data[, lapply(.SD, function(col) length(unique(col)) == 1), by = get(args$idname), .SDcols = args$clustervars]
      # If any cluster variable varies over time within any unit, stop execution
      # (same wording as the slow path and mboot())
      if (!all(unlist(clust_tv[, -1, with = FALSE]))) {
        stop("Time-varying cluster variables are not supported. Please provide a time-invariant cluster variable.")
      }
    }
  }

  # Check if anticipation is numeric using
  if (!is.numeric(args$anticipation)) {
    stop("anticipation must be numeric. Please convert it.")
  }

  # Check if anticipation is positive
  if (args$anticipation < 0) {
    stop("anticipation must be non-negative. Please check your arguments.")
  }

}

#' @title DiD standardization
#' @description A utility function to coerce the data into standard format
#' @param data data.table used in function
#' @param args list of arguments to validate
#'
#' @return A list containing ordered data and arguments
#' @noRd
did_standardization <- function(data, args){
  # keep relevant columns plus the RAW covariate variables (all.vars of xformla), so
  # model.matrix(args$xformla, .) can be rebuilt downstream. Keeping the evaluated
  # model.frame instead would lose the raw variable for transform formulae
  # (~I(X^2), ~poly(X, 2), ~log(X)) and create matrix-valued columns that break the
  # data.table coercion. For bare-variable / factor formulae this is identical.
  xvars <- all.vars(args$xformla)
  cols_to_keep <- unique(c(args$idname, args$tname, args$gname, args$yname,
                           args$weightsname, args$clustervars, xvars))
  # Subset the dataset to keep only the relevant columns
  data <- data[, ..cols_to_keep]

  # Check if any covariates were missing
  n_orig <- data[, .N]
  data <- data[complete.cases(data)]
  # also drop rows whose EVALUATED design is non-finite (e.g. log of a non-positive
  # covariate); safe now that raw-covariate NAs are removed (so poly()/ns()/... will
  # not error on NA input). Use model.frame (NOT model.matrix) with na.action =
  # na.pass: model.frame keeps every row -- including NA/NaN-valued terms -- so the
  # complete.cases() mask stays aligned with `data` (model.matrix would silently drop
  # NaN rows, shortening the mask and letting the offending rows survive). Inf-valued
  # terms are kept, matching the prior behavior.
  if (length(xvars) > 0L && data[, .N] > 0L) {
    mf_check <- suppressWarnings(model.frame(args$xformla, data = data, na.action = na.pass))
    finite_rows <- complete.cases(mf_check)
    if (!all(finite_rows)) data <- data[finite_rows]
  }
  n_new <- data[, .N]
  n_diff <- n_orig - n_new
  if (n_diff != 0) {
    warning(paste0("dropped ", n_diff, " rows from original data due to missing data"))
  }

  # Set weights
  if (is.null(args$weightsname)) weights <- rep(1, n_new) else weights <- data[[args$weightsname]]
  # Validate user-supplied weights: negative weights flip signs and a non-positive
  # mean divides by ~0 during normalization, silently producing NA/NaN ATTs.
  if (!is.null(args$weightsname) && (any(weights < 0, na.rm = TRUE) || isTRUE(mean(weights, na.rm = TRUE) <= 0)))
    stop("The weights variable '", args$weightsname, "' must be non-negative with a positive mean.")
  # enforcing normalization of weights. At this point we already drop any missing values in weights.
  weights <- weights/mean(weights)
  if (".w" %in% names(data)) stop("Your data already contains a column named '.w', which is reserved for internal use by `did`. Please rename this column before calling att_gt().")
  data.table::set(data, j = ".w", value = weights)

  # Check for time-varying weights in panel data
  if (!is.null(args$weightsname) && args$panel) {
    w_range <- data[, .(w_range = max(.w) - min(.w)), by = c(args$idname)]
    if (any(w_range$w_range > .Machine$double.eps^0.5, na.rm = TRUE)) {
      message(
        "Time-varying weights detected. For balanced panel data, the default ",
        "behavior uses the weight from the earlier of the two time periods in ",
        "each 2x2 comparison (the base period for post-treatment cells). ",
        "Use the 'fix_weights' argument to control this behavior. ",
        "See ?att_gt for details."
      )
    }
  }

  # get a list of dates from min to max
  tlist <- data[, sort(unique(get(args$tname)))]

  # Ensure gname column is numeric (double) so it can hold Inf
  if (is.integer(data[[args$gname]])) {
    data[, (args$gname) := as.numeric(get(args$gname))]
  }

  # Coerce control group identified with zero as Inf
  data[get(args$gname) == 0, (args$gname) := Inf]

  # Working out the dates
  # Identify groups with treatment time bigger than the maximum treatment time + anticipation.
  # We account for anticipation because units treated shortly after the last observed period
  # may already exhibit anticipatory effects during the observed time window, and should
  # not be used as controls.
  max_treatment_time <- max(tlist, na.rm = TRUE)
  # gname has no NAs here (complete cases were enforced above), so the comparison
  # cannot yield NA; keep it as a local vector instead of a temporary column
  asif_never_treated <- data[[args$gname]] > max_treatment_time + args$anticipation
  # set gname to Inf for those groups considered as never treated
  if (any(asif_never_treated)) {
    data.table::set(data, i = which(asif_never_treated), j = args$gname, value = Inf)
  }

  # get list of treated groups (by time) from min to max
  glist <- sort(unique(data[[args$gname]]))

  # Check if there is a never treated group in the data
  # if (!(Inf %in% glist)) {
  #   if (args$control_group == "nevertreated") {
  #     stop("There is no available never-treated group")
  #   } else {
  #     # Drop all time periods with time periods >= latest treated
  #     max_treated_time <- max(glist[is.finite(glist)], na.rm = TRUE)
  #     latest_treated_time <- max_treated_time - args$anticipation
  #     data <- data[get(args$tname) < latest_treated_time]
  #
  #     tlist <- sort(unique(data[[args$tname]]))
  #     glist <- sort(unique(data[[args$gname]]))
  #
  #     # don't compute ATT(g,t) for groups that are only treated at end
  #     # and only play a role as a comparison group
  #     glist <- glist[glist < max_treated_time]
  #   }
  # }

  if (!(Inf %in% glist)) {
    # Compute latest treated cohort once, and the cutoff time
    latest_g <- max(glist[is.finite(glist)], na.rm = TRUE)
    cutoff_t  <- latest_g - args$anticipation

    if (args$control_group == "nevertreated") {
      # Warn the user (same wording as the slow path, pre_process_did: the
      # filtering sentence matters -- periods >= cutoff_t are dropped just below)
      warning(
        "No never-treated group is available. ",
        "The last treated cohort is being coerced as 'never-treated' units, and data from periods after that is being filtered out (no available comparison groups)."
      )

      # Drop all periods ≥ (latest_g - anticipation)
      data <- data[get(args$tname) < cutoff_t]

      # For any row where gname == latest_g, set gname := Inf
      data[, (args$gname) := as.numeric(get(args$gname))] # Convert the column to numeric so Inf can be stored
      data[get(args$gname) == latest_g, (args$gname) := Inf]
    } else {
      # 3. If not "nevertreated", we simply drop those periods and leave gnames alone
      data <- data[get(args$tname) < cutoff_t]
    }

    # Recompute tlist and glist from the filtered/modified data
    tlist <- sort(unique(data[[args$tname]]))
    glist <- sort(unique(data[[args$gname]]))

    # 5. If control_group != "nevertreated", drop the max cohort from glist
    if (args$control_group != "nevertreated") {
      glist <- glist[glist < latest_g]
    }
  }


  # get only the first period
  first_period <- tlist[1]

  # drop groups treated in the first period or before; keep only treated groups
  glist <- glist[glist != Inf & glist > first_period + args$anticipation]

  # Check for groups treated in the first period (accounting for anticipation) and drop them
  # identify groups treated in the first period; gname has no NAs here, so keep a
  # local vector instead of a temporary column (it is only read before the subset below)
  gvec <- data[[args$gname]]
  treated_first_period <- is.finite(gvec) & gvec <= first_period + args$anticipation

  # count the number of units treated in the first period
  if (args$panel) {
    nfirstperiod <- length(unique(data[[args$idname]][treated_first_period]))
  } else {
    nfirstperiod <- sum(treated_first_period)
  }

  # handle units treated in the first period
  if (nfirstperiod > 0) {
    warning(paste0("Dropped ", nfirstperiod, " units that were already treated in the first period",
                    if (args$anticipation > 0) paste0(" (accounting for anticipation = ", args$anticipation, ")") else "",
                    "."))
    data <- data[get(args$gname) %in% c(glist, Inf)]

    # update tlist and glist
    tlist <- data[, sort(unique(get(args$tname)))]
    glist <- data[, sort(unique(get(args$gname)))]

    # Drop groups treated in the first period or before
    first_period <- tlist[1]
    glist <- glist[glist != Inf & glist > first_period + args$anticipation]
  }

  # If user specifies repeated cross sections,
  # set that it really is repeated cross sections
  args$true_repeated_cross_sections <- FALSE
  if (!args$panel) {
    args$true_repeated_cross_sections <- TRUE
  }

  #-----------------------------------------------------------------------------
  # setup data in panel case
  #-----------------------------------------------------------------------------
  # Check if data is a balanced panel if panel = TRUE and allow_unbalanced_panel = TRUE
  bal_panel_test <- (args$panel)*(args$allow_unbalanced_panel)

  if (bal_panel_test) {
    # First, focus on complete cases. Rows with missing data were already dropped
    # in did_standardization() above, so this is a guarded no-op in the common case.
    data_comp <- if (anyNA(data)) data[complete.cases(data)] else data

    # Check if the panel is balanced
    is_balanced <- check_balance(data_comp, id_col = args$idname, time_col = args$tname)

    args$allow_unbalanced_panel <- !is_balanced
    message(if (!is_balanced)
      "You have an unbalanced panel. Proceeding as such."
      else
        "You have a balanced panel. Setting allow_unbalanced_panel = FALSE.")
  }

  if (args$panel) {
    # Check for unbalanced panel
    if (args$allow_unbalanced_panel) {
      # Flag for true repeated cross sections
      args$panel <- FALSE
      args$true_repeated_cross_sections <- FALSE
    } else {
      # Coerce balanced panel

      # Focus on complete cases. Rows with missing data were already dropped in
      # did_standardization() above, so this is a guarded no-op in the common case.
      if (anyNA(data)) {
        keepers <- complete.cases(data)
        id_col <- data[[args$idname]]
        n <- length(unique(id_col))
        n_keep <- length(unique(id_col[keepers]))
        if (n_keep < n) {
          warning(paste0("Dropped ", (n - n_keep), " observations that had missing data."))
          data <- data[keepers, ]
        }
      }

      # Make balanced panel. validate_args() guarantees at most one row per (id, t),
      # so the panel is balanced iff nrow == n_ids * n_periods; the grouped count is
      # only needed in the unbalanced case to identify the offending units.
      raw_time_size <- length(unique(data[[args$tname]]))
      if (nrow(data) != length(unique(data[[args$idname]])) * raw_time_size) {
        unit_count <- data[, .(count = .N), by = c(args$idname)]
        mis_unit_ids <- unit_count[[1L]][unit_count[[2L]] < raw_time_size]
        warning(length(mis_unit_ids), " units are missing in some periods. Converting to balanced panel by dropping them.")
        data <- data[!(data[[args$idname]] %in% mis_unit_ids)]
      }

      # If all data is dropped, stop execution
      if (nrow(data) == 0) {
        stop("All observations dropped while converting data to balanced panel. Consider setting `panel = FALSE` and/or revisiting 'idname'.")
      }

      n <- sum(data[[args$tname]] == tlist[1])

      # Note: treatment irreversibility was already checked in validate_args()
    }
  }

  #-----------------------------------------------------------------------------
  # setup data in repeated cross section
  #-----------------------------------------------------------------------------

  if (!args$panel) {
    # Focus on complete cases
    data <- data[complete.cases(data)]

    if (nrow(data) == 0) {
      stop("All observations were dropped due to missing data. Check your outcome, group, time, and covariate variables for missing values.")
    }

    # n-row data.frame to hold the influence function
    if (args$true_repeated_cross_sections) {
      data[, .rowid := .I] # Create row index
      args$idname <- ".rowid"
    } else {
      # Set rowid to idname for repeated cross section/unbalanced
      data[, .rowid := get(args$idname)]
    }

    # Count unique number of cross section observations
    n <- length(unique(data[[args$idname]]))
  }

  # Check if groups is empty (usually a problem with the way people defined groups)
  if(length(glist)==0){
    stop("No valid groups. The variable in 'gname' should be expressed as the time a unit is first treated (0 if never-treated).")
  }

  # if there are only two time periods, then uniform confidence
  # bands are the same as pointwise confidence intervals
  if (length(tlist)==2) {
    args$cband <- FALSE
  }

  #-----------------------------------------------------------------------------
  # more error handling after we have balanced the panel

  # Check against very small groups
  # Calculate group sizes, dividing the count of each group by the length of tlist
  gsize <- data[, .N / length(tlist), by = get(args$gname)]

  # How many in each group before giving a warning
  reqsize <- length(BMisc::rhs_vars(args$xformla)) + 5

  # Filter groups smaller than reqsize
  gsize <- gsize[V1 < reqsize]

  # Warn if some groups are small
  if (nrow(gsize) > 0) {
    small_groups <- gsize[[1L]]
    gpaste <- paste(ifelse(is.infinite(small_groups), 0, small_groups), collapse = ",")
    warning(paste0("Some groups in your dataset have very few observations, which may cause estimation problems.\n  Check groups: ", gpaste, "."))

    # Check if the never treated group is too small
    if ((Inf %in% small_groups) & (args$control_group == "nevertreated")) {
      stop("The never-treated group is too small to serve as a reliable control. Try setting `control_group = 'notyettreated'` to include not-yet-treated units as controls.")
    }
  }

  #-----------------------------------------------------------------------------
  # Sort the data for easy access later on
  if (args$panel){
    setorderv(data, c(args$tname, args$gname, args$idname), c(1,1,1))
  } else {
    setorderv(data, c(args$tname, args$gname), c(1,1))
  }

  # Assign new args regarding number of time periods and groups
  # How many time periods
  args$time_periods_count <- length(tlist)
  # time periods
  args$time_periods <- tlist
  # How many treated groups
  args$treated_groups_count <- length(glist)
  # treated groups
  args$treated_groups <- glist
  # id size
  args$id_count <- n

  return(list(data = data, args = args))
}

#' @title DiD tensors
#' @description A utility function that splits the data into a list of outcome tensors and a list of arguments.
#' Tensors are objects of dimension id_count x 1 x time_periods_count and are used for faster filtering in the computation of the DiD estimator.
#' @param data data.table used in function
#' @param args list of arguments to validate
#'
#' @return A list containing outcome tensors, time invariant data, cohort counts, covariates matrix, cluster vector, and weights vector.
#' @noRd
get_did_tensors <- function(data, args){

  # Drop globally-empty factor levels from the covariates referenced by xformla
  # before any design matrix is built: model.matrix() emits a dummy column for
  # every declared level, so a level absent from the ENTIRE estimation sample
  # (common after users subset their data) would produce an all-zero column in
  # every (g,t) cell and NA-fail the whole estimation. Done before invariant_data
  # is derived below so both model.matrix() calls see the same levels. Levels
  # absent only from a particular 2x2 cell are unaffected. Mirrors the identical
  # fix in compute.att_gt() (slow path).
  for (v in all.vars(args$xformla)) {
    if (is.factor(data[[v]])) set(data, j = v, value = droplevels(data[[v]]))
  }

  # Getting the outcome tensor: a vector of outcome variables per time period of dimension id_count x 1 x time_periods_count
  #if(!args$allow_unbalanced_panel){
  if(args$panel){
    # Extract the outcome column once as a plain vector, then split by time period
    y_vec <- data[[args$yname]]
    n <- args$id_count
    nT <- args$time_periods_count
    outcomes_tensor <- vector("list", nT)
    for(time in seq_len(nT)){
      start <- (time - 1L) * n + 1L
      outcomes_tensor[[time]] <- y_vec[start:(start + n - 1L)]
    }
    # Build weights tensor: one weight vector per time period
    w_vec <- data[[".w"]]
    weights_tensor <- vector("list", nT)
    for(time in seq_len(nT)){
      start <- (time - 1L) * n + 1L
      weights_tensor[[time]] <- w_vec[start:(start + n - 1L)]
    }
  } else {
    # for(time in args$time_periods){
    #   outcome_vector_time <- rep(NA, args$id_count)  # Initialize vector with NAs
    #   # Create a new column with Y values if tname == current_time, otherwise NA
    #   data[, outcome_vector_time := fifelse(get(args$tname) == time, get(args$yname), NA_real_)]
    #   # Store the vector in the outcomes_tensor
    #   outcomes_tensor[[time]] <- data$outcome_vector_time
    #   # drop the column created
    #   data[, outcome_vector_time := NULL]
    # }
    outcomes_tensor <- NULL
    weights_tensor <- NULL
  }

  # Getting the time invariant data
  time_invariant_cols <-  c(args$idname, args$gname, ".w", args$clustervars)
  #if(!args$allow_unbalanced_panel){
  if(args$panel){
    # We can do this filtering because the data is already sorted appropriately
    invariant_data <- data[1:args$id_count]
  } else {
    if (!args$allow_unbalanced_panel) {
      # get the first observation for each unit
      invariant_data <- data[data[, .I[1], by = get(args$idname)]$V1]
      # # order by idname
      # setorderv(invariant_data, c(args$idname), c(1))
    } else {
      invariant_data <- copy(data)
    }
  }

  # Get cohort, period and crosstable counts. In the balanced-panel case
  # invariant_data is exactly the first-period slice of the (t, g, id)-sorted data,
  # so the by-period and by-(period, cohort) tables are degenerate one-period shapes
  # derivable from cohort_counts without re-grouping the table; in the RCS/unbalanced
  # case the cohort and period margins are rollups of a single (t, g) pass.
  cohorts <- c(args$treated_groups, Inf)
  if (args$panel) {
    # Get cohort counts
    cohort_counts <- invariant_data[, .(cohort_size = .N) , by = get(args$gname)]
    names(cohort_counts)[1] <- "cohort" # changing the name

    # Get period counts: all invariant_data rows share the first time period
    period_counts <- data.table(period = args$time_periods[1],
                                period_size = nrow(invariant_data))

    # Get crosstable counts (long form): one row per cohort, all in the first period
    crosstable_counts <- data.table(period = args$time_periods[1],
                                    cohort = cohort_counts$cohort,
                                    N = cohort_counts$cohort_size)
  } else {
    # Get crosstable counts (long form): count occurrences by tname and gname
    crosstable_counts <- invariant_data[, .N, by = .(get(args$tname), get(args$gname))]
    names(crosstable_counts)[1] <- "period" # changing the name
    names(crosstable_counts)[2] <- "cohort" # changing the name

    # Get cohort and period counts as rollups of the (t, g) pass
    cohort_counts <- crosstable_counts[, .(cohort_size = sum(N)), by = cohort]
    period_counts <- crosstable_counts[, .(period_size = sum(N)), by = period]
  }
  crosstable_counts <- dcast(crosstable_counts, period ~ cohort, value.var = "N", fill = 0)  # Reshape


  # # Get covariates if any
  # if(args$xformla == ~1){
  #   covariates <- rep(1, args$id_count)
  # } else {
  #   covariates <- as.data.table(model.matrix(args$xformla, data = invariant_data, na.action = na.pass))
  # }

  # if (args$xformla == ~1) {
  #   covariates_tensor <- list(rep(1, args$id_count))
  # } else {
  #   if (args$panel){
  #     covariates_tensor <- vector("list", args$time_periods_count)
  #     for (tt in seq_along(args$time_periods)) {
  #       rng <- ((tt - 1) * args$id_count + 1):(tt * args$id_count)
  #       covariates_tensor[[tt]] <-
  #         model.matrix(args$xformla, data = data[rng], na.action = na.pass)
  #     }
  #   } else {
  #     covariates_tensor <- as.data.table(model.matrix(args$xformla, data = invariant_data, na.action = na.pass))
  #   }
  #
  # }
  if (args$xformla == ~1) {
    ones_vec <- rep(1, args$id_count)
    covariates_tensor <- replicate(args$time_periods_count, ones_vec,
                                   simplify = FALSE)
    covariates_matrix <- matrix(1, nrow = nrow(invariant_data), ncol = 1)
  } else {
    if (args$panel) {
      # Compute model.matrix once on the full stacked data, then split by time period
      full_mm <- model.matrix(args$xformla, data = data, na.action = na.pass)
      n <- args$id_count
      nT <- args$time_periods_count
      covariates_tensor <- vector("list", nT)
      for (tt in seq_len(nT)) {
        start <- (tt - 1L) * n + 1L
        covariates_tensor[[tt]] <- full_mm[start:(start + n - 1L), , drop = FALSE]
      }
      covariates_matrix <- NULL        # not needed in balanced panel
    } else {                             # RCS / unbalanced
      covariates_tensor <- NULL        # slices not needed
      covariates_matrix <- model.matrix(args$xformla,
                                        data = invariant_data,
                                        na.action = na.pass)
    }
  }

  # Get the cluster variable only (direct column access is faster than .SD/.SDcols)
  if(!is.null(args$clustervars)){
    cluster <- invariant_data[[args$clustervars]]
  } else {
    cluster <- NA
  }

  # Get the weights only
  weights <- invariant_data[[".w"]]

  # Gather all the arguments to return
  return(list(outcomes_tensor = outcomes_tensor,
              data = data,
              time_invariant_data = invariant_data,
              cohort_counts = cohort_counts,
              period_counts = period_counts,
              crosstable_counts = crosstable_counts,
              # covariates = covariates,
              covariates_matrix  = covariates_matrix,
              covariates_tensor = covariates_tensor,
              cluster = cluster,
              weights = weights,
              weights_tensor = weights_tensor))
}

#' @title Process `did` Function Arguments
#'
#' @description Function to process arguments passed to the main methods in the
#'  `did` package as well as conducting some tests to make sure
#'  data is in proper format / try to throw helpful error messages.
#'
#' @inheritParams att_gt
#' @param call Function call to att_gt
#'
#' @return a [`DIDparams`] object
#'
#' @export
pre_process_did2 <- function(yname,
                            tname,
                            idname,
                            gname,
                            xformla = NULL,
                            data,
                            panel = TRUE,
                            allow_unbalanced_panel,
                            control_group = c("nevertreated","notyettreated"),
                            anticipation = 0,
                            weightsname = NULL,
                            fix_weights = NULL,
                            alp = 0.05,
                            bstrap = FALSE,
                            cband = FALSE,
                            biters = 1000,
                            clustervars = NULL,
                            est_method = "dr",
                            base_period = "varying",
                            print_details = TRUE,
                            faster_mode=FALSE,
                            pl = FALSE,
                            cores = 1,
                            call = NULL) {


  # coerce data to data.table first, keeping only the columns the pipeline uses
  # (id/time/group/outcome/weights/cluster plus the raw xformla variables) so wide
  # inputs do not pay a full-table copy. If any referenced name is missing, fall
  # back to coercing everything so the downstream checks (dreamerr/validate_args/
  # xformla) fire with their usual messages and precedence.
  keep_cols <- unique(c(idname, tname, gname, yname, weightsname, clustervars,
                        if (!is.null(xformla)) all.vars(xformla)))
  if (is.data.frame(data) && is.character(keep_cols) && all(keep_cols %in% names(data))) {
    if (is.data.table(data)) {
      data <- data[, ..keep_cols]
    } else {
      # selecting columns from a data.frame is shallow; only kept columns are copied
      data <- as.data.table(data[keep_cols])
    }
  } else if (!is.data.table(data)) {
    data <- as.data.table(data)
  }

  # gathering all the arguments except data
  args_names <- setdiff(names(formals()), "data")
  args <- mget(args_names, sys.frame(sys.nframe()))

  # pick a control_group by default
  args$control_group <- control_group[1]
  if (!(args$control_group %in% c("nevertreated", "notyettreated"))) {
    stop("control_group must be either 'nevertreated' or 'notyettreated'")
  }
  args$base_period <- base_period[1]
  if (!(args$base_period %in% c("universal", "varying"))) {
    stop("base_period must be either 'universal' or 'varying'.")
  }
  check_reserved_did_names(yname = args$yname, tname = args$tname,
                           idname = args$idname, gname = args$gname,
                           xformla = args$xformla,
                           weightsname = args$weightsname,
                           clustervars = args$clustervars)

  # run error checking on arguments
  validate_args(args, data)

  # Strip idname from clustervars in the caller's copy of args.
  # validate_args works on a local copy, so we must repeat this here.
  # Clustering at the unit level (idname) is already done implicitly.
  if (!is.null(args$clustervars) && !is.null(args$idname) && (args$idname %in% args$clustervars)) {
    args$clustervars <- setdiff(args$clustervars, args$idname)
    if (length(args$clustervars) == 0L) args$clustervars <- NULL
  }

  # put in blank xformla if no covariates
  if (is.null(args$xformla)) {
    args$xformla <- ~1
  } else {
    # extract variable names from the formula
    formula_vars <- all.vars(args$xformla)

    # identify variables in xformla not in data
    missing_vars <- setdiff(formula_vars, names(data))

    # error checking for missing variables in data
    if (length(missing_vars) > 0) {
      stop(paste("The following variables are not in data:", paste(missing_vars, collapse = ", ")), call. = FALSE)
    }
  }

  # Put the data in a standard format after some validation
  cleaned_did <- did_standardization(data, args)

  # Partition staggered DiD into a 2x2 DiD for faster implementation
  did_tensors <- get_did_tensors(cleaned_did$data, cleaned_did$args)

  # store parameters for passing around later
  dp <- DIDparams2(did_tensors, cleaned_did$args, call=call)

  return(dp)
}
