R/vi_model.R

In TengMCing/visage: Visual Inference for Linear Regression Diagnostics

# VI_MODEL ----------------------------------------------------------------

# nocov start

class_VI_MODEL <- function(env = new.env(parent = parent.frame())) {

  # Pass CMD check
  .fitted <- .resid <- self <- NULL

  bandicoot::new_class(bandicoot::BASE, env = env, class_name = "VI_MODEL")

  env$prm <- list()
  env$prm_type <- list()

# init --------------------------------------------------------------------

  init_ <- function(prm = list(), prm_type = list(), formula = self$formula, null_formula = NULL, alt_formula = NULL) {

    self$prm <- prm
    self$prm_type <- prm_type

    if (is.null(formula)) {
      stop("`formula` is missing! Unable to create `y`.")
    } else {

      # Update the formula
      self$set_formula(formula = formula,
                       null_formula = null_formula,
                       alt_formula = alt_formula)
    }

    # Get the quoted formula, and eval in the current environment
    formula <- eval(self$formula)

    # Assign values of `prm` into current environment
    list2env(prm, envir = environment())

    # Set the y variable
    self$prm$y <- visage::closed_form(formula, env = new.env(parent = parent.env(self)))

    return(invisible(self))
  }

# set_formula -------------------------------------------------------------

  set_formula_ <- function(...) {

    for_list <- list(...)
    name_for_list <- names(for_list)

    if (length(name_for_list) == 0) stop("No formula name is provided!")

    for (i in 1:length(for_list)) {
      if (name_for_list[i] == "") stop("Formula needs to be provided with a name")

      self[[name_for_list[i]]] <- for_list[[i]]
      attributes(self[[name_for_list[i]]]) <- NULL
    }

    return(invisible(self))
  }


# set_parameter -----------------------------------------------------------

  set_prm_ <- function(prm_name, prm_value) {

    for (i in 1:length(prm_name)) {
      pname <- prm_name[[i]]
      pval <- prm_value[[i]]

      self$prm[[pname]] <- pval
      self$prm$y$set_sym(list(pname), list(pval))
    }

    return(invisible(self))

  }

# gen ---------------------------------------------------------------------

  gen_ <- function(n, fit_model = TRUE, test = FALSE, computed = NULL) {

    # Generate the data frame from the expression
    dat <- visage::CLOSED_FORM$as_dataframe(self$prm$y$gen(n, rhs_val = TRUE, computed = computed), "y")

    # If requires residuals and fitted values, fit the model
    if (fit_model) {
      mod <- self$fit(dat)
      dat$.resid <- mod$residuals
      dat$.fitted <- mod$fitted.values
    }

    # If requires test statistics, get the test results
    if (test) {
      test_results <- self$test(dat)
      dat$test_name <- test_results$name
      dat$statistic <- test_results$statistic
      dat$p_value <- test_results$p_value
    }

    return(dat)
  }

# gen_lineup --------------------------------------------------------------

  gen_lineup_ <- function(n, k = 20, pos = NULL, computed = NULL, test = TRUE) {

    # Generate data with fitted values and residuals and test results
    dat <- self$gen(n, fit_model = TRUE, test = test, computed = computed)

    # Get the null model
    mod <- self$fit(dat)

    # Generate a random position for the true plot
    if (is.null(pos)) pos <- sample(1:k, 1)

    # Record the true position and mark it as true plot
    dat$k <- pos
    dat$null <- FALSE

    # Get a copy of the data frame
    newdat <- dat

    for (i in 1:k) {

      # Skip the true position
      if (i == pos) next

      # Combine the data frame generated by null_resid with the main data frame
      dat <- dplyr::bind_rows(dat, `[[<-`(self$null_resid(newdat, mod, test = test), "k", value = i))
    }

    return(dat)
  }


# null_resid --------------------------------------------------------------

  null_resid_ <- function(dat, mod, test = FALSE) {

    # Get the fitted values
    .fitted <- mod$fitted.values

    # Get the original RSS
    old_rss <- self$rss(mod)

    # Get the number of observations
    n <- length(.fitted)

    # Replace the regressand with random values
    dat$y <- stats::rnorm(n)

    # Update the model
    mod <- stats::update(mod, data = dat)

    # Update the residuals of the data frame
    dat$.resid <- mod$residuals * sqrt(old_rss/self$rss(mod))

    # Update the random y
    dat$y <- dat$.resid + .fitted

    # Get the test result from the updated data
    if (test)
    {
      test_results <- self$test(dat)
      dat$test_name <- test_results$name
      dat$statistic <- test_results$statistic
      dat$p_value <- test_results$p_value
    }

    # Label the data as null data
    dat$null <- TRUE

    return(dat)
  }


# rss ---------------------------------------------------------------------

  rss_ <- function(mod) sum(mod$residuals^2)


# fit ---------------------------------------------------------------------

  fit_ <- function(dat, formula = self$null_formula, ...) {

    if (is.null(formula)) stop("`formula` is missing! Can't fit the model.")

    # Get rid of the environment and class attributes in case it is a formula
    attributes(formula) <- NULL

    # Fit the model with the provided formula and data
    lm_call <- substitute(stats::lm(formula = formula, data = dat, ...))
    mod <- eval(lm_call, envir = parent.frame())

    return(mod)
  }


# test --------------------------------------------------------------------

  test_ <- function(dat, null_formula = self$null_formula, alt_formula = self$alt_formula) {

    # Define the null model and the alternative model
    null_mod <- self$fit(dat, null_formula)
    alt_mod <- self$fit(dat, alt_formula)

    # Use F-test to test the null hypothesis
    F_test <- stats::anova(null_mod, alt_mod, test = "F")

    list(name = "F-test", statistic = F_test$F[2], p_value = F_test$`Pr(>F)`[2])
  }


# average_effect_size -----------------------------------------------------

  average_effect_size_ <- function(n = 50, tol = 1e-2, window_size = 1000, ...) {
    effect_size_series <- c()

    while (TRUE) {
      effect_size_series <- c(effect_size_series, self$sample_effect_size(self$gen(n = n), ...))
      if (length(effect_size_series) < 2 * window_size) next
      if (any(is.na(effect_size_series))) return(NA)

      delta <- abs(mean(effect_size_series) - mean(effect_size_series[1:(length(effect_size_series) - window_size)]))

      if (delta < tol) break
    }

    mean(effect_size_series)
  }


# sample_effect_size ------------------------------------------------------

  sample_effect_size_ <- function(...) {NA}

# plot_resid --------------------------------------------------------------

  plot_resid_ <- function(dat, alpha = 1, size = 0.5, stroke = 0.5) {
    ggplot2::ggplot(dat) +
      ggplot2::geom_point(ggplot2::aes(.fitted, .resid), alpha = alpha, size = size, stroke = stroke)
  }


# plot_qq -----------------------------------------------------------------

  plot_qq_ <- function(dat) {
    ggplot2::ggplot(dat) +
      ggplot2::geom_qq_line(ggplot2::aes(sample = .resid)) +
      ggplot2::geom_qq(ggplot2::aes(sample = .resid))
  }


# plot --------------------------------------------------------------------

  plot_ <- function(dat,
                    type = "resid",
                    theme = ggplot2::theme_grey(),
                    alpha = 1,
                    size = 0.5,
                    stroke = 0.5,
                    remove_axis = FALSE,
                    remove_legend = FALSE,
                    remove_grid_line = FALSE,
                    add_zero_line = TRUE) {
    p <- switch(type,
                "resid" = self$plot_resid(dat, alpha, size, stroke),
                "qq" = self$plot_qq(dat)) +
      theme

    if (add_zero_line) {
      p <- p + ggplot2::geom_hline(yintercept = 0, col = "red")

      # Pop the last layers
      tmp <- p$layers[[length(p$layers)]]
      p$layers[[length(p$layers)]] <- NULL

      # Append the element back to the list
      p$layers <- append(list(tmp), p$layers)
    }

    if (remove_axis) p <- p + ggplot2::theme(axis.line = ggplot2::element_blank(),
                                             axis.ticks = ggplot2::element_blank(),
                                             axis.text.x = ggplot2::element_blank(),
                                             axis.text.y = ggplot2::element_blank(),
                                             axis.title.x = ggplot2::element_blank(),
                                             axis.title.y = ggplot2::element_blank())

    if (remove_legend) p <- p + ggplot2::theme(legend.position = "none")

    if (remove_grid_line) p <- p + ggplot2::theme(panel.grid.major = ggplot2::element_blank(),
                                                  panel.grid.minor = ggplot2::element_blank())

    return(p)
  }


# plot_lineup -------------------------------------------------------------

  plot_lineup_ <- function(dat, type = "resid", ...) {
    single_plot <- self$plot(dat, type, ...)
    single_plot +
      ggplot2::facet_wrap(~k)
  }

# str ---------------------------------------------------------------------


  str_ <- function() {

    if (!self$..instantiated..) {
      return(paste0("<", self$..type.., " class>"))
    }

    results <- bandicoot::use_method(self$prm$y, visage::CLOSED_FORM$..str..)()
    results <- paste0("<", self$..type.., " object>\n ",
                      gsub("<CLOSED_FORM object>\n EXPR", "y", results, fixed = TRUE))

    if (length(names(self$prm_type)[self$prm_type == "o"]) > 0) results <- paste0(results, "\n Parameters:")

    for (i in names(self$prm_type)[self$prm_type == "o"]) results <- paste0(results, "\n  - ", i, ": ", self$prm[[i]])

    results
  }


# register_method ---------------------------------------------------------

  bandicoot::register_method(env,
                             ..init.. = init_,
                             ..str.. = str_,
                             set_formula = set_formula_,
                             set_prm = set_prm_,
                             gen = gen_,
                             test = test_,
                             fit = fit_,
                             average_effect_size = average_effect_size_,
                             sample_effect_size = sample_effect_size_,
                             plot_resid = plot_resid_,
                             plot_qq = plot_qq_,
                             plot = plot_,
                             plot_lineup = plot_lineup_,
                             rss = rss_,
                             null_resid = null_resid_,
                             gen_lineup = gen_lineup_)

  return(env)
}


