#  STK2100 / V25
#  Solution to modified Extra exercise 10
#  Vera Kvisgaard 06.05.25

# a) import data ----
filepath <- "https://www.uio.no/studier/emner/matnat/math/STK2100/data/heart.csv"
df <- read.csv(filepath, stringsAsFactors = TRUE)
sapply(df, class)
head(df)

# exclude rows with NAs
df <- df[!rowSums(is.na(df)) > 0, ]


# draw train-test split
N <- nrow(df)
train <- sample(N, round(N/2))
df_train <- df[train, ]
df_test  <- df[-train, ]


# convert to matrices for methods that prefer matrices to data frames...
x.tr <- vapply(df_train[, -ncol(df_train)], as.numeric, numeric(nrow(df_train)))
x.test <- vapply(df_test[, -ncol(df_train)], as.numeric, numeric(nrow(df_test)))
y.tr <- df_train$AHD
y.test <- df_test$AHD

# b) reproduce Figure 8.6 in the ISLR book ----
library(tree)
set.seed(007)

## initial tree ----
# define hyperparameters of initial tree
?tree.control
control <- tree.control(nobs = length(train),
                        mincut = 1,
                        mindev = .000001)

# fit regression tree and plot
fit <- tree::tree(AHD~., data = df_train, control = control)
plot(fit)
text(fit, digits=2, pretty=3,cex=0.5)


# define function that returns a sequence of pruned trees
prune_tree_sequence <- function(tree) {

  trees  <- list(tree)
  ntrees <- 1
  for (size in seq(summary(tree)$size, 2)) {

    # prune tree
    tree <- prune.tree(tree,
                       best = size, # number of terminal nodes
                       method = "misclass")


    # check if tree is changed
    # - prune.tree() returns the larger tree if no subtree of size `size` is found
    if (identical(trees[[ntrees]], tree)) next

    ntrees <- ntrees +1
    trees[[ntrees]] <- tree
  }

  return(trees)
}

# define function that computes misclasstion error of tree
compute_errorrate  <- function(fit, df) {
  # fit = fitted tree
  # df  = data.frame with predictors

  # compute predicted class probs
  probs  <- predict(fit, df)

  # return mean misclassification error
  mean( (probs[, 2] > 1/2) != (df$AHD == "Yes"))
}


err <- matrix(NA, nrow = summary(fit)$size, ncol = 3)
colnames(err) <- c("Train", "Test", "CV")

## Test- and train error for tree-sequence ----
fits <- prune_tree_sequence(fit)
for (fit in fits) {
  size <- summary(fit)$size
  err[size, "Train"] <- compute_errorrate(fit, df_train)
  err[size, "Test"]  <- compute_errorrate(fit, df_test)
}


## CV estimate of test error for tree-sequence ----
## - partition the data into K folds
## - for each fold, iterate over the tree-sizes, prune tree and compute test error

# draw partition of training data into K folds
K <- 10
fold <- sample(K, length(train), replace = TRUE)

err_cv <- matrix(nrow = maxsize, ncol = K)
for (k in 1:K) {

  # define training data for current fold
  train_k <- which(fold != k)
  df_train_k <- df_train[train_k, ]
  df_test_k  <- df_test[-train_k, ]

  # initiate a deep tree on current training data
  control <- tree.control(nobs = length(train_k),
                          mincut = 1,
                          mindev = .000001)
  fit_k  <- tree::tree(AHD~.,
                       data = df_train_k,
                       control = )

  # compute the sequence of pruned trees
  fits_k <-  prune_tree_sequence(fit_k)

  # estimate test error on hold-out sample for each tree
  for (fit in fits_k) {
    size <- summary(fit)$size
    err_cv[size, k] <- compute_errorrate(fit, df_test_k)
  }
}

# average test-error over all folds
err[, "CV"] <- rowMeans(err_cv, na.rm = TRUE)[1:nrow(err)]


## plot ----
sizes <- 1:nrow(err)

j  <- 1
indx <- !is.na(err[, j])
plot(sizes[indx], err[indx, j], col = j, pch = j, ylim = c(0, .5), ylab = "misclassification error", xlab = "size")
lines(sizes[indx], err[indx, j], col = j, pch = j)

for (j in 2:ncol(err)) {
  indx <- !is.na(err[, j])
  points(sizes[indx], err[indx, j], col = j, pch = j)
  lines(sizes[indx], err[indx, j], col = j, pch = j)
}

legend("topright", legend = colnames(err), col = 1:ncol(err), lty = 1)


# c)-d) Bagging ----

fit <-  randomForest::randomForest(x = x.tr,
                                   y = df_train$AHD,
                                   xtest = x.te,
                                   ytest = df_test$AHD,
                                   ntree = B,
                                   mtry = p)

err <- list(Test = fit$test$err.rate[, 1],  # error on test data
            OOB = fit$err.rate[, 1])        # error computed on OOB samples

plot(err[[1]], col = "black", type = "l", ylim = c(0.1, .3))
lines(err[[2]], col = "green", type = "l", ylim = c(0.1, .3))
legend("topright", legend = names(err), col = c("black", "green"), lty = 1)


# e) Boosting (xgboost) ----

## Fit boosting trees to the data, first with stumps (J=2)
## as the weak learner
library(xgboost)

# Use 10-fold cross-validation to select M and the learning rate is set to 0.1
fit.boost.1.cv <- xgb.cv(data=xgb.DMatrix(data=x.tr,label=y.tr == "Yes"),
                         nrounds = 1000,
                         eta=0.1,
                         nfold = 10,
                         metrics=list("error"),
                         max_depth=1,
                         objective="binary:logistic")
best.boost.1 <- which.min(fit.boost.1.cv$evaluation_log$test_error_mean)
plot(fit.boost.1.cv$evaluation_log$train_error_mean,type="l",lwd=2,xlab="Stage",ylab="Misclassification error")
lines(fit.boost.1.cv$evaluation_log$test_error_mean,col=2,lwd=2)
abline(v=best.boost.1,lty=2,lwd=2)
legend("topright",c("training","test"),col=c(1,2),lty=rep(1,2))