# =============================================================================
# Tests for ANALYTICAL cluster-robust standard errors in the MULTI-PERIOD path.
#
# Feature under test (working tree, multi-period only):
#   ddd(..., cluster = "cl", boot = FALSE) should return ANALYTICAL
#   cluster-robust SEs = cluster-sum CRVE on the influence function with 1/n
#   scaling. The ddd object exposes $cluster_vector (per-unit cluster ids aligned
#   to inf_func_mat rows), $cluster_var, $inf_func_mat, $se, $n.
#
# NOTE on reachability:
#   An earlier revision of run_preprocess_multPeriods() FORCED boot = TRUE (and
#   cband = TRUE) whenever a `cluster` was supplied with boot = FALSE, which made
#   the analytic cluster-robust branch in att_gt() (gated on `!boot`) unreachable
#   in the multi-period path. That override has since been REMOVED, so the
#   analytic path is now reached directly by ddd(..., cluster = "cl",
#   boot = FALSE). The analytic_cluster_reachable() guard below is kept as a
#   REGRESSION CHECK: it asserts boot stayed FALSE and no bootstrap ran, and will
#   fail if that override is ever re-introduced.
# =============================================================================

# ---- shared helper: attach a time-invariant cluster column to mult-period data
make_mult_data_with_clusters <- function(seed = 123, size = 600, n_clusters = 30,
                                         dgp_type = 1, balanced = TRUE) {
  set.seed(seed)
  d <- gen_dgp_mult_periods(size = size, dgp_type = dgp_type)[["data"]]
  ids <- sort(unique(d$id))
  if (balanced) {
    # equal-sized clusters: ids spread round-robin so each cluster has ~same size
    cl <- rep(seq_len(n_clusters), length.out = length(ids))
  } else {
    set.seed(seed + 1)
    cl <- sample(seq_len(n_clusters), size = length(ids), replace = TRUE)
  }
  clmap <- data.frame(id = ids, cl = cl)
  merge(d, clmap, by = "id")
}

# Returns TRUE if ddd() actually reached the analytic cluster path (i.e. it did
# NOT silently flip to bootstrap). Asserted below as a regression guard against
# the old preprocessing override that forced boot = TRUE under clustering.
analytic_cluster_reachable <- function(ddd_obj) {
  isFALSE(ddd_obj$argu$boot) && is.null(ddd_obj$bT)
}

# -----------------------------------------------------------------------------
# CHECK 2: Analytic clustered == manual cluster-sum CRVE (tight tolerance).
# -----------------------------------------------------------------------------
test_that("analytic clustered (g,t) SEs equal the manual cluster-sum CRVE", {
  d <- make_mult_data_with_clusters(seed = 123, size = 600, n_clusters = 30)

  o <- suppressWarnings(ddd(
    yname = "y", tname = "time", idname = "id", gname = "state",
    pname = "partition", xformla = ~cov1 + cov2 + cov3 + cov4,
    data = d, control_group = "nevertreated", base_period = "varying",
    est_method = "dr", cluster = "cl", boot = FALSE))

  # Sanity on exposed fields / alignment
  expect_equal(o$cluster_var, "cl")
  expect_false(is.null(o$cluster_vector))
  expect_equal(length(o$cluster_vector), o$n)

  # Analytic cluster path must be reachable: boot stays FALSE, no bootstrap artifacts.
  expect_true(analytic_cluster_reachable(o))

  IF <- as.matrix(o$inf_func_mat)
  cv <- o$cluster_vector
  n  <- o$n
  # cluster-sum CRVE, exactly as specified in the porting guide check 2
  manual <- sqrt(diag(crossprod(rowsum(IF, cv))) / n) / sqrt(n)

  expect_equal(as.numeric(o$se), as.numeric(manual), tolerance = 1e-8)
})

# -----------------------------------------------------------------------------
# CHECK 4: No clustering -> byte-identical to the i.i.d. analytic formula.
# This path is NOT affected by the bug and must pass.
# -----------------------------------------------------------------------------
test_that("no-cluster (g,t) SEs equal the i.i.d. analytic formula from IF", {
  d <- make_mult_data_with_clusters(seed = 321, size = 500, n_clusters = 25)

  o <- ddd(
    yname = "y", tname = "time", idname = "id", gname = "state",
    pname = "partition", xformla = ~cov1 + cov2 + cov3 + cov4,
    data = d, control_group = "nevertreated", base_period = "varying",
    est_method = "dr", cluster = NULL, boot = FALSE)

  # the new cluster branch must be gated OFF when cluster is NULL
  expect_true(is.null(o$cluster_vector))
  expect_false(o$argu$boot)   # analytic path, no bootstrap
  expect_true(is.null(o$bT))

  IF <- as.matrix(o$inf_func_mat)
  n  <- o$n
  iid_a <- sqrt(colSums(IF^2)) / n
  iid_b <- sqrt(diag(crossprod(IF)) / n) / sqrt(n)  # algebraically identical

  expect_equal(as.numeric(o$se), as.numeric(iid_a), tolerance = 1e-10)
  expect_equal(as.numeric(o$se), as.numeric(iid_b), tolerance = 1e-10)
})

