#' Stride length calculation

#' 

#' Compute the optimal length for Stride.

#'

#' @param maxlen Length of the input sequence.

#' @param plen Length of a patch.

#' @returns Numerical value.

#' @noRd

stridecalc <- function(maxlen, plen) {

  vec <- c()

  for (i in ceiling(plen / 3):(floor(plen / 2) - 1)) {

    if ((maxlen - plen) %% i == 0) {

      vec <- c(vec, i)

    }

  }

  return(vec)

}

#' Number of Patches calculation

#' 

#' Compute the Number of Patches.

#'

#' @param plen Length of a patch.

#' @param maxlen Length of the input sequence.

#' @param stride Stride.

#' @returns Numerical value.

#' @noRd

nopatchescalc <- function(plen, maxlen, stride) {

  ((maxlen - plen)/stride) + 1

}

maxlencalc <- function(plen, nopatches, stride) {

  (nopatches - 1) * stride + plen

}

#' Checkpoints saving function

#'

#' @param cp Type of the checkpoint.

#' @param runname Name of the run. Name will be used to identify output from callbacks.

#' @param model A keras model.

#' @param optimizer A keras optimizer.

#' @param history A keras history object.

#' @returns None. Saves object to file.

#' @noRd

savechecks <- function(cp, runname, model, optimizer, history, path_checkpoint) {

  np = reticulate::import("numpy", convert = FALSE)

  ## define path for saved objects

  modpath <- file.path(path_checkpoint, runname, cp)

  ## save model object

  model %>% keras::save_model_hdf5(paste0(modpath, "mod_temp.h5"))

  file.rename(paste0(modpath, "mod_temp.h5"),

              paste0(modpath, "mod.h5"))

  ## save optimizer object

  np$save(

    paste0(modpath, "opt.npy"),

    np$array(keras::backend(FALSE)$batch_get_value(optimizer$weights),

             dtype = "object"),

    allow_pickle = TRUE

  )

  ## save history object

  saveRDS(history, paste0(modpath, "history_temp.rds"))

  file.rename(paste0(modpath, "history_temp.rds"),

              paste0(modpath, "history.rds"))

  ## print when finished

  message(paste0("---------- New ", cp, " model saved\n"))

}

#' Tensorboard Writer

#' 

#' Writes the loss and the accuracy for a given epoch to the tensorboard.

#'

#' @param writer Name of the tensorboard writer function.

#' @param loss Computed loss for a given epoch.

#' @param acc Computed accracy for a given epoch.

#' @param epoch Epoch, for which the values shall be written to the tensorboard.

#' @returns None. Saves object to file.

#' @noRd

TB_loss_acc <- function(writer, loss, acc, epoch) {

  with(writer$as_default(), {

    tensorflow::tf$summary$scalar('epoch_loss',

                                  loss$result(),

                                  step = tensorflow::tf$cast(epoch, "int64"))

    tensorflow::tf$summary$scalar('epoch_accuracy',

                                  acc$result(),

                                  step = tensorflow::tf$cast(epoch, "int64"))

  })

}

#' Step function 

#'

#' @param trainvaldat A data generator.

#' @param model A keras model.

#' @param train_type Either `"cpc"`, `"Self-GenomeNet"`.

#' @param training Boolean. Whether this step is a training step.

#' @noRd

modelstep <-

  function(trainvaldat,

           model,

           train_type = "cpc",

           training = FALSE) {

    ## get batch

    a <- trainvaldat$x %>% tensorflow::tf$convert_to_tensor()

    if (train_type == "Self-GenomeNet") {

      ## get complement 

      a_complement <-

        tensorflow::tf$convert_to_tensor(array(as.array(a)[, (dim(a)[2]):1, 4:1], dim = c(dim(a)[1], dim(a)[2], dim(a)[3])))

      a <- tensorflow::tf$concat(list(a, a_complement), axis = 0L)

    }

    ## insert data in model

    model(a, training = training)

  }

#' Reading Pretrained Model function

#'

#' @param pretrained_model The path to a saved keras model.

#' @noRd

ReadOpt <- function(pretrained_model) {

  ## Read configuration

  optconf <-

    readRDS(paste(sub("/[^/]+$", "", pretrained_model),

                  "optconfig.rds",

                  sep = "/"))

  ## Read optimizer

  optimizer <- tensorflow::tf$optimizers$Adam$from_config(optconf)

  # Initialize optimizer

  with(

    keras::backend()$name_scope(optimizer$`_name`),

    with(tensorflow::tf$python$framework$ops$init_scope(), {

      optimizer$iterations

      optimizer$`_create_hypers`()

      optimizer$`_create_slots`(model$trainable_weights)

    })

  )

  # Read optimizer weights

  wts2 <-

    np$load(paste(

      sub("/[^/]+$", "", pretrained_model),

      "/",

      utils::tail(stringr::str_remove(

        strsplit(pretrained_model, "/")[[1]], "mod.h5"

      ), 1),

      "opt.npy",

      sep = ""

    ), allow_pickle = TRUE)

  # Set optimizer weights

  optimizer$set_weights(wts2)

  return(optimizer)

}

