##########################
# Extra exercise 13      #
# STK2100, spring 2026   #
# Berit Omli Øksnes      #
# Last modified 16.04.26 #
##########################

# Clean workspace
rm(list=ls())

# Import packages
library(ISLR)      # Wage data set
library(nnet)      # nnet function
library(RSNNS)     # mlp function

# Set seed for reproducibility
set.seed(2)


# a) Read and prep data --------------------------------------------------- ####
data(Wage)
n <- nrow(Wage)

# Train-test split
train      <- sample(1:n, n/2, replace=FALSE)
Wage.train <- Wage[train,]
Wage.test  <- Wage[-train,]

# Scale numerical features using training data
scaler <- scale(Wage.train[,1:2])

Wage.train[,1:2] <- scale(Wage.train[,1:2],
                          center=attr(scaler, "scaled:center"),
                          scale=attr(scaler, "scaled:scale"))
Wage.test[,1:2]  <- scale(Wage.test[,1:2],
                          center=attr(scaler, "scaled:center"),
                          scale=attr(scaler, "scaled:scale"))


# b) Fit single hidden layer neural networks and compute errors ----------- ####
m     <- 10
decay <- 0.1
fmla  <- wage~year+age+maritl+race+education+jobclass+health+health_ins
results.nnet <- data.frame(m=integer(), decay=numeric(), error=numeric())

settings <- list(c(decay=0.1, m=10),
                 c(decay=0.1, m=5),
                 c(decay=0.1, m=20),
                 c(decay=0.01, m=10),
                 c(decay=0.05, m=10))

for (k in seq_along(settings)) {
  decay     <- settings[[k]][1]
  m         <- settings[[k]][2]
  wage.nnet <- nnet(fmla, data=Wage.train, size=m, decay=decay,
                    linout=TRUE, MaxNWts=10000, maxit=300)
  pred.wage <- predict(wage.nnet, Wage.test)
  resid     <- pred.wage - Wage.test$wage
  err.nnet  <- mean(resid^2)
  results.nnet[nrow(results.nnet)+1,] <- list(m, decay, err.nnet)
}

show(results.nnet)
# smallest error for neural network with (m=5, decay=0.1)


# c) Scale response and redo previous exercise ---------------------------- ####
maxWage          <- max(Wage.train$wage)
Wage.train$rwage <- Wage.train$wage / maxWage

fmla_new <- rwage~year+age+maritl+race+education+jobclass+health+health_ins

results.new.nnet <- data.frame(m=integer(), decay=numeric(), error=numeric())

for (k in seq_along(settings)) {
  decay      <- settings[[k]][1]
  m          <- settings[[k]][2]
  rwage.nnet <- nnet(fmla_new, data=Wage.train, size=m, decay=decay,
                     linout=TRUE, MaxNWts=10000, maxit=300)
  pred.rwage <- predict(rwage.nnet, Wage.test)
  pred.wage  <- pred.rwage * maxWage
  resid      <- pred.wage - Wage.test$wage
  err.nnet   <- mean(resid^2)
  results.new.nnet[nrow(results.new.nnet)+1,] <- c(m, decay, err.nnet)
}

show(results.new.nnet)
# overall smaller errors compared to not scaling response
# new best (m=20, decay=0.1)


# d) Deep learning -------------------------------------------------------- ####
X <- model.matrix(~year+age+maritl+race+education+jobclass+health+health_ins,
                  data=Wage)
# change column names to avoid warning when calling mlp ...
colnames(X)[-1] <- sapply(colnames(X)[-1], function(x) gsub(" ", "_", x))

X.train <- X[train,]
X.test  <- X[-train,]

sizes <- list(c(10),
              c(10,10),
              c(10,10,10),
              c(10,10,10,10))

results.dnet <- data.frame(size=character(), error=numeric())

set.seed(2100)
for (size in sizes) {
  wage.dnet <- mlp(X.train, Wage.train$rwage, size=size,
                   linout=TRUE, learnFunc="BackpropWeightDecay",
                   learnFuncParams=c(0.3,0,0.1), maxit=300)

  pred.dnet <- maxWage*predict(wage.dnet, X.test)
  resid     <- pred.dnet - Wage.test$wage
  err.dnet  <- mean(resid^2)

  results.dnet[nrow(results.dnet)+1,] <- list(
    paste0("(", paste0(size, collapse=","), ")"),
    err.dnet)
}

print(results.dnet, digits=7)


# e) Refit best model several times --------------------------------------- ####
best.ind  <- which.min(results.dnet[,"error"])
size      <- sizes[[best.ind]]

# Store predictions and errors from each fitted model
fits      <- 10
pred.fits <- matrix(NA, nrow=fits, ncol=nrow(X.test))
err.fits  <- numeric(fits)

for (k in 1:fits) {
  wage.dnet <- mlp(X.train, Wage.train$rwage, size=size,
                   linout=TRUE, learnFunc="BackpropWeightDecay",
                   learnFuncParams=c(0.3,0,0.1), maxit=300)
  pred.dnet <- maxWage*predict(wage.dnet, X.test)
  resid     <- pred.dnet - Wage.test$wage
  err.dnet  <- mean(resid^2)

  err.fits[k]   <- err.dnet
  pred.fits[k,] <- pred.dnet
}

# See variability across model fits
boxplot(err.fits)


# f) Average predictions from 10 fitted models ---------------------------- ####
pred.dnet.10fits <- colMeans(pred.fits)
resid            <- pred.dnet.10fits - Wage.test$wage
err.dnet.10fits  <- mean(resid^2)

print(err.dnet.10fits)


# e) Summary of results --------------------------------------------------- ####

# For single hidden layer neural networks:
# Improvement by using scaled response
{
  cat("Results for single hidden layer NN: \n")
  cat("Wage as response: \n")
  print(results.nnet[which.min(results.nnet[,"error"]),],
        digits=7, row.names=FALSE)
  cat("Scaled wage as response: \n")
  print(results.new.nnet[which.min(results.new.nnet[,"error"]),],
        digits=7, row.names=FALSE)
}

# For deep neural networks:
# Improvement by averaging over several model outputs
{
  cat("Results for deep NN: \n")
  cat("Predict with single model: \n")
  print(results.dnet[which.min(results.dnet[,"error"]),],
        digits=7, row.names=FALSE)
  cat("Predict with average over 10 models: \n")
  cat(round(err.dnet.10fits, 3))
}