# Tests for the bpbounds package
# 2018-10-23 Tom Palmer

if (!requireNamespace("tidyr", quietly = TRUE)) {
  install.packages("tidyr", repos = "https://cloud.r-project.org/")
}
context("Tests for bpbounds package")

# Balke and Pearl, JASA, 1997 examples ----

## Table 1 - note 9665 should be 9663 from paper ----
tab1dat <- data.frame(
  z = c(0, 0, 1, 1, 1, 1, 0, 0),
  x = c(0, 0, 0, 0, 1, 1, 1, 1),
  y = c(0, 1, 0, 1, 0, 1, 0, 1),
  freq = c(74, 11514, 34, 2385, 12, 9665, 0, 0)
)

tab1inddat <- tidyr::uncount(tab1dat, freq)
xt <- xtabs(~ x + y + z, data = tab1inddat)
p <- prop.table(xt, margin = 3)

## Error checks
test_that("Errors", {
  expect_error(bpbounds(p, fmt = "bivariate"))
  expect_error(bpbounds(p, fmt = "anything-you-like"))
})

## Analyses
test_that("Balke and Pearl Table 1 example: trivariate data with 2 category instrument",
          {
            # Using conditional probabilities
            bpres <- bpbounds(p)

            expect_equal(class(bpres), "bpbounds")
            expect_equal(bpres$fmt, "trivariate")
            expect_equal(bpres$nzcats, 2)

            expect_true(bpres$inequality)
            expect_equal(bpres$bplb, -0.1946, tol = 1e-4)
            expect_equal(bpres$bpub, 0.0054, tol = 1e-4)
            expect_equal(bpres$p10low, 0.9936, tol = 1e-4)
            expect_equal(bpres$p10upp, 0.9936, tol = 1e-4)
            expect_equal(bpres$p11low, 0.7990, tol = 1e-4)
            expect_equal(bpres$p11upp, 0.9990, tol = 1e-4)
            expect_equal(bpres$crrlb, 0.8042, tol = 1e-4)
            expect_equal(bpres$crrub, 1.0054, tol = 1e-4)

            expect_true(bpres$monoinequality)
            expect_equal(bpres$monobplb, -0.1946, tol = 1e-4)
            expect_equal(bpres$monobpub, 0.0054, tol = 1e-4)
            expect_equal(bpres$monop10low, 0.9936, tol = 1e-4)
            expect_equal(bpres$monop10upp, 0.9936, tol = 1e-4)
            expect_equal(bpres$monop11low, 0.7990, tol = 1e-4)
            expect_equal(bpres$monop11upp, 0.9990, tol = 1e-4)
            expect_equal(bpres$monocrrlb, 0.8042, tol = 1e-4)
            expect_equal(bpres$monocrrub, 1.0054, tol = 1e-4)

            print(bpres)
            print(bpres, digits = 4)
            bpres

            sbp = summary(bpres)
            expect_equal(class(sbp), "summary.bpbounds")

            print(sbp)
            sbp
            print(sbp, digits = 2)
            print(sbp, digits = 4)

            # Using cell counts
            bpres = bpbounds(xt)
            expect_equal(class(bpres), "bpbounds")
            expect_equal(bpres$fmt, "trivariate")
            expect_equal(bpres$nzcats, 2)

            expect_true(bpres$inequality)
            expect_equal(bpres$bplb, -0.1946, tol = 1e-4)
            expect_equal(bpres$bpub, 0.0054, tol = 1e-4)
            expect_equal(bpres$p10low, 0.9936, tol = 1e-4)
            expect_equal(bpres$p10upp, 0.9936, tol = 1e-4)
            expect_equal(bpres$p11low, 0.7990, tol = 1e-4)
            expect_equal(bpres$p11upp, 0.9990, tol = 1e-4)
            expect_equal(bpres$crrlb, 0.8042, tol = 1e-4)
            expect_equal(bpres$crrub, 1.0054, tol = 1e-4)

            expect_true(bpres$monoinequality)
            expect_equal(bpres$monobplb, -0.1946, tol = 1e-4)
            expect_equal(bpres$monobpub, 0.0054, tol = 1e-4)
            expect_equal(bpres$monop10low, 0.9936, tol = 1e-4)
            expect_equal(bpres$monop10upp, 0.9936, tol = 1e-4)
            expect_equal(bpres$monop11low, 0.7990, tol = 1e-4)
            expect_equal(bpres$monop11upp, 0.9990, tol = 1e-4)
            expect_equal(bpres$monocrrlb, 0.8042, tol = 1e-4)
            expect_equal(bpres$monocrrub, 1.0054, tol = 1e-4)
          })