#' Learning Rate Schedule - Parameter Check

#' 

#' Checks, whether all necessary parameters for a defined learning rate schedule are given.

#'

#' @param lr_schedule The name of a learning rate schedule.

#' @noRd

LRstop <- function(lr_schedule) {

  # cosine annealing

  if ("cosine_annealing" %in% lr_schedule) {

    if (!isTRUE(all.equal(sort(names(lr_schedule)), sort(

      c("schedule", "lrmin", "lrmax", "restart", "mult")

    )))) {

      stop(

        "Please define lrmin, lrmax, restart, and mult within the list to use cosine annealing"

      )

    }

    # step decay

  } else if ("step_decay" %in% lr_schedule) {

    if (!isTRUE(all.equal(sort(names(lr_schedule)), sort(

      c("schedule", "lrmax", "newstep", "mult")

    )))) {

      stop("Please define lrmax, newstep, and mult within the list to use step decay")

    }

    # exponential decay

  } else if ("exp_decay" %in% lr_schedule) {

    if (!isTRUE(all.equal(sort(names(lr_schedule)), sort(c(

      "schedule", "lrmax", "mult"

    ))))) {

      stop("Please define lrmax, and mult within the list to use exponential decay")

    }

  }

}

#' Learning Rate Calculator

#' 

#' Computes the learning rate for a given epoch.

#'

#' @param lr_schedule The name of a learning rate schedule.

#' @param epoch Epoch, for which the learning rate shall be calculated.

#' @noRd

getEpochLR <- function(lr_schedule, epoch) {

  if (lr_schedule$schedule == "cosine_annealing") {

    # cosine annealing

    sgdr(

      lrmin = lr_schedule$lrmin,

      restart = lr_schedule$restart,

      lrmax = lr_schedule$lrmax,

      mult = lr_schedule$mult,

      epoch = epoch

    )

  } else if (lr_schedule$schedule == "step_decay") {

    # step decay

    stepdecay(

      newstep = lr_schedule$newstep,

      lrmax = lr_schedule$lrmax,

      mult = lr_schedule$mult,

      epoch = epoch

    )

  } else if (lr_schedule$schedule == "exp_decay") {

    # exponential decay

    exp_decay(

      lrmax = lr_schedule$lrmax,

      mult = lr_schedule$mult,

      epoch = epoch

    )

  }

}

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

########################################### Parameter Lists ############################################

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

#

GenParams <- function(maxlen,

                      batch_size,

                      step,

                      proportion_per_seq,

                      max_samples) {

  checkmate::assertInt(maxlen, lower = 1)

  checkmate::assertInt(batch_size, lower = 1)

  checkmate::assertInt(step, lower = 1)

  checkmate::assertInt(max_samples, lower = 1, null.ok = TRUE)

  checkmate::assertNumber(

    proportion_per_seq,

    lower = 0,

    upper = 1,

    null.ok = TRUE

  )

  structure(

    list(

      maxlen = maxlen,

      batch_size = batch_size,

      step = step,

      proportion_per_seq = proportion_per_seq,

      max_samples = max_samples

    ),

    class = "Params"

  )

}

GenTParams <- function(path,

                       shuffle_file_orderTrain,

                       path_file_log,

                       seed) {

  checkmate::assertLogical(shuffle_file_orderTrain)

  checkmate::assertInt(seed)

  structure(

    list(

      path_corpus = path,

      shuffle_file_order = shuffle_file_orderTrain,

      path_file_log = path_file_log,

      seed = seed

    ),

    class = "Params"

  )

}

GenVParams <- function(path_val,

                       shuffle_file_orderVal) {

  checkmate::assertLogical(shuffle_file_orderVal)

  structure(list(path_corpus = path_val[[1]],

                 shuffle_file_order = shuffle_file_orderVal),

            class = "Params")

}

# add list of hyperparameters to model

