if (interactive()) library(tinytest)
suppressPackageStartupMessages(library(glmnet))

# --- Linear regression ----------------------------------------------

# Works
X <- matrix(rnorm(500), 50, 10)
y <- X[, 1] + rnorm(50)
cvfit <- cv.ncvreg(X, y)
if (interactive()) print(summary(cvfit))

# Predict
b <- coef(cvfit)
p <- predict(cvfit, X, type = "link")
p <- predict(cvfit, X, type = "response")
p <- predict(cvfit, X, type = "coef")
p <- predict(cvfit, X, type = "vars")
p <- predict(cvfit, X, type = "nvars")

# Integers
X <- matrix(rpois(500, 1), 50, 10)
y <- rpois(50, 1)
cvfit <- cv.ncvreg(X, y)

# Data frame
cvfit <- cv.ncvreg(as.data.frame(X), y)

# LOOCV
X <- matrix(rnorm(25 * 4), 25, 4)
y <- rnorm(25)
cvfit <- cv.ncvreg(X, y, nfolds = 25)
summary(cvfit)
cvfit <- cv.ncvreg(X, y, fold = 1:25)

# --- Logistic regression --------------------------------------------

# Works
X <- matrix(rnorm(500), 50, 10)
y <- rbinom(50, 1, binomial()$linkinv(X[, 1]))
cvfit <- cv.ncvreg(X, y, family = "binomial")
if (interactive()) print(summary(cvfit))

# Predict
b <- coef(cvfit)
p <- predict(cvfit, X, type = "link")
p <- predict(cvfit, X, type = "response")
p <- predict(cvfit, X, type = "class")
p <- predict(cvfit, X, type = "coef")
p <- predict(cvfit, X, type = "vars")
p <- predict(cvfit, X, type = "nvars")

# LOOCV
X <- matrix(rnorm(30 * 2), 30, 2)
y <- rbinom(30, 1, 0.5)
cvfit <- cv.ncvreg(X, y, nfolds = 30, family = "binomial")
summary(cvfit)
cvfit <- cv.ncvreg(X, y, fold = 1:30, family = "binomial")

# --- Poisson regression ---------------------------------------------

# Works
cvfit <- cv.ncvreg(X, y, family = "poisson")
summary(cvfit)

# Predict
b <- coef(cvfit)
p <- predict(cvfit, X, type = "link")
p <- predict(cvfit, X, type = "response")
p <- predict(cvfit, X, type = "coef")
p <- predict(cvfit, X, type = "vars")
p <- predict(cvfit, X, type = "nvars")

# LOOCV
cvfit <- cv.ncvreg(X, y, nfolds = 30, family = "poisson")
summary(cvfit)
cvfit <- cv.ncvreg(X, y, fold = 1:30, family = "poisson")

# --- cv.ncvreg() seems to work --------------------------------------
n <- 40
p <- 10
X <- matrix(rnorm(n * p), ncol = p)
b <- c(2, -2, 1, -1, rep(0, p - 4))
y <- rnorm(n, mean = X %*% b, sd = 2)
yb <- y > .5
yp <- rpois(n, exp(X %*% b / 3))
par(mfrow = c(3, 2))

gcvfit <- cv.glmnet(X, y, nfolds = n, grouped = FALSE)
plot(gcvfit)
ncvfit <- cv.ncvreg(X, y, penalty = "lasso", lambda = gcvfit$lambda, nfolds = n)
plot(ncvfit)
gcvfit <- cv.glmnet(X, yb, family = "binomial", nfolds = n, grouped = FALSE)
plot(gcvfit)
ncvfit <- cv.ncvreg(
  X,
  yb,
  family = "binomial",
  penalty = "lasso",
  lambda = gcvfit$lambda,
  nfolds = n,
  lambda.min = 0.2
)
plot(ncvfit)
cvfit <- cv.glmnet(X, yp, family = "poisson")
plot(cvfit)
cvfit <- cv.ncvreg(X, yp, family = "poisson", penalty = "lasso", lambda = cvfit$lambda)
plot(cvfit)

# --- cv.ncvreg() returning array of linear predictors works ---------
n <- 100
p <- 10
X <- matrix(rnorm(n * p), ncol = p)
b <- c(-3, 3, rep(0, 8))

y <- rnorm(n, mean = X %*% b, sd = 1)
cvfit <- cv.ncvreg(X, y, returnY = TRUE)
cve <- apply(cvfit$Y - y, 2, crossprod) / n
expect_equivalent(cve, cvfit$cve, tol = .001)

y <- rnorm(n, mean = X %*% b) > 0
cvfit <- cv.ncvreg(X, y, family = "binomial", returnY = TRUE)
pe <- apply((cvfit$Y >= 0.5) != y, 2, mean)
expect_equivalent(pe, cvfit$pe, tol = .001)

# --- cv.ncvreg() return Y works -------------------------------------

X <- matrix(rnorm(500), 50, 10)
y <- 100 + X[, 1] + rnorm(50)
cvfit <- cv.ncvreg(X, y, returnY = TRUE)
expect_equivalent(
  cvfit$cve,
  apply((cvfit$Y - y)^2, 2, mean)
)
