#' Create transformer model

#'

#' Creates transformer network for classification. Model can consist of several stacked attention blocks.

#' 

#' @inheritParams keras::layer_multi_head_attention

#' @inheritParams create_model_lstm_cnn

#' @param pos_encoding Either `"sinusoid"` or `"embedding"`. How to add positional information.

#' If `"sinusoid"`, will add sine waves of different frequencies to input.

#' If `"embedding"`, model learns positional embedding.

#' @param embed_dim Dimension for token embedding. No embedding if set to 0. Should be used when input is not one-hot encoded

#' (integer sequence).

#' @param head_size Dimensions of attention key.

#' @param n Frequency of sine waves for positional encoding. Only applied if `pos_encoding = "sinusoid"`.

#' @param ff_dim Units of first dense layer after attention blocks.

#' @param dropout Vector of dropout rates after attention block(s). 

#' @param dropout_dense Dropout for dense layers.

#' @param flatten_method How to process output of last attention block. Can be `"max_ch_first"`, `"max_ch_last"`, `"average_ch_first"`,

#' `"average_ch_last"`, `"both_ch_first"`, `"both_ch_last"`, `"all"`, `"none"` or `"flatten"`.

#' If `"average_ch_last"` /  `"max_ch_last"`  or `"average_ch_first"` / `"max_ch_first"`, will apply global average/max pooling.

#' `_ch_first` / `_ch_last` to decide along which axis. `"both_ch_first"` / `"both_ch_last"` to use max and average together. `"all"` to use all 4 

#' global pooling options together. If `"flatten"`, will flatten output after last attention block. If `"none"` no flattening applied.

#' @examples

#' 

#' maxlen <- 50

#' \donttest{

#' library(keras)

#' model <- create_model_transformer(maxlen = maxlen,

#'                                   head_size=c(10,12),

#'                                   num_heads=c(7,8),

#'                                   ff_dim=c(5,9),

#'                                   dropout=c(0.3, 0.5))

#' }

#' @returns A keras model implementing transformer architecture.

#' @export

create_model_transformer <- function(maxlen,

                                     vocabulary_size = 4,

                                     embed_dim = 64,

                                     pos_encoding = "embedding",

                                     head_size = 4L,

                                     num_heads = 5L,

                                     ff_dim = 8,

                                     dropout=0,

                                     n = 10000, # pos emb frequency

                                     layer_dense = 2,

                                     dropout_dense = NULL,

                                     flatten_method = "flatten",

                                     last_layer_activation = "softmax",

                                     loss_fn = "categorical_crossentropy",

                                     solver = "adam",

                                     learning_rate = 0.01,

                                     label_noise_matrix = NULL,

                                     bal_acc = FALSE,

                                     f1_metric = FALSE,

                                     auc_metric = FALSE,

                                     label_smoothing = 0,

                                     verbose = TRUE,

                                     model_seed = NULL,

                                     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_transformer, argg)

    })

    return(model)

  }

  stopifnot(length(head_size) == length(num_heads))

  stopifnot(length(head_size) == length(dropout))

  stopifnot(length(head_size) == length(ff_dim))

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

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

  stopifnot(pos_encoding %in% c("sinusoid", "embedding"))

  num_dense_layers <- length(layer_dense)

  head_size <- as.integer(head_size)

  num_heads <- as.integer(num_heads)

  maxlen <- as.integer(maxlen)

  num_attention_blocks <- length(num_heads)

  vocabulary_size <-  as.integer(vocabulary_size)

  if (!is.null(model_seed)) tensorflow::tf$random$set_seed(model_seed)

  if (embed_dim == 0) {

    input_tensor <- keras::layer_input(shape = c(maxlen, vocabulary_size))

  } else {

    input_tensor <- keras::layer_input(shape = c(maxlen))

  }

  # positional encoding 

  if (pos_encoding == "sinusoid") { 

    pos_enc_layer <- layer_pos_sinusoid_wrapper(maxlen = maxlen, vocabulary_size = vocabulary_size,

                                                n = n, embed_dim = embed_dim)

  } 

  if (pos_encoding == "embedding") { 

    pos_enc_layer <- layer_pos_embedding_wrapper(maxlen = maxlen, vocabulary_size = vocabulary_size,

                                                 embed_dim = embed_dim)

  }

  output_tensor <- input_tensor %>% pos_enc_layer

  # attention blocks

  for (i in 1:num_attention_blocks) {

    attn_block <- layer_transformer_block_wrapper(

      num_heads = num_heads[i],

      head_size = head_size[i],

      dropout_rate = dropout[i],

      ff_dim = ff_dim[i],

      embed_dim = embed_dim,

      vocabulary_size = vocabulary_size,

      load_r6 = FALSE)

    output_tensor <- output_tensor %>% attn_block

  }

  if (flatten_method != "none") {

    output_tensor <- pooling_flatten(global_pooling = flatten_method, output_tensor = output_tensor)

  }

  # dense layers

  if (num_dense_layers > 1) {

    for (i in 1:(num_dense_layers - 1)) {

      output_tensor <- output_tensor %>% keras::layer_dense(units = layer_dense[i], activation = "relu")

      if (!is.null(dropout_dense)) {

        output_tensor <- output_tensor %>% keras::layer_dropout(rate = dropout_dense[i])

      }

    }

  }  

  output_tensor <- output_tensor %>% keras::layer_dense(units = layer_dense[length(layer_dense)], activation = last_layer_activation, dtype = "float32")

  # create model

  model <- keras::keras_model(inputs = input_tensor, outputs = output_tensor)

  model <- compile_model(model = model, label_smoothing = label_smoothing, layer_dense = layer_dense,

                         solver = solver, learning_rate = learning_rate, loss_fn = loss_fn, 

                         num_output_layers = 1, label_noise_matrix = label_noise_matrix,

                         bal_acc = bal_acc, f1_metric = f1_metric, auc_metric = auc_metric)

  if (verbose) print(model)

  argg <- c(as.list(environment()))

  model <- add_hparam_list(model, argg)

  reticulate::py_set_attr(x = model, name = "hparam", value = model$hparam)

  model

}