add_hparam_list <- function(model, argg) {

  argg["model_metrics"] <- NULL

  argg["model"] <- NULL

  argg["i"] <- NULL

  argg["optimizer"] <- NULL

  argg["layer_lstm"] <- paste(as.character(argg$layer_lstm), collapse = " ")

  argg["filters"] <- paste(as.character(argg$filters), collapse = " ")

  argg["kernel_size"] <- paste(as.character(argg$kernel_size), collapse = " ")

  argg["pool_size"] <- paste(as.character(argg$pool_size), collapse = " ")

  argg["strides"] <- paste(as.character(argg$strides), collapse = " ")

  argg["residual_block"] <- paste(as.character(argg$residual_block), collapse = " ")

  argg["residual_block_length"] <- paste(as.character(argg$residual_block_length), collapse = " ")

  argg["size_reduction_1Dconv"] <- paste(as.character(argg$size_reduction_1Dconv), collapse = " ")

  argg["layer_dense"] <- paste(as.character(argg$layer_dense), collapse = " ")

  argg["padding"] <- paste(as.character(argg$padding), collapse = " ")

  argg["use_bias"] <- paste(as.character(argg$use_bias), collapse = " ")

  argg["input_label_list"] <- paste(as.character(argg$layer_dense), collapse = " ")

  argg["num_heads"] <- paste(as.character(argg$num_heads), collapse = " ")

  argg["head_size"] <- paste(as.character(argg$head_size), collapse = " ")

  argg["dropout"] <- paste(as.character(argg$dropout), collapse = " ")

  argg["input_tensor"] <- NULL

  argg["label_inputs"] <- NULL

  argg["f1"] <- NULL

  argg["multi_acc"] <- NULL

  argg[["number_model_params"]] <- model$count_params()

  for (i in 1:length(argg$label_input)) {

    argg[paste0("input_tensor_", i)] <- NULL

    argg[paste0("label_input_layer_", i)] <- NULL

  }

  argg["output_tensor"] <- NULL

  argg["output_list"] <- NULL

  argg["residual_layer"] <- NULL

  argg["label_noise_matrix"] <- NULL

  argg["smooth_loss"] <- NULL

  argg["noisy_loss"] <- NULL

  argg["col_sums"] <- NULL

  argg["auc"] <- NULL

  argg["multi_label"] <- NULL

  argg["macro_average_cb"] <- NULL

  argg["verbose"] <- NULL

  argg["embedded_indices"] <- NULL

  argg["position_indices"] <- NULL

  argg["optimizer"] <- NULL

  argg["residual_blocks"] <- paste(as.character(argg$residual_blocks), collapse = " ")

  argg["model_metrics"] <- NULL

  argg["i"] <- NULL

  argg["optimizer"] <- NULL

  argg["model"] <- NULL

  argg["input_tensor_1"] <- NULL

  argg["input_tensor_2"] <- NULL

  argg["input_label_list"] <- NULL

  for (i in 1:length(argg$label_input)) {

    argg[paste0("input_tensor_", i)] <- NULL

    argg[paste0("label_input_layer_", i)] <- NULL

  }

  argg[["number_model_params"]] <- model$count_params()

  argg["label_input"] <- NULL

  argg["label_inputs"] <- NULL

  argg["maxlen_1"] <- NULL

  argg["maxlen_2"] <- NULL

  argg["f1"] <- NULL

  argg["output_tensor"] <- NULL

  argg["output_tensor_1"] <- NULL

  argg["output_tensor_2"] <- NULL

  argg[["attn_block"]] <- NULL

  argg["feature_ext_model"] <- NULL

  argg["ff_dim"] <- NULL

  argg["pos_enc_layer"] <- NULL

  argg["number_of_cnn_layers"] <- paste(as.character(argg$number_of_cnn_layers), collapse = " ")

  argg["feature_ext_model"] <- NULL

  argg["pe_matrix"] <- NULL

  argg["position_embedding_layer"] <- NULL 

  argg["layer_add_td"] <- NULL 

  model$hparam <- argg

  model

}

get_maxlen <- function(model, set_learning, target_middle, read_data, return_int = FALSE,

                       n_gram = NULL) {

  if (is.null(set_learning)) {

    num_in_layers <- length(model$inputs)

    if (num_in_layers == 1) {

      maxlen <- model$input$shape[[2]]

    } else {

      if (!target_middle & !read_data) {

        maxlen <- model$input[[num_in_layers]]$shape[[2]]

      } else {

        maxlen <- model$inputs[[num_in_layers - 1]]$shape[[2]] + model$inputs[[num_in_layers]]$shape[[2]]

      }

    }

    if (!is.null(n_gram)) {

      maxlen <- maxlen + n_gram - 1

    }

  } else {

    maxlen <- set_learning$maxlen

  }

  return(maxlen)

}

# combine lists containing x, y and sample weight subsets

reorder_masked_lm_lists <- function(array_lists, include_sw = NULL) {

  if (is.null(include_sw)) include_sw <- FALSE

  x <- list()

  y <- list()

  sw <- list()

  for (i in 1:length(array_lists)) {

    x[[i]] <- array_lists[[i]]$x

    y[[i]] <- array_lists[[i]]$y

    if (include_sw) sw[[i]] <- array_lists[[i]]$sample_weight

  }

  x <- abind::abind(x, along = 1)

  y <- abind::abind(y, along = 1)

  if (include_sw) sw <- abind::abind(sw, along = 1)

  if (include_sw) {

    return(list(x=x, y=y, sw=sw))

  } else {

    return(list(x=x, y=y))

  }

}

# stack 