# CUBIC_MODEL -------------------------------------------------------------

class_CUBIC_MODEL <- function(env = new.env(parent = parent.frame())) {

  # Pass CMD check
  self <- NULL

  bandicoot::new_class(VI_MODEL, env = env, class_name = "CUBIC_MODEL")

  # Run the `set_formula` method for the class
  env$set_formula(formula = y ~ 1 + (2-c) * x + c * z + a * (((2-c) * x)^2 + (c * z)^2) + b * (((2-c) * x)^3 + (c * z)^3) + e,
                  null_formula = y ~ x + z,
                  alt_formula = y ~ x + I(x^2) + I(x^3) + z + I(z^2) + I(z^3))

# init --------------------------------------------------------------------

  init_ <- function(a = 1, b = 1, c = 1, sigma = 1,
                    x = visage::rand_uniform(-1, 1, env = new.env(parent = parent.env(self))),
                    z = visage::rand_uniform(-1, 1, env = new.env(parent = parent.env(self))),
                    e = visage::rand_normal(0, sigma, env = new.env(parent = parent.env(self)))) {

    # Use the init method from the VI_MODEL class
    bandicoot::use_method(self, visage::VI_MODEL$..init..)(
      prm = list(a = a, b = b, c = c, sigma = sigma, x = x, z = z, e = e),
      prm_type = list(a = "o", b = "o", c = "o", sigma = "o", x = "r", z = "r", e = "r"),
      formula = self$formula,
      null_formula = self$null_formula,
      alt_formula = self$alt_formula
      )

    return(invisible(self))
  }


# set_prm -----------------------------------------------------------------

  set_prm_ <- function(prm_name, prm_value) {

    for (i in 1:length(prm_name)) {
      pname <- prm_name[[i]]
      pval <- prm_value[[i]]

      if (pname == "sigma") {
        self$prm[[pname]] <- pval
        self$prm$e$set_prm(list("sigma"), list(pval))
      } else {
        bandicoot::use_method(self, visage::VI_MODEL$set_prm)(list(pname), list(pval))
      }
    }

    return(invisible(self))

  }

# E -----------------------------------------------------------------------

  E_ <- function(dat,
                 a = self$prm$a,
                 b = self$prm$b,
                 c = self$prm$c) {

    Xa <- as.matrix(data.frame(1, dat$x, dat$z))
    Ra <- diag(nrow(dat)) - Xa %*% solve(t(Xa) %*% Xa) %*% t(Xa)
    Xb <- as.matrix(data.frame(dat$x^2, dat$z^2, dat$x^3, dat$z^3))
    beta_b <- matrix(c(a*(2-c)^2, a*c^2, b*(2-c)^3, b*c^3))

    c(Ra %*% Xb %*% beta_b)
  }


# sample_effect_size ------------------------------------------------------

  sample_effect_size_ <- function(dat,
                                  a = self$prm$a,
                                  b = self$prm$b,
                                  c = self$prm$c,
                                  sigma = self$prm$sigma) {

    Xa <- as.matrix(data.frame(1, dat$x, dat$z))
    Ra <- diag(nrow(dat)) - Xa %*% solve(t(Xa) %*% Xa) %*% t(Xa)
    Xb <- as.matrix(data.frame(dat$x^2, dat$z^2, dat$x^3, dat$z^3))
    beta_b <- matrix(c(a*(2-c)^2, a*c^2, b*(2-c)^3, b*c^3))

    (1/nrow(dat)) * (1/sigma^2) * sum((diag(sqrt(diag(Ra))) %*% Xb %*% beta_b)^2)
  }

# register_method ---------------------------------------------------------

  bandicoot::register_method(env,
                             ..init.. = init_,
                             E = E_,
                             sample_effect_size = sample_effect_size_,
                             set_prm = set_prm_)

  return(env)
}

