# This is a basic utility function to plot the top 20 feasture importance scores based on BART posterior samples.

# Depends on the library bartMachine. 

VarImp.learner <- function(fit,

                           num.var = 20,

                           layer.names = NULL){

  ################################################

  # Extract required elements from the IL object #

  ################################################

  if(fit$meta_learner=="SL.nnls.auc"){

    VIMP_stack<- cbind.data.frame(fit$weights)

    colnames(VIMP_stack)<-c('mean')

    VIMP_stack$sd <- NA

    VIMP_stack$type<-'stack'

  }else{

    VIMP_stack <- NULL

  }

  if(fit$base_learner=="SL.BART"){

    if(is.null(layer.names)){

      layer.names <- names(fit$model_fits$model_layers)  

    }

    if(all(layer.names %in% names(fit$model_fits$model_layers))==FALSE){

      stop(paste(layer.names[!(layer.names %in% names(fit$model_fits$model_layers))],

                             "is not a valid layer in the fit object."))

    }

    #######################################################

    # Extract per-layer feature importance scores (VIMPs) #

    #######################################################

    VIMP_list <- list()

    for( i in 1:length(layer.names)){

      qq<-bartMachine::investigate_var_importance(

        fit$model_fits$model_layers[[layer.names[i]]],plot = FALSE)

      VIMP_layer<-cbind.data.frame(qq$avg_var_props, qq$sd_var_props)

      colnames(VIMP_layer)<-c('mean', 'sd')

      VIMP_layer$type<-layer.names[i]

      VIMP_list[[i]] <- VIMP_layer[1:num.var, ]

    }

    VIMP <- do.call(rbind,VIMP_list)

  }else{

    stop("This functionality is currently available only for BART base learner")

  }

  ###########################

  # Feature importance plot #

  ###########################

  if(!is.null(VIMP_stack)){ VIMP <- as.data.frame(rbind.data.frame(VIMP_stack,VIMP))}

  VIMP<-rownames_to_column(VIMP, 'ID')

  p<-VIMP %>% 

    filter(type %in% layer.names) %>% 

    arrange(mean) %>% 

    mutate(ID = str_replace_all(ID, fixed("_"), " ")) %>% 

    mutate(type = factor(type, 

                         levels = layer.names,

                         labels = layer.names)) %>% 

    ggplot(aes(reorder(ID, -mean), mean, fill = type)) +

    facet_wrap(.~ type, scale = 'free') + 

    geom_bar(stat = "identity", fill = "lightsalmon") + 

    geom_errorbar(aes(ymin=ifelse(mean-sd>0,mean-sd,0), ymax=mean+sd), width=.2, position=position_dodge(.9)) +

    theme_bw() + 

    coord_flip() + 

    omicsEye_theme() + 

    theme (strip.background = element_blank()) + 

    ylab('Inclusion proportion') + 

    xlab('') 

  return(p)  

}

# This is a ggplot theme from the Rahnavard lab at GWU.

omicsEye_theme <- function() {

  # set default text format based on categorical and length

  angle = 45

  hjust = 1

  size = 6

  return (ggplot2::theme_bw() + ggplot2::theme(

    axis.text.x = ggplot2::element_text(size = 8, vjust = 1, hjust = hjust, angle = angle),

    axis.text.y = ggplot2::element_text(size = 8, hjust = 1),

    axis.title = ggplot2::element_text(size = 10),

    plot.title = ggplot2::element_text(size = 10),

    plot.subtitle = ggplot2::element_text(size = 8),

    legend.title = ggplot2::element_text(size = 6, face = 'bold'),

    legend.text = ggplot2::element_text(size = 7),

    axis.line = ggplot2::element_line(colour = 'black', size = .25),

    ggplot2::element_line(colour = 'black', size = .25),

    axis.line.x = ggplot2::element_line(colour = 'black', size = .25),

    axis.line.y = ggplot2::element_line(colour = 'black', size = .25),

    panel.border = ggplot2::element_blank(),

    panel.grid.major = ggplot2::element_blank(),

    panel.grid.minor = ggplot2::element_blank())

  )

}