#' @title Create GenomeNet Model with Given Architecture Parameters
#'
#' @param maxlen (integer `numeric(1)`)\cr
#'   Input sequence length.
#' @param learning_rate (`numeric(1)`)\cr
#'   Used by the `keras` optimizer that is specified by `optimizer`.
#' @param number_of_cnn_layers (integer `numeric(1)`)\cr
#'   Target number of CNN-layers to use in total. If `number_of_cnn_layers` is
#'   greater than `conv_block_count`, then the effective number of CNN layers
#'   is set to the closest integer that is divisible by `conv_block_count`.
#' @param conv_block_count (integer `numeric(1)`)\cr
#'   Number of convolutional blocks, into which the CNN layers are divided.
#'   If this is greater than `number_of_cnn_layers`, then it is set to
#'   `number_of_cnn_layers` (the convolutional block size will then be 1).\cr
#'   Convolutional blocks are used when `model_type` is `"gap"` (the output of
#'   the last `conv_block_count * (1 - skip_block_fraction)` blocks is
#'   fed to global average pooling and then concatenated), and also when
#'   `residual_block` is `TRUE` (the number of filters is held constant within
#'   blocks). If neither of these is the case, `conv_block_count` has little
#'   effect besides the fact that `number_of_cnn_layers` is set to the closest
#'   integer divisible by `conv_block_count`.
#' @param kernel_size_0 (`numeric(1)`)\cr
#'   Target CNN kernel size of the first CNN-layer. Although CNN kernel size is
#'   always an integer, this value can be non-integer, potentially affecting
#'   the kernel-sizes of intermediate layers (which are geometrically
#'   interpolated between `kernel_size_0` and `kernel_size_end`).
#' @param kernel_size_end (`numeric(1)`)\cr
#'   Target CNN kernel size of the last CNN-layer; ignored if only one
#'   CNN-layer is used (i.e. if `number_of_cnn_layers` is 1). Although CNN
#'   kernel size is always an integer, this value can be non-integer,
#'   potentially affecting the kernel-sizes of intermediate layers (which are
#'   geometrically interpolated between `kernel_size_0` and `kernel_size_end`).
#' @param filters_0 (`numeric(1)`)\cr
#'   Target filter number of the first CNN-layer. Although CNN filter number is
#'   always an integer, this value can be non-integer, potentially affecting
#'   the filter-numbers of intermediate layers (which are geometrically
#'   interpolated between `filters_0` and `filters_end`).\cr
#'   Note that filters are constant within convolutional blocks when
#'   `residual_block` is `TRUE`.
#' @param filters_end (`numeric(1)`)\cr
#'   Target filter number of the last CNN-layer; ignored if only one CNN-layer
#'   is used (i.e. if `number_of_cnn_layers` is 1). Although CNN filter number
#'   is always an integer, this value can be non-integer, potentially affecting
#'   the filter-numbers of intermediate dilation_rates layers (which are geometrically
#'   interpolated between `kernel_size_0` and `kernel_size_end`).\cr
#'   Note that filters are constant within convolutional blocks when
#'   `residual_block` is `TRUE`.
#' @param dilation_end (`numeric(1)`)\cr
#'   Dilation of the last CNN-layer *within each block*. Dilation rates within
#'   each convolutional block grows exponentially from 1 (no dilation) for the
#'   first CNN-layer to each block, to this value. Set to 1 (default) to
#'   disable dilation.
#' @param max_pool_end (`numeric(1)`)\cr
#'   Target total effective pooling of CNN part of the network. "Effective
#'   pooling" here is the product of the pooling rates of all previous
#'   CNN-layers. A network with three CNN-layers, all of which are followed
#'   by pooling layers of size 2, therefore has effective pooling of 8, with
#'   the effective pooling at intermediate positions being 1 (beginning), 2,
#'   and 4. Effective pooling after each layer is set to the power of 2 that is,
#'   on a logarithmic scale, closest to
#'   `max_pool_end ^ (<CNN layer number> / <total number of CNN layers>)`.
#'   Therefore, even though the total effective pooling size of the whole
#'   CNN part of the network will always be a power of 2, having different,
#'   possibly non-integer values of `max_pool_end`, will still lead to
#'   different networks.
#' @param dense_layer_num (integer `numeric(1)`)\cr
#'   number of dense layers at the end of the network, not counting the output
#'   layer.
#' @param dense_layer_units (integer `numeric(1)`)\cr
#'   Number of units in each dense layer, except for the output layer.
#' @param dropout (`numeric(1)`)\cr
#'   Dropout rate of dense layers, except for the output layer.
#' @param batch_norm_momentum (`numeric(1)`)\cr
#'   `momentum`-parameter of `layer_batch_normalization` layers used in the
#'   convolutional part of the network.
#' @param leaky_relu_alpha (`numeric(1)`)\cr
#'   `alpha`-parameter of the `layer_activation_leaky_relu` activation layers
#'   used in the convolutional part of the network.
#' @param dense_activation (`character(1)`)\cr
#'   Which activation function to use for dense layers. Should be one of
#'   `"relu"`, `"sigmoid"`, or `"tanh"`.
#' @param skip_block_fraction (`numeric(1)`)\cr
#'   What fraction of the first convolutional blocks to skip.
#'   Only used when `model_type` is `"gap"`.
#' @param residual_block (`logical(1)`)\cr
#'   Whether to use residual layers in the convolutional part of the network.
#' @param reverse_encoding (`logical(1)`)\cr
#'   Whether the network should have a second input for reverse-complement
#'   sequences.
#' @param optimizer (`character(1)`)\cr
#'   Which optimizer to use. One of `"adam"`, `"adagrad"`, `"rmsprop"`, or `"sgd"`.
#' @param model_type (`character(1)`)\cr
#'   Whether to use the global average pooling (`"gap"`) or recurrent
#'   (`"recurrent"`) model type.
#' @param recurrent_type (`character(1)`)\cr
#'   Which recurrent network type to use. One of `"lstm"` or `"gru"`.
#'   Only used when `model_type` is `"recurrent"`.
#' @param recurrent_layers (integer `numeric(1)`)\cr
#'   Number of recurrent layers.
#'   Only used when `model_type` is `"recurrent"`.
#' @param recurrent_bidirectional (`logical(1)`)\cr
#'   Whether to use bidirectional recurrent layers.
#'   Only used when `model_type` is `"recurrent"`.
#' @param recurrent_units (integer `numeric(1)`)\cr
#'   Number of units in each recurrent layer.
#'   Only used when `model_type` is `"recurrent"`.
#' @param vocabulary_size (integer `numeric(1)`)\cr
#'   Vocabulary size of (one-hot encoded) input strings. This determines the
#'   input tensor shape, together with `maxlen`.
#' @param last_layer_activation Either `"sigmoid"` or `"softmax"`.
#' @param loss_fn Either `"categorical_crossentropy"` or `"binary_crossentropy"`. If `label_noise_matrix` given, will use custom `"noisy_loss"`.
#' @param auc_metric Whether to add AUC metric.
#' @param num_targets (integer `numeric(1)`)\cr
#'   Number of output units to create.
#' @return A keras model.
#' @inheritParams create_model_lstm_cnn
#' @examplesIf reticulate::py_module_available("tensorflow")
#' model <- create_model_genomenet()
#' model
#' 
#' @returns A keras model implementing genomenet architecture.
#' @export
create_model_genomenet <- function(
    maxlen = 300,
    learning_rate = 0.001,
    number_of_cnn_layers = 1,
    conv_block_count = 1,
    kernel_size_0 = 16,
    kernel_size_end = 16,
    filters_0 = 256,
    filters_end = 512,
    dilation_end = 1,
    max_pool_end = 1,
    dense_layer_num = 1,
    dense_layer_units = 100,
    dropout_lstm = 0,
    dropout = 0,
    batch_norm_momentum = 0.8,
    leaky_relu_alpha = 0,
    dense_activation = "relu",
    skip_block_fraction = 0,
    residual_block = FALSE,
    reverse_encoding = FALSE,
    optimizer = "adam",
    model_type = "gap",
    recurrent_type = "lstm",
    recurrent_layers = 1,
    recurrent_bidirectional = FALSE,
    recurrent_units = 100,
    vocabulary_size = 4,
    last_layer_activation = "softmax",
    loss_fn = "categorical_crossentropy",
    auc_metric = FALSE,
    num_targets = 2,
    model_seed = NULL,
    bal_acc = FALSE,
    f1_metric = FALSE,
    mixed_precision = FALSE,
    mirrored_strategy = NULL) {
  
  if (mixed_precision) tensorflow::tf$keras$mixed_precision$set_global_policy("mixed_float16")
  
  if (is.null(mirrored_strategy)) mirrored_strategy <- ifelse(count_gpu() > 1, TRUE, FALSE)
  if (mirrored_strategy) {
    mirrored_strategy <- tensorflow::tf$distribute$MirroredStrategy()
    with(mirrored_strategy$scope(), { 
      argg <- as.list(environment())
      argg$mirrored_strategy <- FALSE
      model <- do.call(create_model_genomenet, argg)
    })
    return(model)
  }
  
  if (!is.null(model_seed)) tensorflow::tf$random$set_seed(model_seed)
  stopifnot(maxlen > 0 & maxlen %% 1 == 0)
  stopifnot(learning_rate > 0)
  stopifnot(number_of_cnn_layers >= 1 & number_of_cnn_layers %% 1 == 0)
  stopifnot(conv_block_count >= 1 & conv_block_count %% 1 == 0)
  stopifnot(kernel_size_0 > 0)
  stopifnot(kernel_size_end > 0)
  stopifnot(filters_0 > 0)
  stopifnot(filters_end > 0)
  stopifnot(dilation_end >= 1)
  stopifnot(max_pool_end >= 1)
  stopifnot(dense_layer_num >= 0 & dense_layer_num %% 1 == 0)
  stopifnot(dense_layer_units >= 0 & dense_layer_units %% 1 == 0)
  stopifnot(0 <= dropout_lstm & dropout_lstm <= 1)
  stopifnot(0 <= batch_norm_momentum & batch_norm_momentum <= 1)
  
  stopifnot(0 <= leaky_relu_alpha& leaky_relu_alpha <= 1)
  dense_activation <- match.arg(dense_activation, c("relu", "sigmoid", "tanh"))
  stopifnot(0 <= skip_block_fraction & skip_block_fraction <= 1)
  
  model_type = match.arg(model_type, c("gap", "recurrent"))
  
  stopifnot(isTRUE(residual_block) || isFALSE(residual_block))
  stopifnot(isTRUE(residual_block) || isFALSE(residual_block))
  
  optimizer <- match.arg(optimizer, c("adam", "adagrad", "rmsprop", "sgd"))
  recurrent_type <- match.arg(recurrent_type, c("lstm", "gru"))
  stopifnot(recurrent_layers >= 1 & recurrent_layers %% 1 == 0)
  stopifnot(isTRUE(recurrent_bidirectional) || isFALSE(recurrent_bidirectional))
  stopifnot(recurrent_units >= 1 && recurrent_units %% 1 == 0)
  
  stopifnot(vocabulary_size >= 2 & vocabulary_size %% 1 == 0)
  stopifnot(num_targets >= 2 & num_targets %% 1 == 0)
  
  if (number_of_cnn_layers < conv_block_count){
    conv_block_size <- 1
    conv_block_count <- number_of_cnn_layers
  } else {
    conv_block_size <- round(number_of_cnn_layers / conv_block_count)
    number_of_cnn_layers <- conv_block_size * conv_block_count
  }
  
  if (residual_block) {
    filters_exponent <- rep(seq(from = 0, to = 1, length.out = conv_block_count), each = conv_block_size)
  } else {
    filters_exponent <- seq(from = 0, to = 1, length.out = number_of_cnn_layers)
  }
  
  filters <- ceiling(filters_0 * (filters_end / filters_0) ^ filters_exponent)
  
  kernel_size <- ceiling(kernel_size_0 * (kernel_size_end / kernel_size_0) ^ seq(from = 0, to = 1, length.out = number_of_cnn_layers))
  
  dilation_rates <- round(dilation_end ^ seq(0, 1, length.out = conv_block_size))
  dilation_rates <- rep(dilation_rates, conv_block_count)
  
  max_pool_divider <- round(log2(max_pool_end) * seq(0, 1, length.out = number_of_cnn_layers + 1))
  max_pool_array <- 2 ^ diff(max_pool_divider)
  
  input_tensor <- keras::layer_input(shape = c(maxlen, vocabulary_size))
  output_tensor <- input_tensor
  
  output_collection <- list()
  
  for (i in seq_len(number_of_cnn_layers)) {
    layer <- keras::layer_conv_1d(kernel_size = kernel_size[i],
                                  padding = "same",
                                  activation = "linear",
                                  filters = filters[i],
                                  dilation_rate = dilation_rates[i])
    
    output_tensor <- output_tensor %>% layer
    
    if (model_type == "gap" && i %% conv_block_size == 0) {
      output_collection[[length(output_collection) + 1]] <- keras::layer_global_average_pooling_1d(output_tensor)
    }
    
    if (max_pool_array[i] > 1) {
      layer <- keras::layer_max_pooling_1d(pool_size = max_pool_array[i], padding = "same")
      output_tensor <- output_tensor %>% layer
    }
    
    if (residual_block) {
      if (i > 1) {
        if (max_pool_array[i] > 1) {
          layer <- keras::layer_average_pooling_1d(pool_size = max_pool_array[i], padding = "same")
          residual_layer <- residual_layer %>% layer
        }
        if (filters[i - 1] != filters[i]) {
          layer <- keras::layer_conv_1d(kernel_size = 1,
                                        padding = "same",
                                        activation = "linear",
                                        filters = filters[i]
          )
          residual_layer <- residual_layer %>% layer
        }
        
        output_tensor <- keras::layer_add(list(output_tensor, residual_layer))
      }
      
      residual_layer <- output_tensor
    }
    
    layer <- keras::layer_batch_normalization(momentum = batch_norm_momentum)
    output_tensor <- output_tensor %>% layer
    
    layer <- keras::layer_activation_leaky_relu(alpha = leaky_relu_alpha)
    
    output_tensor <- output_tensor %>% layer
  }
  
  if (model_type == "gap") {
    # skip 'skip_block_fraction' of outputs we collected --> use the last (1 - skip_block_fraction) part of them
    use_blocks <- ceiling((1 - skip_block_fraction) * length(output_collection))
    use_blocks <- max(use_blocks, 1)
    
    output_collection <- utils::tail(output_collection, use_blocks)
    
    # concatenate outputs from blocks (that we are using)
    if (length(output_collection) > 1) {
      output_tensor <- keras::layer_concatenate(output_collection)
    } else {
      output_tensor <- output_collection[[1]]
    }
  } else {
    # recurrent model
    
    recurrent_layer_constructor = switch(recurrent_type,
                                         lstm = keras::layer_lstm,
                                         gru = keras::layer_gru
    )
    
    for (i in seq_len(recurrent_layers)) {
      if (recurrent_bidirectional) {
        layer <- keras::bidirectional(
          layer = recurrent_layer_constructor(units = recurrent_units, return_sequences = (i != recurrent_layers)))
      } else {
        layer <- recurrent_layer_constructor(units = recurrent_units, return_sequences = (i != recurrent_layers))
      }
      output_tensor <- output_tensor %>% layer
    }
  }
  
  for (i in seq_len(dense_layer_num)) {
    layer <- keras::layer_dropout(rate = dropout)
    output_tensor <- output_tensor %>% layer
    layer <- keras::layer_dense(units = dense_layer_units, activation = dense_activation)
    output_tensor <- output_tensor %>% layer
  }
  
  output_tensor <- output_tensor %>%
    keras::layer_dense(units = num_targets, activation = last_layer_activation, dtype = "float32")
  
  # define "model" as the mapping from input_tensor to output_tensor
  model <- keras::keras_model(inputs = input_tensor, outputs = output_tensor)
  
  if (reverse_encoding) {
    input_tensor_reversed <- keras::layer_input(shape = c(maxlen, vocabulary_size))
    
    # define "output_tensor_reversed" as what comes out of input_tensor_reversed when model() is applied to it
    output_tensor_reversed <- model(input_tensor_reversed)
    
    # define a new model: model from above (with input_tensor, output_tensor), and
    model <- keras::keras_model(
      inputs = c(input_tensor, input_tensor_reversed),
      outputs = keras::layer_average(c(output_tensor, output_tensor_reversed))
    )
  }
  
  # assign optimization method
  keras_optimizer <- set_optimizer(optimizer, learning_rate) 
  
  #add metrics
  if (loss_fn == "binary_crossentropy") {
    model_metrics <- c(tensorflow::tf$keras$metrics$BinaryAccuracy(name = "acc"))
  } else {
    model_metrics <- c("acc")
  } 
  
  cm_dir <-
    file.path(tempdir(), paste(sample(letters, 7), collapse = ""))
  while (dir.exists(cm_dir)) {
    cm_dir <-
      file.path(tempdir(), paste(sample(letters, 7), collapse = ""))
  }
  dir.create(cm_dir)
  model$cm_dir <- cm_dir
  
  if (loss_fn == "categorical_crossentropy") {
    if (bal_acc) {
      macro_average_cb <- balanced_acc_wrapper(num_targets, cm_dir)
      model_metrics <- c(macro_average_cb, "acc")
    }
    
    if (f1_metric) {
      f1 <- f1_wrapper(num_targets)
      model_metrics <- c(model_metrics, f1)
    }
  }
  
  if (auc_metric) {
    auc <-
      auc_wrapper(model_output_size = dense_layer_units[length(dense_layer_units)],
                  loss = loss_fn)
    model_metrics <- c(model_metrics, auc)
  }
  
  model %>% keras::compile(loss = loss_fn, optimizer = keras_optimizer, metrics = model_metrics)
  
  argg <- c(as.list(environment()))
  model <- add_hparam_list(model, argg)
  
  model
}


