test_that("check_data fails if not given a data.table with unique variables.", {
  data <- c("A", "B")
  expect_error(check_data(data), "'data' must be a data.table.")
  baseline_data <- c("A", "B")
  expect_error(check_data(baseline_data), "'baseline_data' must be a data.table.")

  data <- data.table(id = 1:10, id = 10:1)
  expect_error(check_data(data), "'data' has duplicated variable names.")
})

test_that("policy_data checks inputs",{
  d <- sim_single_stage(10, seed=1)

  expect_error(
    policy_data(d, action = 1),
    "'action' must be a vector or a list of type character."
  )
  expect_error(
    policy_data(d, action = TRUE),
    "'action' must be a vector or a list of type character."
  )
  expect_error(
    policy_data(d, action = "a"),
    "Action variables a not found in data."
  )
  expect_error(
    policy_data(d, action = c("A", "a")),
    "Action variables a not found in data."
  )
  expect_error(
    policy_data(d, action = c("a1", "a")),
    "Action variables a1,a not found in data."
  )

  expect_error(
    policy_data(d, action = c("A"), covariates = 1),
    "'covariates' must be a character vector or a list of character vectors."
  )

  expect_error(
    policy_data(d,
                action = c("A"),
                covariates = list(X="a",Y="b"),
                utility = "U"),
    "Variables not found in data: \"a\", \"b\"."
  )
  expect_error(
    policy_data(d,
                action = c("A"),
                covariates = "a",
                utility = "U"),
    "Variables not found in data: \"a\"."
  )
  expect_error(
    policy_data(d,
                action = c("A"),
                covariates = c("Z"),
                utility = "U"),
    NA
  )
  expect_error(
    policy_data(c(1,2,3),
                action = c("A"),
                covariates = c("Z"),
                utility = "U"),
    "'data' must be a data.table."
  )
  expect_error(
    policy_data(d,
                type = "test",
                action = c("A"),
                covariates = c("Z"),
                utility = "U"),
    "'type' must be either \"wide\" or \"long\"."
  )
  expect_error(
    suppressWarnings({
      policy_data(d,
                  type = "long",
                  action = c("A"),
                  covariates = c("Z"),
                  utility = "U")
    })
  )
})

test_that("policy_data() checks id for type = 'wide'", {
  ## long data
  ## 3 cases: no right censoring, right censored before/at stage 1 action,
  ## right censored utility outcome (at stage 2)
  ld <- data.table(
    id = c(1,1,2,3,3),
    stage = c(1,2,1,1,2),
    event = c(0,1,2,0,2),
    A = c("0", NA, NA, "1", NA),
    B = c("gr1","gr1", "gr2", "gr3", "gr3"),
    Z = c("A", NA, "A", "B", NA),
    L = c(1, 2, 2, 1, 3),
    time = c(1, 2, 0.5, 1, 1.5),
    U = c(0, 10, NA, 0, NA),
    U_A0 = c(0,0,NA,0,NA),
    U_A1 = c(0,0,NA,0,NA)
  )
  setkey(ld, id, stage)
  setindex(ld, event)

  expect_error(
    policy_data(data = ld,
                action = "A",
                utility = "U",
                covariates = c("Z"),
                id = "id",
                type = "wide"),
    "'id' column in data must be unique."
  )
})

test_that("policy_data handles varying actions set",{
  d <- sim_two_stage_multi_actions(n = 1e4)
  expect_error(
    pd <- policy_data(data = d,
                      action = c("A_1", "A_2"),
                      baseline = c("B", "BB"),
                      covariates = list(L = c("L_1", "L_2"),
                                        C = c("C_1", "C_2")),
                      utility = c("U_1", "U_2", "U_3")),
    NA
  )

  expect_equal(
    pd$action_set,
    sort(c("yes", "no", "default"))
  )

  sas <- list(c("no", "yes"), c("default","no", "yes"))
  names(sas) <- c("stage_1", "stage_2")
  expect_equal(
    pd$stage_action_sets,
    sas
  )

  expect_equal(
    partial(pd, K=2)$stage_action_sets,
    sas
  )

  expect_equal(
    partial(pd, K=1)$stage_action_sets,
    sas[1]
  )

  expect_equal(
    get_stage_action_sets(subset_id(pd, id = 1:2)),
    sas
  )

})