create_x_y_tensors_lm <- function(sequence_list, nuc_dist_list, target_middle,

                                  maxlen, vocabulary, ambiguous_nuc,

                                  start_index_list, quality_list, target_len,

                                  coverage_list, use_coverage, max_cov,n_gram,

                                  n_gram_stride, output_format, wavenet_format) {

  if (!wavenet_format) {

    array_list <- purrr::map(1:length(sequence_list),

                             ~seq_encoding_lm(sequence_list[[.x]], nuc_dist = nuc_dist_list[[.x]], adjust_start_ind = TRUE,

                                              maxlen = maxlen, vocabulary = vocabulary, ambiguous_nuc = ambiguous_nuc,

                                              start_ind =  start_index_list[[.x]], 

                                              quality_vector = quality_list[[.x]], target_len = target_len,

                                              cov_vector = coverage_list[[.x]], use_coverage = use_coverage, max_cov = max_cov,

                                              n_gram = n_gram, n_gram_stride = n_gram_stride, output_format = output_format)

    )

    if (!is.list(array_list[[1]][[2]])) {

      if (!target_middle) {

        x <- array_list[[1]][[1]]

        y <- array_list[[1]][[2]]

        if (length(array_list) > 1) {

          for (i in 2:length(array_list)) {

            x <- abind::abind(x, array_list[[i]][[1]], along = 1)

            y <- rbind(y, array_list[[i]][[2]])

          }

        }

        # coerce y type to matrix

        if (dim(x)[1] == 1) {

          if (is.null(n_gram)) {

            dim(y) <-  c(1, length(vocabulary))

          } else {

            dim(y) <-  c(1, length(vocabulary)^n_gram)

          }

        }

      } else {

        x_1 <- array_list[[1]][[1]][[1]]

        x_2 <- array_list[[1]][[1]][[2]]

        y <- array_list[[1]][[2]]

        if (length(array_list) > 1) {

          for (i in 2:length(array_list)) {

            x_1 <- abind::abind(x_1, array_list[[i]][[1]][[1]], along = 1)

            x_2 <- abind::abind(x_2, array_list[[i]][[1]][[2]], along = 1)

            y <- rbind(y, array_list[[i]][[2]])

          }

        }

        x <- list(x_1, x_2)

        # coerce y type to matrix

        if (dim(x_1)[1] == 1) {

          if (is.null(n_gram)) {

            dim(y) <-  c(1, length(vocabulary))

          } else {

            dim(y) <-  c(1, length(vocabulary)^n_gram)

          }

        }

      }

    } else {

      if (!target_middle) {

        x <- array_list[[1]][[1]]

        y <- array_list[[1]][[2]]

        if (length(array_list) > 1) {

          for (i in 2:length(array_list)) {

            x <- abind::abind(x, array_list[[i]][[1]], along = 1)

            for (j in 1:length(y)) {

              y[[j]] <- rbind(y[[j]], array_list[[i]][[2]][[j]] )

            }

          }

        }

        # coerce y type to matrix

        if (dim(x)[1] == 1) {

          for (i in 1:length(y)) {

            if (is.null(n_gram)) {

              dim(y[[i]]) <-  c(1, length(vocabulary))

            } else {

              dim(y[[i]]) <-  c(1, length(vocabulary)^n_gram)

            }

          }

        }

      } else {

        x_1 <- array_list[[1]][[1]][[1]]

        x_2 <- array_list[[1]][[1]][[2]]

        y <- array_list[[1]][[2]]

        if (length(array_list) > 1) {

          for (i in 2:length(array_list)) {

            x_1 <- abind::abind(x_1, array_list[[i]][[1]][[1]], along = 1)

            x_2 <- abind::abind(x_2, array_list[[i]][[1]][[2]], along = 1)

            for (j in 1:length(y)) {

              y[[j]] <- rbind(y[[j]], array_list[[i]][[2]][[j]] )

            }

          }

        }

        x <- list(x_1, x_2)

        # coerce y type to matrix

        if (dim(x_1)[1] == 1) {

          for (i in 1:length(y)) {

            if (is.null(n_gram)) {

              dim(y[[i]]) <-  c(1, length(vocabulary))

            } else {

              dim(y[[i]]) <-  c(1, length(vocabulary)^n_gram)

            }

          }

        }

      }

    }

    # wavenet format

  } else {

    if (target_len > 1) {

      stop("target_len must be 1 when using wavenet_format")

    }

    # one hot encode strings collected in sequence_list and connect arrays

    array_list <- purrr::map(1:length(sequence_list),

                             ~seq_encoding_lm(sequence_list[[.x]], ambiguous_nuc = ambiguous_nuc, adjust_start_ind = TRUE,

                                              maxlen = maxlen, vocabulary = vocabulary, nuc_dist = nuc_dist_list[[.x]],

                                              start_ind =  start_index_list[[.x]], 

                                              quality_vector = quality_list[[.x]], n_gram = n_gram,

                                              cov_vector = coverage_list[[.x]], use_coverage = use_coverage, max_cov = max_cov,

                                              output_format = output_format)

    )

    x <- array_list[[1]][[1]]

    y <- array_list[[1]][[2]]

    if (length(array_list) > 1) {

      for (i in 2:length(array_list)) {

        x <- abind::abind(x, array_list[[i]][[1]], along = 1)

        y <- abind::abind(y, array_list[[i]][[2]], along = 1)

      }

    }

  }

  return(list(x, y))

}