# simple_cubic_model ------------------------------------------------------

class_SIMPLE_CUBIC_MODEL <- function(env = new.env(parent = parent.frame())) {

  # Pass CMD check
  self <- NULL

  bandicoot::new_class(VI_MODEL, env = env, class_name = "SIMPLE_CUBIC_MODEL")

  # Run the `set_formula` method for the class
  env$set_formula(formula = y ~ 1 + x + a * x^2 + b * x^3 + e,
                  null_formula = y ~ x,
                  alt_formula = y ~ x + I(x^2) + I(x^3))

# init --------------------------------------------------------------------

  init_ <- function(a = 1, b = 1, sigma = 1,
                    x = visage::rand_uniform(-1, 1, env = new.env(parent = parent.env(self))),
                    e = visage::rand_normal(0, sigma, env = new.env(parent = parent.env(self)))) {

    # Use the init method from the VI_MODEL class
    bandicoot::use_method(self, visage::VI_MODEL$..init..)(
      prm = list(a = a, b = b, sigma = sigma, x = x, e = e),
      prm_type = list(a = "o", b = "o", sigma = "o", x = "r", e = "r"),
      formula = self$formula,
      null_formula = self$null_formula,
      alt_formula = self$alt_formula
    )

    return(invisible(self))
  }


# set_prm -----------------------------------------------------------------

  set_prm_ <- function(prm_name, prm_value) {

    # Reuse the CUBIC_MODEL$set_prm method
    bandicoot::use_method(self, visage::CUBIC_MODEL$set_prm)(prm_name, prm_value)

    return(invisible(self))
  }

# E -----------------------------------------------------------------------

  E_ <- function(dat, a = self$prm$a, b = self$prm$b) {

    Xa <- as.matrix(data.frame(1, dat$x))
    Ra <- diag(nrow(dat)) - Xa %*% solve(t(Xa) %*% Xa) %*% t(Xa)
    Xb <- as.matrix(data.frame(dat$x^2, dat$x^3))
    beta_b <- matrix(c(a, b))

    c(Ra %*% Xb %*% beta_b)
  }


# effect_size -------------------------------------------------------------

  sample_effect_size_ <- function(dat,
                                  a = self$prm$a,
                                  b = self$prm$b,
                                  sigma = self$prm$sigma) {

    Xa <- as.matrix(data.frame(1, dat$x))
    Ra <- diag(nrow(dat)) - Xa %*% solve(t(Xa) %*% Xa) %*% t(Xa)
    Xb <- as.matrix(data.frame(dat$x^2, dat$x^3))
    beta_b <- matrix(c(a, b))

    (1/nrow(dat)) * (1/sigma^2) * sum((diag(sqrt(diag(Ra))) %*% Xb %*% beta_b)^2)
  }

# register_method ---------------------------------------------------------

  bandicoot::register_method(env,
                             ..init.. = init_,
                             E = E_,
                             sample_effect_size = sample_effect_size_,
                             set_prm = set_prm_)

  return(env)
}