## Test the bivariate formulation ----
g  = xtabs(~ y + z, data = tab1inddat)
gp = prop.table(g, margin = 2)

t  = xtabs(~ x + z, data = tab1inddat)
tp = prop.table(t, margin = 2)

test_that("Balke and Pearl Table 1 example treated as bivariate data", {
  bpres = bpbounds(p = gp, t = tp, fmt = "bivariate")
  expect_true(bpres$inequality)
  expect_equal(bpres$bplb, -0.1974, tol = 1e-4)
  expect_equal(bpres$bpub, 0.0064, tol = 1e-4)
  expect_equal(bpres$p10low, 0.9936, tol = 1e-4)
  expect_equal(bpres$p10upp, 0.9936, tol = 1e-4)
  expect_equal(bpres$p11low, 0.7962, tol = 1e-4)
  expect_equal(bpres$p11upp, 1, tol = 1e-4)
  expect_equal(bpres$crrlb, 0.8013, tol = 1e-4)
  expect_equal(bpres$crrub, 1.0064, tol = 1e-4)

  expect_true(bpres$monoinequality)
  expect_equal(bpres$monobplb, -0.1974, tol = 1e-4)
  expect_equal(bpres$monobpub, 0.0064, tol = 1e-4)
  expect_equal(bpres$monop10low, 0.9936, tol = 1e-4)
  expect_equal(bpres$monop10upp, 0.9936, tol = 1e-4)
  expect_equal(bpres$monop11low, 0.7962, tol = 1e-4)
  expect_equal(bpres$monop11upp, 1, tol = 1e-4)
  expect_equal(bpres$monocrrlb, 0.8013, tol = 1e-4)
  expect_equal(bpres$monocrrub, 1.0064, tol = 1e-4)

  sbp = summary(bpres)
  print(sbp, digits = 3)
  print(sbp)
})

test_that("Balke and Pearl, bivariate data using cell counts",
          {
            bpres <- bpbounds(p = g, t = t, fmt = "bivariate")
            expect_true(bpres$inequality)
            expect_equal(bpres$bplb, -0.1974, tol = 1e-4)
            expect_equal(bpres$bpub, 0.0064, tol = 1e-4)
            expect_equal(bpres$p10low, 0.9936, tol = 1e-4)
            expect_equal(bpres$p10upp, 0.9936, tol = 1e-4)
            expect_equal(bpres$p11low, 0.7962, tol = 1e-4)
            expect_equal(bpres$p11upp, 1, tol = 1e-4)
            expect_equal(bpres$crrlb, 0.8013, tol = 1e-4)
            expect_equal(bpres$crrub, 1.0064, tol = 1e-4)

            expect_true(bpres$monoinequality)
            expect_equal(bpres$monobplb, -0.1974, tol = 1e-4)
            expect_equal(bpres$monobpub, 0.0064, tol = 1e-4)
            expect_equal(bpres$monop10low, 0.9936, tol = 1e-4)
            expect_equal(bpres$monop10upp, 0.9936, tol = 1e-4)
            expect_equal(bpres$monop11low, 0.7962, tol = 1e-4)
            expect_equal(bpres$monop11upp, 1, tol = 1e-4)
            expect_equal(bpres$monocrrlb, 0.8013, tol = 1e-4)
            expect_equal(bpres$monocrrub, 1.0064, tol = 1e-4)

            sbp <- summary(bpres)
            print(sbp, digits = 3)
            print(sbp)
          })

