library(automatedRecLin)
library(data.table)
library(xgboost)

df_1 <- data.frame(
  "name" = c("James", "Emma", "William", "Olivia", "Thomas",
             "Sophie", "Harry", "Amelia", "George", "Isabella"),
  "surname" = c("Smith", "Johnson", "Brown", "Taylor", "Wilson",
                "Davis", "Clark", "Harris", "Lewis", "Walker")
)
df_2 <- data.frame(
  "name" = c("James", "Ema", "Wimliam", "Olivia", "Charlotte",
             "Henry", "Lucy", "Edward", "Alice", "Jack"),
  "surname" = c("Smith", "Johnson", "Bron", "Tailor", "Moore",
                "Evans", "Hall", "Wright", "Green", "King")
)
matches <- data.frame("a" = 1:4, "b" = 1:4)
df_new_1 <- data.frame(
  "name" = c("John", "Emily", "Mark", "Anna", "David"),
  "surname" = c("Smith", "Johnson", "Taylor", "Williams", "Brown")
)
df_new_2 <- data.frame(
  "name" = c("John", "Emely", "Mark", "Michael"),
  "surname" = c("Smitth", "Johnson", "Tailor", "Henders")
)
true_matches <- data.frame("a" = 1:3, "b" = 1:3)

comparators = list("name" = jarowinkler_complement(),
                   "surname" = jarowinkler_complement())
methods_cpar <- list("name" = "continuous_parametric",
                     "surname" = "continuous_parametric")
methods_cnonpar <- list("name" = "continuous_nonparametric",
                        "surname" = "continuous_nonparametric")

confusion_cpar <- matrix(c(2, 0, 1, 17), nrow = 2, ncol = 2)
rownames(confusion_cpar) <- c("Actual Positive", "Actual Negative")
colnames(confusion_cpar) <- c("Predicted Positive", "Predicted Negative")

expect_rate_bounds <- function(result) {
  expect_true(is.finite(result$flr_est))
  expect_true(result$flr_est >= 0 && result$flr_est <= 1)
  expect_true(is.finite(result$mmr_est))
  expect_true(result$mmr_est >= 0 && result$mmr_est <= 1)
}

expect_silent(
  train_rec_lin(A = df_1, B = df_2,
                variables = c("name", "surname"),
                matches = matches)
)

expect_warning(
  train_rec_lin(A = df_1, B = df_2,
                variables = c("name", "surname"),
                matches = matches,
                prob_ratio = "2")
)

expect_equal(
  train_rec_lin(A = df_1, B = df_2,
                variables = c("name", "surname"),
                matches = matches)$b_params,
  data.table("variable" = c("gamma_name", "gamma_surname"),
             "theta" = c(0.5, 0.5),
             "eta" = c(0.02, 0.02))
)

expect_equal(
  train_rec_lin(A = df_1, B = df_2,
                variables = c("name", "surname"),
                matches = matches)$pi_est,
  0.04
)

expect_equal(
  train_rec_lin(A = df_1, B = df_2,
                variables = c("name", "surname"),
                matches = matches)$match_prop,
  0.4
)

expect_silent(
  train_rec_lin(A = df_1, B = df_2,
                variables = c("name", "surname"),
                matches = matches,
                comparators = comparators,
                methods = methods_cpar)
)

expect_warning(
  train_rec_lin(A = df_1, B = df_2,
                variables = c("name", "surname"),
                matches = matches,
                comparators = comparators,
                methods = methods_cpar,
                prob_ratio = "2")
)

expect_equal(
  train_rec_lin(A = df_1, B = df_2,
                variables = c("name", "surname"),
                matches = matches,
                comparators = comparators,
                methods = methods_cpar)$cpar_params,
  data.table("variable" = c("gamma_name", "gamma_surname"),
             "p_0_M" = c(0.5, 0.5),
             "p_0_Omega" = c(0.02000000000000000041633, 0.02000000000000000041633),
             "alpha_M" = c(224.66614858286149569722, 15.65946723172959487158),
             "alpha_Omega" = c(5.407152306652986517577, 6.105019794944631428280),
             "beta_M" = c(2516.2608641280485244351, 176.1690063569579081104),
             "beta_Omega" = c(7.975415092480718115553, 9.363290038376906210260))
)

expect_warning(
  train_rec_lin(A = df_1, B = df_2,
                variables = c("name", "surname"),
                matches = matches,
                comparators = comparators,
                methods = methods_cnonpar)
)

expect_equal(
  train_rec_lin(A = df_1, B = df_2,
                variables = c("name", "surname"),
                matches = matches,
                comparators = comparators,
                methods = methods_cnonpar)$cnonpar_params,
  data.table("variable" = c("gamma_name", "gamma_surname"),
             "p_0_M_cnonpar" = c(0.5, 0.5),
             "p_0_U_cnonpar" = c(0.02, 0.02))
)

