## ItMLiHSmar2022
## regular_fun.R, child script
## Regularisation model building function
## Andreas Gammelgaard Damsbo, agdamsbo@clin.au.dk
##
## Now modified to use in publication
##

regular_fun <- function(X, y, K, lambdas, alpha) {
  n <- nrow(X)
  set.seed(321)

  # Using caret function to ensure both levels represented in all folds
  c <- caret::createFolds(y = y, k = K, list = FALSE, returnTrain = TRUE)

  B <- yhatTestProbKeep <- list()
  accTrain <- accTest <- err_train <- err_test <- auc_train <- auc_test <- matrix(nrow = K, ncol = length(lambdas))

  TrainProb <- TestProb <- list()
  
  catinfo <- levels(y)

  cMatTrain <- cMatTest <- table(true = factor(c(0, 0), levels = catinfo), pred = factor(c(0, 0), levels = catinfo))


  ## Iterate over partitions
  for (idx1 in 1:K) {
    # Status
    cat("Processing fold", idx1, "of", K, "\n")

    # idx1=1
    # Get training- and test sets
    I_train <- c != idx1 ## Creating selection vector of TRUE/FALSE
    I_test <- !I_train

    Xtrain <- X[I_train, ]
    ytrain <- y[I_train]
    Xtest <- X[I_test, ]
    ytest <- y[I_test]


    ## Model matrices for glmnet
    ## Using the complicated approach not to include first level.
    # Xmat.train<-model.matrix(~ .-1, data=Xtrain,
    #              contrasts.arg = lapply(Xtrain[,sapply(Xtrain, is.factor)],
    #                                     contrasts, contrasts=T))
    # Xmat.test<-model.matrix(~ .-1, data=Xtest,
    #                contrasts.arg = lapply(Xtest[,sapply(Xtest, is.factor)],
    #                                       contrasts, contrasts=T))

    # Xmat.train<-model.matrix(~.-1,Xtrain)
    # Xmat.test<-model.matrix(~.-1,Xtest)

    # Weights
    ytrain_weight <- as.vector(1 - (table(ytrain)[ytrain] / length(ytrain)))
    # ytest_weight<-as.vector(1 / (table(ytest)[ytest] / length(ytest)))

    # Fit regularized linear regression model
    mod <- glmnet::glmnet(Xtrain, ytrain,
      alpha = alpha, ## Alpha = 1 for lasso
      lambda = lambdas, ## Setting lambdas
      standardize = TRUE, ## Scales and centers
      weights = ytrain_weight,
      family = "binomial"
    )

    # Keep coefficients for plot
    B[[idx1]] <- as.matrix(coef(mod))

    # TrainProb[[idx1]] <- list()
    TestProb[[idx1]] <- list()
    
    # Iterate over regularization strengths to compute training- and test
    # errors for individual regularization strengths.
    for (idx2 in 1:length(lambdas)) {
      # idx2=1

      # Predict
      yhatTrainProb <- predict(mod,
        s = lambdas[idx2],
        newx = data.matrix(Xtrain),
        type = "response"
      )

      yhatTestProb <- predict(mod,
        s = lambdas[idx2],
        newx = data.matrix(Xtest),
        type = "response"
      )
      
      # Compute training and test error
      yhatTrain <- round(yhatTrainProb)
      yhatTest <- round(yhatTestProb)

      TestProb[[idx1]][[idx2]] <- dplyr::bind_cols(data.matrix(Xtest),y=ytest,pred=yhatTestProb,.name_repair = "unique_quiet")
      
      # Make predictions categorical again (instead of 0/1 coding)
      yhatTrainCat <- factor(round(yhatTrainProb), levels = c("0", "1"), labels = catinfo, ordered = TRUE)
      yhatTestCat <- factor(round(yhatTestProb), levels = c("0", "1"), labels = catinfo, ordered = TRUE)

      # Evaluate classifier performance
      # Accuracy
      # accTrain[idx1,idx2] <- sum(yhatTrainCat==ytrain)/length(ytrain)
      # accTest [idx1,idx2] <- sum(yhatTestCat==ytest)/length(ytest)
      # #
      # # Error rate
      # err_train[idx1,idx2] =  1 - accTrain[idx1,idx2]
      # err_test [idx1,idx2] =   1 - accTest[idx1,idx2]

      # AUROC
      suppressMessages(
        auc_train[idx1, idx2] <- pROC::auc(ytrain, yhatTrainCat)
      )
      suppressMessages(
        auc_test[idx1, idx2] <- pROC::auc(ytest, yhatTestCat)
      )

      # Compute confusion matrices
      cMatTrain <- cMatTrain + table(true = ytrain, pred = yhatTrainCat)
      cMatTest <- cMatTest + table(true = ytest, pred = yhatTestCat)
    }
  }
  list(mod = mod, B = B, auc_train = auc_train, auc_test = auc_test, cMatTrain = cMatTrain, cMatTest = cMatTest, 
       TrainProb=TrainProb, 
       TestProb=TestProb)
}