test_that("Bivariate data with cell counts for one and cond probs other", {
  bpres <- bpbounds(p = gp, t = t, fmt = "bivariate")
  sbp <- summary(bpres)

  expect_equal(class(bpres), "bpbounds")
  expect_equal(bpres$fmt, "bivariate")
  expect_equal(bpres$nzcats, 2)

  expect_true(bpres$inequality)
  expect_equal(bpres$bplb, -0.1974, tol = 1e-4)
  expect_equal(bpres$bpub, 0.0064, tol = 1e-4)
  expect_equal(bpres$p10low, 0.9936, tol = 1e-4)
  expect_equal(bpres$p10upp, 0.9936, tol = 1e-4)
  expect_equal(bpres$p11low, 0.7962, tol = 1e-4)
  expect_equal(bpres$p11upp, 1, tol = 1e-4)
  expect_equal(bpres$crrlb, 0.8013, tol = 1e-4)
  expect_equal(bpres$crrub, 1.0064, tol = 1e-4)

  expect_true(bpres$inequality)
  expect_equal(bpres$monobplb, -0.1974, tol = 1e-4)
  expect_equal(bpres$monobpub, 0.0064, tol = 1e-4)
  expect_equal(bpres$monop10low, 0.9936, tol = 1e-4)
  expect_equal(bpres$monop10upp, 0.9936, tol = 1e-4)
  expect_equal(bpres$monop11low, 0.7962, tol = 1e-4)
  expect_equal(bpres$monop11upp, 1, tol = 1e-4)
  expect_equal(bpres$monocrrlb, 0.8013, tol = 1e-4)
  expect_equal(bpres$monocrrub, 1.0064, tol = 1e-4)
})

## Balke and Pearl, 1997, Table 2 - 0.001 was 0 in published table ----
tab2cp <- c(.0064, 0, .9936, 0, .0028, 0.001, .1972, .799)
p2 <- array(tab2cp,
           dim = c(2, 2, 2),
           dimnames = list(
             x = c(0, 1),
             y = c(0, 1),
             z = c(0, 1)
           ))
p2 <- as.table(p2)
sum(p2)

test_that("Balke and Pearl Table 2 example: trivariate data with 2 category instrument",
          {
            bpres <- bpbounds(p2, fmt = "trivariate")
            sbp <- summary(bpres)
            print(sbp)

            expect_equal(class(bpres), "bpbounds")
            expect_equal(bpres$fmt, "trivariate")
            expect_equal(bpres$nzcats, 2)

            expect_true(bpres$inequality)
            expect_equal(bpres$bplb, -0.1946, tol = 1e-4)
            expect_equal(bpres$bpub, 0.0054, tol = 1e-4)
            expect_equal(bpres$p10low, 0.9936, tol = 1e-4)
            expect_equal(bpres$p10upp, 0.9936, tol = 1e-4)
            expect_equal(bpres$p11low, 0.7990, tol = 1e-4)
            expect_equal(bpres$p11upp, 0.9990, tol = 1e-4)
            expect_equal(bpres$crrlb, 0.8042, tol = 1e-4)
            expect_equal(bpres$crrub, 1.0054, tol = 1e-4)

            expect_true(bpres$monoinequality)
            expect_equal(bpres$monobplb, -0.1946, tol = 1e-4)
            expect_equal(bpres$monobpub, 0.0054, tol = 1e-4)
            expect_equal(bpres$monop10low, 0.9936, tol = 1e-4)
            expect_equal(bpres$monop10upp, 0.9936, tol = 1e-4)
            expect_equal(bpres$monop11low, 0.7990, tol = 1e-4)
            expect_equal(bpres$monop11upp, 0.9990, tol = 1e-4)
            expect_equal(bpres$monocrrlb, 0.8042, tol = 1e-4)
            expect_equal(bpres$monocrrub, 1.0054, tol = 1e-4)
          })


# Meleady AJCN 2003; 3 category instrument - Table 3 of paper ----
dat <- data.frame(
  count = c(341, 47, 297, 17, 63, 18, 272, 41, 269, 38, 56, 35),
  z = c(0, 0, 1, 1, 2, 2, 0, 0, 1, 1, 2, 2),
  y = c(0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1),
  x = c(0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1)
)
longdat <- tidyr::uncount(dat, weights = count)

