context("predict")

test_that("Sucessful prediction", {
   
   #testthat::skip_if_not_installed("tensorflow")
   testthat::skip_if_not(reticulate::py_module_available("tensorflow"))
   
   sequence <- "AAACCNGGGTTT"
   maxlen <- 8
   filename <- tempfile(fileext = ".h5")
   
   model <- create_model_lstm_cnn(
      maxlen = maxlen,
      verbose = FALSE,
      layer_dense = 4,
      layer_lstm = 8)
   
   # test h5 output
   pred <- predict_model(layer_name = NULL, sequence = sequence,
                         filename = filename, step = 1,
                         batch_size = 1, 
                         return_states = TRUE,
                         verbose = FALSE,
                         output_type = "h5",
                         model = model,
                         mode = "label", 
                         include_seq = TRUE)
   
   expect_true(all(pred$states >= 0))
   expect_true(all(pred$states <= 1))
   expect_equal(pred$sample_end_position, 8:12)
   
   pred_h5 <- load_prediction(filename, get_sample_position = TRUE)
   expect_equal(pred_h5$states, pred$states)
   expect_equal(pred_h5$sample_end_position, pred$sample_end_position)
   
   # batch size bigger than number of samples
   pred2 <- predict_model(layer_name = NULL, sequence = sequence,
                          filename = NULL, step = 1,
                          batch_size = 100, 
                          return_states = TRUE,
                          verbose = FALSE,
                          output_type = "h5",
                          model = model,
                          mode = "label", 
                          include_seq = TRUE)
   
   expect_true(all(abs(pred$states - pred2$states) < 1e-06))
   expect_equal(pred$sample_end_position, pred2$sample_end_position)
   
   # test csv + padding maxlen + ... (nuc_dist)
   filename <- tempfile(fileext = ".csv")
   pred <- predict_model(layer_name = NULL, sequence = sequence,
                         filename = filename, step = 1,
                         batch_size = 2, 
                         return_states = TRUE,
                         padding = "maxlen",
                         verbose = FALSE,
                         output_type = "csv", model = model,
                         mode = "label", ambiguous_nuc = "empirical",
                         nuc_dist = c(0.1,0.4,0.4,0.1),
                         include_seq = TRUE)
   
   expect_true(all(pred$states >= 0))
   expect_true(all(pred$states <= 1))
   expect_equal(pred$sample_end_position, 0:12)
   
   pred_csv <- read.csv(filename)
   expect_equal(as.matrix(pred_csv), pred$states)
   
   # padding 
   pred <- predict_model(layer_name = NULL, sequence = "AAA",
                         filename = NULL, step = 2,
                         batch_size = 2, 
                         return_states = TRUE,
                         padding = "standard",
                         verbose = FALSE,
                         output_type = "csv", model = model,
                         mode = "label", 
                         include_seq = TRUE)
   
   expect_true(all(pred$states >= 0))
   expect_true(all(pred$states <= 1))
   expect_equal(pred$sample_end_position, 3)
   expect_equal(nrow(pred$states), length(pred$sample_end_position))
   
   # step
   pred <- predict_model(layer_name = NULL, sequence = "AAAAACCCCC",
                         filename = NULL, step = 2,
                         batch_size = 2, 
                         return_states = TRUE,
                         padding = "standard",
                         verbose = FALSE,
                         output_type = "csv", model = model,
                         mode = "label", 
                         include_seq = TRUE)
   
   expect_equal(pred$sample_end_position, c(8, 10))
   expect_equal(nrow(pred$states), length(pred$sample_end_position))
   
   # fasta file by_entry
   Sequence <- c("AAAACCCC", "TT", "AAACCCGGGTTT")
   Header <- letters[1:3]   
   df <- data.frame(Sequence, Header)
   fasta_path <- tempfile(fileext = ".fasta")
   microseq::writeFasta(df, fasta_path)
   output_path <- tempfile()
   dir.create(output_path)
   
   expect_message(
      predict_model(layer_name = NULL, 
                    path_input = fasta_path,
                    output_format = "by_entry",
                    output_dir = output_path,
                    filename = "states.h5",
                    step = 2,
                    batch_size = 2, 
                    padding = "none",
                    verbose = TRUE,
                    output_type = "h5",
                    model = model,
                    mode = "label", 
                    include_seq = TRUE)
   )
   
   h5_files <- list.files(output_path, full.names = TRUE)
   expect_true(basename(h5_files[1]) == "states_nr_1.h5")
   expect_true(basename(h5_files[2]) == "states_nr_3.h5")
   
   output_list_1 <- load_prediction(h5_files[1], get_sample_position = TRUE)
   expect_equal(output_list_1$sample_end_position, 8)
   output_list_2 <- load_prediction(h5_files[2], get_sample_position = TRUE)
   expect_equal(output_list_2$sample_end_position, c(8,10,12))
   
   # fasta file, by_entry
   h5_file <- tempfile(fileext = ".h5")
   pred <- predict_model(layer_name = NULL, 
                         path_input = fasta_path,
                         output_format = "by_entry_one_file",
                         filename = h5_file,
                         step = 2,
                         batch_size = 2, 
                         padding = "none",
                         verbose = FALSE,
                         output_type = "h5",
                         model = model,
                         mode = "label")
   
   output_list <- load_prediction(h5_file, get_sample_position = TRUE)
   expect_true(all(output_list[[1]]$states == output_list_1$states))
   expect_true(all(output_list[[1]]$sample_end_position == output_list_1$sample_end_position))
   expect_true(all(output_list[[2]]$states == output_list_2$states))
   expect_true(all(output_list[[2]]$sample_end_position == output_list_2$sample_end_position))
   
   # one pred per entry
   h5_file <- tempfile(fileext = ".h5")
   pred <- predict_model(layer_name = NULL, 
                         path_input = fasta_path,
                         output_format = "one_pred_per_entry",
                         filename = h5_file,
                         step = 2,
                         batch_size = 2, 
                         verbose = FALSE,
                         output_type = "h5",
                         model = model,
                         mode = "label")
   
   output_list <- load_prediction(h5_file)
   expect_equal(nrow(output_list$states),  nrow(df))
   
   # lm, target middle
   model <- create_model_lstm_cnn_target_middle(
      maxlen = maxlen,
      verbose = FALSE,
      layer_dense = 4,
      layer_lstm = 8)
   
   h5_file <- tempfile(fileext = ".h5")
   pred <- predict_model(layer_name = NULL, 
                         path_input = fasta_path,
                         output_format = "by_entry_one_file",
                         filename = h5_file,
                         step = 2,
                         target_len = 1,
                         batch_size = 2, 
                         padding = "standard",
                         verbose = FALSE,
                         output_type = "h5",
                         lm_format = "target_middle_lstm",
                         model = model,
                         mode = "lm")
   
   output_list <- load_prediction(h5_file, get_sample_position = TRUE)
   expect_equal(output_list[[1]]$sample_end_position, 8)
   expect_equal(output_list[[2]]$sample_end_position, 2)
   expect_equal(output_list[[3]]$sample_end_position[1], 9)
   
})