# HETER_MODEL -------------------------------------------------------------


class_HETER_MODEL <- function(env = new.env(parent = parent.frame())) {

  # Pass CMD check
  self <- NULL

  bandicoot::new_class(VI_MODEL, env = env, class_name = "HETER_MODEL")

  # Run the `set_formula` method for the class
  env$set_formula(formula = y ~ 1 + x + sqrt(1 + (2 - abs(a)) * (x - a)^2 * b) * e,
                  null_formula = y ~ x,
                  alt_formula = NULL)

# init --------------------------------------------------------------------

  init_ <- function(a = 0, b = 1,
                    x = visage::rand_uniform(-1, 1, env = new.env(parent = parent.env(self))),
                    e = visage::rand_normal(0, 1, env = new.env(parent = parent.env(self)))) {

    # Use the init method from the VI_MODEL class
    bandicoot::use_method(self, visage::VI_MODEL$..init..)(
      prm = list(a = a, b = b, x = x, e = e),
      prm_type = list(a = "o", b = "o", x = "r", e = "r"),
      formula = self$formula,
      null_formula = self$null_formula,
      alt_formula = self$alt_formula
    )

    return(invisible(self))
  }


# test --------------------------------------------------------------------

  test_ <- function(dat) {

    # Get the null model
    mod <- self$fit(dat)

    # construct proxy variables
    tmp_data <- data.frame(x = dat$x)
    tmp_data$xs <- tmp_data$x^2

    # Run the BP test
    bp_test <- lmtest::bptest(mod, varformula = ~ x + xs, data = tmp_data)

    return(list(name = "BP-test",
                statistic = unname(bp_test$statistic),
                p_value = unname(bp_test$p.value)))
  }


# effect_size -------------------------------------------------------------

  sample_effect_size_ <- function(dat, a = self$prm$a, b = self$prm$b, type = NULL) {

    if (!is.null(type) && type == "simple") {
      # return(sqrt(nrow(dat)) * b)
      return(sqrt(nrow(dat)) * b / (abs(a) + 1))
    }

    if (!is.null(type) && type == "kl") {

      if (b == 0) return(0)

      n <- nrow(dat)
      Xa <- as.matrix(data.frame(1, dat$x))
      Ra <- diag(nrow(dat)) - Xa %*% solve(t(Xa) %*% Xa) %*% t(Xa)
      V <- diag(1 + b * (2 - abs(a)) * (dat$x - a)^2)

      Ra_V_Ra <- Ra %*% V %*% t(Ra)
      diag_Ra_V_Ra <- diag(Ra_V_Ra)

      # sigma2_assumed <- mean(diag(V))

      # Correct one:
      # diag_Ra_sigma2 <- diag(Ra) * sigma2_assumed
      diag_Ra_sigma2 <- diag(Ra)

      log_det_s2_div_det_s1 <- sum(log(diag_Ra_V_Ra)) - sum(log(diag_Ra_sigma2))
      tr_inv_s2_s1 <- sum(1/diag_Ra_V_Ra * diag_Ra_sigma2)

      return((log_det_s2_div_det_s1 - n + tr_inv_s2_s1)/2)
    }

    return(sqrt(nrow(dat)) * b)
  }


# register_method ---------------------------------------------------------

  bandicoot::register_method(env,
                             ..init.. = init_,
                             test = test_,
                             sample_effect_size = sample_effect_size_)

  return(env)
}


# QUARTIC_MODEL -----------------------------------------------------------

