library(Biostrings)
library(parallel)

parse_data <- function(file) {
  reversed_sequences <- Biostrings::readQualityScaledDNAStringSet(file)
  sequences <- Biostrings::reverseComplement(reversed_sequences)
  vj_segments <- union(
    readRDS("data/v_segments.rds"),
    readRDS("data/j_segments_phe.rds")
  )
  return(list(sequences, vj_segments))
}

parse_metadata <- function(metadata) {
  id_elements <- unlist(strsplit(metadata, split = " "))
  v_identifier <- id_elements[2]
  j_identifier <- id_elements[3]
  return(list(v_id = v_identifier, j_id = j_identifier))
}

match_id_sequence <- function(names, vdj_segments, id) {
  matches <- grep(names, pattern = id)
  row <- matches[1]
  return(as.character(vdj_segments[row]))
}

get_vj_sequence <- function(metadata, names, vdj_segments) {
  identifiers <- parse_metadata(metadata)
  v_sequence <- match_id_sequence(names, vdj_segments, id = identifiers["v_id"])
  j_sequence <- match_id_sequence(names, vdj_segments, id = identifiers["j_id"])
  return(list(v_seq = v_sequence, j_seq = j_sequence))
}

fetch_vj_sequences <- function(sequences, vdj_segments) {
  vj_sequences <- sapply(names(sequences),
    names(vdj_segments),
    vdj_segments,
    FUN = get_vj_sequence
  )
  results <- data.frame(t(vj_sequences))
  return(results)
}

align_sequence <- function(sequence, vdj_segment) {
  return(Biostrings::pairwiseAlignment(
    subject = sequence,
    pattern = vdj_segment,
    type = "global-local",
    gapOpening = 1
  ))
}

handle_indels <- function(insertion, deletion, cys, alignment) {
  ins_start <- sum(Biostrings::width(deletion[start(deletion) <= cys$start]))
  ins_end <- sum(Biostrings::width(deletion[end(deletion) <= cys$end]))
  shift_num <- c(0, cumsum(Biostrings::width(insertion))[-length(ins_start)])
  shifted_ins <- IRanges::shift(insertion, shift_num)
  gaps <- sum(width(shifted_ins[end(shifted_ins) < cys$start + ins_start])) +
    nchar(stringr::str_extract(alignedSubject(alignment), "^-*"))
  return(list("start" = ins_start - gaps, "end" = ins_end - gaps))
}

get_cys_coordinates <- function(alignment) {
  cys <- list("start" = 310, "end" = 312)
  insertion <- unlist(Biostrings::insertion(alignment))
  deletion <- unlist(Biostrings::deletion(alignment))
  delta_coordinates <- handle_indels(insertion, deletion, cys, alignment)
  cys_start <- cys$start + delta_coordinates$start
  cys_end <- cys$end + delta_coordinates$end
  return(list("start" = cys_start, "end" = cys_end))
}

get_hvr_sequences <- function(sequences, vdj_segments, cores = detectCores()) {
  df <- fetch_vj_sequences(sequences, vdj_segments)
  v_alignment <- parallel::mcmapply(sequences,
    df$v_seq,
    FUN = align_sequence,
    mc.cores = cores
  )
  cys_coordinates <- parallel::mclapply(v_alignment, FUN = get_cys_coordinates)
  cys_df <- as.data.frame(do.call(rbind, cys_coordinates))
  remaining <- Biostrings::subseq(sequences, start = unlist(cys_df$end))
  j_alignment <- parallel::mcmapply(remaining,
    df$j_seq,
    FUN = align_sequence,
    mc.cores = cores
  )
  j_start <- parallel::mclapply(
    j_alignment,
    function(x) start(Biostrings::Views(x)),
    mc.cores = cores
  )
  hvr_start <- unlist(cys_df$start)
  hvr_end <- unlist(cys_df$start) + unlist(j_start) + 2
  hvr <- Biostrings::subseq(sequences, start = hvr_start, end = hvr_end)
  return(hvr)
}

data <- parse_data(file = "data/curesim_sequence.fastq")
hvr <- get_hvr_sequences(sequences = data[[1]], vdj_segments = data[[2]])
Biostrings::writeXStringSet(hvr, "data/CuReSim-HVR.fastq", format = "fastq")