# Zvolení Working Directory
setwd("C:/VUT/BAKALARKA/RfilesBc")

################################
# Instalace a nacteni knihoven
BiocManager::install("Biostrings")
BiocManager::install("Rsamtools")
#
library(Biostrings)
library(Rsamtools)
################################


MyPileup <- function(bam_file) {
  # Funkce pro vtvoreni pileupu s urcitymi parametry z namapovanych dat
  # Vstup: BAM soubor s namapovanymi daty
  # Vystup: pileup
  
  my_pileup <- pileup(bam_file, pileupParam = PileupParam(max_depth = 5000, 
                      min_base_quality=13, distinguish_strands = TRUE, 
                      distinguish_nucleotides=TRUE, min_nucleotide_depth = 20, 
                      min_minor_allele_depth = 50))
  return(my_pileup)
}
############

CompletePileup <- function(p1_bam, p2_bam, p3_bam, p4_bam, tprC_bam, tprE_bam, 
                           ind, ref_seq) {
  # Funkce pro vytvoreni pileupu pro cely genom
  # Vstup1-6: BAM soubory pro jednotlive pooly
  # Vstup7: csv soubor s indexy zacatku a koncu poolu
  # Vstup8: referencni sekvence genomu
  # Vystup: csv soubor s pileupem
  
  # Nacteni referencni sekvence a csv souboru s indexy
  ref_seq <- readDNAStringSet(ref_seq)
  ref_seq <- ref_seq[[1]]
  indices <- read.csv(ind, header = FALSE, sep = " ")
  
  # Vytvoreni pileupu pro kazdy pool pomoci funkce MyPileup
  pileup_p1 <- MyPileup(p1_bam)
  pileup_p2 <- MyPileup(p2_bam)
  pileup_p3 <- MyPileup(p3_bam)
  pileup_p4 <- MyPileup(p4_bam)
  pileup_tprC <- MyPileup(tprC_bam)
  pileup_tprE <- MyPileup(tprE_bam)

  # Rozdeleni pool1 na konec a zacatek referencni sekvence 
  length_p1_1 <- length(c(indices[1,1]:length(ref_seq)))
  pileup_p1_1 <- subset(pileup_p1, pos < (length_p1_1 + 1))
  pileup_p1_2 <- subset(pileup_p1, pos > length_p1_1)

  # Pridani hodnot pozic puvodni sekvence do pileupu
  pileup_p1_1$pos[] <- pileup_p1_1$pos[] + indices[1,1] - 1
  pileup_p1_2$pos[] <- pileup_p1_2$pos[] - length_p1_1
  pileup_p2$pos[] <- pileup_p2$pos[] + indices[2,1] - 1
  pileup_p3$pos[] <- pileup_p3$pos[] + indices[3,1] - 1
  pileup_p4$pos[] <- pileup_p4$pos[] + indices[4,1] - 1
  pileup_tprC$pos[] <- pileup_tprC$pos[] + indices[5,1] - 1
  pileup_tprE$pos[] <- pileup_tprE$pos[] + indices[6,1] - 1

  # Propojeni vsech pileupu do jednoho souboru a serazeni podle pozice 
  # v sekvenci
  all_pileups <- rbind(pileup_p1_1, pileup_p1_2, pileup_p2, pileup_p3,pileup_p4, 
                       pileup_tprC, pileup_tprE)
  all_pileups <- all_pileups[with(all_pileups, order(all_pileups$pos, 
                                                     all_pileups$nucleotide)),]
  row.names(all_pileups) <- c(1:nrow(all_pileups))

  # Nalezeni stejnych nukleotidu na stejne pozici a secteni jejich poctu
  for(i in c(1:(nrow(all_pileups) - 1))) {
    
    if  (all_pileups$pos[i] == all_pileups$pos[i+1] & 
         all_pileups$nucleotide[i] == all_pileups$nucleotide[i+1] & 
         all_pileups$strand[i] == all_pileups$strand[i+1]) {
    
      all_pileups$count[i] <- all_pileups$count[i] + all_pileups$count[i+1]
      all_pileups$pos[i+1] <- 0
    }
  }

  final_pileup <- all_pileups[all_pileups$pos != 0, ]

  # Vytvoreni souboru csv s pileupem
  write.csv(final_pileup, file = "final_pileup.csv")
}
############

StrandsRate <- function(my_pileup) {
  # Funkce pro selekci pozic v pileupu, kde je pomer zastoupeni nukleotidu
  # v obou vlaknech minimalne 30:70
  # Vstup: pileup
  # Vystup: pileup s dostatecnym pomerem zastoupeni nukleotidu ve vlaknech
  
  # Vytvoreni vektoru pozic vyskytujicich se v pileupu
  u_pos <- unique(my_pileup$pos)
  
  # Vytvoreni promenne pro novy pileup
  s_pileup <- my_pileup[1,]
  
  # Cyklus for v kazde behu vybere data z pileupu pro danou pozici, pokud je na
  # teto pozici nukleotid zastoupen v obou vlaknech, vstoupi do dalsiho cyklu
  # for, kde jsou selektovany hodnoty s pomerem v obou vlaknech minimalne 30:70,
  # a jsou zapsany do promenne s_pileup
  for (i in u_pos) {
    position <- my_pileup[my_pileup$pos == i,]
    nuc_table <- table(position$nucleotide)
    nuc_both <- names(nuc_table[nuc_table == 2])
    l <- length(nuc_both)
    if (l != 0) {
      for (j in 1:l) {
        nuc_counts <- position$count[position$nucleotide == nuc_both[j]]
        s_counts <- sum(nuc_counts)
        if ((nuc_counts[1] / s_counts) > 0.3 &
            (nuc_counts[2] / s_counts) > 0.3) {
          s_pileup <-
            rbind(s_pileup, position[position$nucleotide == nuc_both[j],])
        }
      }
    }
  }
  s_pileup <- s_pileup[-1, ]
  return(s_pileup)
}
############

