devtools::load_all("../mlr3")
devtools::load_all(".")
set.seed(8)

# load data
tsk_adult_train = tsk("adult_train")
tsk_adult_test = tsk("adult_test")
tsk_adult_train

# train and predict
learner = lrn("classif.rpart")
learner$train(tsk_adult_train)
prediction = learner$predict(tsk_adult_test)

# set protected attribute
tsk_adult_train$set_col_roles("race", add_to = "pta")
tsk_adult_test$set_col_roles("race", add_to = "pta")

# create groupwise metrics
msr_3 = groupwise_metrics(msr("classif.fomr"),  tsk_adult_train)
unname(sapply(msr_3, function(x) x$id))

# subset test task
adult_subset = tsk_adult_test$clone()
df = adult_subset$data()
rows = seq_len(nrow(df))[df$race %in% c("Black", "White") & df$sex %in% c("Female")]
adult_subset$filter(rows)
adult_subset$set_col_roles("race", add_to = "pta")

# adult subset contains 4731 rows
# tsk_adult_test and prediction contains 15315 rows

# in get_pta() we call task$data(rows = prediction$row_ids) (helpers.R Line 132)
# which ignores the filter on the task
# in mlr3 1.0.0 this is an error
#
# in MeasureSubgroup we assume that the return of get_pta() is of length 15315
# see groups[, row_ids := prediction$row_ids] in MeasureSubgroup.R Line 75
# do we need both tasks?
prediction$score(msr_3, task = adult_subset)