library(distfreereg)
all.equal.distfreereg <- distfreereg:::all.equal.distfreereg
test_dfr_functions <- distfreereg:::test_dfr_functions

n <- 1e2

### Poisson

set.seed(20250302)
p_pois <- 2
X_pois <- matrix(rexp(n*p_pois), ncol = p_pois)
theta_pois <- c(2,1)
means <- exp(colSums(theta_pois * t(X_pois)) + 1)
Y_pois <- rpois(n, means)
df_pois <- data.frame(y = Y_pois, x = X_pois[,1], z = X_pois[,2],
                      g = rep(1:10, 10))
form_pois <- y ~ x + z
m_pois <- glm(form_pois, data = df_pois, family = "poisson", x = TRUE, y = TRUE)

set.seed(20250302)
dfr_pois <- distfreereg(test_mean = m_pois, verbose = FALSE,
                        control = list(return_on_error = FALSE))

dfr_pois

set.seed(20250302)
dfr_pois_x_false <- distfreereg(test_mean = update(m_pois, x = FALSE),
                                control = list(return_on_error = FALSE))

set.seed(20250302)
dfr_pois_J <- distfreereg(test_mean = m_pois,
                          override = list(J = dfr_pois[["J"]]),
                          control = list(return_on_error = FALSE))

set.seed(20250302)
dfr_pois_fitted <- distfreereg(test_mean = m_pois,
                               override = list(fitted_values = dfr_pois[["fitted_values"]]),
                               control = list(return_on_error = FALSE))

stopifnot(all.equal(dfr_pois, dfr_pois_x_false))
stopifnot(all.equal(dfr_pois, dfr_pois_J))
stopifnot(all.equal(dfr_pois, dfr_pois_fitted))

set.seed(20250516)
dfr_form_glm_verbose <- distfreereg(test_mean = form_pois, data = df_pois,
                                    method = "glm",
                                    method_args = list(family = "poisson"),
                                    control = list(return_on_error = FALSE))

newdata_pois <- data.frame(y = rpois(10, lambda = 1), x = rexp(10), z = rexp(10))
test_dfr_functions(dfr_pois, newdata = newdata_pois)

set.seed(20250302)
dfr_form_pois <- distfreereg(test_mean = form_pois, data = df_pois,
                             method = "glm",
                             method_args = list(family = "poisson"),
                             verbose = FALSE,
                             control = list(return_on_error = FALSE))

dfr_form_pois
test_dfr_functions(dfr_form_pois, newdata = newdata_pois)

stopifnot(all.equal(dfr_pois, dfr_form_pois))

set.seed(20250225)
cdfr_form_pois <- asymptotics(dfr_form_pois, reps = 5)
set.seed(20250225)
cdfr_pois <- asymptotics(dfr_pois, reps = 5)

signif(rejection(cdfr_form_pois, alpha = c(0.1, 0.5))[,2:3], digits = 3)
signif(rejection(cdfr_pois, alpha = c(0.1, 0.5))[,2:3], digits = 3)


# Orderings

set.seed(20250516)
dfr_pois_asis <- update(dfr_pois, ordering = "asis")
set.seed(20250516)
dfr_form_pois_asis <- update(dfr_form_pois, ordering = "asis")
stopifnot(all.equal(dfr_pois_asis, dfr_form_pois_asis))

set.seed(20250516)
dfr_pois_optimal <- update(dfr_pois, ordering = "optimal")
set.seed(20250516)
dfr_form_pois_optimal <- update(dfr_form_pois, ordering = "optimal")
stopifnot(all.equal(dfr_pois_optimal, dfr_form_pois_optimal))

set.seed(20250516)
dfr_pois_natural <- update(dfr_pois, ordering = "natural")
set.seed(20250516)
dfr_form_pois_natural <- update(dfr_form_pois, ordering = "natural")
stopifnot(all.equal(dfr_pois_natural, dfr_form_pois_natural))

set.seed(20250516)
dfr_pois_g <- update(dfr_pois, ordering = list("g"))
set.seed(20250516)
dfr_form_pois_g <- update(dfr_form_pois, ordering = list("g"))
stopifnot(all.equal(dfr_pois_g, dfr_form_pois_g))

df_pois[dfr_pois_g[["res_order"]],][["g"]]
df_pois[dfr_form_pois_g[["res_order"]],][["g"]]