test_that("policy_data() handles varying sets of/missing covariates in a given stage",{
  n <- 20
  set.seed(1)
  library("polle")
  library("data.table")
  d <- data.table(id = 1:n,
                  Y_1 = rnorm(n),
                  X_1 = rnorm(n),
                  Z_1 = rnorm(n),
                  W_1 = rep("test", n),
                  A_1 = rbinom(n = n, size = 1, prob = .5),
                  X_2 = rnorm(n),
                  Z_2 = rnorm(n),
                  Y_2 = rnorm(n),
                  A_2 = rbinom(n = n, size = 1, prob = .5),
                  U = rnorm(n))

  expect_error(
    pd <- policy_data(data = d,
                      id = "id",
                      action = c("A_1"),
                      covariates = c("X_1", NA),
                      utility = "U"),
    "covariate NA is invalid."
  )

  pd <- policy_data(data = d,
                    id = "id",
                    action = c("A_1", "A_2"),
                    covariates = list(X = c("X_1", "X_2"),
                                      Y = c(NA, "Y_2")),
                    utility = "U")
  expect_equal(
    polle:::get_stage_data.policy_data(pd)$Y,
    unlist(lapply(d$Y_2, function(x) c(NA, x, NA)))
  )

  pd <- policy_data(data = d,
                    id = "id",
                    action = c("A_1", "A_2"),
                    covariates = list(X = c("X_1", "X_2"),
                                      Y = c("Y_1", NA)),
                    utility = "U")
  expect_equal(
    polle:::get_stage_data.policy_data(pd)$Y,
    unlist(lapply(d$Y_1, function(x) c(x, NA, NA)))
  )

  pd <- policy_data(data = d,
                    id = "id",
                    action = c("A_1", "A_2"),
                    covariates = list(X = c("X_1", "X_2"),
                                      W = c("W_1", NA)),
                    utility = "U")
  expect_equal(
    polle:::get_stage_data.policy_data(pd)$W,
    unlist(lapply(d$W_1, function(x) c(x, NA, NA)))
  )

  expect_error(
    policy_data(data = d,
                id = "id",
                action = c("A_1", "A_2"),
                covariates = list(X = c("X_1", "X_2"),
                                  W = c(NA, NA)),
                utility = "U"),
    "'covariates' must be a character vector or a list of character vectors."
  )

  pd <- policy_data(data = d,
                    id = "id",
                    action = c("A_1", "A_2"),
                    covariates = list(X = c("X_1", "X_2"),
                                      Y = c(NA, "Y_2"),
                                      W = c("W_1", NA)),
                    utility = "U")
  expect_equal(
    polle:::get_stage_data.policy_data(pd)$W,
    unlist(lapply(d$W_1, function(x) c(x, NA, NA)))
  )
  expect_equal(
    polle:::get_stage_data.policy_data(pd)$Y,
    unlist(lapply(d$Y_2, function(x) c(NA, x, NA)))
  )

})