class_QUARTIC_MODEL <- function(env = new.env(parent = parent.frame())) {

  # Pass CMD check
  self <- NULL

  bandicoot::new_class(VI_MODEL, env = env, class_name = "QUARTIC_MODEL")

  # Run the `set_formula` method for the class
  env$set_formula(formula = y ~ 1 + x + a * x^2 + b * x^3 + c * x^4 + e,
                  null_formula = y ~ x,
                  alt_formula = y ~ x + I(x^2) + I(x^3) + I(x^4))

# init --------------------------------------------------------------------

  init_ <- function(a = 1, b = 1, c = 1, sigma = 1,
                    x = visage::rand_uniform(-1, 1, env = new.env(parent = parent.env(self))),
                    e = visage::rand_normal(0, sigma, env = new.env(parent = parent.env(self)))) {

    # Use the init method from the VI_MODEL class
    bandicoot::use_method(self, visage::VI_MODEL$..init..)(
      prm = list(a = a, b = b, c = c, sigma = sigma, x = x, e = e),
      prm_type = list(a = "o", b = "o", c = "o", sigma = "o", x = "r", e = "r"),
      formula = self$formula,
      null_formula = self$null_formula,
      alt_formula = self$alt_formula
    )

    return(invisible(self))
  }


# set_prm -----------------------------------------------------------------

  set_prm_ <- function(prm_name, prm_value) {

    # Reuse the CUBIC_MODEL$set_prm method
    bandicoot::use_method(self, visage::CUBIC_MODEL$set_prm)(prm_name, prm_value)

    return(invisible(self))
  }

# E -----------------------------------------------------------------------

  E_ <- function(dat,
                 a = self$prm$a,
                 b = self$prm$b,
                 c = self$prm$c) {

    Xa <- as.matrix(data.frame(1, dat$x))
    Ra <- diag(nrow(dat)) - Xa %*% solve(t(Xa) %*% Xa) %*% t(Xa)
    Xb <- as.matrix(data.frame(dat$x^2, dat$x^3, dat$x^4))
    beta_b <- matrix(c(a, b, c))

    c(Ra %*% Xb %*% beta_b)
  }


# sample_effect_size ------------------------------------------------------

  sample_effect_size_ <- function(dat,
                                  a = self$prm$a,
                                  b = self$prm$b,
                                  c = self$prm$c,
                                  sigma = self$prm$sigma) {

    Xa <- as.matrix(data.frame(1, dat$x))
    Ra <- diag(nrow(dat)) - Xa %*% solve(t(Xa) %*% Xa) %*% t(Xa)
    Xb <- as.matrix(data.frame(dat$x^2, dat$x^3, dat$x^4))
    beta_b <- matrix(c(a, b, c))

    (1/nrow(dat)) * (1/sigma^2) * sum((diag(sqrt(diag(Ra))) %*% Xb %*% beta_b)^2)
  }

# register_method ---------------------------------------------------------

  bandicoot::register_method(env,
                             ..init.. = init_,
                             E = E_,
                             sample_effect_size = sample_effect_size_,
                             set_prm = set_prm_)

  return(env)
}


# POLY_MODEL --------------------------------------------------------------