# 

slice_tensor_lm <- function(xy, output_format, target_len, n_gram,

                            n_gram_stride, 

                            # maxlen_n_gram = NULL,

                            # target_len_n_gram = NULL, 

                            total_seq_len, return_int) {

  xy_dim <- dim(xy)

  if (!is.null(n_gram)) {

    target_len <- floor(target_len/n_gram)

  }

  if (output_format == "target_right") {

    x_index <- get_x_index(xy_dim, output_format, target_len)

    if (return_int) {

      x <- xy[ , x_index, drop=FALSE]

      y <- xy[ , -x_index]

    } else {

      x <- xy[ , x_index, , drop=FALSE]

      y <- xy[ , -x_index, ]

    }

  }

  if (output_format == "wavenet") {

    if (target_len != 1) {

      stop("Target length must be 1 for wavenet model")

    }

    x_index <- 1:(xy_dim[2] - target_len)

    y_index <- 2:dim(xy)[2]

    if (return_int) {

      x <- xy[ , x_index, drop=FALSE]

      y <- xy[ , y_index, drop=FALSE]

    } else {

      x <- xy[ , x_index, , drop=FALSE]

      y <- xy[ , y_index, , drop=FALSE]

    }

  }

  if (output_format == "target_middle_cnn") {

    seq_middle <- ceiling(xy_dim[2]/2)

    y_index <- (1:target_len) + (seq_middle - ceiling(target_len/2))

    if (return_int) {

      x <- xy[ , -y_index, drop=FALSE]

      y <- xy[ , y_index]

    } else {

      x <- xy[ , -y_index, , drop=FALSE]

      y <- xy[ , y_index, ]

    }

  }

  if (output_format == "target_middle_lstm") {

    seq_middle <- ceiling(xy_dim[2]/2)

    y_index <- (1:target_len) + (seq_middle - ceiling(target_len/2))

    if (return_int) {

      x1 <- xy[ , 1:(min(y_index) - 1), drop=FALSE]

      # reverse order of x2

      x2 <- xy[ , xy_dim[2] : (max(y_index) + 1), drop=FALSE]

      y <- xy[ , y_index]

    } else {

      x1 <- xy[ , 1:(min(y_index) - 1), , drop=FALSE]

      # reverse order of x2

      x2 <- xy[ ,  xy_dim[2] : (max(y_index) + 1), , drop=FALSE]

      y <- xy[ , y_index, ]

    }

    x <- list(x1, x2)

  }

  if (target_len == 1 & xy_dim[1] == 1 & output_format != "wavenet") {

    y <- matrix(y, nrow = 1)

  }

  if (target_len > 1 & xy_dim[1] == 1) {

    y <- array(y, dim = c(1, dim(y)))

  }

  return(list(x=x, y=y))

}

get_x_index <- function(xy_dim, output_format, target_len) {

  if (output_format == "target_right") {

    x_index <- 1:(xy_dim[2] - target_len)

  }

  #TODO: subset for other formats with n_gram/stride

  return(x_index)

}

add_dim <- function(x) {

  if (is.null(dim(x))) {

    return(matrix(x, nrow = 1))

  } else {

    return(array(x, dim = c(1, dim(x))))

  }

}

shuffle_batches <- function(x, shuffle_index) {

  if (!is.list(x)) {

    dim_len <- length(dim(x))

    x <- shuffle_sample(x, dim_len, shuffle_index)

  } else {

    dim_len <- length(dim(x[[1]]))

    for (i in 1:length(x)) {

      x[[i]] <- shuffle_sample(x[[i]], dim_len, shuffle_index)

    }

  }

}  

shuffle_sample <- function(x, dim_len, shuffle_index) {

  if (is.null(dim_len) | dim_len == 1) {

    x <- x[shuffle_index]

  }

  if (dim_len == 2) {

    x <- x[shuffle_index, ]

  }

  if (dim_len == 3) {

    x <- x[shuffle_index, , ]

  }

  return(x)

}

count_gpu <- function() {

  pd <- tensorflow::tf$config$list_physical_devices()

  count <- 0

  for (i in 1:length(pd)) {

    if (pd[[i]]$device_type == "GPU") count <- count + 1

  }

  return(count)

}

to_time_dist <- function(x, samples_per_target) {

  x_dim <- dim(x)

  x_dim_td <- c(x_dim[1], samples_per_target, x_dim[2]/samples_per_target, x_dim[3])

  x_td <- keras::k_reshape(x, shape = x_dim_td)

  keras::k_eval(x_td)

}

#' Plot confusion matrix

#' 

#' Plot confusion matrix, either with absolute numbers or percentages per column (true labels).

#' 

#' @param cm A confusion matrix

