fit_GAN <- function(prot_seq,
                    label = NULL,
                    length_seq,
                    embedding_dim,
                    embedding_args = list(),
                    latent_dim = NULL,
                    intermediate_generator_layers,
                    intermediate_discriminator_layers,
                    prot_seq_val = NULL,
                    label_val = NULL,
                    epochs,
                    batch_size,
                    preprocessing = list(
                        x_train = NULL,
                        x_val = NULL,
                        y_train = NULL,
                        y_val = NULL,
                        lenc = NULL,
                        length_seq = NULL,
                        num_seq = NULL,
                        embedding_dim = NULL,
                        embedding_matrix = NULL,
                        removed_prot_seq = NULL,
                        removed_prot_seq_val = NULL,
                        latent_dim = NULL),
                    optimizer = "adam",
                    validation_split = 0) {
    result <- NULL
    result$preprocessing <- NULL
    
    ### pre-processing
    tryCatch({
        if (all(unlist(lapply(preprocessing, Negate(is.null))))) {
            result$preprocessing <- preprocessing
            x_train <- preprocessing$x_train
            if (is.null(x_train)) stop("check embedded sequence data for train in the preprocessed result")
            y_train <- preprocessing$y_train
            x_val <- preprocessing$x_val
            y_val <- preprocessing$y_val
            length_seq <- preprocessing$length_seq
            if (is.null(length_seq)) stop("check length of sequence in the preprocessed result")
            num_seq <- preprocessing$num_seq
            if (is.null(num_seq)) stop("check number of sequences for train in the preprocessed result")
            num_seq_val <- preprocessing$num_seq_val
            embedding_dim <- preprocessing$embedding_dim
            if (is.null(embedding_dim)) stop("check dimension of the dense embedding in the preprocessed result")
            embedding_matrix <- preprocessing$embedding_matrix
            if (is.null(embedding_matrix)) stop("check embedding matrix in the preprocessed result")
            if (is.null(latent_dim)) latent_dim <- preprocessing$latent_dim
        }
        
        if (any(unlist(lapply(preprocessing, is.null)))) {
            message("pre-processing...")
            
            # check
            checked_prot_seq <- prot_seq_check(prot_seq = prot_seq, label = label)
            prot_seq <- checked_prot_seq$prot_seq
            label <- checked_prot_seq$label
            result$preprocessing$removed_prot_seq <- checked_prot_seq$removed_prot_seq
            
            # sequence to vector
            prot2vec_train <- do.call(prot2vec,
                                    c(list(prot_seq = prot_seq, embedding_dim = embedding_dim),
                                        embedding_args))
            x_train <- prot2vec_train$prot_vec
            num_seq <- dim(x_train)[1]
            y_train <- label
            embedding_matrix <- prot2vec_train$embedding_matrix
            
            x_val <- NULL
            y_val <- NULL
            num_seq_val <- NULL
            if (!is.null(prot_seq_val)) {
                # check
                checked_prot_seq_val <- prot_seq_check(prot_seq = prot_seq_val, label = label_val)
                prot_seq_val <- checked_prot_seq_val$prot_seq
                label_val <- checked_prot_seq_val$label
                result$preprocessing$removed_prot_seq_val <- checked_prot_seq_val$removed_prot_seq
                
                # sequence to vector
                prot2vec_val <- prot2vec(prot_seq = prot_seq_val,
                                        embedding_matrix = embedding_matrix)
                x_val <- prot2vec_val$prot_vec
                y_val <- label_val
                num_seq_val <- dim(x_val)[1]
            }
            
            if (is.null(prot_seq_val) & validation_split) {
                x <- x_train
                idx <- sample(seq_len(nrow(x)))
                train_idx <- seq_len(nrow(x)) %in%
                    idx[seq_len(round(nrow(x) * (1 - validation_split)))]
                x_train <- x[train_idx,,]
                x_val <- x[!train_idx,,]
                if (!is.null(y_train) & is.null(y_val)) {
                    y <- y_train
                    y_train <- y[train_idx]
                    y_val <- y[!train_idx]
                }
                num_seq <- dim(x_train)[1]
                num_seq_val <- dim(x_val)[1]
            }
        }
    },
    error = function(e)
        message("In pre-processing : ", e)
    )
    
    
    ### building model
    tryCatch({
        # GAN
        g_model <- function(latent_dim, label) {
            generator_input <- layer_input(shape = c(latent_dim))
            if (!is.null(label)) {
                input_latent <- layer_input(shape = c(latent_dim))
                input_label <- layer_input(shape = c(1))
                embedding_label <- input_label %>%
                    layer_embedding(input_dim = length(unique(label)), output_dim = latent_dim) %>%
                    layer_flatten()
                generator_input <- layer_multiply(list(input_latent, embedding_label))
            }
            
            x <- generator_input
            for (i in seq_len(length(intermediate_generator_layers))) {
                x <- intermediate_generator_layers[[i]](x)
            }
            
            generator_output <- x %>%
                layer_dense(units = length_seq * embedding_dim) %>%
                layer_reshape(c(length_seq, embedding_dim))
            
            if (is.null(label)) {
                generator <- keras_model(generator_input, generator_output)
            } else {
                generator <- keras_model(list(input_latent, input_label), generator_output)
            }
        }
        
        d_model <- function(length_seq, embedding_dim, label, optimizer) {
            discriminator_input <- layer_input(shape = c(length_seq, embedding_dim))
            discriminator_output <- discriminator_input %>%
                layer_flatten(input_shape = c(length_seq, embedding_dim))
            
            x <- discriminator_output
            for (i in seq_len(length(intermediate_discriminator_layers))) {
                x <- intermediate_discriminator_layers[[i]](x)
            }
            
            if (is.null(label)) {
                discriminator_output <- x %>%
                    layer_dense(units = 1, activation = "sigmoid")
                discriminator <- keras_model(discriminator_input, discriminator_output)
            } else {
                fake <- discriminator_output %>% 
                    layer_dense(units = 1, activation = "sigmoid", name = "generation")
                aux <- discriminator_output %>%
                    layer_dense(units = length(unique(label)), activation = "softmax", name = "auxiliary")
                discriminator <- keras_model(discriminator_input, list(fake, aux))
            }
            
            if (is.null(label)) {
                discriminator %>% compile(
                    optimizer = optimizer,
                    loss = "binary_crossentropy"
                )        
            } else {
                discriminator %>% compile(
                    optimizer = optimizer,
                    loss = list("binary_crossentropy", "sparse_categorical_crossentropy")
                )
            }
        }
        
        combined_model <- function(g_model, d_model, label, optimizer) {
            freeze_weights(d_model)
            
            if (is.null(label)) {
                gan <- keras_model_sequential() %>% g_model %>% d_model
            } else {
                input_latent <- layer_input(shape = c(latent_dim))
                input_label <- layer_input(shape = c(1))
                results <- g_model(list(input_latent, input_label)) %>% d_model
                gan <- keras_model(list(input_latent, input_label), results)
            }
            
            gan_optimizer <- optimizer_adam()
            
            if (is.null(label)) {
                gan %>% compile(
                    optimizer = optimizer, 
                    loss = list("binary_crossentropy")
                )
            } else {
                gan %>% compile(
                    optimizer = optimizer,
                    loss = list("binary_crossentropy", "sparse_categorical_crossentropy")
                )
            }
        }
        
        generator <- g_model(latent_dim, y_train)
        discriminator <- d_model(length_seq, embedding_dim, y_train, optimizer)
        combined <- combined_model(generator, discriminator, y_train, optimizer)
    },
        error = function(e)
        message("In building model : ", e)
    )
    
    
    ### training
    tryCatch({
        message("training...")
        if (!is.null(y_train)) {
            if (any(unlist(lapply(preprocessing, is.null)))) {
                lenc <- CatEncoders::LabelEncoder.fit(y_train)
                y_train <- CatEncoders::transform(lenc, y_train) - 1
                result$preprocessing$lenc <- lenc
                if (!is.null(y_val)) {
                    y_val <- CatEncoders::transform(lenc, y_val) - 1
                }
            }
        }
        
        # epoch
        for (epoch in seq_len(epochs)) {
            epoch_gen_loss <- NULL
            epoch_disc_loss <- NULL
            
            num_batches <- trunc(num_seq/batch_size)
            possible_indexes <- seq_len(num_seq)
            
            cat(sprintf("Epoch %s/%s \n", epoch, epochs))
            progbar <- tensorflow::tf$keras$utils$Progbar(target = num_batches)
            
            # batch
            for (index in seq_len(num_batches)) {
                batch <- sample(possible_indexes, size = batch_size)
                possible_indexes <- possible_indexes[!possible_indexes %in% batch]
                
                
                # train discriminator
                noise <- matrix(stats::rnorm(batch_size * latent_dim), 
                                nrow = batch_size, ncol = latent_dim)
                x_train_batch <- x_train[batch,,,drop = FALSE]
                
                if (!is.null(y_train)) {
                    y_train_batch <- y_train[batch]
                    y_train_sample <- sample(unique(y_train), batch_size, replace = TRUE)
                    aux_y <- matrix(c(y_train_sample, y_train_batch), ncol = 1)
                    x_gen_batch <- predict(generator, list(noise, y_train_sample))
                } else {
                    x_gen_batch <- predict(generator, noise)
                }
                
                X <- array(0, dim = c(2 * batch_size, length_seq, embedding_dim))
                X[seq(batch_size),,] <- x_gen_batch
                X[seq(batch_size + 1, 2 * batch_size),,] <- x_train_batch
                
                fake_real <- c(rep(0, batch_size), rep(1, batch_size))
                
                if (!is.null(y_train)) {
                    disc_loss <- train_on_batch(
                        discriminator,
                        x = X,
                        y = list(fake_real, aux_y)
                    )
                } else {
                    disc_loss <- train_on_batch(
                        discriminator,
                        x = X,
                        y = fake_real
                    )
                }
                epoch_disc_loss <- rbind(epoch_disc_loss, unlist(disc_loss))
                
                
                # train generator
                noise <- matrix(stats::rnorm(2 * batch_size * latent_dim), 
                                nrow = 2 * batch_size, ncol = latent_dim)
                
                if (!is.null(y_train)) {
                    y_train_sample <- matrix(
                        sample(unique(y_train), size = 2 * batch_size, replace = TRUE),
                        ncol = 1)
                }
                
                misleading_real <- rep(1, 2 * batch_size)
                
                if (!is.null(y_train)) {
                    combined_loss <- train_on_batch(
                        combined,
                        list(noise, y_train_sample),
                        list(misleading_real, y_train_sample)
                    )
                } else {
                    combined_loss <- train_on_batch(
                        combined,
                        noise,
                        misleading_real)
                }
                epoch_gen_loss <- rbind(epoch_gen_loss, unlist(combined_loss))
                progbar$add(1)
            }
            
            discriminator_train_loss <- apply(epoch_disc_loss, 2, mean)
            generator_train_loss <- apply(epoch_gen_loss, 2, mean)
            
            if (!is.null(x_val)) {
                
                # evaluate discriminator
                noise <- matrix(stats::rnorm(num_seq_val * latent_dim), 
                                nrow = num_seq_val, ncol = latent_dim)
                
                if (!is.null(y_val)) {
                    y_val_sample <- sample(unique(y_val), size = num_seq_val, replace = TRUE)
                    aux_y <- matrix(c(y_val_sample, y_val), ncol = 1)
                    x_gen <- predict(generator, list(noise, y_val_sample))
                } else {
                    x_gen <- predict(generator, noise)
                }
                
                X <- array(0, dim = c(2 * num_seq_val, length_seq, embedding_dim))
                X[seq(num_seq_val),,] <- x_gen
                X[seq(num_seq_val + 1, 2 * num_seq_val),,] <- x_val
                
                fake_real <- c(rep(0, num_seq_val), rep(1, num_seq_val))
                
                if (!is.null(y_val)) {
                    discriminator_test_loss <- evaluate(
                        discriminator,
                        x = X,
                        y = list(fake_real, aux_y),
                        verbose = FALSE
                    ) %>% unlist()
                } else {
                    discriminator_test_loss <- evaluate(
                        discriminator,
                        x = X,
                        y = fake_real,
                        verbose = FALSE
                    ) %>% unlist()
                }
                
                
                # evaluate generator
                noise <- matrix(stats::rnorm(2 * num_seq_val * latent_dim), 
                                nrow = 2 * num_seq_val, ncol = latent_dim)
                
                if (!is.null(y_val)) {
                    y_val_sample <- matrix(
                        sample(unique(y_val), size = 2 * num_seq_val, replace = TRUE),
                        ncol = 1)
                }
                
                misleading_real <- rep(1, 2 * num_seq_val)
                
                if (!is.null(y_val)) {
                    generator_test_loss <- combined %>% evaluate(
                        list(noise, y_val_sample),
                        list(misleading_real, y_val_sample),
                        verbose = FALSE
                    )
                } else {
                    generator_test_loss <- combined %>% evaluate(
                        noise,
                        misleading_real,
                        verbose = FALSE
                    )
                }
            }
            
            
            # performance
            if (length(discriminator$metrics_names) == 1) {
                row_fmt <- "%s : loss %f \n"
            } else {
                row_fmt <- "%s : loss %f | generation_loss %f | auxiliary_loss %f \n"
            }
            
            cat(do.call(sprintf, c(list(row_fmt, "generator (train)"), generator_train_loss)))
            if (!is.null(x_val)) {
                cat(do.call(sprintf, c(list(row_fmt, "generator (test)"), generator_test_loss)))
            }
            cat(do.call(sprintf, c(list(row_fmt, "discriminator (train)"), discriminator_train_loss)))
            if (!is.null(x_val)) {
                cat(do.call(sprintf, c(list(row_fmt, "discriminator (test)"), discriminator_test_loss)))
            }
            cat("\n")
        }
    },
    error = function(e)
        message("In training : ", e)
    )
    
    result$model <- combined
    result$generator <- generator
    result$discriminator <- discriminator
    result$preprocessing$x_train <- x_train
    result$preprocessing$y_train <- y_train
    result$preprocessing$x_val <- x_val
    result$preprocessing$y_val <- y_val
    result$preprocessing$length_seq <- length_seq
    result$preprocessing$num_seq <- num_seq
    result$preprocessing$num_seq_val <- num_seq_val
    result$preprocessing$embedding_dim <- embedding_dim
    result$preprocessing$embedding_matrix <- embedding_matrix
    result$preprocessing$latent_dim <- latent_dim
    result
}