set.seed(20250516)
dfr_pois_g_grouped <- update(dfr_pois_g, group = TRUE)
set.seed(20250516)
dfr_form_pois_g_grouped <- update(dfr_form_pois_g, group = TRUE)
stopifnot(all.equal(dfr_pois_g_grouped, dfr_form_pois_g_grouped))



### Partial output

dfr_pois_partial <- distfreereg(test_mean = m_pois, verbose = FALSE,
                                control = list(orth_tol = 1e-100))
names(dfr_pois_partial)



### Binomial

theta_binom <- c(1,2,-1,-3)/30
p_binom <- length(theta_binom) - 1
set.seed(20250225)
X_binom <- cbind(1, matrix(rexp(n*p_binom, rate = 1/5), ncol = p_binom))
probs <- 1/(1 + exp(-colSums(theta_binom * t(X_binom)) - 1))
Y_binom <- rbinom(n, size = 1, prob = probs)

df_binom <- as.data.frame(cbind(Y_binom, X_binom[,-c(1,4)], rep(1:10, 10)))# omit intercept and one covariate
colnames(df_binom) <- c("y", letters[1:(length(theta_binom)-2)], "g")
df_binom$a <- df_binom$a^2
form_binom <- reformulate(termlabels = colnames(df_binom)[-1], response = "y")
m_binom <- glm(form_binom, data = df_binom, family = "binomial", x = TRUE, y = TRUE)

set.seed(20250302)
dfr_binom <- distfreereg(test_mean = m_binom, verbose = FALSE,
                         control = list(return_on_error = TRUE))
dfr_binom
newdata_binom <- data.frame(y = rbinom(10, size = 1, prob = runif(10)),
                            a = rexp(10, rate = 1/5), b = rexp(10, rate = 1/5))
test_dfr_functions(dfr_binom, newdata = newdata_binom)

set.seed(20250302)
dfr_form_binom <- distfreereg(test_mean = form_binom, data = df_binom,
                              method = "glm",
                              method_args = list(family = "binomial"),
                              verbose = FALSE,
                              control = list(return_on_error = FALSE))
dfr_form_binom
test_dfr_functions(dfr_form_binom, newdata = newdata_binom)

stopifnot(all.equal(dfr_binom, dfr_form_binom))

cdfr_form_binom <- asymptotics(dfr_form_binom, reps = 5)
cdfr_binom <- asymptotics(dfr_binom, reps = 5)

signif(rejection(cdfr_form_binom, alpha = c(0.1, 0.5))[,2:3], digits = 3)
signif(rejection(cdfr_binom, alpha = c(0.1, 0.5))[,2:3], digits = 3)

# Orderings

set.seed(20250516)
dfr_binom_asis <- update(dfr_binom, ordering = "asis")
set.seed(20250516)
dfr_form_binom_asis <- update(dfr_form_binom, ordering = "asis")
stopifnot(all.equal(dfr_binom_asis, dfr_form_binom_asis))

set.seed(20250516)
dfr_binom_optimal <- update(dfr_binom, ordering = "optimal")
set.seed(20250516)
dfr_form_binom_optimal <- update(dfr_form_binom, ordering = "optimal")
stopifnot(all.equal(dfr_binom_optimal, dfr_form_binom_optimal))

set.seed(20250516)
dfr_binom_natural <- update(dfr_binom, ordering = "natural")
set.seed(20250516)
dfr_form_binom_natural <- update(dfr_form_binom, ordering = "natural")
stopifnot(all.equal(dfr_binom_natural, dfr_form_binom_natural))

set.seed(20250516)
dfr_binom_g <- update(dfr_binom, ordering = list("g"))
set.seed(20250516)
dfr_form_binom_g <- update(dfr_form_binom, ordering = list("g"))
stopifnot(all.equal(dfr_binom_g, dfr_form_binom_g))

df_binom[dfr_binom_g[["res_order"]],][["g"]]
df_binom[dfr_form_binom_g[["res_order"]],][["g"]]

set.seed(20250516)
dfr_binom_g_grouped <- update(dfr_binom_g, group = TRUE)
set.seed(20250516)
dfr_form_binom_g_grouped <- update(dfr_form_binom_g, group = TRUE)
stopifnot(all.equal(dfr_binom_g_grouped, dfr_form_binom_g_grouped))




### Failures

tryCatch(distfreereg(test_mean = form_pois, data = df_pois,
                     method = "glm",
                     method_args = list(family = "poisson",
                                        weights = rep(1, n)),
                     verbose = FALSE,
                     control = list(return_on_error = FALSE)),
         error = function(e) warning(e))