#' @param perc Whether to use absolute numbers or percentages.

#' @param cm_labels Labels corresponding to confusion matrix entries.

#' @param round_dig How to round numbers.

#' @param text_size Size of text annotations.

#' @param highlight_diag Whether to highlight entries in diagonal.

#' @examplesIf reticulate::py_module_available("tensorflow")

#' cm <- matrix(c(90, 1, 0, 2, 7, 1, 8, 3, 1), nrow = 3, byrow = TRUE)

#' plot_cm(cm, perc = TRUE, cm_labels = paste0('label_', 1:3), text_size = 8)

#' 

#' @returns A ggplot of a confusion matrix.

#' @export

plot_cm <- function(cm, perc = FALSE, cm_labels, round_dig = 2, text_size = 1, highlight_diag = TRUE) {

  if (perc) cm <- cm_perc(cm, round_dig)

  cm <- create_conf_mat_obj(cm, cm_labels)

  cm_plot <- ggplot2::autoplot(cm, type = "heatmap") +

    ggplot2::scale_fill_gradient(low="#D6EAF8", high = "#2E86C1")  +

    ggplot2::theme(axis.text.x =

                     ggplot2::element_text(angle=90, hjust=1, size = text_size)) +

    ggplot2::theme(axis.text.y =

                     ggplot2::element_text(size = text_size))

  if (highlight_diag) {

    #diagonal_data <- data.frame(x = levels(y_true), y = levels(y_pred))

    diagonal_data <- data.frame(x = cm_labels, y = cm_labels)

    cm_plot <- cm_plot + ggplot2::geom_tile(data = diagonal_data, ggplot2::aes(x = x, y = y),

                                            fill = "red", colour = "white", size = 1,

                                            alpha = 0.000001) 

  }

  # TODO: add conf mat with ComplexHeatmap for bigger sizes

  cm_plot

}

#' Remove checkpoints

#' 

#' Remove all but n 'best' checkpoints, based on some condition. Condition can be 

#' accuracy, loss or epoch number.

#' 

#' @param cp_dir Directory containing checkpoints.

#' @param metric Either `"acc"`, `"loss"` or `"last_ep"`. Condition which checkpoints to keep.

#' @param best_n Number of checkpoints to keep.

#' @param ask_before_remove Whether to show files to keep before deleting rest.

#' @examplesIf reticulate::py_module_available("tensorflow")

#' model <- create_model_lstm_cnn(layer_lstm = 8)

#' checkpoint_folder <- tempfile()

#' dir.create(checkpoint_folder)

#' keras::save_model_hdf5(model, file.path(checkpoint_folder, 'Ep.007-val_loss11.07-val_acc0.6.hdf5'))

#' keras::save_model_hdf5(model, file.path(checkpoint_folder, 'Ep.019-val_loss8.74-val_acc0.7.hdf5'))

#' keras::save_model_hdf5(model, file.path(checkpoint_folder, 'Ep.025-val_loss0.03-val_acc0.8.hdf5'))

#' remove_checkpoints(cp_dir = checkpoint_folder, metric = "acc", best_n = 2,

#'                    ask_before_remove = FALSE)

#' list.files(checkpoint_folder)

#'  

#' @returns None. Deletes certain files.

#' @export 

remove_checkpoints <- function(cp_dir, metric = "acc", best_n = 1, ask_before_remove = TRUE) {

  stopifnot(metric %in% c("acc", "loss", "last_ep"))

  stopifnot(dir.exists(cp_dir))

  stopifnot(best_n >= 1)

  files <- list.files(cp_dir, full.names = TRUE)

  if (length(files) == 0) {

    stop("Directory is empty")

  }

  files_basename <- basename(files)

  num_cp <- length(files)

  if (metric == "acc") {

    if (!all(stringr::str_detect(files_basename, "acc"))) {

      stop("No accuracy information in checkpoint names ('acc' string), use other metric.")

    }

    acc_scores <- files_basename %>% stringr::str_extract("acc\\d++\\.\\d++") %>% 

      stringr::str_remove("acc") %>% as.numeric()

    rank_order <- rank(acc_scores, ties.method = "last")

    index <- rank_order > (num_cp - best_n)

  }

  if (metric == "loss") {

    if (!all(stringr::str_detect(files_basename, "loss"))) {

      stop("No loss information in checkpoint names ('loss' string), use other metric.")

    }

    loss_scores <- files_basename %>% stringr::str_extract("loss\\d++\\.\\d++") %>% 

      stringr::str_remove("loss") %>% as.numeric()

    rank_order <- rank(loss_scores, ties.method = "last")

    index <- rank_order <= best_n

  }

  if (metric == "last_ep") {

    ep_scores <- files_basename %>% stringr::str_extract("Ep\\.\\d++") %>% 

      stringr::str_remove("Ep\\.") %>% as.numeric()

    rank_order <- rank(ep_scores)

    index <- rank_order > (num_cp - best_n)

  }

  if (ask_before_remove) {

    message("Deleting", sum(!index), paste0(ifelse(sum(!index) == 1, "file", "files") , "."),

            "Only keep \n", paste0(basename(files[index]), collapse = ",\n "), "\n")

    remove_cps <- utils::askYesNo("")

  } else {

    remove_cps <- TRUE

  }

  if (is.na(remove_cps)) return(NULL)

  if (remove_cps) {

    invisible(file.remove(files[!index]))

  }

}

