source('http://myweb.uiowa.edu/pbreheny/7600/s16/notes/fun.R')

# Slide 4 (degrees of freedom in forward selection)
set.seed(1)
n <- 25
p <- 100
xnam <- paste0("V", 1:p)
form <- as.formula(paste("y ~ ", paste(xnam, collapse= "+")))
N <- 100
Y <- Yhat <- matrix(NA, N, n)
pb <- txtProgressBar(1, N, style=3)
for (i in 1:N) {
  X <- matrix(runif(n*p), n, p)
  Data <- as.data.frame(X)
  y <- rnorm(n)
  fit0 <- lm(y~1, data=Data)
  fit <- step(fit0, scope=form, direction="forward", trace=0, k=log(n), steps=5)
  #fit <- lm(y~V1+V2+V3+V4+V5, Data)  # You can check that this one really does have 5 df
  Y[i,] <- y
  Yhat[i,] <- fit$fitted
  setTxtProgressBar(pb, i)
}
sum(sapply(1:n, function(i) cov(Y[,i], Yhat[,i])))

# Slide 7
load(url('http://myweb.uiowa.edu/pbreheny/data/pollution.RData'))
require(ncvreg)
fit <- ncvreg(X, y, penalty="lasso")
ll <- log(fit$lambda)
IC <- cbind(AIC(fit), BIC(fit))
col <- c("#FF4E37FF", "#008DFFFF")
matplot(ll, IC, type="l", lwd=3, lty=1, xlab=expression(lambda), xaxt="n", bty="n", xlim=rev(range(ll)), col=col, las=1, ylab="Information criterion")
at <- c(40, 4, 0.4, 0.04)
axis(1, at=log(at), labels=at)
toplegend(legend=c("AIC", "BIC"), lwd=3, col=col)
abline(v=ll[which.min(IC[,1])], lty=2, lwd=2, col=col[1])
abline(v=ll[which.min(IC[,2])], lty=2, lwd=2, col=col[2])

# Basic glmnet usage
require(glmnet)
fit <- glmnet(X, y)
plot(fit, label=TRUE)
cvfit <- cv.glmnet(X, y)
cvfit <- cv.glmnet(X, y, nfolds=length(y))  # For leave-one-out CV

# Slide 16
ll <- log(fit$lambda)
plot(cvfit, xlim=rev(range(ll)), xlab=expression(log(lambda)), las=1, bty="n", xaxt="n", ylab=expression(CV(lambda)))
at <- c(40, 4, 0.4, 0.04)
axis(1, at=log(at), labels=at)
mtext("Number of nonzero coefficients", 3, 3)

# Slide 17 (range of lambda's within 1 se of min)
ind <- which(cvfit$lambda==cvfit$lambda.min)
range(cvfit$lambda[which(cvfit$cvm < cvfit$cvup[ind])])

# Slide 25
cvfit <- cv.glmnet(X, y)
1-cvfit$cvm/var(y)
cvfit <- cv.ncvreg(X, y, penalty="lasso")
plot(cvfit, type="rsq", bty="n", xlab=expression(lambda), xaxt="n")
at <- c(40, 4, 0.4, 0.04)
axis(1, at=log(at), labels=at)
summary(cvfit)

##########################################################
#### Remaining code is for the comparison on slide 22 ####
##########################################################
set.seed(1)

###############
## Null case ##
###############

require(glmnet)
require(ncvreg)
n <- 100
p <- 1000
X <- matrix(rnorm(n*p), n, p)
X <- std(X)
b <- rep(0, p)
y <- rnorm(n, X%*%b)

## Fit
fit <- ncvreg(X, y, penalty="lasso", lambda.min=0.06)
fit <- glmnet(X, y, lambda=fit$lambda)
cvfit <- cv.glmnet(X, y, lambda=fit$lambda)
n <- nrow(X)
l <- fit$lambda
ll <- log(l)
ll1 <- ll
nv <- sapply(predict(fit, type="nonzero"), length)+1

## sigma
Sigma <- array(NA, dim=c(length(l), 3, 2))
Y <- predict(fit, X)
RSS <- apply(y-Y, 2, crossprod)
Sigma[,1,1] <- sqrt(RSS/(n-nv))
Sigma[,2,1] <- cvfit$cvm