test_that("policy_data() melts wide data correctly in a two stage case.", {
  wide_data <- data.table(
    B = c("gr1", "gr2"),
    Z_1 = c("A", "B"),
    L_1 = c(1,2),
    Z_2 = c("C", "D"),
    L_2 = c(3,4),
    treat_1 = c(0,1),
    treat_2 = c(1,0),
    outcome = c(10, 5)
  )

  target_stage_data <- data.table(
    id = c(1,1,1,2,2,2),
    stage = c(1,2,3,1,2,3),
    event = c(0,0,1,0,0,1),
    time = as.numeric(c(NA,NA,NA,NA,NA,NA)),
    time2 = as.numeric(c(NA,NA,NA,NA,NA,NA)),
    A = as.character(c(0,1,NA,1,0,NA)),
    Z = c("A","C",NA,"B","D",NA),
    L = c(1,3,NA,2,4,NA),
    U = c(0,0,10,0,0,5),
    U_A0 = c(0,0,NA,0,0,NA),
    U_A1 = c(0,0,NA,0,0,NA)
  )
  setkey(target_stage_data, id, stage)
  setindex(target_stage_data, event)

  pd <- policy_data(
    data = wide_data,
    action = c("treat_1", "treat_2"),
    covariates = list(
      Z = c("Z_1", "Z_2"),
      L = c("L_1", "L_2")
    ),
    baseline = "B",
    utility = "outcome"
  )
  setindex(pd$stage_data, event)
  expect_equal(pd$stage_data, target_stage_data, check.attributes = TRUE)

  # including deterministic rewards:
  wide_data <- data.table(
    B = c("gr1", "gr2"),
    Z_1 = c("A", "B"),
    L_1 = c(1,2),
    Z_2 = c("C", "D"),
    L_2 = c(3,4),
    treat_1 = c(0,1),
    treat_2 = c(1,0),
    outcome_1 = c(10, 5),
    outcome_2 = c(7, 3),
    outcome_3 = c(-2, 1),
    outcome_1_A0 = c(1.5, 2.5),
    outcome_2_A0 = c(3.5, 4.5)
  )

  target_stage_data <- data.table(
    id = c(1,1,1,2,2,2),
    stage = c(1,2,3,1,2,3),
    event = c(0,0,1,0,0,1),
    time = rep(as.numeric(NA), 6),
    time2 = rep(as.numeric(NA), 6),
    A = as.character(c(0,1,NA,1,0,NA)),
    Z = c("A","C",NA,"B","D",NA),
    L = c(1,3,NA,2,4,NA),
    U = c(10,7,-2,5,3,1),
    U_A0 = c(1.5,3.5,NA,2.5,4.5,NA),
    U_A1 = c(0,0,NA,0,0,NA)
  )
  setkey(target_stage_data, id, stage)
  setindex(target_stage_data, event)

  pd <- policy_data(
    data = wide_data,
    action = c("treat_1", "treat_2"),
    covariates = list(
      Z = c("Z_1", "Z_2"),
      L = c("L_1", "L_2")
    ),
    baseline = "B",
    utility = c("outcome_1", "outcome_2", "outcome_3"),
    deterministic_rewards = list(U_A0 = c("outcome_1_A0", "outcome_2_A0"))
  )
  setindex(pd$stage_data, event)
  expect_equal(pd$stage_data, target_stage_data, check.attributes = TRUE)

  # invalid inputs:
  args <- list(
    data = wide_data,
    action = c("treat_1", "treat_2"),
    covariates = list(
      Z = c("Z_1", "Z_2"),
      L = c("L_1", "L_2")
    ),
    baseline = "B",
    utility = c("outcome_1", "outcome_2", "outcome_3"),
    deterministic_rewards = list(U_A0 = c("outcome_1_A0", "outcome_2_A0"))
  )

  args_copy <- args
  args_copy$covariates$Z <- c("Z_1", "Z_1")
  expect_error(do.call(what = "policy_data", args_copy), "Duplicated variables: \"Z_1\".")

  args_copy <- args
  args_copy$covariates$Z <- c("Z_1")
  expect_error(do.call(what = "policy_data", args_copy), "Each element in 'covariates' must have length 2.")
})

