#' Doubly robust DDD estimator for ATT, with panel data and multiple periods
#'
#' This function implements a doubly robust estimator for assessing the average
#' treatment effect on the treated (ATT) using a triple differences (DDD) approach
#' in panel data settings with staggered treatment adoption. The function takes preprocessed
#' data structured specifically for this analysis.
#'
#' @importFrom Matrix Matrix
#' @param did_preprocessed A list containing preprocessed data and specifications for the DDD estimation.
#'        Expected elements include:
#'        - `preprocessed_data`: A data table containing the data with variables needed for the analysis.
#'        - `xformula`: The formula for the covariates to be included in the model. It should be of the form \code{~ x1 + x2}.
#'        Default is \code{xformla = ~1} (no covariates).
#'        - `control_group`: A character string indicating the control group. Default is \code{"notyettreated"}.
#'        - `n`: The number of unique ID observations in the data.
#'        - `cohorts`: A vector containing the number of treated cohorts
#'        - `time_periods`: A vector containing the number of time periods
#'        - `tlist`: A vector containing the time periods
#'        - `glist`: A vector containing the groups
#'        - `base_period`: A character string indicating the base period. Default is \code{"universal"}.
#'        - `cohort_size`: A vector containing the size of each cohort
#'        - `boot`: Logical. If \code{TRUE}, the function computes the bootstrap standard errors. Default is \code{FALSE}.
#'        - `nboot`: The number of bootstrap samples to be used. Default is \code{NULL}. If \code{boot = TRUE}, the default is \code{nboot = 999}.
#'        - `use_parallel`: Boolean of whether or not to use parallel processing in the multiplier bootstrap, default is \code{use_parallel=FALSE}
#'        - `cores`: the number of cores to use with parallel processing, default is \code{cores=1}
#'        - `cband`: Boolean of whether or not to compute simultaneous confidence bands, default is \code{cband=FALSE}
#'        - `alpha` The level of significance for the confidence intervals.  Default is \code{0.05}.
#'
#' @keywords internal
#' @return A list with the estimated ATT, standard error, upper and lower confidence intervals, and influence function.
#' @noRd
NULL
# ------------------------------------------------------------------------------

