args <- commandArgs(trailingOnly = FALSE)
script_arg <- grep("^--file=", args, value = TRUE)
if (length(script_arg) > 0) {
  script_path <- normalizePath(sub("^--file=", "", script_arg[1]))
  setwd(normalizePath(file.path(dirname(script_path), "..")))
}

library(caret)

pep_mat_local <- function(pep) {
  do.call(rbind, strsplit(as.character(pep), split = ""))
}

one_hot <- function(y, num_classes = 2) {
  y <- as.integer(as.character(y))
  out <- matrix(0, nrow = length(y), ncol = num_classes)
  out[cbind(seq_along(y), y + 1)] <- 1
  out
}

load("model/BLOSUM62_with_deletion.Rdata")

pep_encode_blosum <- function(pep) {
  bl62_prob <- BLOSUM62
  p_mat <- pep_mat_local(pep)
  n_peps <- length(pep)
  l_peps <- nchar(pep[1])
  l_enc <- ncol(bl62_prob)
  o_tensor <- array(data = NA, dim = c(n_peps, l_peps, l_enc))
  for (i in 1:n_peps) {
    pep_i_residues <- p_mat[i, ]
    pep_img <- bl62_prob[pep_i_residues, ]
    o_tensor[i, , ] <- pep_img
  }
  o_tensor
}

all <- read.delim("model/CNN/all_ab_pre_post.txt", sep = "\t", header = TRUE, stringsAsFactors = FALSE)

set.seed(42)
all <- all[sample(nrow(all)), ]

max_k <- max(nchar(all$CDR3K))
max_h <- max(nchar(all$CDR3H))

all$paddedh <- gsub("\\s", "-", format(all$CDR3H, width = max_h))
all$paddedk <- gsub("\\s", "-", format(all$CDR3K, width = max_k))
all$lengthh <- nchar(all$CDR3H)
all$lengthk <- nchar(all$CDR3K)
all$padded <- paste0(all$paddedk, all$paddedh)

all <- all[order(all$padded, -all$post), ]
all2 <- all[!duplicated(all$padded), ]

all2$enriched <- ifelse((all2$fc >= 1.8 & all2$post >= 0.01), 1,
                        ifelse(all2$fc < 1 & all2$pre >= 0.01, 0, "ambi"))
all2 <- all2[all2$enriched != "ambi", ]
all2$enriched <- as.integer(all2$enriched)

split_train_test <- function(antigenx) {
  mini <- all2[all2$antigen == antigenx, ]
  set.seed(22)
  trainIndex <- createDataPartition(mini$enriched, p = .8, list = FALSE, times = 1)
  list(train = mini[trainIndex, ], test = mini[-trainIndex, ])
}

c1 <- split_train_test("CTLA-4")
p1 <- split_train_test("PD-1")

saveRDS(c1, file = "model/CNN/c1.RDS")
saveRDS(p1, file = "model/CNN/p1.RDS")

c1_train <- pep_encode_blosum(c1$train$padded)
c1_test  <- pep_encode_blosum(c1$test$padded)
p1_train <- pep_encode_blosum(p1$train$padded)
p1_test  <- pep_encode_blosum(p1$test$padded)

c1_train <- array(c1_train, dim = c(dim(c1_train), 1))
c1_test  <- array(c1_test,  dim = c(dim(c1_test), 1))
p1_train <- array(p1_train, dim = c(dim(p1_train), 1))
p1_test  <- array(p1_test,  dim = c(dim(p1_test), 1))

c1_train_y <- one_hot(c1$train$enriched)
c1_test_y  <- one_hot(c1$test$enriched)
p1_train_y <- one_hot(p1$train$enriched)
p1_test_y  <- one_hot(p1$test$enriched)

saveRDS(p1_train,   file = "model/CNN/p1_train.RDS")
saveRDS(p1_test,    file = "model/CNN/p1_test.RDS")
saveRDS(p1_train_y, file = "model/CNN/p1_train_y.RDS")
saveRDS(p1_test_y,  file = "model/CNN/p1_test_y.RDS")
saveRDS(c1_train,   file = "model/CNN/c1_train.RDS")
saveRDS(c1_test,    file = "model/CNN/c1_test.RDS")
saveRDS(c1_train_y, file = "model/CNN/c1_train_y.RDS")
saveRDS(c1_test_y,  file = "model/CNN/c1_test_y.RDS")

seq_all <- readRDS("model/GAN/seq_all.RDS")
encoded <- vector("list", 15)
names(encoded) <- names(seq_all)

for (i in 1:15) {
  seq_encoded <- pep_encode_blosum(seq_all[[i]])
  encoded[[i]] <- array(seq_encoded, dim = c(dim(seq_encoded), 1))
}
saveRDS(encoded, file = "model/GAN/seq_all_encoded.RDS")

print("prepare data OK")
print(dim(c1_train))
print(dim(p1_train))
print(names(encoded))