class_POLY_MODEL <- function(env = new.env(parent = parent.frame())) {

  # Pass CMD check
  self <- NULL

  bandicoot::new_class(VI_MODEL, env = env, class_name = "POLY_MODEL")

  # Run the `set_formula` method for the class
  env$set_formula(formula = y ~ 1 + x + include_z * z + e,
                  z_formula = z ~ (raw_z - min(raw_z))/max(raw_z - min(raw_z)) * 2 - 1,
                  raw_z_formula = raw_z ~ hermite(shape)((x - min(x))/max(x - min(x)) * 4 - 2),
                  null_formula = y ~ x,
                  alt_formula = y ~ x + z)

# init --------------------------------------------------------------------

  init_ <- function(shape = 2, sigma = 1, include_z = TRUE,
                    x = visage::rand_uniform(-1, 1, env = new.env(parent = parent.env(self))),
                    e = visage::rand_normal(0, sigma, env = new.env(parent = parent.env(self)))) {

    if (!shape %in% 1:4) stop("Parameter `shape` out of range [1, 4].")

    hermite <- self$hermite
    raw_z <- visage::closed_form(eval(self$raw_z_formula), env = new.env(parent = parent.env(self)))
    z <- visage::closed_form(eval(self$z_formula), env = new.env(parent = parent.env(self)))

    # Use the init method from the VI_MODEL class
    bandicoot::use_method(self, visage::VI_MODEL$..init..)(
      prm = list(shape = shape, include_z = include_z, sigma = sigma, x = x, raw_z = raw_z, z = z, e = e),
      prm_type = list(shape = "o", include_z = "o", sigma = "o", x = "r", raw_z = "r", z = "r", e = "r"),
      formula = self$formula,
      null_formula = self$null_formula,
      alt_formula = self$alt_formula
    )

    return(invisible(self))
  }

# test --------------------------------------------------------------------

  test_ <- function(dat,
                    null_formula = self$null_formula,
                    alt_formula = self$alt_formula,
                    test = "F",
                    power = 2:3,
                    power_type = "fitted") {

    if (test == "F") {
      # Use the test method (F-test) from VI_MODEL class
      return(bandicoot::use_method(self, visage::VI_MODEL$test)(dat, null_formula, alt_formula))
    }

    if (test == "RESET") {
      null_mod <- self$fit(dat, null_formula)
      RESET_test <- lmtest::resettest(null_mod, power = power, type = power_type, data = dat)
      return(list(name = "RESET-test",
                  statistic = unname(RESET_test$statistic),
                  p_value = unname(RESET_test$p.value)))
    }

    stop("Unknown test type!")
  }


# set_prm -----------------------------------------------------------------

  set_prm_ <- function(prm_name, prm_value) {

    for (i in 1:length(prm_name)) {
      pname <- prm_name[[i]]
      pval <- prm_value[[i]]

      if (pname == "shape") {
        self$prm[[pname]] <- pval
        self$prm$raw_z$set_sym(list(pname), list(pval))
        next
      }

      if (pname == "raw_z") {
        self$prm[[pname]] <- pval
        self$prm$z$set_sym(list(pname), list(pval))
        next
      }

      # Reuse the CUBIC_MODEL$set_prm method
      bandicoot::use_method(self, visage::CUBIC_MODEL$set_prm)(list(pname), list(pval))
    }

    return(invisible(self))
  }

# E -----------------------------------------------------------------------

  E_ <- function(dat, include_z = self$prm$include_z) {

    Xa <- as.matrix(data.frame(1, dat$x))
    Ra <- diag(nrow(dat)) - Xa %*% solve(t(Xa) %*% Xa) %*% t(Xa)
    Xb <- as.matrix(data.frame(dat$z))
    beta_b <- matrix(as.numeric(include_z))

    c(Ra %*% Xb %*% beta_b)
  }


# hermite -----------------------------------------------------------------

  hermite_ <- function(shape) {
    if (!shape %in% 1:4) stop("Parameter `shape` out of range [1, 4].")
    suppressMessages(as.function(mpoly::hermite(c(2, 3, 6, 18)[shape])))
  }


# sample_effect_size ------------------------------------------------------

  sample_effect_size_ <- function(dat,
                                  sigma = self$prm$e$prm$sigma,
                                  include_z = self$prm$include_z,
                                  shape = self$prm$shape,
                                  type = "kl") {

    if (type == "simple") {
      n <- nrow(dat)
      return(sqrt(n) / (sigma * c(1, 2, 3, 5)[shape]))
      # return(sqrt(n) / sigma)
    }

    if (type == "kl") {

      if (include_z == 0) return(0)

      n <- nrow(dat)
      Xa <- as.matrix(data.frame(1, dat$x))
      Ra <- diag(nrow(dat)) - Xa %*% solve(t(Xa) %*% Xa) %*% t(Xa)
      Xb <- as.matrix(data.frame(dat$z))
      beta_b <- matrix(as.numeric(include_z))

      inv_sigma_2 <- diag(1/(sigma^2 * diag(Ra)))
      mu_2 <- Ra %*% Xb %*% beta_b

      # sum(log(diag(sigma_2)/diag(sigma_1))) - n + sum(diag(solve(sigma_2) %*% sigma_1))
      # log(det A/ det A) = log(1) = 0
      # solve(sigma_2) %*% sigma_1 = I_n
      # sum(diag(I_n)) = n
      # -n + n = 0
      # mu_1 = 0
      return(c((t(mu_2) %*% inv_sigma_2 %*% mu_2)/2))
    }

    Xa <- as.matrix(data.frame(1, dat$x))
    Ra <- diag(nrow(dat)) - Xa %*% solve(t(Xa) %*% Xa) %*% t(Xa)
    Xb <- as.matrix(data.frame(dat$z))
    beta_b <- matrix(as.numeric(include_z))

    return((1/nrow(dat)) * (1/sigma^2) * sum((diag(sqrt(diag(Ra))) %*% Xb %*% beta_b)^2))
  }

# register_method ---------------------------------------------------------

  bandicoot::register_method(env,
                             ..init.. = init_,
                             test = test_,
                             E = E_,
                             sample_effect_size = sample_effect_size_,
                             set_prm = set_prm_,
                             hermite = hermite_)

  return(env)
}


# AR1_MODEL ---------------------------------------------------------------