att_gt <- function(did_preprocessed){
  data <- did_preprocessed$preprocessed_data
  xformla <- did_preprocessed$xformula
  control_group <- did_preprocessed$control_group
  n <- did_preprocessed$n
  cohorts <- did_preprocessed$cohorts
  time_periods <- did_preprocessed$time_periods
  tlist <- did_preprocessed$tlist
  glist <- did_preprocessed$glist
  panel <- did_preprocessed$panel
  allow_unbalanced_panel <- did_preprocessed$allow_unbalanced_panel
  base_period <- did_preprocessed$base_period
  boot <- did_preprocessed$boot
  alpha <- did_preprocessed$alpha
  cohort_size <- did_preprocessed$cohort_size
  use_parallel <- did_preprocessed$use_parallel # to perform bootstrap
  cores <- did_preprocessed$cores # to perform bootstrap
  cband <- did_preprocessed$cband # to perform bootstrap + simult. conf. band

  orig_data <- copy(data)
  id_index <- unique(data$id)
  attgt_list <- list()

  counter <- 1
  tlist_length <- length(tlist)
  tfac <- 0
  if (base_period != "universal") {
    tlist_length <- tlist_length - 1
    tfac <- 1
  }

  # influence function matrix to be populated
  inf_func_mat <- Matrix::Matrix(data=0,nrow=n, ncol=cohorts*(time_periods - tfac), sparse=TRUE)
  se_gt_ddd_nyt <- rep(NA, cohorts * (time_periods - tfac)) # standard errors for not yet treated control group
  # If "nevertreated", create the control column that gonna indicate the comparison group for each iteration
  if (control_group == "nevertreated") {
    data[, control := as.integer(first_treat == 0)]
  }


  for (g in 1:cohorts){
    data[, treat := as.integer(first_treat == glist[g])]

    for (t in 1:tlist_length){
      pret <- t
      if (base_period == "universal"){
        pret <- tail(which(tlist < glist[g]), 1)
      }

      # creating control in case of "not yet treated option"
      if (control_group == "notyettreated") {
        # max_period <- tlist[max(t, pret) + tfac] + anticipation
        max_period <- tlist[max(t, pret) + tfac]
        data[, control := as.integer((first_treat == 0) | (first_treat > max_period & first_treat != glist[g]))]
      }


      # check if in post-treatment period
      if ((glist[g] <= tlist[(t + tfac)])) {

        # update pre-period if in post-treatment period to be period (g-delta-1)
        pret <- tail(which(tlist < glist[g]), 1)

        # print a warning message if there are no pre-treatment period
        if (length(pret) == 0) {
          warning(paste0("There are no pre-treatment periods for the group first treated at ", glist[g], "\nUnits from this group are dropped"))

          # if there are not pre-treatment periods, move on to next iteration
          break
        }
      }

      # normalize in period (g-1) to be equal to 0 when based period is universal and move on to next iteration
      if (base_period == "universal") {
        if (tlist[pret] == tlist[t + tfac]) {
          attgt_list[[counter]] <- list(att = 0, group = glist[g], year = tlist[t + tfac], post = 0)
          inf_func_mat[, counter] <- rep(0, n)
          counter <- counter + 1
          next
        }
      }

      # post treatment dummy variable
      post_treat <- 1*(glist[g] <= tlist[t + tfac])

      # filter data with only 2 periods
      cohort_data <- data[period %in% c(tlist[t + tfac], tlist[pret])]

      # ============================================
      # BRANCH: PANEL vs RCS/UNBALANCED
      # ============================================

      if (panel && !allow_unbalanced_panel) {
          # ==========================================
          # PANEL DATA PATH (balanced panels only)
          # ==========================================
        
          # get total number of units
          # n_size = uniqueN(cohort_data[, id])
        
          # filter data for treated and control groups in each (g,t) cell. Save index

          # index of unit in current cell (g,t) when treat = 1 and control = 1
          index_units_in_gt <- cohort_data[, treat == 1 | control == 1]

          # filter by units selected in current cell (g,t)
          cohort_data <- cohort_data[index_units_in_gt]

          # save number of units after filtering
          size_gt = uniqueN(cohort_data[, id])

          # Identify all available control states in this cohort.
          available_controls <- sort(unique(cohort_data$first_treat[cohort_data$control == 1 ]))
          available_controls <- available_controls[available_controls != glist[g]]

        if (length(available_controls) == 1) {
          # creating subgroup variable
          # 4 if (eligible ==1 & enabled == 1); 3 if (eligible ==0 & enabled == 1); 2 if (eligible ==1 & enabled == 0); 1 if (eligible ==0 & enabled == 0)
          # Create subgroup variable based on the control_group option
          cohort_data[, subgroup := NA_integer_]
          cohort_data[(treat == 1 & partition == 1), subgroup := 4]
          cohort_data[(treat == 1 & partition == 0), subgroup := 3]
          cohort_data[(treat == 0 & partition == 1), subgroup := 2]
          cohort_data[(treat == 0 & partition == 0), subgroup := 1]

          # add post treatment dummy variable if period is equal to t + tfac
          cohort_data[, post := 0]

          # Ypre is always the baseline, even in pre-treatment periods were baseline could be in later periods with respect to t.
          pseudo_post = ifelse(post_treat == 1, tlist[max(t, pret) + tfac], tlist[t + tfac])
          cohort_data[(period == pseudo_post), post := 1]

          # cohort_data[(period == tlist[max(t, pret) + tfac]), post := 1]
          # Calculate the size of each subgroup in the 'subgroup' column
          subgroup_counts <- cohort_data[, .(count = .N/2), by = subgroup][order(-subgroup)]

          # Reassign dp object and run att_dr
          did_preprocessed$preprocessed_data <- cohort_data
          did_preprocessed$subgroup_counts <- subgroup_counts
          did_preprocessed$boot <- FALSE # forcing false to avoid bootstrap inside att_dr() function.
          did_preprocessed$inffunc <- TRUE # forcing true to recover influence function

          # run att_dr based on 2 time periods subdata
          attgt_inf_func <- att_dr(did_preprocessed)

          # adjust influence function
          attgt_inf_func$inf_func <- (n/size_gt) * attgt_inf_func$inf_func
          # save results in a list
          attgt_list[[counter]] <- list(att = attgt_inf_func$ATT, group = glist[g], year = tlist[t + tfac], post = post_treat)

          # recover influence function
          inff <- rep(0, n)
          # avoid repetition in index
          inff[index_units_in_gt[seq(1, length(index_units_in_gt), by = 2)]] <- attgt_inf_func$inf_func
          # save in influence function matrix
          inf_func_mat[, counter] <- inff
          # update counter
          counter  <- counter + 1
        } else {
          # Here we have more than one control group available -> notyettreated

          # Initialize local containers.
          ddd_over_controls_res <- list()
          inf_mat_local <- NULL

          # Loop over each available control state.
          for (ctrl in available_controls) {
            # Subset data: keep treated units (first_treat == glist[g]) and control units (first_treat == ctrl)

            index_units_in_gt_ctrl <- cohort_data[, first_treat == glist[g] | first_treat == ctrl]
            subset_data <- cohort_data[index_units_in_gt_ctrl]
            size_gt_ctrl <- uniqueN(subset_data[, id])

            # Create subgroup variable based on the control_group option
            subset_data[, subgroup := NA_integer_]
            subset_data[(treat == 1 & partition == 1), subgroup := 4]
            subset_data[(treat == 1 & partition == 0), subgroup := 3]
            subset_data[(treat == 0 & partition == 1), subgroup := 2]
            subset_data[(treat == 0 & partition == 0), subgroup := 1]

            # add post treatment dummy variable if period is equal to t + tfac
            subset_data[, post := 0]

            # Ypre is always the baseline, even in pre-treatment periods were baseline could be in later periods with respect to t.
            pseudo_post = ifelse(post_treat == 1, tlist[max(t, pret) + tfac], tlist[t + tfac])
            subset_data[(period == pseudo_post), post := 1]

            # subset_data[(period == tlist[max(t, pret) + tfac]), post := 1]
            # Calculate the size of each subgroup in the 'subgroup' column
            subgroup_counts <- subset_data[, .(count = .N/2), by = subgroup][order(-subgroup)]

            # Reassign dp object and run att_dr
            did_preprocessed$preprocessed_data <- subset_data
            did_preprocessed$subgroup_counts <- subgroup_counts
            did_preprocessed$boot <- FALSE # forcing false to avoid bootstrap inside att_dr() function.
            did_preprocessed$inffunc <- TRUE # forcing true to recover influence function


            # run att_dr based on 2 time periods subdata
            ddd_out <- triplediff::att_dr(did_preprocessed)

            # computing the aggregate influence function for every g'>t: RIF_{g,g',t}
            # adjust influence function (rescale + recentered)
            ddd_out$inf_func <- ((size_gt/size_gt_ctrl)*ddd_out$inf_func) #+ 1*ddd_out$ATT

            # Save the ddd output in the local list using the control state as name.
            ddd_over_controls_res[[as.character(ctrl)]] <- ddd_out

            # recover influence function
            inff <- rep(0, n)
            # avoid repetition in index
            inff[index_units_in_gt_ctrl[seq(1, length(index_units_in_gt_ctrl), by = 2)]] <- ddd_out$inf_func

            # Append this vector as a new column in the local matrix.
            if (is.null(inf_mat_local)) {
              inf_mat_local <- matrix(inff, ncol = 1)
              colnames(inf_mat_local) <- as.character(ctrl)
            } else {
              inf_mat_local <- cbind(inf_mat_local, inff)
              colnames(inf_mat_local)[ncol(inf_mat_local)] <- as.character(ctrl)
            }
          } # end loop over available_controls

          # If no ddd estimations were obtained, skip this (g,t) cell.
          if (length(ddd_over_controls_res) == 0) next

          # Compute the naive ATT as the average of all ddd ATT estimates.
          att_vals <- sapply(ddd_over_controls_res, function(res) res$ATT)

          # -----------------------------------------
          # Compute the GMM-based estimator for ATT(g,t)
          # -----------------------------------------

          # computing the aggregate influence function for every g'>t: RIF_{g,g',t}
          OMEGA <- cov(inf_mat_local, use = "complete.obs")

          ones <- rep(1, length(ddd_over_controls_res))
          inv_OMEGA <- solve(OMEGA)
          w <- colSums(inv_OMEGA) / sum(inv_OMEGA)

          # Compute the GMM estimator for ATT(g,t)
          ATT_gmm <- sum(w * att_vals) / sum(w)
          # IF_gmm
          IF_gmm <- (inf_mat_local) %*% w
          # gmm_se
          gmm_se <- sqrt(1/ (n * sum(inv_OMEGA)))

          # save the results for this (g,t) cell
          attgt_list[[counter]] <- list(att = ATT_gmm, group = glist[g], year = tlist[t + tfac], post = post_treat)

          # save in influence function matrix
          inf_func_mat[, counter] <- IF_gmm

          # save the standard errors from GMM
          se_gt_ddd_nyt[counter] <- gmm_se

          # update counter
          counter  <- counter + 1

        } # end of GMM-based procedure

      } else { # else of if (panel)

        # ==========================================
        # RCS/UNBALANCED PANEL PATH
        # ==========================================

        # Get total number of observations (NOT unique IDs)
        # n_size = uniqueN(cohort_data[, id])

        # Filter for treated and control groups
        index_units_in_gt <- cohort_data[, treat == 1 | control == 1]
        cohort_data <- cohort_data[index_units_in_gt]

        # Save number of observations after filtering
        size_gt <- if (allow_unbalanced_panel) uniqueN(cohort_data$id) else nrow(cohort_data)
        # size_gt <- nrow(cohort_data)

        # Identify available control states
        available_controls <- sort(unique(cohort_data$first_treat[cohort_data$control == 1]))
        available_controls <- available_controls[available_controls != glist[g]]

        # ----------------------------------------
        # SINGLE CONTROL GROUP CASE (RCS)
        # ----------------------------------------
        if (length(available_controls) == 1) {

          # Get total sample size for influence function (number of unique IDs in original data)
          n_total <- uniqueN(orig_data$id)

          # Create subgroup variable
          cohort_data[, subgroup := NA_integer_]
          cohort_data[(treat == 1 & partition == 1), subgroup := 4]
          cohort_data[(treat == 1 & partition == 0), subgroup := 3]
          cohort_data[(treat == 0 & partition == 1), subgroup := 2]
          cohort_data[(treat == 0 & partition == 0), subgroup := 1]

          # Add post treatment dummy
          cohort_data[, post := 0]
          pseudo_post <- ifelse(post_treat == 1, tlist[max(t, pret) + tfac], tlist[t + tfac])
          cohort_data[(period == pseudo_post), post := 1]

          # Calculate subgroup counts - count unique IDs (handles both RCS and unbalanced panel)
          # For unbalanced panel, we need to count unique IDs in the cohort_data, not the full data
          subgroup_counts <- cohort_data[, .(count = uniqueN(id)), by = subgroup][order(-subgroup)]
          

          # Prepare did_preprocessed for att_dr_rc
          did_preprocessed$preprocessed_data <- cohort_data
          did_preprocessed$subgroup_counts <- subgroup_counts
          did_preprocessed$boot <- FALSE
          did_preprocessed$inffunc <- TRUE

          # Call att_dr_rc (RCS version)
          attgt_inf_func <- att_dr_rc(did_preprocessed)

          # Save ATT results
          attgt_list[[counter]] <- list(att = attgt_inf_func$ATT,
                                         group = glist[g],
                                         year = tlist[t + tfac],
                                         post = post_treat)

          # Populate influence function
          # Following did package approach: rescale BEFORE aggregating
          # This is important for unbalanced panels where IDs have different observation counts
          inff <- rep(0, n_total)  # Use total sample size

          # Rescale influence function (still at observation level)
          attgt_inf_func$inf_func <- (n_total / nrow(cohort_data)) * attgt_inf_func$inf_func

          # Aggregate influence function by ID (sum for repeated IDs)
          ids_in_cohort <- cohort_data$id
          aggte_inffunc <- suppressWarnings(stats::aggregate(attgt_inf_func$inf_func,
                                                              list(ids_in_cohort),
                                                              sum))

          # Match aggregated IDs to positions in unique ID index
          id_positions <- match(aggte_inffunc[, 1], id_index)
          inff[id_positions] <- aggte_inffunc[, 2]
          inf_func_mat[, counter] <- inff

          counter <- counter + 1

        # ----------------------------------------
        # MULTIPLE CONTROL GROUPS (GMM) CASE (RCS)
        # ----------------------------------------
        } else {

          # Get total sample size for influence function (number of unique IDs in original data)
          n_total <- uniqueN(orig_data$id)

          # Initialize containers
          ddd_over_controls_res <- list()
          inf_mat_local <- NULL

          # Loop over each available control state
          for (ctrl in available_controls) {

            # Subset data for this control group
            index_units_in_gt_ctrl <- cohort_data[, first_treat == glist[g] | first_treat == ctrl]
            subset_data <- cohort_data[index_units_in_gt_ctrl]
            size_gt_ctrl <- if (allow_unbalanced_panel) uniqueN(subset_data$id) else nrow(subset_data)

            # Create subgroup variable
            subset_data[, subgroup := NA_integer_]
            subset_data[(treat == 1 & partition == 1), subgroup := 4]
            subset_data[(treat == 1 & partition == 0), subgroup := 3]
            subset_data[(treat == 0 & partition == 1), subgroup := 2]
            subset_data[(treat == 0 & partition == 0), subgroup := 1]

            # Add post indicator
            subset_data[, post := 0]
            pseudo_post <- ifelse(post_treat == 1, tlist[max(t, pret) + tfac], tlist[t + tfac])
            subset_data[(period == pseudo_post), post := 1]

            # Calculate subgroup counts - count unique IDs (handles both RCS and unbalanced panel)
            subgroup_counts <- subset_data[, .(count = uniqueN(id)), by = subgroup][order(-subgroup)]

            # Prepare and run att_dr_rc
            did_preprocessed$preprocessed_data <- subset_data
            did_preprocessed$subgroup_counts <- subgroup_counts
            did_preprocessed$boot <- FALSE
            did_preprocessed$inffunc <- TRUE

            ddd_out <- att_dr_rc(did_preprocessed)

            # Save results
            ddd_over_controls_res[[as.character(ctrl)]] <- ddd_out

            # Populate influence function
            inff <- rep(0, n_total)  # Use total sample size

            # Rescale influence function (still at observation level)
            ddd_out$inf_func <- (n_total / nrow(subset_data)) * ddd_out$inf_func

            # Aggregate influence function by ID (sum for repeated IDs)
            ids_in_subset <- subset_data$id
            aggte_inffunc_ctrl <- suppressWarnings(stats::aggregate(ddd_out$inf_func,
                                                                     list(ids_in_subset),
                                                                     sum))

            # Match aggregated IDs to positions in unique ID index
            id_positions_ctrl <- match(aggte_inffunc_ctrl[, 1], id_index)
            inff[id_positions_ctrl] <- aggte_inffunc_ctrl[, 2]

            if (is.null(inf_mat_local)) {
              inf_mat_local <- matrix(inff, ncol = 1)
              colnames(inf_mat_local) <- as.character(ctrl)
            } else {
              inf_mat_local <- cbind(inf_mat_local, inff)
              colnames(inf_mat_local)[ncol(inf_mat_local)] <- as.character(ctrl)
            }
          } # end loop over available_controls

          # Skip if no estimations obtained
          if (length(ddd_over_controls_res) == 0) next

          # Compute GMM estimator
          att_vals <- sapply(ddd_over_controls_res, function(res) res$ATT)
          OMEGA <- cov(inf_mat_local, use = "complete.obs")
          inv_OMEGA <- solve(OMEGA)
          w <- colSums(inv_OMEGA) / sum(inv_OMEGA)

          ATT_gmm <- sum(w * att_vals) / sum(w)
          IF_gmm <- (inf_mat_local) %*% w
          gmm_se <- sqrt(1 / (n_total * sum(inv_OMEGA)))

          # Save results
          attgt_list[[counter]] <- list(att = ATT_gmm,
                                         group = glist[g],
                                         year = tlist[t + tfac],
                                         post = post_treat)
          inf_func_mat[, counter] <- IF_gmm
          se_gt_ddd_nyt[counter] <- gmm_se

          counter <- counter + 1

        } # end of GMM-based procedure (RCS)

      } # end of RCS/unbalanced path

    } # end of tlist loop
  } # end of glist loop

  # recover original arguments in did_preprocessed
  did_preprocessed$preprocessed_data <- orig_data
  did_preprocessed$boot <- boot # restoring boot argument

  # PREPROCESS attgt_list AND inf_func_mat
  attgt_res <- process_attgt(attgt_list)
  groups <- attgt_res$group
  periods <- attgt_res$periods
  att_gt_ddd <- attgt_res$att
  
  # ---- per-unit cluster vector aligned to the rows of inf_func_mat -------------
  # inf_func_mat rows are in id_index = unique(data$id) order (see top of att_gt()).
  # Map one cluster id per unit onto that exact ordering. `data` still carries the
  # `cluster` column from preprocessing.
  cluster      <- did_preprocessed$cluster           # cluster variable NAME (or NULL)
  cluster_vec  <- NULL
  cluster_analytic <- FALSE
  if (!is.null(cluster) && length(cluster) > 0) {
    cl_map      <- unique(data[, .(id, cluster)])     # data.table: one row per id
    cluster_vec <- cl_map$cluster[match(id_index, cl_map$id)]
    # only usable if it aligns with inf_func_mat and has no missing ids
    if (length(cluster_vec) == n && !anyNA(cluster_vec)) {
      cluster_analytic <- !boot                       # analytic branch only when not bootstrapping
    } else {
      cluster_vec <- NULL
    }
  }

  # Get analytical errors: This is analogous to cluster robust standard errors at the unit level
  n <- did_preprocessed$n
  if (cluster_analytic) {
    # Cluster-robust analytical variance: cluster sums of the influence function,
    # mirroring the cluster-sum aggregation of the multiplier bootstrap
    # (Callaway & Sant'Anna 2021, Remark 10). Same 1/n scaling as the i.i.d. V.
    Sc <- rowsum(as.matrix(inf_func_mat), cluster_vec)   # n_clusters x k cluster sums
    V  <- Matrix::Matrix(crossprod(Sc) / n)
  } else {
    V <- Matrix::t(inf_func_mat) %*% inf_func_mat / n
  }
  se_gt_ddd <- sqrt(Matrix::diag(V) / n)

  # If clustering was requested without the bootstrap but the cluster-robust V could
  # not be formed (e.g. cluster ids unavailable), warn that SEs are NOT cluster-robust.
  if (!is.null(cluster) && length(cluster) > 0 && !boot && !cluster_analytic) {
    warning("Clustered standard errors could not be computed analytically; the reported standard errors do NOT account for clustering. Set boot = TRUE for the cluster-robust multiplier bootstrap.")
  }

  # Zero standard error replaced by NA
  se_gt_ddd[se_gt_ddd <= sqrt(.Machine$double.eps)*10] <- NA

  # if control group is "notyettreated", we need to adjust the standard errors
  if (control_group == "notyettreated" && !cluster_analytic) {
    # adjust standard errors for not yet treated control group
    se_gt_ddd[!is.na(se_gt_ddd_nyt)] <- se_gt_ddd_nyt[!is.na(se_gt_ddd_nyt)]
  }
  # When cluster_analytic, the GMM (not-yet-treated) cells take their cluster-robust SE
  # from the cluster-sum V above (IF_gmm is already stored in inf_func_mat). This is
  # exact: without clustering sum(IF_gmm^2)/n == 1/sum(inv_OMEGA) == the old gmm_se.

  # Identify entries of main diagonal V that are zero or NA
  zero_na_sd_entry <- unique(which(is.na(se_gt_ddd)))

  # compute bootstrap standard errors
  bT <- NULL
  if (boot){
    # perform multiplier bootstrap
    boot_result <- mboot(inf_func_mat, did_preprocessed=did_preprocessed, use_parallel=use_parallel, cores=cores)
    bT <- boot_result$bT # sup-t confidence band
    # save bootstrap standard errors
    if(length(zero_na_sd_entry)>0) {
      se_gt_ddd[-zero_na_sd_entry] <- boot_result$se[-zero_na_sd_entry]
    } else {
      se_gt_ddd <- boot_result$se
    }
  }

  # Zero standard error replaced by NA
  se_gt_ddd[se_gt_ddd <= sqrt(.Machine$double.eps)*10] <- NA

  #-----------------------------------------------------------------------------
  # compute confidence intervals / bands
  #-----------------------------------------------------------------------------

  # compute critical values for point-wise interval
  cv <- qnorm(1 - alpha/2)

  # in case uniform confidence bands are requested
  if (boot){
    if (cband){
      # get critical value to compute uniform confidence bands
      cv <- boot_result$unif_crit_val
      if(cv >= 7){
        warning("Simultaneous critical value is arguably `too large' to be realible. This usually happens when number of observations per group is small and/or there is no much variation in outcomes.")
      }
    }
  }

  # compute confidence intervals/bands
  ci_upper <- att_gt_ddd + cv * se_gt_ddd
  ci_lower <- att_gt_ddd - cv * se_gt_ddd

  # ------------------------------------------------------------------------------
  # Return results
  # ------------------------------------------------------------------------------

  # we need this for aggregation
  # For balanced panel: use first period data (all IDs are present in first period)
  # For unbalanced panel or RCS: aggregate by ID to get one observation per unique ID
  if (panel && !allow_unbalanced_panel) {
    # Balanced panel: all IDs present in first period
    first_period_dta <- orig_data[period == tlist[1]]
  } else {
    # Unbalanced panel or RCS: take one observation per unique ID
    # This ensures we have one row per ID for aggregation weights
    first_period_dta <- orig_data[, .SD[1], by = id]
  }
  ret <- (list(ATT = att_gt_ddd,
               se = se_gt_ddd,
               uci = ci_upper,
               lci = ci_lower,
               groups = groups,
               periods = periods,
               tlist = tlist,
               glist = glist,
               cohort_size = cohort_size,
               n = n,
               bT = bT,
               inf_func_mat = inf_func_mat,
               first_period_dta = first_period_dta,
               cluster_vector = cluster_vec,   # per-unit clusters aligned to inf_func_mat rows (or NULL)
               cluster_var    = if (!is.null(cluster_vec)) cluster else NULL
              ))

  return(ret)
}
