# Parity tests for the TabICL in-context learning stage (R/tabicl-learning.R)
# and the full model forward (R/tabicl-model.R) against the Python reference.
# Fixtures are generated by dev/tabicl/dump_primitives.py, classification and
# regression variants each.

# ------------------------------------------------------------------------------
# Stage 3: ICLearning

for (fixture in c("icl_learning", "icl_learning_reg")) {
  local({
    fixture_name <- fixture
    test_that(
      paste0(
        "tabicl_icl_learning 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)

        icl <- brulee:::tabicl_icl_learning(
          max_classes = meta$max_classes,
          out_dim = meta$out_dim,
          d_model = meta$d_model,
          num_blocks = meta$num_blocks,
          nhead = meta$nhead,
          dim_feedforward = 2 * meta$d_model,
          bias_free_ln = isTRUE(meta$bias_free_ln),
          ssmax = "qassmax-mlp-elementwise"
        )
        icl$eval()
        tabicl_copy_icl_learning(icl, f)

        out <- icl(f$R, f$y_train)

        expect_equal(dim(out), dim(f$out))
        expect_lt(tabicl_max_abs_diff(out, f$out), 1e-5)
      }
    )
  })
}

# ------------------------------------------------------------------------------
# Full model forward (col -> row -> icl)

for (fixture in c("full_model", "full_model_reg")) {
  local({
    fixture_name <- fixture
    test_that(
      paste0(
        "tabicl_model full forward 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)

        model <- brulee:::tabicl_model(meta$config)
        model$eval()
        tabicl_copy_model(model, f)

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

        expect_equal(dim(out), dim(f$out))
        expect_lt(tabicl_max_abs_diff(out, f$out), 1e-5)
      }
    )
  })
}