## ItMLiHSmar2022
## regularisation_steps.R, child script
## Regularised model building and analysation for assignment
## Andreas Gammelgaard Damsbo, agdamsbo@clin.au.dk
##
## Now modified to use in publication
##



#' Title
#'
#' @param data
#' @param outcome.var
#' @param weighted
#'
#' @return
#' @export
#'
#' @examples
#' data <- targets::tar_read(df_pred_data) |> 
#' pred_ls_split(excluded.vars = "reg_bmi")
#' 
#' data <- data[[1]]
#' mod <- data |> regularisation_steps(auto.l=TRUE)
regularisation_steps <- function(data, outcome.var = "pase_bin", weighted = FALSE, auto.l = FALSE) {
  n <- nrow(data)

  y <- data |> dplyr::select({{ outcome.var }})
  X <- data |> dplyr::select(-{{ outcome.var }})

  ## ====================================================================
  ## Step 0: data import and wrangling
  ## ====================================================================

  # setwd("/Users/au301842/PhysicalActivityandStrokeOutcome/1 PA Decline/")

  # source("data_format.R")
  y1 <- factor(as.integer(y[[1]])) ## Outcome is required to be factor of 0 or 1.
  # summary(y1)

  ## ====================================================================
  ## Step 1: settings
  ## ====================================================================

  ## Folds
  K <- 10
  set.seed(3)
  c <- caret::createFolds(
    y = y1,
    k = K,
    list = FALSE,
    returnTrain = TRUE
  ) # Fold IDs for tuning

  ## Defining tuning parameters
  if (auto.l){
    lambdas <- NULL
  } else {
    lambdas <- 2^seq(-10, 20, 1)
  }
  
  alphas <- seq(0, 1, .1)

  ## Weights for models
  if (weighted) {
    wght <- as.vector(1 - (table(y1)[y1] / length(y1)))
  } else {
    wght <- rep(1, length(y1))
  }


  ## Standardise numeric
  ## Centered and



  ## ====================================================================
  ## Step 2: all cross validations for each alpha
  ## ====================================================================

  # library(furrr)
  # library(purrr)
  # library(doMC)
  # registerDoMC(cores=6)

  future::plan(strategy = "multisession", workers = 2)

  # Nested CVs with analysis for all lambdas for each alpha
  #
  set.seed(3)
  cvs <- furrr::future_map(alphas, .options = furrr::furrr_options(seed = 3), function(a) {
    glmnet::cv.glmnet(model.matrix(~ . - 1, X),
      y1,
      weights = wght,
      lambda = lambdas,
      type.measure = "deviance", # This is standard measure and recommended for tuning
      foldid = c, # Per recommendation the folds are kept for alpha optimisation
      alpha = a,
      standardize = TRUE,
      family = quasibinomial, # Same as binomial, but not as picky
      keep = TRUE
    )
  })

  ## ====================================================================
  # Step 3: optimum lambda for each alpha
  ## ====================================================================


  # For each alpha, lambda is chosen for the lowest meassure (deviance)
  each_alpha <- sapply(seq_along(alphas), function(id) {
    each_cv <- cvs[[id]]
    alpha_val <- alphas[id]
    index_lmin <- match(
      each_cv$lambda.min,
      each_cv$lambda
    )
    c(
      lamb = each_cv$lambda.min,
      alph = alpha_val,
      cvm = each_cv$cvm[index_lmin]
    )
  })

  if (auto.l){
    # Best (min) lambda
    best_lamb <- min(each_alpha["lamb", ])
    
    # Alpha is chosen for best lambda with lowest model deviance, each_alpha["cvm",]
    best_alph <- each_alpha["alph", ][each_alpha["cvm", ] == min(each_alpha["cvm", ]
                                                                 [each_alpha["lamb", ] %in% best_lamb])]
    
    # BEst lamb is set to NULL to allow glm.net to use the optimal method, which is bult in.
    # best_lamb <- NULL
    
  } else {
    # Best (min) lambda
    best_lamb <- min(each_alpha["lamb", ])
    
    # Alpha is chosen for best lambda with lowest model deviance, each_alpha["cvm",]
    best_alph <- each_alpha["alph", ][each_alpha["cvm", ] == min(each_alpha["cvm", ]
                                                                 [each_alpha["lamb", ] %in% best_lamb])]
  }
  

  ## https://stackoverflow.com/questions/42007313/plot-an-roc-curve-in-r-with-ggplot2
  # df_roc <- glmnet::roc.glmnet(cvs[[match(best_alph, alphas)]]$fit.preval, newy = y1)[match(best_lamb, lambdas)]# |> # Plots performance from model with best alpha
  #
  #   df_roc |> plot_roc_curve()

  ## ====================================================================
  # Step 4: Creating the final model
  ## ====================================================================

  # source(here::here("R/regular_fun.R")) # Custom function
  optimised_model <- regular_fun(X = X, y = y1, K = K, lambdas = best_lamb, alpha = best_alph)
  # With lambda and alpha specified, the function is just a k-fold cross-validation wrapper,
  # but keeps model performance figures from each fold.
  
  # list2env(optimised_model, .GlobalEnv)
  # Function outputs a list, which is unwrapped to Env.
  # See source script for reference.

  ## ====================================================================
  # Step 5: creating table of coefficients for inference
  ## ====================================================================

  # reg_coef_tbl <- optimised_model$B |> purrr::reduce(cbind)

  # Bmatrix <- optimised_model$B |> purrr::reduce(cbind)
  # Bmedian <- apply(Bmatrix, 1, median)
  # Bmean <- apply(Bmatrix, 1, mean)
  #
  # reg_coef_tbl <- dplyr::tibble(
  #   name = rownames(Bmatrix),
  #   medianX = round(Bmedian, 5),
  #   ORmed = round(exp(Bmedian), 5),
  #   meanX = round(Bmean, 5),
  #   ORmea = round(exp(Bmean), 5)
  # ) # |>
  # arrange(desc(abs(medianX)))%>%
  # gt::gt()

  ## ====================================================================
  # Step 6: plotting predictive performance
  ## ====================================================================

  # reg_cfm <- caret::confusionMatrix(optimised_model$cMatTest)
  # reg_cfm <- optimised_model$cMatTest
  # reg_auc_sum <- optimised_model$auc_test[, 1]

  ## ====================================================================
  # Step 7: Packing list to save in loop
  ## ====================================================================

  list(
    "IncludedN" = n,
    "model" = optimised_model,
    "alphas" = alphas,
    "bestA" = best_alph,
    "lambdas" = lambdas,
    "bestL" = best_lamb,
    # "TestTable" = reg_cfm,
    # "AUROC" = reg_auc_sum,
    # "ROC curve" = df_roc,
    "y1" = y1,
    "X" = X,
    "data" = data,
    "cvs" = cvs
  )
}