# -----------------------------------------------------------------------------
# CHECK 5: NYT / GMM cells.
#   (a) cluster = NULL : single-control cells use the analytic i.i.d. SE
#       sqrt(colSums(IF^2))/n EXACTLY; multi-control GMM cells instead report the
#       per-cell gmm_se = sqrt(1/(n*sum(inv_OMEGA))) override (se_gt_ddd_nyt).
#       We regression-test BOTH facts and confirm NO bootstrap is triggered
#       (this cluster = NULL NYT path is NOT affected by the preprocessing bug).
#   (b) WITH clustering : GMM cells should report the cluster-sum CRVE (matching
#       check 2). Depends on the analytic cluster path -> SKIPPED if unreachable.
#
# Note: the source comment in att_gt() asserts sum(IF_gmm^2)/n == 1/sum(inv_OMEGA),
# i.e. that the i.i.d. formula reproduces gmm_se exactly. In practice the two
# differ slightly for GMM cells (the GMM weights make IF_gmm's raw second moment
# depart from 1/sum(inv_OMEGA)); the reported SE is the gmm_se override, so we
# test against gmm_se on GMM cells and against the i.i.d. formula elsewhere.
# -----------------------------------------------------------------------------
test_that("NYT/GMM cells without clustering: i.i.d. on single-control, gmm_se on GMM", {
  d <- make_mult_data_with_clusters(seed = 99, size = 700, n_clusters = 35)

  o <- ddd(
    yname = "y", tname = "time", idname = "id", gname = "state",
    pname = "partition", xformla = ~cov1 + cov2 + cov3 + cov4,
    data = d, control_group = "notyettreated", base_period = "varying",
    est_method = "dr", cluster = NULL, boot = FALSE)

  expect_false(o$argu$boot)        # analytic, no bootstrap (bug does not touch this path)
  expect_true(is.null(o$bT))
  expect_true(is.null(o$cluster_vector))

  IF <- as.matrix(o$inf_func_mat)
  n  <- o$n
  iid <- sqrt(colSums(IF^2)) / n
  iid[iid <= sqrt(.Machine$double.eps) * 10] <- NA

  ok <- !is.na(o$se) & !is.na(iid)
  # GMM cells = cells whose reported SE differs from the plain i.i.d. formula
  # (these are the not-yet-treated cells with >1 control, carrying the gmm_se
  # override). All other reported SEs must equal the i.i.d. formula EXACTLY.
  is_gmm <- ok & (abs(o$se - iid) > 1e-7)
  is_iid_cell <- ok & !is_gmm

  # single-control cells reproduce the i.i.d. analytic formula exactly
  expect_equal(as.numeric(o$se[is_iid_cell]),
               as.numeric(iid[is_iid_cell]), tolerance = 1e-8)

  # GMM cells carry a distinct, finite, positive override SE that is close to
  # (but not identical to) the i.i.d. value (regression-test of the gmm_se path)
  if (any(is_gmm)) {
    expect_true(all(is.finite(o$se[is_gmm])))
    expect_true(all(o$se[is_gmm] > 0))
    rel <- abs(o$se[is_gmm] - iid[is_gmm]) / iid[is_gmm]
    expect_true(all(rel < 0.05))   # gmm_se override stays in the same ballpark
  } else {
    skip("No multi-control GMM cells present in this design; nothing to assert for the gmm_se override.")
  }
})

test_that("NYT/GMM cells WITH clustering match the manual cluster-sum CRVE", {
  d <- make_mult_data_with_clusters(seed = 99, size = 700, n_clusters = 35)

  o <- suppressWarnings(ddd(
    yname = "y", tname = "time", idname = "id", gname = "state",
    pname = "partition", xformla = ~cov1 + cov2 + cov3 + cov4,
    data = d, control_group = "notyettreated", base_period = "varying",
    est_method = "dr", cluster = "cl", boot = FALSE))

  expect_equal(o$cluster_var, "cl")
  expect_false(is.null(o$cluster_vector))

  # Analytic cluster path must be reachable for GMM cells too.
  expect_true(analytic_cluster_reachable(o))

  IF <- as.matrix(o$inf_func_mat)
  cv <- o$cluster_vector
  n  <- o$n
  manual <- sqrt(diag(crossprod(rowsum(IF, cv))) / n) / sqrt(n)

  ok <- !is.na(o$se) & !is.na(manual)
  expect_equal(as.numeric(o$se[ok]), as.numeric(manual[ok]), tolerance = 1e-8)
})