test_that("policy_data() melts wide data correctly in a single stage case.", {

  wide_data <- data.table(
    B = c("gr1", "gr2"),
    Z = c("A", "B"),
    L = c(1, 2),
    treat = c(0,1),
    outcome = c(10, 5)
  )

  target_stage_data <- data.table(
    id = c(1,1,2,2),
    stage = c(1,2,1,2),
    event = c(0,1,0,1),
    time = rep(as.numeric(NA), 4),
    time2 = rep(as.numeric(NA), 4),
    A = c("0", NA, "1", NA),
    B = c("gr1", NA, "gr2", NA),
    Z = c("A", NA, "B", NA),
    L = c(1, NA, 2, NA),
    U = c(0, 10, 0, 5),
    U_A0 = c(0,NA,0,NA),
    U_A1 = c(0,NA,0,NA)
  )
  setkey(target_stage_data, id, stage)
  setindex(target_stage_data, event)

  ## except equal:
  pd <- policy_data(data = wide_data,
                    action = "treat",
                    covariates = c("B", "Z", "L"),
                    utility = "outcome")
  setindex(pd$stage_data, event)
  expect_equal(pd$stage_data, target_stage_data, check.attributes = TRUE)
  rm(pd)

  ## duplicate variable names:
  expect_error(
    policy_data(data = wide_data,
                action = "treat",
                covariates = c("B", "Z", "Z"),
                utility = "outcome"),
    "Duplicated variables: \"Z\"."
  )
  expect_error(
    policy_data(data = wide_data,
                action = "treat",
                covariates = c("B", "Z", "treat"),
                utility = "outcome"),
    "Duplicated variables: \"treat\"."
  )
  expect_error(
    policy_data(data = wide_data,
                action = "treat",
                covariates = c("B", "Z", "L"),
                utility = "treat"),
    "Duplicated variables: \"treat\"."
  )

  # invalid variable names:
  wide_data_copy <- copy(wide_data)
  wide_data_copy$stage <- 1
  wide_data_copy$event <- 0
  expect_error(
    policy_data(data = wide_data_copy,
                action = "treat",
                covariates = c("B", "Z", "stage"),
                utility = "outcome"),
    "'covariates' can not have named elements \"event\", \"stage\", \"time\", or \"time2\"."
  )
  expect_error(
    policy_data(data = wide_data_copy,
                action = "treat",
                covariates = c("B", "Z", "event"),
                utility = "outcome"),
    "'covariates' can not have named elements \"event\", \"stage\", \"time\", or \"time2\"."
  )

  ## id input
  wide_data_copy <- copy(wide_data)
  wide_data_copy$id <- c(1,2)
  expect_error(
    policy_data(data = wide_data_copy,
                action = "treat",
                covariates = c("B", "Z", "L"),
                utility = "outcome"),
    "'data' has a variable id, but 'id' = NULL. Please set 'id' = \"id\" or change the name of the id variable."
  )
  pd <- policy_data(data = wide_data_copy,
                    action = "treat",
                    covariates = c("B", "Z", "L"),
                    utility = "outcome",
                    id = "id")
  setindex(pd$stage_data, event)
  expect_equal(
    pd$stage_data,
    target_stage_data,
    check.attributes = TRUE
  )
  wide_data_copy$id <- NULL
  wide_data_copy$ID <- c(1,2)

  pd <- policy_data(data = wide_data_copy,
                action = "treat",
                covariates = c("B", "Z", "L"),
                utility = "outcome",
                id = "ID")
  setindex(pd$stage_data, event)
  expect_equal(
    pd$stage_data,
    target_stage_data,
    check.attributes = TRUE
  )

  pd <- policy_data(data = wide_data,
                    action = "treat",
                    covariates = c("B", "Z", "L"),
                    utility = "outcome")
  expect_equal(
    pd$colnames$state_names,
    c("B", "Z", "L")
  )

  ## expected baseline
  target_baseline <- data.table(id = c(1,2), B = c("gr1", "gr2"))
  setkey(target_baseline, id)
  pd <- policy_data(data = wide_data,
                    action = "treat",
                    covariates = c("Z", "L"),
                    utility = "outcome",
                    baseline = "B")
  expect_equal(
    pd$baseline_data,
    pd$baseline_data
  )
  expect_equal(
    pd$colnames$baseline_names,
    c("B")
  )
  rm(pd)

  ## deterministic rewards
  wide_data_copy <- copy(wide_data)
  wide_data_copy$reward_A0 <- 2
  wide_data_copy$reward_A1 <- 1
  expect_error(
    policy_data(data = wide_data_copy,
                action = "treat",
                covariates = c("B", "Z", "L"),
                utility = "outcome",
                deterministic_rewards = list("reward_A0", "reward_A1")),
    "'deterministic_rewards' must be a named list with names in the set 'U_A0', 'U_A1'."
  )
  expect_error(
    policy_data(data = wide_data_copy,
                action = "treat",
                covariates = c("B", "Z", "L"),
                utility = "outcome", deterministic_rewards = list(U_0 = "reward_A0", U_1 = "reward_A1")),
    "'deterministic_rewards' must be a named list with names in the set 'U_A0', 'U_A1'."
  )
  pd <- policy_data(data = wide_data_copy,
                    action = "treat",
                    covariates = c("B", "Z", "L"),
                    utility = "outcome",
                    deterministic_rewards = list(U_A0 = "reward_A0", U_A1 = "reward_A1"))
  expect_identical(
    pd$stage_data[, c("U_A0", "U_A1"), with = FALSE],
    data.table(U_A0 = c(2,NA,2,NA), U_A1 = c(1,NA,1,NA))
  )
})

