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

################################
# Instalace a nacteni knihoven
BiocManager::install("Biostrings")
BiocManager::install("genbankr")
install.packages('pheatmap')
#
library(Biostrings)
library(genbankr)
library(pheatmap)
################################


RefCount <- function(align_seq) {
  # Funkce pro vytvoreni vektoru s pozicemi zarovnane sekvence na indexech 
  # referencni sekvence
  # Vstup: zarovnana sekvence
  # Vystup: vektor s pozicemi zarovnane sekvence
  
  # Nalezeni pozic, na nichz jsou v zarovnane sekvenci delece
  del_pos <- gregexpr('-', align_seq)[[1]]
  
  # Vytvoreni vektoru o delce zarovnane sekvence a odstraneni pozic s delecemi
  align_pos <- 1:nchar(align_seq)
  align_pos <- align_pos[-(del_pos)]
  return(align_pos)
}
############

VarPercent <- function(var_pileup) {
  # Funkce pro vytvoreni pileupu s pozicemi, na kterych je pomer jednotlivych 
  # Vstup1: csv soubor s pileupem
  # Vystup: list obsahujici vektor s procenty a vektor s unikatnimi pozicemi
  
  # Nacteni dat z csv souboru
  var_pileup <- read.csv(var_pileup)
  
  # Vytvoreni vektoru pozic vyskytujicich se v pileupu
  u_pos <- unique(var_pileup$pos)
  
  # Cyklus for v kazde behu vybere data z pileupu pro danou pozici, vypocita
  # procentuelni zastoupeni alternativni alely a zapise ho do vektoru percents
  percents <- c()
  j <- 1
  for (i in u_pos) {
    position <- var_pileup[var_pileup$pos == i,]
    p_sum <- sum(position$count)
    p_ref <- position$ref[1]
    ref_sum <- sum(position$count[position$nucleotide == p_ref])
    percents[j] <- ((p_sum - ref_sum) / p_sum) * 100
    j <- j + 1
  }
  return(list(percents, u_pos))
}
###########

VarAlign <- function(var_pileup, strain, complete_align) {
  # Funkce pro vytvoreni vektoru delky zarovnane sekvence s vyznacenymi 
  # variabilnimi pozicemi
  # Vstup1: pileup s variabilnimi pozicemi
  # Vstup2: retezec se jmenem kmenu
  # Vstup3: zarovnane sekvence
  # Vystup: list obsahujici vektor s hodnotami procentuelniho zastoupeni
  # variabilit a vektor s binarne vyznacenymi variabilnimi pozicemi
  
  # Vyber zarovnane sekvence pro konkretni kmen 
  align_seq <- complete_align[[strain]]
  
  # Pouziti funkce RefCount pro vytvoreni vektoru s pozicemi zarovnane sekvence 
  # na indexech referencni sekvence
  align_pos <- RefCount(align_seq)
  
  # Nacteni pileupu, ziskani vektoru variabilnich pozic pomoci funkce 
  # VarPercent a nalezeni techto pozic v zarovnane sekvenci
  per_pos <- VarPercent(var_pileup)
  var_pos <- per_pos[[2]]
  var_align <- align_pos[var_pos]
  
  # Vytvoreni vektoru s nulami o delce zarovnane sekvence a zapsani 
  # variabilnich pozic do tohoto vektoru jako jednicky
  var_bin <- integer(nchar(align_seq))
  var_bin[var_align] <- 1
  
  # Vytvoreni vektoru s nulami o delce zarovnane sekvence a zapsani 
  # procentualniho zastoupeni variability na pozicich
  percents <- per_pos[[1]]
  var_per <- integer(nchar(align_seq))
  var_per[var_align] <- percents
  
  return(list(var_per, var_bin))
}
############

GetCDS <- function(gb_file, gene_names) {
  # Funkce pro nalezeni vsech pozic referencniho genomu na kterych jsou 
  # kodujici sekvence
  # Vstup1: soubor genbank
  # Vstup2: csv soubor se jmeny genu zapsanymi do prvniho sloupce tabulky
  
  # Nacteni gb souboru a vektoru se jmeny genu
  gb <- readGenBank(gb_file)
  gene_names <- read.csv(gene_names, header = FALSE)
  gene_names <- gene_names$V1
  
  # Extrakce startu a delek CDS a vypocet koncu CDS
  cds_starts <- gb@cds@ranges@start
  cds_widths <- gb@cds@ranges@width
  cds_ends <- cds_starts + (cds_widths - 1)
  
  # Vytvoreni vektoru se vsemi pozicemi CDS a vektoru o stejne delce s indexy
  # a jmeny genu
  cds_pos <- c()
  cds_ind <- c()
  g_names <- c()
  for (i in 1:length(cds_starts)) {
    cds_pos <- append(cds_pos, cds_starts[i]:cds_ends[i])
    cds_ind <- append(cds_ind, rep(i, cds_widths[i]))
  }
  g_names <- gene_names[cds_ind]
  return(list(cds_pos, cds_ind, g_names))
}
##############################