## Fan estimator
N <- 25
V <- matrix(NA, N, length(l))
lmvar <- function(y, X, s) {
  if (!length(s)) return(var(y))
  ols <- lm(y~X[,s])
  summary(ols)$sigma^2
}
pb <- txtProgressBar(1, N, style=3)
for (i in 1:N) {
  ind <- sample(1:n, n/2)
  X1 <- X[ind,]
  X2 <- X[-ind,]
  y1 <- y[ind]
  y2 <- y[-ind]
  fit1 <- glmnet(X1, y1, lambda=fit$lambda)
  fit2 <- glmnet(X2, y2, lambda=fit$lambda)
  s1 <- predict(fit1, type="nonzero")
  s2 <- predict(fit2, type="nonzero")
  v1 <- v2 <- numeric(length(l))
  for (j in 1:length(l)) {
    v1[j] <- lmvar(y1, X1, s2[[j]])
    v2[j] <- lmvar(y2, X2, s1[[j]])
  }
  V[i,] <- 0.5*v1 + 0.5*v2
  setTxtProgressBar(pb, i)
}
Sigma[,3,1] <- sqrt(apply(V, 2, mean, na.rm=TRUE))

#################
## Signal case ##
#################

b <- numeric(p); b[1:5] <- 1
y <- rnorm(n, X%*%b)

## Fit
fit <- ncvreg(X, y, penalty="lasso", lambda.min=0.02)
fit <- glmnet(X, y, lambda=fit$lambda)
cvfit <- cv.glmnet(X, y, lambda=fit$lambda)
n <- nrow(X)
l <- fit$lambda
ll <- log(l)
ll2 <- ll
nv <- sapply(predict(fit, type="nonzero"), length)+1

## sigma
Y <- predict(fit, X)
RSS <- apply(y-Y, 2, crossprod)
Sigma[,1,2] <- sqrt(RSS/(n-nv))
Sigma[1:length(cvfit$cvm),2,2] <- cvfit$cvm

## Fan estimator
N <- 25
V <- matrix(NA, N, length(l))
pb <- txtProgressBar(1, N, style=3)
for (i in 1:N) {
  ind <- sample(1:n, n/2)
  X1 <- X[ind,]
  X2 <- X[-ind,]
  y1 <- y[ind]
  y2 <- y[-ind]
  fit1 <- glmnet(X1, y1, lambda=fit$lambda)
  fit2 <- glmnet(X2, y2, lambda=fit$lambda)
  s1 <- predict(fit1, type="nonzero")
  s2 <- predict(fit2, type="nonzero")
  v1 <- v2 <- numeric(length(l))
  for (j in 1:length(l)) {
    v1[j] <- lmvar(y1, X1, s2[[j]])
    v2[j] <- lmvar(y2, X2, s1[[j]])
  }
  V[i,] <- 0.5*v1 + 0.5*v2
  setTxtProgressBar(pb, i)
}
Sigma[,3,2] <- sqrt(apply(V, 2, mean, na.rm=TRUE))

## Plot
col <- c("black", "#FF4E37FF", "#008DFFFF")
matplot(ll1, Sigma[,,1], type="l", col=col, lty=1, lwd=3, xlim=rev(range(ll1)), ylab=expression(hat(sigma)), las=1, xaxt="n", xlab=expression(lambda), bty="n")
abline(h=1, lty=2, col="gray", lwd=2)
at <- seq(max(ll1), min(ll1), length=5)
axis(1, at=at, labels=round(exp(at), 2))
toplegend(legend=c("Plug-in", "CV", "RCV"), col=col, lwd=3)
matplot(ll2, Sigma[,,2], type="l", col=col, lty=1, lwd=3, xlim=rev(range(ll2)), ylab=expression(hat(sigma)), las=1, xaxt="n", xlab=expression(lambda), bty="n", ylim=c(0, 3))
abline(h=1, lty=2, col="gray", lwd=2)
at <- seq(max(ll2), min(ll2), length=5)
axis(1, at=at, labels=round(exp(at), 2))
toplegend(legend=c("Plug-in", "CV", "RCV"), col=col, lwd=3)