class_AR1_MODEL <- function(env = new.env(parent = parent.frame())) {

  # pass CMD check
  self <- NULL

  bandicoot::new_class(VI_MODEL, env = env, class_name = "AR1_MODEL")

  # Run the `set_formula` method for the class
  env$set_formula(formula = y ~ ar1(1 + x + e, phi),
                  null_formula = y ~ x,
                  alt_formula = y ~ x + lag(y))


# init_ -------------------------------------------------------------------

  init_ <- function(phi = 0.5, sigma = 1,
                    x = visage::rand_uniform(-1, 1, env = new.env(parent = parent.env(self))),
                    e = visage::rand_normal(0, sigma, env = new.env(parent = parent.env(self)))) {

    ar1 <- self$ar1

    self$prm <- list(phi = phi, sigma = sigma, x = x, e = e)
    self$prm_type <- list(phi = "o", sigma = "o", x = "r", e = "r")

    # Get the quoted formula, and eval in the current environment
    formula <- eval(self$formula)

    # Set the y variable
    self$prm$y <- visage::closed_form(formula, env = new.env(parent = parent.env(self)))

    return(invisible(self))
  }

# test --------------------------------------------------------------------

  test_ <- function(dat,
                    null_formula = self$null_formula,
                    type = "Ljung-Box",
                    lag = 1) {

      null_mod <- self$fit(dat, null_formula)
      BOX_test <- stats::Box.test(stats::resid(null_mod),
                                  lag = lag,
                                  type = type)
      return(list(name = "Box-test",
                  statistic = unname(BOX_test$statistic),
                  p_value = unname(BOX_test$p.value)))

  }

# set_prm -----------------------------------------------------------------

  set_prm_ <- function(prm_name, prm_value) {

    # Reuse the CUBIC_MODEL$set_prm method
    bandicoot::use_method(self, visage::CUBIC_MODEL$set_prm)(prm_name, prm_value)

    return(invisible(self))
  }

# ar1 ---------------------------------------------------------------------

  ar1_ <- function(lag_diff, phi) {
    Reduce(function(lag_y_t, y_t) phi * lag_y_t + y_t,
           lag_diff,
           accumulate = TRUE)
  }

# register_method ---------------------------------------------------------

  bandicoot::register_method(env,
                             ..init.. = init_,
                             test = test_,
                             set_prm = set_prm_,
                             ar1 = ar1_)

}


# NON_NORMAL_MODEL --------------------------------------------------------

class_NON_NORMAL_MODEL <- function(env = new.env(parent = parent.frame())) {

  # pass CMD check
  self <- NULL

  bandicoot::new_class(VI_MODEL, env = env, class_name = "NON_NORMAL_MODEL")

  # Run the `set_formula` method for the class
  env$set_formula(formula = y ~ 1 + x + e,
                  null_formula = y ~ x,
                  alt_formula = y ~ x)


# init_ -------------------------------------------------------------------

  init_ <- function(x = visage::rand_uniform(-1, 1, env = new.env(parent = parent.env(self))),
                    e = visage::rand_lognormal(0, 1, env = new.env(parent = parent.env(self)))) {

    # Use the init method from the VI_MODEL class
    bandicoot::use_method(self, visage::VI_MODEL$..init..)(
      prm = list(x = x, e = e),
      prm_type = list(x = "r", e = "r"),
      formula = self$formula,
      null_formula = self$null_formula,
      alt_formula = self$alt_formula
    )

    return(invisible(self))
  }

# test --------------------------------------------------------------------

  test_ <- function(dat,
                    null_formula = self$null_formula) {

    null_mod <- self$fit(dat, null_formula)
    SHAPIRO_test <- stats::shapiro.test(stats::resid(null_mod))
    return(list(name = "Shapiro-test",
                statistic = unname(SHAPIRO_test$statistic),
                p_value = unname(SHAPIRO_test$p.value)))

  }

# register_method ---------------------------------------------------------

  bandicoot::register_method(env,
                             ..init.. = init_,
                             test = test_)

}


# PHN_MODEL ---------------------------------------------------------------