# #' Load pretrained Genomenet model
# #'
# #' Classification model with labels "bacteria", "virus-no-phage","virus-phage".
# #' TODO: add link to paper
# #'
# #' @inheritParams create_model_lstm_cnn
# #' @param maxlen Model input size. Either 150 or 10000.
# #' @param learning_rate Learning rate for optimizer. If compile is TRUE and learning_rate is NULL,
# #' will use learning rate from previous training.
# #' @export
# load_model_self_genomenet <- function(maxlen, compile = FALSE, optimizer = "adam",
#                                       learning_rate = NULL) {
#   
#   stopifnot(any(maxlen == c(150,10000)))
#   
#   if (maxlen == 150) {
#     load(model_self_genomenet_maxlen_150)
#     model <- keras::unserialize_model(model_self_genomenet_maxlen_150, compile = FALSE)
#   }
#   
#   if (maxlen == 10000) {
#     load("data/self_genomenet_model_maxlen_10k.rda")
#     #data(model_self_genomenet_maxlen_10k)
#     model <- keras::unserialize_model(model_self_genomenet_maxlen_10k, compile = FALSE)
#   }
#   
#   if (is.null(learning_rate)) {
#     if (maxlen == 150) learning_rate <- 0.00039517784549691
#     if (maxlen == 10000) learning_rate <- 8.77530464905713e-05
#   }
#   
#   if (compile) {
#     keras_optimizer <- set_optimizer(optimizer, learning_rate)
#     model %>% keras::compile(loss = "categorical_crossentropy", optimizer = keras_optimizer, metrics = "acc")
#   }
#   
#   return(model)
# }