# -----------------------------------------------------------------------------
# CHECK 1: Analytic == bootstrap (large nboot), at (g,t) level and through
# agg_ddd() (eventstudy + simple). Bootstrap is Monte-Carlo noisy, so a loose
# relative tolerance is used. Depends on the analytic cluster path -> SKIPPED if
# unreachable.
# -----------------------------------------------------------------------------
test_that("analytic clustered SEs approx bootstrap clustered SEs (g,t and agg)", {
  skip_on_cran()
  d <- make_mult_data_with_clusters(seed = 2024, size = 800, n_clusters = 40)

  o_analytic <- suppressWarnings(ddd(
    yname = "y", tname = "time", idname = "id", gname = "state",
    pname = "partition", xformla = ~cov1 + cov2 + cov3 + cov4,
    data = d, control_group = "nevertreated", base_period = "varying",
    est_method = "dr", cluster = "cl", boot = FALSE))

  # Analytic cluster path must be reachable so we can compare against bootstrap.
  expect_true(analytic_cluster_reachable(o_analytic))

  o_boot <- suppressWarnings(ddd(
    yname = "y", tname = "time", idname = "id", gname = "state",
    pname = "partition", xformla = ~cov1 + cov2 + cov3 + cov4,
    data = d, control_group = "nevertreated", base_period = "varying",
    est_method = "dr", cluster = "cl", boot = TRUE, nboot = 5000, cband = TRUE))

  # (g,t)-level comparison on non-missing cells
  ok <- !is.na(o_analytic$se) & !is.na(o_boot$se)
  rel <- abs(o_analytic$se[ok] - o_boot$se[ok]) / o_analytic$se[ok]
  expect_lt(max(rel), 0.10)

  # through agg_ddd(): eventstudy overall + simple overall
  es_a <- agg_ddd(o_analytic, type = "eventstudy")$aggte_ddd
  es_b <- agg_ddd(o_boot,     type = "eventstudy")$aggte_ddd
  expect_lt(abs(es_a$overall.se - es_b$overall.se) / es_a$overall.se, 0.10)

  s_a <- agg_ddd(o_analytic, type = "simple")$aggte_ddd
  s_b <- agg_ddd(o_boot,     type = "simple")$aggte_ddd
  expect_lt(abs(s_a$overall.se - s_b$overall.se) / s_a$overall.se, 0.10)
})

# -----------------------------------------------------------------------------
# CHECK 7: Alignment. Shuffling input rows must not change clustered (g,t) SEs
# (guards the match(id_index, ...) alignment of the cluster vector).
#
# The i.i.d. (cluster = NULL) variant of this check is order-invariant via the
# same alignment code and is NOT affected by the bug, so it is asserted directly.
# The clustered variant is also asserted: even though the reported SEs currently
# come from the bootstrap, the cluster_vector alignment + bootstrap (with a fixed
# seed) must still be invariant to input row order.
# -----------------------------------------------------------------------------
test_that("clustered (g,t) SEs are invariant to input row permutation", {
  d <- make_mult_data_with_clusters(seed = 55, size = 500, n_clusters = 25)

  run <- function(dd) {
    set.seed(1)  # fix seed so any internal bootstrap is reproducible across runs
    suppressWarnings(ddd(
      yname = "y", tname = "time", idname = "id", gname = "state",
      pname = "partition", xformla = ~cov1 + cov2 + cov3 + cov4,
      data = dd, control_group = "nevertreated", base_period = "varying",
      est_method = "dr", cluster = "cl", boot = FALSE))
  }

  o_ordered <- run(d)

  set.seed(777)
  d_shuf <- d[sample(seq_len(nrow(d))), ]
  o_shuffled <- run(d_shuf)

  ok <- !is.na(o_ordered$se) & !is.na(o_shuffled$se)
  expect_equal(as.numeric(o_ordered$se[ok]),
               as.numeric(o_shuffled$se[ok]),
               tolerance = 1e-6)
})

test_that("no-cluster (g,t) SEs are invariant to input row permutation (analytic)", {
  d <- make_mult_data_with_clusters(seed = 56, size = 500, n_clusters = 25)

  run <- function(dd) {
    ddd(yname = "y", tname = "time", idname = "id", gname = "state",
        pname = "partition", xformla = ~cov1 + cov2 + cov3 + cov4,
        data = dd, control_group = "nevertreated", base_period = "varying",
        est_method = "dr", cluster = NULL, boot = FALSE)
  }

  o_ordered <- run(d)
  set.seed(888)
  d_shuf <- d[sample(seq_len(nrow(d))), ]
  o_shuffled <- run(d_shuf)

  ok <- !is.na(o_ordered$se) & !is.na(o_shuffled$se)
  expect_equal(as.numeric(o_ordered$se[ok]),
               as.numeric(o_shuffled$se[ok]),
               tolerance = 1e-9)
})