pooling_flatten <- function(global_pooling = NULL, output_tensor) {

  if (!is.null(global_pooling)) {

    stopifnot(global_pooling %in% c("max_ch_first", "max_ch_last", "average_ch_first",

                                    "average_ch_last", "both_ch_first", "both_ch_last", "all", "none", "flatten"))

  } else {

    out <- output_tensor %>% keras::layer_flatten()

    return(out)

  }

  #if (!is.null(global_pooling) & global_pooling != "flatten") {

  if (stringr::str_detect(global_pooling, "_ch_")) {

    if (global_pooling == "max_ch_first") {

      out <- output_tensor %>% keras::layer_global_max_pooling_1d(data_format="channels_first")

    } 

    if (global_pooling == "max_ch_last") {

      out <- output_tensor %>% keras::layer_global_max_pooling_1d(data_format="channels_last")

    } 

    if (global_pooling ==  "average_ch_first") {

      out <- output_tensor %>% keras::layer_global_average_pooling_1d(data_format="channels_first")

    } 

    if (global_pooling ==  "average_ch_last") { 

      out <- output_tensor %>% keras::layer_global_average_pooling_1d(data_format="channels_last")

    } 

    if (global_pooling ==  "both_ch_last") { 

      out1 <- output_tensor %>% keras::layer_global_average_pooling_1d(data_format="channels_last")

      out2 <- output_tensor %>% keras::layer_global_max_pooling_1d(data_format="channels_last")

      out <- keras::layer_concatenate(list(out1, out2))

    } 

    if (global_pooling ==  "both_ch_first") {

      out1 <- output_tensor %>% keras::layer_global_average_pooling_1d(data_format="channels_first")

      out2 <- output_tensor %>% keras::layer_global_max_pooling_1d(data_format="channels_first")

      out <- keras::layer_concatenate(list(out1, out2))

    } 

    if (global_pooling == "all") {

      out1 <- output_tensor %>% keras::layer_global_average_pooling_1d(data_format="channels_first")

      out2 <- output_tensor %>% keras::layer_global_max_pooling_1d(data_format="channels_first")

      out3 <- output_tensor %>% keras::layer_global_average_pooling_1d(data_format="channels_last")

      out4 <- output_tensor %>% keras::layer_global_max_pooling_1d(data_format="channels_last")

      out <- keras::layer_concatenate(list(out1, out2, out3, out4))

    }       

  }

  if (global_pooling == "flatten") {

    out <- output_tensor %>% keras::layer_flatten()

  }

  if (global_pooling == "none") {

    return(output_tensor)

  }

  return(out)

}

pooling_flatten_time_dist <- function(global_pooling = NULL, output_tensor) {

  if (!is.null(global_pooling)) {

    stopifnot(global_pooling %in% c("max_ch_first", "max_ch_last", "average_ch_first",

                                    "average_ch_last", "both_ch_first", "both_ch_last", "all", "none", "flatten"))

  } else {

    out <- output_tensor %>% keras::layer_flatten()

    return(out)

  }

  #if (!is.null(global_pooling) & global_pooling != "flatten") {

  if (stringr::str_detect(global_pooling, "_ch_")) {

    if (global_pooling == "max_ch_first") {

      out <- output_tensor %>% keras::time_distributed(keras::layer_global_max_pooling_1d(data_format="channels_first"))

    } 

    if (global_pooling == "max_ch_last") {

      out <- output_tensor %>% keras::time_distributed(keras::layer_global_max_pooling_1d(data_format="channels_last"))

    } 

    if (global_pooling ==  "average_ch_first") {

      out <- output_tensor %>% keras::time_distributed(keras::layer_global_average_pooling_1d(data_format="channels_first"))

    } 

    if (global_pooling ==  "average_ch_last") { 

      out <- output_tensor %>% keras::time_distributed(keras::layer_global_average_pooling_1d(data_format="channels_last"))

    } 

    if (global_pooling ==  "both_ch_last") { 

      out1 <- output_tensor %>% keras::time_distributed(keras::layer_global_average_pooling_1d(data_format="channels_last"))

      out2 <- output_tensor %>% keras::time_distributed(keras::layer_global_max_pooling_1d(data_format="channels_last"))

      out <- keras::layer_concatenate(list(out1, out2))

    } 

    if (global_pooling ==  "both_ch_first") {

      out1 <- output_tensor %>% keras::time_distributed(keras::layer_global_average_pooling_1d(data_format="channels_first"))

      out2 <- output_tensor %>% keras::time_distributed(keras::layer_global_max_pooling_1d(data_format="channels_first"))

      out <- keras::layer_concatenate(list(out1, out2))

    } 

    if (global_pooling == "all") {

      out1 <- output_tensor %>% keras::time_distributed(keras::layer_global_average_pooling_1d(data_format="channels_first"))

      out2 <- output_tensor %>% keras::time_distributed(keras::layer_global_max_pooling_1d(data_format="channels_first"))

      out3 <- output_tensor %>% keras::time_distributed(keras::layer_global_average_pooling_1d(data_format="channels_last"))

      out4 <- output_tensor %>% keras::time_distributed(keras::layer_global_max_pooling_1d(data_format="channels_last"))

      out <- keras::layer_concatenate(list(out1, out2, out3, out4))

    }       

  }

  if (global_pooling == "flatten") {

    out <- output_tensor %>% keras::time_distributed(keras::layer_flatten())

  }

  if (global_pooling == "none") {

    return(output_tensor)

  }

  return(out)

}