gen_GAN <- function(x, label = NULL, num_seq, remove_gap = TRUE) {
    result <- NULL
    latent_dim <- x$preprocessing$latent_dim
    if (is.null(latent_dim)) stop("check dimension of latent vector in the fitted result")
    
    ### generating
    tryCatch({
        message("generating...")
        if (is.null(x$preprocessing$y_train)) {
            gen_vec <- predict(x$generator,
                            matrix(stats::rnorm(num_seq * latent_dim),
                                    nrow = num_seq, ncol = latent_dim))
        } else {
            lenc <- x$preprocessing$lenc
            lable_enc <- CatEncoders::transform(lenc, label) - 1
            gen_vec <- predict(x$generator,
                            list(matrix(stats::rnorm(num_seq * latent_dim),
                                        nrow = num_seq, ncol = latent_dim),
                                        lable_enc))
        }
    },
        error = function(e)
        message("In generating : ", e)
    )
    
    
    ### post-processing
    tryCatch({
        message("post-processing...")
        # vector to sequence
        gen_seq <- vec2prot(gen_vec, x$preprocessing$embedding_matrix)
        
        if (remove_gap) {
            gen_seq <- gsub("-", "", gen_seq)
        } else {
            gen_seq <- gen_seq
        }
    },
        error = function(e)
        message("In post-processing : ", e)
    )
    
    result$gen_seq <- gen_seq
    result$label <- label
    result
}