class_PHN_MODEL <- function(env = new.env(parent = parent.frame())) {

  # Pass CMD check
  self <- NULL

  bandicoot::new_class(VI_MODEL, env = env, class_name = "PHN_MODEL")

  # Run the `set_formula` method for the class
  env$set_formula(formula = y ~ 1 + x1 + include_x2 * x2 + include_z * (z + include_x2 * w) + k * e,
                  z_formula = z ~ (raw_z - min(raw_z))/max(raw_z - min(raw_z)) * 2 - 1,
                  raw_z_formula = raw_z ~ hermite(j)((x1 - min(x1))/max(x1 - min(x1)) * 4 - 2),
                  w_formula = w ~ (raw_w - min(raw_w))/max(raw_w - min(raw_w)) * 2 - 1,
                  raw_w_formula = raw_w ~ hermite(j)((x2 - min(x2))/max(x2 - min(x2)) * 4 - 2),
                  k_formula = sigma ~ sqrt(1 + (2 - abs(a)) * ((x1 + include_x2 * x2) - a)^2 * b),
                  null_formula = y ~ x1 + x2,
                  alt_formula = y ~ x1 + x2 + z)


# ..init.. ----------------------------------------------------------------


  init_ <- function(j = 2,
                    include_x2 = FALSE,
                    include_z = TRUE,
                    sigma = 1,
                    a = 0,
                    b = 0,
                    x1 = visage::rand_uniform(-1, 1, env = new.env(parent = parent.env(self))),
                    x2 = visage::rand_uniform(-1, 1, env = new.env(parent = parent.env(self))),
                    e = visage::rand_normal(0, sigma, env = new.env(parent = parent.env(self)))) {

    hermite <- self$hermite
    raw_z <- visage::closed_form(eval(self$raw_z_formula), env = new.env(parent = parent.env(self)))
    z <- visage::closed_form(eval(self$z_formula), env = new.env(parent = parent.env(self)))
    raw_w <- visage::closed_form(eval(self$raw_w_formula), env = new.env(parent = parent.env(self)))
    w <- visage::closed_form(eval(self$w_formula), env = new.env(parent = parent.env(self)))
    k <- visage::closed_form(eval(self$k_formula), env = new.env(parent = parent.env(self)))

    if (!include_x2) {
      self$set_formula(null_formula = y ~ x1,
                       alt_formula = y ~ x1 + z)
    }

    # Use the init method from the VI_MODEL class
    bandicoot::use_method(self, visage::VI_MODEL$..init..)(
      prm = list(j = j,
                 include_x2 = include_x2,
                 include_z = include_z,
                 sigma = sigma,
                 a = a,
                 b = b,
                 x1 = x1,
                 x2 = x2,
                 raw_z = raw_z,
                 z = z,
                 raw_w = raw_w,
                 w = w,
                 k = k,
                 e = e),
      prm_type = list(j = "o",
                      include_x2 = "o",
                      include_z = "o",
                      sigma = "o",
                      a = "o",
                      b = "o",
                      x1 = "r",
                      x2 = "r",
                      raw_z = "r",
                      z = "r",
                      raw_w = "r",
                      w = "r",
                      k = "r",
                      e = "r"),
      formula = self$formula,
      null_formula = self$null_formula,
      alt_formula = self$alt_formula
    )

    return(invisible(self))
  }


# hermite -----------------------------------------------------------------

  hermite_ <- function(j) {
    suppressMessages(as.function(mpoly::hermite(j)))
  }


# sample_effect_size ------------------------------------------------------

  sample_effect_size_ <- function(dat, times = 1000L) {

    include_z <- self$prm$include_z
    include_x2 <- self$prm$include_x2
    sigma <- self$prm$sigma
    a <- self$prm$a
    b <- self$prm$b
    e_dist <- self$prm$e$..class..[1]

    if (include_z == 0 && b == 0 && e_dist == "RAND_NORMAL") return(0)

    n <- nrow(dat)

    if (include_x2) {
      Xa <- as.matrix(data.frame(1, dat$x1, dat$x2))
      Xb <- as.matrix(data.frame(dat$z, dat$w))
      beta_b <- matrix(c(as.numeric(include_z),
                         as.numeric(include_z) * as.numeric(include_x2)))
    } else {
      Xa <- as.matrix(data.frame(1, dat$x1))
      Xb <- as.matrix(data.frame(dat$z))
      beta_b <- matrix(as.numeric(include_z))
    }

    Ra <- diag(nrow(dat)) - Xa %*% solve(t(Xa) %*% Xa) %*% t(Xa)

    if (e_dist == "RAND_NORMAL") {

      e <- self$prm$e

      V <- diag(1 + (2 - abs(a)) * ((dat$x1 + include_x2 * dat$x2) - a)^2 * b) * e$prm$sigma^2

      Ra_V_Ra <- Ra %*% V %*% t(Ra)
      diag_Ra_V_Ra <- diag(Ra_V_Ra)

      mu_z <- Ra %*% Xb %*% beta_b
      mean_diff <- t(mu_z) %*% diag(1/diag_Ra_V_Ra) %*% mu_z

      if (b == 0) {
        return(c(mean_diff/2))
      } else {

        # The variance of a uniform mixture distribution with mean zero is the mean
        # of each individual variance
        diag_Ra_sigma <- diag(Ra) * mean(diag(V))

        log_det_s2_div_det_s1 <- sum(log(diag_Ra_V_Ra)) - sum(log(diag_Ra_sigma))

        tr_inv_s2_s1 <- sum(1/diag_Ra_V_Ra * diag_Ra_sigma)

        return(c((log_det_s2_div_det_s1 - n + tr_inv_s2_s1 + mean_diff)/2))
      }
    } else {

      k <- dat$k
      e <- self$prm$e
      Xb_beta_b <- c(Xb %*% beta_b)

      sample_e <- e$gen(n * times)

      # Mean correction in case the provided error distribution has mean other
      # than zero
      sample_e <- sample_e - mean(sample_e)

      # This is our assumed standard deviation for the error
      # distribution (after scaling by k)
      sd_vector <- sqrt(diag(Ra)) * sd(sample_e) * mean(k)

      before_ra <- matrix(sample_e, ncol = times) * k + Xb_beta_b

      final_result <- 0

      for (i in 1:n) {
        this_row <- Ra[i, ]
        after_ra <- c(this_row %*% before_ra)

        this_pdf <- function(x) {
          result <- approxfun(density(after_ra))(x)
          ifelse(is.na(result), 0, result)
        }

        current_sd <- sd_vector[i]

        kl_integrand <- function(x) {
          p_x <- this_pdf(x)
          # q_x <- assumed_pdf(x, sd = current_sd)
          log_qx <- -0.5 * log(2 * pi * current_sd^2) - x^2 / (2 * current_sd^2)

          result <- log(p_x) - log_qx

          result <- ifelse(p_x == 0, 0, result)
          # result <- ifelse(q_x == 0 & p_x != 0, Inf, result)

          return(result)
        }

        final_result <- final_result + mean(kl_integrand(after_ra))
      }

      return(final_result)
    }
  }

  bandicoot::register_method(env,
                             ..init.. = init_,
                             hermite = hermite_,
                             sample_effect_size = sample_effect_size_)

}


# nocov end