context("lfe")

skip_on_cran()

skip_if_not_installed("modeltests")
library(modeltests)

skip_if_not_installed("lfe")

set.seed(27)
n <- 100
df <- data.frame(
  id = sample(5, n, TRUE),
  v1 = sample(5, n, TRUE),
  v2 = sample(1e6, n, TRUE),
  v3 = sample(round(runif(100, max = 100), 4), n, TRUE),
  v4 = sample(round(runif(100, max = 100), 4), n, TRUE)
)

# no FE or clus
fit <- lfe::felm(v2 ~ v3, df)
# with FE
fit2 <- lfe::felm(v2 ~ v3 | id + v1, df, na.action = na.exclude)

# with clus
fit3 <- lfe::felm(v2 ~ v3 | 0 | 0 | id + v1, df, na.action = na.exclude)

## with multiple outcomes
fit_multi <- lfe::felm(v1 + v2 ~ v3 , df)
fit_Y2 <- lfe::felm(v1 ~ v3 , df)

form <- v2 ~ v4
fit_form <- lfe::felm(form, df) # part of a regression test

test_that("felm tidier arguments", {
  check_arguments(tidy.felm)
  check_arguments(glance.felm)
  check_arguments(augment.felm)
})

test_that("tidy.felm", {
  td1 <- tidy(fit)
  td2 <- tidy(fit2, conf.int = TRUE, fe = TRUE, fe.error = FALSE)
  td3 <- tidy(fit2, conf.int = TRUE, fe = TRUE)
  td4 <- tidy(fit_form)
  td5 <- tidy(fit, se = "robust")
  td6 <- tidy(fit2, se = "robust")
  td7 <- tidy(fit2, se = "robust", fe = TRUE)
  td8 <- tidy(fit3)
  td9 <- tidy(fit3, se = "iid")
  
  
  td_multi <- tidy(fit_multi)
  td_multi_CI <- tidy(fit_multi, conf.int = TRUE)
  
  check_tidy_output(td1)
  check_tidy_output(td2)
  check_tidy_output(td3)
  check_tidy_output(td4)
  check_tidy_output(td5)
  check_tidy_output(td6)
  check_tidy_output(td7)
  check_tidy_output(td8)
  check_tidy_output(td9)
  check_tidy_output(td_multi)
  check_tidy_output(td_multi_CI)
  
  check_dims(td1, 2, 5)
  check_dims(td_multi_CI, 4, 8)
  
  expect_equal(tidy(fit_multi)[3:4, -1],
               tidy(fit))
  expect_equal(tidy(fit_multi, conf.int = TRUE)[3:4, -1],
               tidy(fit, conf.int = TRUE))
  expect_equal(tidy(fit_multi, conf.int = TRUE)[1:2, -1],
               tidy(fit_Y2, conf.int = TRUE))
  
  expect_equal(dplyr::pull(td5, std.error),
               as.numeric(lfe:::summary.felm(fit, robust = TRUE)$coef[, "Robust s.e"]))
  expect_equal(dplyr::pull(td6, std.error),
               as.numeric(lfe:::summary.felm(fit2, robust = TRUE)$coef[, "Robust s.e"]))
  expect_equal(dplyr::pull(td8, std.error),
               as.numeric(lfe:::summary.felm(fit3)$coef[, "Cluster s.e."]))
  expect_equal(dplyr::pull(td9, std.error),
               as.numeric(lfe:::summary.felm(fit3, robust = FALSE)$coef[, "Std. Error"]))

  # check for deprecation warning from 0.7.0.9001
  expect_warning(
    tidy(fit, robust = TRUE),
    '"robust" argument has been deprecated'
  )
})

test_that("glance.felm", {
  gl <- glance(fit)
  gl2 <- glance(fit2)

  check_glance_outputs(gl, gl2)
  check_dims(gl, expected_cols = 8)

  expect_error(glance(fit_multi), "Glance does not support linear models with multiple responses.")
})

test_that("augment.felm", {
  check_augment_function(
    aug = augment.felm,
    model = fit,
    data = df
  )

  check_augment_function(
    aug = augment.felm,
    model = fit2,
    data = df
  )

  check_augment_function(
    aug = augment.felm,
    model = fit_form,
    data = df
  )
  expect_error(augment(fit_multi), 
               "Augment does not support linear models with multiple responses.")

  # Ensure that the .resid and .fitted columns are basic columns, not matrix
  aug <- augment(fit)
  expect_false(inherits(aug$.resid, "matrix"))
  expect_false(inherits(aug$.fitted, "matrix"))
  expect_null(c(colnames(aug$.resid), colnames(aug$.fitted)))
})