## Trivariate data
xt3 <- xtabs(count ~ x + y + z, data = dat)
p3  <- prop.table(xt3, margin = 3)

test_that("Mendelian randomization with 3 category instrument, trivariate data",
          {
            bpres <- bpbounds(p3)
            expect_true(bpres$inequality)
            expect_equal(bpres$bplb, -0.3101, tol = 1e-4)
            expect_equal(bpres$bpub, 0.4622, tol = 1e-4)
            expect_equal(bpres$p10low, 0.4332, tol = 1e-4)
            expect_equal(bpres$p10upp, 0.5136, tol = 1e-4)
            expect_equal(bpres$p11low, 0.2035, tol = 1e-4)
            expect_equal(bpres$p11upp, 0.8953, tol = 1e-4)
            expect_equal(bpres$crrlb, 0.3962, tol = 1e-4)
            expect_equal(bpres$crrub, 2.0670, tol = 1e-4)
            expect_false(bpres$monoinequality)
            sbp <- summary(bpres)
            print(sbp)
          })

## Bivariate data
gtab <- xtabs(~ y + z, data = longdat)
gp <- prop.table(gtab, margin = 2)
gp

ttab <- xtabs(~ x + z, data = longdat)
tp <- prop.table(ttab, margin = 2)
tp

test_that("Mendelian randomization with 3 category instrument, bivariate data",
          {
            bpres <- bpbounds(p = gp, t = tp, fmt = "bivariate")
            print(bpres)
            expect_true(bpres$inequality)
            expect_equal(bpres$bplb, -0.5720, tol = 1e-4)
            expect_equal(bpres$bpub, 0.5942, tol = 1e-4)
            expect_equal(bpres$p10low, 0.4058, tol = 1e-4)
            expect_equal(bpres$p10upp, 0.5720, tol = 1e-4)
            expect_equal(bpres$p11low, 0, tol = 1e-4)
            expect_equal(bpres$p11upp, 1, tol = 1e-4)
            expect_equal(bpres$crrlb, 0, tol = 1e-4)
            expect_equal(bpres$crrub, 2.4643, tol = 1e-4)
            expect_false(bpres$monoinequality)

            sbp <- summary(bpres)
            print(sbp)
          })

# GitHub issue #3: IV inequality with 3 category instrument ----
# The reporter applied rowSums(apply(tabp, c(1,2), max)) as the IV inequality check,
# which is only valid for a binary instrument. The correct inequalities for a
# 3-category instrument are the Ramsahai constraints implemented in A_tri_x2y2z3.

test_that("Issue 3, example 1: inequality correctly detected as satisfied", {
  # This table is compatible with a valid IV model: a joint distribution over
  # the 8 compliance types and 4 response types exists that reproduces it
  # exactly (verified by linear programming), so all Ramsahai constraints
  # hold. Versions <= 0.1.7 wrongly reported a violation because the panels
  # of p were fed to A_tri_x2y2z3 in the wrong within-panel order.
  tabp <- as.table(array(
    data = c(0.5279183, 0.02208171, 0.1220817, 0.3279183,
             0.0849975, 0.3150025, 0.4650025, 0.1349975,
             0.1132796, 0.3867204, 0.1867204, 0.3132796),
    dim = c(2, 2, 3),
    dimnames = list(x = c(0, 1), y = c(0, 1), z = c(0, 1, 2))
  ))
  bpres <- bpbounds(tabp)
  expect_true(bpres$inequality)
})

test_that("Issue 3, example 2 (vignette MR data): inequality correctly not violated", {
  # Reporter expected inequality = FALSE based on rowSums(apply(p3, c(1,2), max)) = 1.08,
  # but this is not a valid IV inequality check for a 3-category instrument.
  # All Ramsahai constraints are satisfied for this dataset.
  bpres <- bpbounds(p3)
  expect_true(bpres$inequality)
})

