setwd("T:/Lincoln/Projects A-E/DNA/pimeleaDNA/final run")
library(dplyr)

samples <- list.files(path = ".", pattern = "merged.fastq_tssv.txt")
loci <- c("pim4", "pim8", "pim10", "pim11", "pim12", "pim13", "pim15", "pim16", "pim17", "pim20")

#-------------------------------------------
# SAFE FILE READER
#-------------------------------------------
safe_read <- function(file_name, min_cols = 3) {
  df <- tryCatch(
    read.table(file_name, skipNul = TRUE, fill = TRUE),  
    error = function(e) return(NULL)
  )
  
  # If completely unreadable
  if (is.null(df)) {
    message("Skipping unreadable file: ", file_name)
    return(NULL)
  }
  
  # If too few columns
  if (ncol(df) < min_cols) {
    message("Skipping malformed file (wrong number of columns): ", file_name)
    return(NULL)
  }
  
  df
}

#-------------------------------------------
# PROCESS A SINGLE LOCUS
#-------------------------------------------
process_locus <- function(seqs, locus, file_name) {
  
  seqsX <- seqs %>% filter(V1 == locus)
  if (nrow(seqsX) == 0) return(NULL)
  
  seqsX <- seqsX %>%
    mutate(
      len = nchar(V2),
      len_scaled = len / max(len)
    ) %>%
    filter(len_scaled > 0.6)
  
  if (nrow(seqsX) < 2) return(NULL)
  if (max(as.numeric(seqsX$V3)) < 50) return(NULL)
  
  seqsX <- seqsX %>%
    mutate(
      cov_scaled = as.numeric(V3) / max(as.numeric(V3))
    ) %>%
    filter(cov_scaled > 0.3)
  
  seqsX <- seqsX %>%
    mutate(final_len = nchar(V2)) %>%
    filter(final_len > 50)
  
  if (nrow(seqsX) < 1) return(NULL)
  
  letters_vec <- letters[1:nrow(seqsX)]
  
  seqsX <- seqsX %>%
    mutate(
      label = paste(">", file_name, locus, letters_vec, V3, sep = "_")
    )
  
  data.frame(seqsX$label, seqsX$V2, stringsAsFactors = FALSE)
}

#-------------------------------------------
# PROCESS A FILE SAFELY
#-------------------------------------------
process_file <- function(file_name, loci) {
  
  seqs <- safe_read(file_name)
  
  # Skip malformed files
  if (is.null(seqs)) return(NULL)
  
  for (locus in loci) {
    out <- process_locus(seqs, locus, file_name)
    
    if (!is.null(out)) {
      write.table(out,
                  file = paste0(locus, ".fas"),
                  quote = FALSE,
                  append = TRUE,
                  row.names = FALSE,
                  col.names = FALSE,
                  sep = "\r")
    }
  }
}

#-------------------------------------------
# MAIN LOOP
#-------------------------------------------
for (file in samples) {
  process_file(file, loci)
}