test_that("policy_data() formats long data correctly for a single stage case.", {

  ## long data:
  ld <- data.table(
    id = c(1,1,2,2),
    stage = c(1,2,1,2),
    event = c(0,1,0,1),
    time = rep(as.numeric(NA), 4),
    time2 = rep(as.numeric(NA), 4),
    A = c("0", NA, "1", NA),
    B = c("gr1", NA, "gr2", NA),
    Z = c("A", NA, "B", NA),
    L = c(1, NA, 2, NA),
    U = c(0, 10, 0, 5),
    U_A0 = rep(0, 4),
    U_A1 = rep(0, 4)
  )
  setkey(ld, id, stage)
  setindex(ld, event)

  ## baseline data:
  bd <- data.table(
    id = c(1,2),
    W = c("blue", "red")
  )
  setkey(bd, id)

  ## correct
  expect_error(
    policy_data(data = ld, baseline_data = bd, type = "long"),
    NA
  )
  pd <- policy_data(data = ld, baseline_data = bd, type = "long")
  expect_equal(ld, pd$stage_data)
  expect_equal(bd, pd$baseline_data)

  ## invalid variable names
  expect_error(
    policy_data(data = ld, type = "long", id = 1),
    "'id' must be a character string."
  )
  expect_error(
    policy_data(data = ld, type = "long", id = "ID"),
    "'id' is invalid."
  )
  expect_error(
    policy_data(data = ld, type = "long", action = "AA"),
    "'action' is invalid."
  )

  # setting new names:
  ld_copy <- copy(ld)
  setnames(ld_copy, "id", "ID")
  setnames(ld_copy, "stage", "k")
  setnames(ld_copy, "event", "evt")
  setnames(ld_copy, "A", "action")
  setnames(ld_copy, "U", "reward")

  bd_copy <- copy(bd)
  setnames(bd_copy, "id", "ID")

  ## correct variable names:
  expect_error(
    policy_data(
      data = ld_copy,
      baseline_data = bd_copy,
      type = "long",
      action = "action",
      id = "ID",
      stage = "k",
      event = "evt",
      utility = "reward"
    ),
    NA
  )
  ## comparison with original:
  expect_identical(
    policy_data(
      data = ld_copy,
      baseline_data = bd_copy,
      type = "long",
      action = "action",
      id = "ID",
      stage = "k",
      event = "evt",
      utility = "reward"
    ),
    pd
  )

  # preventing duplicated id variable:
  ld_copy$id <- c(3,3,4,4)
  expect_error(
    policy_data(data = ld_copy, type = "long", id = "ID"),
    "'data' has a variable called \"id\", but 'id' = \"ID\". Please remove or rename 'id'."
  )
  ld_copy$id <- NULL

  # preventing duplicated action variable:
  ld_copy$A <- c(0,NA,0,NA)
  expect_error(
    policy_data(
      data = ld_copy,
      type = "long",
      action = "action",
      id = "ID",
      stage = "k",
      event = "evt",
      utility = "U"
    ),
    "'data' has a variable called \"A\", but 'action' = \"action\". Please remove or rename 'action'."
  )
  ld_copy$A <- NULL

  # missing variable with non-matching default\
  expect_error(
    policy_data(
      data = ld_copy,
      type = "long",
      action = "action",
      id = "ID",
      stage = "k",
      event = "evt"
    ),
    "'utility' is invalid."
  )

})

test_that("policy_data() handles missing values.", {
  # long data:
  ld <- data.table(
    id = c(1,1,2,2),
    stage = c(1,2,1,2),
    event = c(0,1,0,1),
    A = c(0, NA, 1, NA),
    B = c("gr1", NA, "gr2", NA),
    Z = c("A", NA, "B", NA),
    L = c(1, NA, 2, NA),
    U = c(0, 10, 0, 5),
    U_A0 = rep(0, 4),
    U_A1 = rep(0, 4)
  )
  setkey(ld, id, stage)
  setindex(ld, event)

  # baseline data:
  bd <- data.table(
    id = c(1,2),
    W = c("blue", "red")
  )
  setkey(bd, id)

  # correct
  expect_error(
    policy_data(data = ld, baseline_data = bd, type = "long"),
    NA
  )

  # missing values:
  ld$A <- c(NA , NA, 1, NA)
  expect_error(policy_data(data = ld, baseline_data = bd, type = "long"))
  ld$A <- c(0, NA, 1, NA)

  # missing utility
  ld$U <- c(NA, 10, 0, 5)
  expect_error(
    policy_data(data = ld, baseline_data = bd, type = "long"),
    "The utility variable U has missing values"
  )
  ld$U <- c(0, 10, 0, 5)

  # allowing for missing covariates
  ld$Z <- c(NA, NA, "B", NA)
  expect_error(
    policy_data(data = ld, baseline_data = bd, type = "long"),
    NA
  )
})

test_that("policy_data() runs without covariates.", {

  ## long data without covariates:
  ld <- data.table(
    id = c(1,1,2,2),
    stage = c(1,2,1,2),
    event = c(0,1,0,1),
    A = c(0, NA, 1, NA),
    U = c(0, 10, 0, 5),
    U_A0 = rep(0, 4),
    U_A1 = rep(0, 4)
  )
  setkey(ld, id, stage)
  setindex(ld, event)

  expect_error(
    pd <- policy_data(data = ld,  type = "long"),
    NA
  )

  ## wide data without covariates:
  wide_data <- data.table(
    treat = c(0,1),
    outcome = c(10, 5)
  )

  expect_error(
    pd <- policy_data(data = wide_data, action = "treat", utility = "outcome", covariates = list()),
    NA
  )

})