VarPositions <- function(my_pileup, ref_seq) {
  # Funkce pro vytvoreni pileupu s pozicemi, na kterych je pomer jednotlivych 
  # nukleotidu na obou vlaknech minimalne 30:70 a zastoupeni alternativni
  # alely minimalne 5 %
  # Vstup1: csv soubor s pileupem
  # Vstup2: referencni sekvence ve formatu fasta
  # Vystup: csv soubor s pileupem
  
  # Nacteni dat z csv souboru
  my_pileup <- read.csv(my_pileup)
  
  # Selekce nukleotidu s dostatecnym zastoupenim na obou vlaknech pomoci funkce
  # StrandsRate
  s_pileup <- StrandsRate(my_pileup)
  
  # Nacteni referencni sekvence
  ref_seq <- readDNAStringSet(ref_seq)
  ref_seq <- as.character(ref_seq[[1]])
  ref_seq <- strsplit(ref_seq, split = "")[[1]]
  
  # Pridani sloupce s referencni sekvenci do pileupu
  ref_col <- c()
  p_pos <- s_pileup$pos
  for (i in 1:(length(p_pos))) {
    ref_col[i] <- ref_seq[p_pos[i]]
  }
  s_pileup$ref <- ref_col
 
  # Vytvoreni vektoru pozic vyskytujicich se v pileupu
  u_pos <- unique(s_pileup$pos)
  
  # Vytvoreni promenne pro novy pileup
  var_pileup <- s_pileup[1,]
  
  # Cyklus for v kazde behu vybere data z pileupu pro danou pozici, pokud je na
  # teto pozici zastoupeni alternativni alely minimalne 5 %, jsou hodnoty
  # z dane pozice zapsany do promenne var_pileup
  for (j in u_pos) {
    position <- s_pileup[s_pileup$pos == j,]
    p_sum <- sum(position$count)
    p_ref <- position$ref[1]
    ref_sum <- sum(position$count[position$nucleotide == p_ref])
    var_rate <- (p_sum - ref_sum) / p_sum
    if (var_rate > 0.05) {
      var_pileup <- rbind(var_pileup, position)
    }
  }
  var_pileup <- var_pileup[-1,]
  
  # Vytvoreni souboru csv s pileupem
  write.csv(var_pileup, 'var_pileup.csv')
}
###################################


# Vytvoreni csv souboru s pileupy s variabilnimi pozicemi pomoci funkci 
# CompletePileup a VarPositions pro jednotlive kmeny
setwd("C:/VUT/BAKALARKA/RfilesBc/TPA_Grady")
CompletePileup('pool1.sam.bam_sorted.bam','pool2.sam.bam_sorted.bam',
                    'pool3.sam.bam_sorted.bam','pool4.sam.bam_sorted.bam',
                    'tprC.sam.bam_sorted.bam','tprE.sam.bam_sorted.bam',
                    'indices.csv', 'reference.fasta')
VarPositions('final_pileup.csv', 'reference.fasta')

setwd("C:/VUT/BAKALARKA/RfilesBc/TPA_HaitiB")
CompletePileup('pool1.sam.bam_sorted.bam','pool2.sam.bam_sorted.bam',
                    'pool3.sam.bam_sorted.bam','pool4.sam.bam_sorted.bam',
                    'tprC.sam.bam_sorted.bam','tprE.sam.bam_sorted.bam',
                    'indices.csv', 'reference.fasta')
VarPositions('final_pileup.csv', 'reference.fasta')

setwd("C:/VUT/BAKALARKA/RfilesBc/TPA_Madras")
CompletePileup('pool1.sam.bam_sorted.bam','pool2.sam.bam_sorted.bam',
                    'pool3.sam.bam_sorted.bam','pool4.sam.bam_sorted.bam',
                    'tprC.sam.bam_sorted.bam','tprE.sam.bam_sorted.bam',
                    'indices.csv', 'reference.fasta')
VarPositions('final_pileup.csv', 'reference.fasta')

setwd("C:/VUT/BAKALARKA/RfilesBc/TPA_Philadelphia")
CompletePileup('pool1.sam.bam_sorted.bam','pool2.sam.bam_sorted.bam',
                    'pool3.sam.bam_sorted.bam','pool4.sam.bam_sorted.bam',
                    'tprC.sam.bam_sorted.bam','tprE.sam.bam_sorted.bam',
                    'indices.csv', 'reference.fasta')
VarPositions('final_pileup.csv', 'reference.fasta')

setwd("C:/VUT/BAKALARKA/RfilesBc/TPA_UZ1974")
CompletePileup('pool1.sam.bam_sorted.bam','pool2.sam.bam_sorted.bam',
                    'pool3.sam.bam_sorted.bam','pool4.sam.bam_sorted.bam',
                    'tprC.sam.bam_sorted.bam','tprE.sam.bam_sorted.bam',
                    'indices.csv', 'reference.fasta')
VarPositions('final_pileup.csv', 'reference.fasta')

##################################