get_pooling_flatten_layer <- function(global_pooling = NULL) {

  if (!is.null(global_pooling)) {

    stopifnot(global_pooling %in% c("max_ch_first", "max_ch_last", "average_ch_first",

                                    "average_ch_last", "both_ch_first", "both_ch_last", "all", "none", "flatten"))

  }

  if (stringr::str_detect(global_pooling, "_ch_")) {

    if (global_pooling == "max_ch_first") {

      out <- keras::layer_global_max_pooling_1d(data_format="channels_first")

    } 

    if (global_pooling == "max_ch_last") {

      out <- keras::layer_global_max_pooling_1d(data_format="channels_last")

    } 

    if (global_pooling ==  "average_ch_first") {

      out <- keras::layer_global_average_pooling_1d(data_format="channels_first")

    } 

    if (global_pooling ==  "average_ch_last") { 

      out <- keras::layer_global_average_pooling_1d(data_format="channels_last")

    } 

    if (global_pooling ==  "both_ch_last") { 

      out1 <- keras::layer_global_average_pooling_1d(data_format="channels_last")

      out2 <- keras::layer_global_max_pooling_1d(data_format="channels_last")

      out <- keras::layer_concatenate(list(out1, out2))

    } 

    if (global_pooling == "both_ch_first") {

      out1 <- keras::layer_global_average_pooling_1d(data_format="channels_first")

      out2 <- keras::layer_global_max_pooling_1d(data_format="channels_first")

      out <- keras::layer_concatenate(list(out1, out2))

    } 

    if (global_pooling == "all") {

      out1 <- keras::layer_global_average_pooling_1d(data_format="channels_first")

      out2 <- keras::layer_global_max_pooling_1d(data_format="channels_first")

      out3 <- keras::layer_global_average_pooling_1d(data_format="channels_last")

      out4 <- keras::layer_global_max_pooling_1d(data_format="channels_last")

      out <- keras::layer_concatenate(list(out1, out2, out3, out4))

    }     

    return(out)

  }

  if (is.null(global_pooling) | global_pooling == "flatten") {

    out <- keras::layer_flatten()

    return(out)

  }

  return(keras::layer_flatten())

}

f_reshape <- function(x, y, reshape_xy, reshape_x_bool, reshape_y_bool, reshape_sw_bool = FALSE, sw = NULL) {

  if (is.null(reshape_xy)) {

    return(list(X = x, Y = y, SW = sw))

  }

  if (reshape_sw_bool) {

    if (reshape_x_bool) {

      x_new <- reshape_xy$x(x = x, y = y, sw = sw)

    } else {

      x_new <- x

    }

    if (reshape_y_bool) {

      y_new <- reshape_xy$y(x = x, y = y, sw = sw)

    } else {

      y_new <- y

    }

    if (reshape_sw_bool) {

      sw_new <- reshape_xy$sw(x = x, y = y, sw = sw)

    } else {

      sw_new <- sw

    }

    return(list(X = x_new, Y = y_new, SW = sw_new))

  } else {

    if (reshape_x_bool) {

      x_new <- reshape_xy$x(x = x, y = y)

    } else {

      x_new <- x

    }

    if (reshape_y_bool) {

      y_new <- reshape_xy$y(x = x, y = y)

    } else {

      y_new <- y

    }

    #browser()

    return(list(X = x_new, Y = y_new))

  }

}

bal_acc_from_cm <- function(cm, verbose = TRUE) {

  class_acc <- list()

  names_list <- list()

  count <- 1

  for (i in 1:ncol(cm)) {

    v_col <- cm[,i]

    if (sum(v_col) == 0) next

    class_acc[count] <- v_col[i]/sum(v_col)

    names_list[count] <- colnames(cm)[i]

    count <- count + 1

  }

  class_acc <- unlist(class_acc)

  names(class_acc) <- unlist(names_list)

  if (verbose) print(class_acc)

  return(mean(class_acc))

}