kliep_model <- densityratio::kliep(
  df_numerator = data.table("gamma_name" = c(0.1, 0.2, 0.3)),
  df_denominator = data.table("gamma_name" = c(0.2, 0.4, 0.6)),
  progressbar = FALSE,
  nfold = 2
)
kliep_pred <- as.vector(stats::predict(
  kliep_model,
  data.table("gamma_name" = 0.4)
))

expect_equal(
  automatedRecLin:::kliep_hurdle_ratio(
    df = data.table("gamma_name" = c(0, 0.4)),
    variables = "gamma_name",
    p_0_numerator = 0.4,
    p_0_denominator = 0.2,
    ratio_kliep_list = list(
      "gamma_name" = kliep_model
    )
  ),
  c(2, (1 - 0.4) / (1 - 0.2) * kliep_pred)
)

model_b <- train_rec_lin(A = df_1, B = df_2,
                         variables = c("name", "surname"),
                         matches = matches)

expect_equal(
  predict(model_b, df_new_1, df_new_2)$M_est,
  data.table("a" = 1:3, "b" = 1:3, "ratio" = rep(12.75510204081633602868, 3))
)

expect_equal(
  predict(model_b, df_new_1, df_new_2)$n_M_est,
  3.35401915868387945352
)

model_cpar <- train_rec_lin(A = df_1, B = df_2,
                            variables = c("name", "surname"),
                            matches = matches,
                            comparators = comparators,
                            methods = methods_cpar)

expect_equal(
  predict(model_cpar, df_new_1, df_new_2)$M_est,
  data.table("a" = c(1, 3), "b" = c(1, 3),
             "ratio" = c(59126.609831455622042995, 4201.052699951693284675))
)

expect_equal(
  predict(model_cpar, df_new_1, df_new_2)$n_M_est,
  2.000000561079755989624
)

expect_equal(
  predict(model_cpar, df_new_1, df_new_2)$mmr_est,
  0
)

expect_equal(
  predict(model_cpar, df_new_1, df_new_2, true_matches = true_matches)$eval_metrics,
  c("FLR" = 0, "MMR" = 1/3)
)

expect_equal(
  predict(model_cpar, df_new_1, df_new_2, true_matches = true_matches)$confusion,
  confusion_cpar
)

model_cnonpar <- train_rec_lin(A = df_1, B = df_2,
                               variables = c("name", "surname"),
                               matches = matches,
                               comparators = comparators,
                               methods = methods_cnonpar)

expect_equal(
  predict(model_cnonpar, df_new_1, df_new_2)$M_est,
  data.table("a" = c(2, 3, 1), "b" = c(2, 3, 1),
             "ratio" = c(47.61403316943127350669, 29.29835603699235235808, 26.65717991654741325647))
)

expect_equal(
  predict(model_cnonpar, df_new_1, df_new_2)$mmr_est,
  0
)

pred_cpar_flr <- predict(model_cpar, df_new_1, df_new_2,
                         set_construction = "flr", target_rate = 0.1)
expect_rate_bounds(pred_cpar_flr)

pred_cpar_mmr <- predict(model_cpar, df_new_1, df_new_2,
                         set_construction = "mmr", target_rate = 0.1)
expect_rate_bounds(pred_cpar_mmr)

# custom model

vectors <- comparison_vectors(A = df_1, B = df_2, variables = c("name", "surname"),
                              comparators = comparators, matches = matches)
model_xgb <- xgboost::xgboost(x = as.matrix(vectors$Omega[, c("gamma_name", "gamma_surname")]),
                              y = factor(vectors$Omega$match),
                              objective = "binary:logistic", eval_metric = "logloss",
                              nrounds = 100, verbosity = 0, nthread = 1)

expect_silent(
  custom_rec_lin_model(model_xgb, vectors)
)

expect_equal(
  custom_rec_lin_model(model_xgb, vectors)$pi_est,
  0.04
)

expect_equal(
  custom_rec_lin_model(model_xgb, vectors)$match_prop,
  0.4
)

custom_xgb_model <- custom_rec_lin_model(model_xgb, vectors)

expect_equal(
  predict(custom_xgb_model, df_new_1, df_new_2,
          data_type = "matrix", type = "response")$M_est,
  data.table("a" = 1:3, "b" = 1:3,
             "ratio" = rep(29.94770389694980394779, 3))
)

expect_equal(
  predict(custom_xgb_model, df_new_1, df_new_2,
          data_type = "matrix", type = "response")$flr_est,
  0.1591118016965802350882
)

expect_equal(
  predict(custom_xgb_model, df_new_1, df_new_2,
          data_type = "matrix", type = "response")$mmr_est,
  0.1591118016965802350882
)