# Nacteni zarovnanych sekvenci
complete_align <- readDNAStringSet('TrepAlignComplete.fa')

# Ziskani vektoru se zaznamenanymi variabilnimi pozicemi pomoci funkce VarAlign
setwd("C:/VUT/BAKALARKA/RfilesBc/TPA_Grady")
var_Grady <- VarAlign('var_pileup.csv', 1, complete_align)

setwd("C:/VUT/BAKALARKA/RfilesBc/TPA_HaitiB")
var_HaitiB <- VarAlign('var_pileup.csv', 2, complete_align)

setwd("C:/VUT/BAKALARKA/RfilesBc/TPA_Madras")
var_Madras <- VarAlign('var_pileup.csv', 3, complete_align)

setwd("C:/VUT/BAKALARKA/RfilesBc/TPA_Philadelphia")
var_Philadelphia <- VarAlign('var_pileup.csv',4,complete_align)

setwd("C:/VUT/BAKALARKA/RfilesBc/TPA_UZ1974")
var_UZ1974 <- VarAlign('var_pileup.csv', 5, complete_align)

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

# Spojeni vektoru s procenty na variabilnich pozicich do jedne matice a zapsani
# teto matice do souboru csv
setwd("C:/VUT/BAKALARKA/RfilesBc")
var_matrix <- cbind(var_Grady[[1]], var_HaitiB[[1]], var_Madras[[1]], 
                    var_Philadelphia[[1]], var_UZ1974[[1]])
rownames(var_matrix) <- 1:nrow(var_matrix)
colnames(var_matrix) <- (c('Grady','HaitiB','Madras','Philadelphia','UZ1974'))
write.csv(var_matrix,'var_matrix.csv')

# Selekce pozic, ktere jsou variabilni minimalne pro dva kmeny zaroven
# pomoci matice s binarne zaznamenanymi variabilnimi pozicemi a zapsani
# vysledne matice di souboru csv
bin_matrix <- cbind(var_Grady[[2]], var_HaitiB[[2]], var_Madras[[2]], 
                    var_Philadelphia[[2]], var_UZ1974[[2]])
rownames(bin_matrix) <- 1:nrow(bin_matrix)
bin_var <- bin_matrix[rowSums(bin_matrix)>1,]
var_vec <- rownames(bin_var)
com_matrix <- var_matrix[var_vec,]
write.csv(com_matrix, 'com_matrix.csv')

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

# Ziskani vsech pozic CDS v referencni sekvenci a jmen genu
cds_list <- GetCDS('Nichols.gb', 'gene_names.csv')
cds_pos <- cds_list[[1]]
cds_ind <- cds_list[[2]]
cds_names <- cds_list[[3]]

# Prepocitani pozic CDS na pizice v zarovnani
align_pos <- RefCount(complete_align[[3]])
cds_align <- align_pos[cds_pos]

# Vyber pozic, ktere se nachazeji v CDS z matice s procenty variability,
# pojmenovani radku jmeny danych genu a zapsani do souboru csv
cds_matrix <- var_matrix[cds_align,]
rownames(cds_matrix) <- cds_names
cds_matrix <- cds_matrix[rowSums(cds_matrix)>5,]
write.csv(cds_matrix, 'cds_matrix.csv')

# Vytvoření matice se stejnými hodnotami pro heatmap
cds_hmatrix <- var_matrix[cds_align,]
cds_hmatrix <- cds_hmatrix[rowSums(cds_hmatrix)>5,]
          
#############################

# Vytvoreni heatmap pomoci funkce pheatmap pro matici se spolecnymi 
# variabilitami a pro matici s variabilitami v CDS
rn_com <- rownames(com_matrix)
rn_cds <- rownames(cds_hmatrix)

pheatmap(com_matrix, legend = TRUE, cluster_rows = FALSE,
        cluster_cols = FALSE, color = rev(heat.colors(120)[1:100]), 
        labels_row = rn_com[seq(1, length(rn_com),length(rn_com)/33)])

pheatmap(cds_hmatrix, legend = TRUE, cluster_rows = FALSE,
         cluster_cols = FALSE, color = rev(heat.colors(120)[1:100]),
         labels_row = rn_cds[seq(1, length(rn_cds),length(rn_cds)/33)])

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