suppressPackageStartupMessages({
  library(hdrm)
  library(grpreg)
})

# Geometry of solution
par(mfrow = c(1, 2))

# Generate data
set.seed(2)
b <- seq(-0.5, 2.1, len = 101)
B <- as.matrix(expand.grid(b, b))
n <- 100
x1 <- rnorm(n)
x2 <- rnorm(n, mean = x1 / 2, sd = .25)
X <- cbind(x1, x2)
y <- rnorm(n, mean = X %*% c(1, 1))

# Non-orthogonal
L <- matrix(apply((y - tcrossprod(X, B))^2, 2, sum), ncol = length(b))
P <- matrix(sqrt(apply(B^2, 1, sum)), ncol = length(b))
Q <- L / n + 0 * P
l <- min(Q)
u <- max(Q)
levels <- c(.999*l + .001*u, .993*l + .007*u, .98*l + .02*u)
contour(
  b, b, Q, levels = levels, drawlabels = FALSE, lwd = 2, bty = 'n', las = 1,
  col = 'gray60', xlab = expression(beta[1]), ylab = expression(beta[2])
)
lam <- exp(seq(log(.01), log(20), len = 10))
Z <- matrix(NA, nrow = length(lam), ncol = 2)
rownames(Z) <- lam
for (i in 1:length(lam)) {
  Q <- L/n + lam[i]*P
  ind <- arrayInd(which.min(Q),dim(Q))
  Z[i,] <- c(b[ind[1]], b[ind[2]])
}
z <- coef(lm(y~0+X))
Z <- rbind(z, Z, c(0,0))
fit <- xspline(Z[, 1], Z[, 2], c(0, rep(1, length(lam)), 0), draw = FALSE)
lines(fit$x, fit$y, lwd = 3, col = 'blue')

# Orthogonal
XX <- grpreg:::orthogonalize(std(X), c(1,1))
L <- matrix(apply((y - tcrossprod(XX, B))^2, 2, sum), ncol = length(b))
P <- matrix(sqrt(apply(B^2, 1, sum)), ncol = length(b))
Q <- L / n + 0 * P
l <- min(Q)
u <- max(Q)
levels <- c(.999*l + .001*u, .993*l + .007*u, .985*l + .015*u)
contour(
  b, b, Q, levels = levels, drawlabels = FALSE, lwd = 2, las = 1, bty = 'n',
  col = 'gray60', xlab = expression(beta[1]), ylab = expression(beta[2])
)
z <- coef(lm(y ~ 0 + XX))
lines(c(0, z[1]), c(0, z[2]), lwd = 3, col = 'blue')

attach_data('glcamd')
gene <- group

# Expand X
xx <- matrix(0, nrow(x), 3 * ncol(x))
for (j in 1:ncol(x)) {
  xx[, (j - 1) * 3 + 1] <- x[, j] == 0
  xx[, (j - 1) * 3 + 2] <- x[, j] == 1
  xx[, (j - 1) * 3 + 3] <- x[, j] == 2
}
locus <- rep(1:ncol(x), each = 3)

# cv.grpreg(x, y, group = gene, family = 'binomial')

# cv.grpreg(xx, y, group = locus, family = 'binomial')

# Fit
fold <- assign_fold(y, 10, 2)
cvfit <- list(
  cv.grpreg(x, y, family = 'binomial', fold = fold, lambda.min = 0.2),
  cv.grpreg(x, y, group = gene, family = 'binomial', fold = fold, lambda.min = 0.2),
  cv.grpreg(xx, y, family = 'binomial', fold = fold, lambda.min = 0.2),
  cv.grpreg(xx, y, group = locus, family = 'binomial', fold = fold, lambda.min = 0.2),
  cv.grpreg(x, y, penalty = 'grMCP', family = 'binomial', fold = fold, lambda.min = 0.2),
  cv.grpreg(x, y, penalty = 'grMCP', group = gene, family = 'binomial', fold = fold, lambda.min = 0.2)
)

ylim <- c(0, 0.10)
par(mfrow = c(1, 2), mar = c(5, 5, 5, 1))
plot(cvfit[[1]], type = 'rsq', ylim = ylim)
mtext('Ungrouped (ordinary lasso)', line = 2.5)
plot(cvfit[[2]], type = 'rsq', ylim = ylim)
mtext('Group lasso', line = 2.5)

par(mfrow = c(1, 2), mar = c(5, 5, 5, 1))
plot(cvfit[[1]], type = 'pred', ylim = c(0.35, 0.55))
mtext('Ungrouped (ordinary lasso)', line = 2.5)
plot(cvfit[[2]], type = 'pred', ylim = c(0.35, 0.55))
mtext('Group lasso', line = 2.5)

# Summary / number of genes selected
s <- lapply(cvfit, summary)
b1 <- coef(cvfit[[1]], lambda = cvfit[[1]]$lambda.min)
n1 <- length(table(gene[b1[-1] != 0]))
n2 <- predict(cvfit[[2]], lambda = cvfit[[2]]$lambda.min, type = "ngroups")

par(mfrow = c(1, 2), mar = c(5, 5, 5, 1))
ylim <- c(0, 0.10)
plot(cvfit[[3]], type = 'rsq', ylim = ylim)
mtext('Ungrouped (ordinary lasso)', line = 2.5)
plot(cvfit[[4]], type = 'rsq', ylim = ylim)
mtext('Group lasso', line = 2.5)

par(mfrow = c(1, 2), mar = c(5, 5, 5, 1))
plot(cvfit[[3]], type = 'pred', ylim = c(0.35, 0.55))
mtext('Ungrouped (ordinary lasso)', line = 2.5)
plot(cvfit[[4]], type = 'pred', ylim = c(0.35, 0.55))
mtext('Group lasso', line = 2.5)

# Summary / number of loci selected
b3 <- coef(cvfit[[3]])
n3 <- length(table(locus[b3[-1] != 0]))
n4 <- predict(cvfit[[4]], type = "ngroups")
v3 <- predict(cvfit[[3]], type = "nvars")
v4 <- predict(cvfit[[4]], type = "nvars")