# Trivariate data, 3 category instrument, monotonicity bounds ----
# Constructed from an explicit IV model with monotone compliance:
# compliance types never/(z=2 only)/(z=1,2)/always with probabilities
# 0.4/0.3/0.2/0.1 and response types (y0,y1) = (0,0)/(1,0)/(0,1)/(1,1) with
# probabilities 0.3/0.2/0.4/0.1, independent of each other. Hence
# P(Y=1|do(X=0)) = 0.3, P(Y=1|do(X=1)) = 0.5, and the true ACE = 0.2.
test_that("Trivariate 3 category instrument monotonicity bounds", {
  cpm <- c(0.63, 0.05, 0.27, 0.05,
           0.49, 0.15, 0.21, 0.15,
           0.28, 0.30, 0.12, 0.30)
  tabm <- as.table(array(
    cpm,
    dim = c(2, 2, 3),
    dimnames = list(x = c(0, 1), y = c(0, 1), z = c(0, 1, 2))
  ))
  bpres <- bpbounds(tabm)

  expect_true(bpres$inequality)
  expect_true(bpres$monoinequality)

  # mlow = p112 + p000 - 1; mupp = 1 - p100 - p012
  expect_equal(bpres$monobplb, -0.07, tol = 1e-8)
  expect_equal(bpres$monobpub, 0.43, tol = 1e-8)

  expect_equal(bpres$monop10low, 0.27, tol = 1e-8)
  expect_equal(bpres$monop10upp, 0.37, tol = 1e-8)
  expect_equal(bpres$monop11low, 0.30, tol = 1e-8)
  expect_equal(bpres$monop11upp, 0.70, tol = 1e-8)

  expect_equal(bpres$monocrrlb, 0.30 / 0.37, tol = 1e-8)
  expect_equal(bpres$monocrrub, 0.70 / 0.27, tol = 1e-8)

  # bounds must contain the true causal quantities of the generating model
  expect_true(bpres$monobplb <= 0.2 && 0.2 <= bpres$monobpub)
  expect_true(bpres$monop10low <= 0.3 && 0.3 <= bpres$monop10upp)
  expect_true(bpres$monop11low <= 0.5 && 0.5 <= bpres$monop11upp)
})

# CRR bounds reported as NA when unbounded ----
test_that("CRR bound is NA with a note when P(Y|do(X=0)) lower bound is 0", {
  g <- as.table(array(c(.9, .1, .8, .2), dim = c(2, 2),
                      dimnames = list(y = 0:1, z = 0:1)))
  t <- as.table(array(c(.5, .5, .4, .6), dim = c(2, 2),
                      dimnames = list(x = 0:1, z = 0:1)))
  bpres <- bpbounds(p = g, t = t, fmt = "bivariate")

  expect_true(bpres$inequality)
  expect_true(bpres$monoinequality)
  expect_equal(bpres$p10low, 0)
  expect_true(is.na(bpres$crrub))
  expect_true(is.na(bpres$monocrrub))
  expect_equal(bpres$crrlb, 0)

  expect_output(print(summary(bpres)), "CRR bounds reported as NA are unbounded")
})

## More error checks
test_that("Cond probs and 1 cell count error", {
  cpr <- c(.0064, 0, .9936, 0, .0028, .001, .1972, 20)
  tabpr <- as.table(array(
    cpr,
    dim = c(2, 2, 2),
    dimnames = list(
      x = c(0, 1),
      y = c(0, 1),
      z = c(0, 1)
    )
  ))
  expect_error(bpbounds(tabpr))
})

test_that("Cond probs and 1, giving cond probs sum error", {
  cpr <- c(.0064, 0, .9936, 0, .0028, .001, .1972, 1)
  tabpr <- array(cpr,
                dim = c(2, 2, 2),
                dimnames = list(
                  x = c(0, 1),
                  y = c(0, 1),
                  z = c(0, 1)
                )) |>
    as.table()
  expect_error(bpbounds(tabpr))
})

test_that("Cell counts and one cond prob", {
  cpr <- c(640, 0, 9936, 0, 28, 1, 1972, 0.5)
  tabpr <- array(cpr,
                dim = c(2, 2, 2),
                dimnames = list(
                  x = c(0, 1),
                  y = c(0, 1),
                  z = c(0, 1)
                )) |>
    as.table()
  expect_error(bpbounds(tabpr))
})