library(TauStar)
context("Testing the tStar function.")

# The Bergsma and Dassios (2014) 'a' function.
a <- function(z) {
  sign(round(abs(z[1] - z[2]) +
    abs(z[3] - z[4]) -
    abs(z[1] - z[3]) -
    abs(z[2] - z[4]), 10))
}

# An extremely naive implementation of tStar just to check things work
# correctly in general.
tStarSlow <- function(x, y, vStat = F) {
  if (length(x) != length(y) || length(x) < 4) {
    stop("Input to tStarSlow of invalid length.")
  }
  n <- length(x)
  val <- 0
  for (i in 1:n) {
    for (j in 1:n) {
      for (k in 1:n) {
        for (l in 1:n) {
          inds <- c(i, j, k, l)
          if (length(unique(inds)) == 4 || vStat == TRUE) {
            val <- val + a(x[inds]) * a(y[inds])
          }
        }
      }
    }
  }
  if (vStat) {
    return(val / n^4)
  } else {
    return(val / (n * (n - 1) * (n - 2) * (n - 3)))
  }
}

# A distribution that is a mixture of continuous and discrete, used to check
# the tStar algorithm works on such input.
poissonGaussMix <- function(n) {
  poisOrGaus <- sample(c(0, 1), n, replace = TRUE)
  return(rpois(n, 5) * poisOrGaus + rnorm(n) * (1 - poisOrGaus))
}

test_that("tStar implementations agree", {
  set.seed(283721)

  reps <- 3
  m <- 6
  # Just a sanity check that the R naive version agrees with the C++ naive
  # version
  for (i in reps) {
    x <- rnorm(m)
    y <- rnorm(m)
    expect_equal(tStarSlow(x, y), tStar(x, y, slow = TRUE))
    expect_equal(tStarSlow(x, y, TRUE), tStar(x, y, TRUE, slow = TRUE))

    x <- rpois(m, 5)
    y <- rpois(m, 5)
    expect_equal(tStarSlow(x, y), tStar(x, y, slow = TRUE))
    expect_equal(tStarSlow(x, y, TRUE), tStar(x, y, TRUE, slow = TRUE))

    x <- rnorm(m)
    y <- rpois(m, 5)
    expect_equal(tStarSlow(x, y), tStar(x, y, slow = TRUE))
    expect_equal(tStarSlow(x, y, TRUE), tStar(x, y, TRUE, slow = TRUE))

    x <- poissonGaussMix(m)
    y <- poissonGaussMix(m)
    expect_equal(tStarSlow(x, y), tStar(x, y, slow = TRUE))
    expect_equal(tStarSlow(x, y, TRUE), tStar(x, y, TRUE, slow = TRUE))
  }

  m <- 30
  reps <- 10
  methods <- c("heller", "weihs", "naive")
  areAllEq <- function(x, y, vstat) {
    vals <- numeric(length(methods))
    for (i in 1:length(methods)) {
      vals[i] <- tStar(x, y, method = methods[i], vStatistic = vstat)
    }
    for (i in 1:(length(methods) - 1)) {
      expect_equal(vals[i], vals[i + 1])
    }
  }
  for (i in 1:reps) {
    x <- rnorm(m)
    y <- rnorm(m)
    areAllEq(x, y, F)
    areAllEq(x, y, TRUE)

    x <- rpois(m, 5)
    y <- rpois(m, 5)
    areAllEq(x, y, FALSE)
    areAllEq(x, y, TRUE)

    x <- rnorm(m)
    y <- rpois(m, 5)
    areAllEq(x, y, FALSE)
    areAllEq(x, y, TRUE)

    x <- poissonGaussMix(m)
    y <- poissonGaussMix(m)
    areAllEq(x, y, FALSE)
    areAllEq(x, y, TRUE)
  }

  x <- rnorm(100)
  y <- rnorm(100)
  ts <- tStar(x, y)
  tvs <- tStar(x, y, TRUE)
  expect_equal(ts, tStar(x, y, slow = TRUE))
  expect_true(abs(tStar(x, y,
    resample = TRUE, sampleSize = 10,
    numResamples = 10000
  ) - ts) < 2 * 10^-3)
})

test_that("tStar errors on bad input", {
  x <- list(1, 2, 3, 4)
  y <- c(1, 2, 3, 4)
  expect_error(tStar(x, y))

  expect_error(tStar(numeric(0), numeric(0)))
  for (i in 1:3) {
    expect_error(tStar(1:i, 1:i))
  }

  expect_error(tStar(1:10, 1:9))
  expect_error(tStar(1:9, 1:10))
  expect_error(tStar(1:10, 1:10, resample = TRUE, slow = TRUE))
  expect_error(tStar(1:10, 1:10, resample = TRUE, numResamples = -1))
  expect_error(tStar(1:10, 1:10, resample = TRUE, sampleSize = -1))
  expect_error(tStar(1:10, 1:10, vStatistic = TRUE, resample = TRUE))
})