# Parity test for the TabICL column-embedding stage (R/tabicl-embedding.R)
# against the Python reference (ColEmbedding). Fixtures are generated by
# dev/tabicl/dump_primitives.py in the configuration the released checkpoints
# use (feature_group = "same", target_aware, affine = FALSE, qassmax SSMax), in
# two variants: classification (one-hot target encoder, biased LayerNorms) and
# regression (linear target encoder, bias-free LayerNorms).
#
# A relative tolerance is used: with affine = FALSE the CLS placeholder slots
# keep the -100 skip value through residual connections, so the output contains
# large-magnitude entries for which an absolute float32 tolerance is misleading.

for (fixture in c("col_embedding", "col_embedding_reg")) {
  local({
    fixture_name <- fixture
    test_that(
      paste0(
        "tabicl_col_embedding matches the Python reference (",
        fixture_name,
        ")"
      ),
      {
        skip_if_no_tabicl_fixtures(fixture_name)

        f <- tabicl_load_fixture(fixture_name)
        meta <- tabicl_fixture_meta(fixture_name)

        ce <- brulee:::tabicl_col_embedding(
          embed_dim = meta$embed_dim,
          num_blocks = meta$num_blocks,
          nhead = meta$nhead,
          dim_feedforward = meta$dim_feedforward,
          num_inds = meta$num_inds,
          feature_group_size = meta$feature_group_size,
          target_aware = TRUE,
          max_classes = meta$max_classes,
          reserve_cls_tokens = meta$reserve_cls_tokens,
          bias_free_ln = isTRUE(meta$bias_free_ln),
          ssmax = "qassmax-mlp-elementwise"
        )
        ce$eval()
        tabicl_copy_col_embedding(ce, f)

        out <- ce(f$X, f$y_train)

        expect_equal(dim(out), dim(f$out))
        expect_lt(tabicl_rel_max_diff(out, f$out), 1e-4)
      }
    )
  })
}

test_that("tabicl_col_embedding output shape is (B, T, H + reserve_cls, E)", {
  skip_on_cran()
  skip_if_not_installed("torch")
  if (!torch::torch_is_installed()) {
    skip("libtorch not installed")
  }

  embed_dim <- 16
  reserve_cls <- 4
  ce <- brulee:::tabicl_col_embedding(
    embed_dim = embed_dim,
    num_blocks = 2,
    nhead = 4,
    dim_feedforward = 32,
    num_inds = 8,
    feature_group_size = 3,
    target_aware = TRUE,
    max_classes = 10,
    reserve_cls_tokens = reserve_cls
  )
  ce$eval()

  h <- 5
  x <- torch::torch_randn(1, 7, h)
  y_train <- torch::torch_randint(0, 3, c(1, 4))$to(
    dtype = torch::torch_float()
  )
  out <- ce(x, y_train)

  expect_equal(dim(out), c(1, 7, h + reserve_cls, embed_dim))
})