---
title: "Lecture 3: Prediction"
output: html_document
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
```

```{r}
library(tidyverse)
library(ISLR)
library(glmnet)
```

## Train-Validate-Test with Linear Regression

### Synthetic Data Generation (Training, Validation, Test Sets)

```{r}

# Set seed for reproducibility
set.seed(123)

total_obs = 300

# Generate covariates and outcomes
X1 = rnorm(total_obs,0,1)
X2 = rnorm(total_obs,0,1)
Y = 1 + 2*X1 + 3*X2 + rnorm(total_obs,0,5)

# Split into training, validation, testing
newdata.syn = data.frame(cbind(X1,X2,Y))
newdata.syn$type = factor(c(rep("train",total_obs/3),rep("val",total_obs/3),rep("test",total_obs/3)))
```

### Model Fitting

```{r}
fm1 = lm(formula = Y ~ 1 + X1, data = subset(newdata.syn, type == "train"))
fm2 = lm(formula = Y ~ 1 + X2, data = subset(newdata.syn, type == "train"))
fm3 = lm(formula = Y ~ 1 + X1 + X2 , data = subset(newdata.syn, type == "train"))
fm4 = lm(formula = Y ~ 1 + X1 + X2 + I(X1^2) + I(X2^2), data = subset(newdata.syn, type == "train"))
fm5 = lm(formula = Y ~ 1 + X1 + X2 + I(X1^2) + I(X2^2) + I(X1^3) + I(X2^3)
         + I(X1^4) + I(X2^4) + I(X1^5) + I(X2^5), data = subset(newdata.syn, type == "train"))
```

### Errors on Training/Validation/Test Sets

```{r}
# Define a function to calculate RMSE
calculate_rmse <- function(model, data) {
  sqrt(mean((data$Y - predict(model, data))^2))
}

# Create a list of models
models_list <- list(fm1 = fm1, fm2 = fm2, fm3 = fm3, fm4 = fm4, fm5 = fm5)

# Create a data frame to store the results
results_df <- tibble(
  Model = character(),
  Train_RMSE = numeric(),
  Val_RMSE = numeric(),
  Test_RMSE = numeric()
)

# Loop through each model and calculate RMSE for train, validation, and test sets
for (model_name in names(models_list)) {
  
  model <- models_list[[model_name]]
  
  train_rmse <- calculate_rmse(model, subset(newdata.syn, type == "train"))
  val_rmse <- calculate_rmse(model, subset(newdata.syn, type == "val"))
  test_rmse <- calculate_rmse(model, subset(newdata.syn, type == "test"))
  
  results_df <- results_df %>%
    add_row(
      Model = model_name,
      Train_RMSE = train_rmse,
      Val_RMSE = val_rmse,
      Test_RMSE = test_rmse
    )
  
}

# View the results
results_df
```

## Lasso and Ridge Regression

### Load Data:

```{r}
data(Hitters)
hitters.df = subset(na.omit(Hitters))
```

### Model Matrix and Standardization:

```{r}
X = model.matrix(Salary ~ 0 + ., hitters.df)
Y = hitters.df$Salary

X = scale(X)
Y = scale(Y)
```

### Train-Test Split

```{r}

set.seed(123) # Set seed for reproducibility 

train.ind = sample(nrow(X), round(nrow(X)/2))
X.train = X[train.ind,]
X.test = X[-train.ind,]
Y.train = Y[train.ind]
Y.test = Y[-train.ind]
```

## Ridge and Lasso Regression

### Set Lambda Sequence

```{r}
lambdas = 10^seq(-3,1.5,0.1)
```

### Ridge Regression

```{r}
fm.ridge = glmnet(X.train, Y.train, alpha = 0, lambda = lambdas, thresh = 1e-12)
```

#### Test Error of Ridge Regression at Each Lambda

```{r}
ridge.test <- tibble(lambda = lambdas) %>%
  mutate(
    TestErr = map_dbl(lambda, ~ mean((Y.test - predict(fm.ridge, s = .x, newx = X.test))^2)),
    Model = "Ridge"
  )
colnames(ridge.test) = c("lambda", "TestErr", "Model")
```

### Lasso Regression

```{r}
fm.lasso = glmnet(X.train, Y.train, alpha = 1, lambda = lambdas, thresh = 1e-12)
```

#### Test Error of Lasso at Each Lambda

```{r}
lasso.test <- tibble(lambda = lambdas) %>%
mutate(
  TestErr = map_dbl(lambda, ~ mean((Y.test - predict(fm.lasso, s = .x, newx = X.test))^2)),
  Model = "Lasso"
)
colnames(lasso.test) = c("lambda", "TestErr", "Model")
```

### Combine Test Error of Each Model into One Data Frame

```{r}
test.df = rbind(ridge.test, lasso.test)
test.df$Model = factor(test.df$Model)
```

### Plot Test Error

```{r}
ggplot(data = test.df, aes(x = lambda, y = TestErr, color = Model)) + 
  geom_line()  + scale_x_log10() + ylab("Test set error")
```

## Analysis of Minimum Test Error

```{r}
ridge.min.err = min(ridge.test$TestErr)
lasso.min.err = min(lasso.test$TestErr)

ridge.min.l = lambdas[which.min(ridge.test$TestErr)]
lasso.min.l = lambdas[which.min(lasso.test$TestErr)]
```

## Coefficients at the Test Error Minimizing Models

```{r}
cbind(predict(fm.lasso, s = lasso.min.l, type = "coefficients"), 
predict(fm.ridge, s = ridge.min.l, type = "coefficients"))
```

## Plot Ridge Coefficients

```{r}
plot(fm.ridge, xvar="lambda")
```

## Plot Lasso Coefficients

```{r}
plot(fm.lasso, xvar="lambda")
```