#' Plot the summary curves produced by IntegratedLearner object

#'

#'@description Plots the R^2/AUC curves for the training (and test, if provided) set produced by IntegratedLearner object 

#'

#' @param fit fitted "IntegratedLearner" object 

#' @param label_size (optional) Numerical value indicating the label size. Default is 8.

#' @param label_x (optional) Single value or vector of x positions for plot labels, relative to each subplot. Defaults to 0.3 for all labels. (Each label is placed all the way to the left of each plot.)

#' @param vjust Adjusts the vertical position of each label. More positive values move the label further down on the plot canvas. Can be a single value (applied to all labels) or a vector of values (one for each label). Default is 0.1.

#' @param rowwise_plot If both train and test data is available, should the train and test plots be rowwise_plot. Default is TRUE. If FALSE, plots are aligned column-wise.

#'

#' @return ggplot2 object

#' @export

plot.learner <- function(fit,label_size=8, label_x=0.3,vjust=0.1, rowwise_plot=TRUE){

  clean_base_learner <- str_remove_all(fit$base_learner, 'SL.')

  clean_meta_learner <- str_remove_all(fit$meta_learner, 'SL.')  

  method <- paste(clean_base_learner,clean_meta_learner,sep=' + ')

  if(rowwise_plot) {

    nrow = 2

    ncol = 1

  } else{

    nrow = 1

    ncol = 2

    }

  if(fit$family=='binomial'){

    # Extract ROC plot data 

    list.ROC<-vector("list", length = ncol(fit$yhat.train))

    names(list.ROC)<-colnames(fit$yhat.train)

    y <- fit$Y_train

    # Loop over layers 

    for(k in 1:length(list.ROC)){

      preds<-fit$yhat.train[ ,k]

      pred = ROCR::prediction(preds, y)

      AUC = round(ROCR::performance(pred, "auc")@y.values[[1]], 2)

      perf = ROCR::performance(pred, "sens", "spec") 

      list.ROC[[k]] <- data.frame(sensitivity = methods::slot(perf, "y.values")[[1]],

                                  specificity = 1 - methods::slot(perf, "x.values")[[1]],

                                  AUC = AUC,

                                  layer = names(list.ROC)[k],

                                  method = method)

    }

    # Combine

    ROC_table<-do.call('rbind', list.ROC)

    # Prepare data for plotting

    plot_data<-ROC_table

    plot_data$displayItem<-paste(plot_data$layer, " AUC = ", plot_data$AUC, sep="")

    plot_data$displayItem<-factor(plot_data$displayItem,

                                  levels = unique(plot_data$displayItem))

    # ROC curves

    p1<-ggplot(plot_data,

               aes(x=specificity,

                   y=sensitivity,

                   group=displayItem)) + 

      geom_line(aes(x=specificity,y=sensitivity,color=displayItem)) +

      #ggtitle(paste('Training data: ', method, sep=''))+

      theme(legend.position="bottom", 

            legend.background=element_blank(),

            legend.box.background=element_rect(colour="black")) + 

      theme_bw() +

      xlab("False Positive Rate") +

      ylab("True Positive Rate") +

      theme(legend.position = "right", legend.direction = "vertical") +

      labs(color='') 

    if(fit$test==TRUE){

      # Extract ROC plot data 

      list.ROC.valid<-vector("list", length = ncol(fit$yhat.test))

      names(list.ROC.valid)<-colnames(fit$yhat.test)

      y <- fit$Y_test

      # Loop over layers 

      for(k in 1:length(list.ROC.valid)){

        preds<-fit$yhat.test[ ,k]

        pred = ROCR::prediction(preds, y)

        AUC = round(ROCR::performance(pred, "auc")@y.values[[1]], 2)

        perf = ROCR::performance(pred, "sens", "spec") 

        list.ROC.valid[[k]] <- data.frame(sensitivity = methods::slot(perf, "y.values")[[1]],

                                          specificity = 1 - methods::slot(perf, "x.values")[[1]],

                                          AUC = AUC,

                                          layer = names(list.ROC.valid)[k],

                                          method = method)

      }

      # Combine

      ROC_table_valid<-do.call('rbind', list.ROC.valid)

      # Prepare data for plotting

      plot_data<-ROC_table_valid

      plot_data$displayItem<-paste(plot_data$layer, " AUC = ", plot_data$AUC, sep="")

      plot_data$displayItem<-factor(plot_data$displayItem,

                                    levels = unique(plot_data$displayItem))

      # ROC curves

      p2<-ggplot(plot_data,

                 aes(x=specificity,

                     y=sensitivity,

                     group=displayItem)) + 

        geom_line(aes(x=specificity,y=sensitivity,color=displayItem)) +

        #ggtitle(paste('Test data: ', method, sep=''))+

        theme(legend.position="bottom", 

              legend.background=element_blank(),

              legend.box.background=element_rect(colour="black")) + 

        theme_bw() +

        xlab("False Positive Rate") +

        ylab("True Positive Rate") +

        theme(legend.position = "right", legend.direction = "vertical") +

        labs(color='') 

      p<-plot_grid(p1, 

                   p2, 

                   nrow = 2, 

                   labels = c(paste('A. ', fit$folds,'-fold CV',sep = ''), 

                              'B. Independent Validation'),

                   label_size = label_size, label_x = label_x,vjust = vjust)+

        theme(plot.margin = unit(c(1,1,1,1), "cm"))  

      print(p)

      return(list('plot'=p,'ROC_table'=ROC_table,'ROC_table_valid'=ROC_table_valid))

    }

    p <- plot_grid(p1, 

                   nrow = nrow,

                   ncol = ncol,

                   labels = c(paste('A. ', fit$folds,'-fold CV',sep = '')), 

                   label_size = label_size, label_x = label_x,vjust = vjust)+

      theme(plot.margin = unit(c(1,1,1,1), "cm"))

    print(p)

    return(list('plot'=p,'ROC_table'=ROC_table)) 

  }

  else if(fit$family=='gaussian'){

    # Extract R2 plot data 

    list.R2<-vector("list", length = ncol(fit$yhat.train))

    names(list.R2)<-colnames(fit$yhat.train)

    y <- fit$Y_train

    # Loop over layers 

    for(k in 1:length(list.R2)){

      preds<-fit$yhat.train[ ,k]

      R2<- as.vector(cor(preds, y)^2)

      list.R2[[k]] <- data.frame(R2 = R2,

                                 layer = names(list.R2)[k],

                                 method = method)

    }

    # Combine 

    R2_table<-do.call('rbind', list.R2)

    # Plot

    p1<-ggplot(R2_table, aes(x = method, y = R2)) +

      geom_bar(position="dodge", stat="identity", aes(fill=layer)) +

      xlab("") + 

      ylab(expression(paste("Prediction accuracy (", R^2, ")"))) +

      scale_fill_discrete(name="") + 

      theme(legend.position="bottom", 

            legend.background=element_blank(),

            legend.box.background=element_rect(colour="black")) + 

      theme_bw() +

      guides(fill=guide_legend(title="")) +

      theme(legend.position = "right", legend.direction = "vertical",

            strip.background = element_blank()) +

      labs(fill='') 

    if(fit$test==TRUE){

      # Extract R2 plot data 

      list.R2.valid<-vector("list", length = ncol(fit$yhat.test))

      names(list.R2.valid)<-colnames(fit$yhat.test)

      y <- fit$Y_test

      # Loop over layers 

      for(k in 1:length(list.R2.valid)){

        preds<-fit$yhat.test[ ,k]

        R2<- as.vector(cor(preds, y)^2)

        list.R2.valid[[k]] <- data.frame(R2 = R2,

                                         layer = names(list.R2.valid)[k],

                                         method = method)

      }

      # Combine 

      R2_table_valid<-do.call('rbind', list.R2.valid)

      # Plot

      p2<-ggplot(R2_table_valid, aes(x = method, y = R2)) +

        geom_bar(position="dodge", stat="identity", aes(fill=layer)) +

        xlab("") + 

        ylab(expression(paste("Prediction accuracy (", R^2, ")"))) +

        scale_fill_discrete(name="") + 

        theme(legend.position="bottom", 

              legend.background=element_blank(),

              legend.box.background=element_rect(colour="black")) + 

        theme_bw() +

        guides(fill=guide_legend(title="")) +

        theme(legend.position = "right", legend.direction = "vertical",

              strip.background = element_blank()) +

        labs(fill='') 

      nrow = NULL

      ncol = NULL

      p<-plot_grid(p1, 

                   p2, 

                   nrow = nrow,

                   ncol = ncol,

                   labels = c(paste('A. ', fit$folds,'-fold CV',sep = ''), 

                              'B. Independent Validation'),

                   label_size = label_size, label_x = label_x,vjust = vjust)+

        theme(plot.margin = unit(c(1,1,1,1), "cm"))  

      print(p)

      return(list('plot'=p,'R2_table'=R2_table,'R2_table_valid'=R2_table_valid))

    }

    p <- plot_grid(p1, 

                   ncol = 1, 

                   labels = c(paste('A. ', fit$folds,'-fold CV',sep = '')), 

                   label_size = label_size, label_x = label_x,vjust = vjust)+

      theme(plot.margin = unit(c(1,1,1,1), "cm"))

    print(p)

    return(list('plot'=p,'R2_table'=R2_table)) 

  }

}