#!/usr/bin/env Rscript
# ==============================================================================
# Chapter 3: R for Next Gen Biologist
# Script: r03_functions_control_flow.R
# Description: Control Flow, Logic & Functions in R for Biology:
#              Conditionals, vectorised ifelse, apply family operations,
#              and building modular bioinformatics functions (Codon Translation & Ti/Tv).
# Usage: Rscript r03_functions_control_flow.R
# ==============================================================================

cat("\n=======================================================\n")
cat(" OMICS.CO.IN - R FOR NEXT GEN BIOLOGIST: CHAPTER 3\n")
cat(" Control Flow, Vectorization & Custom Functions\n")
cat("=======================================================\n\n")

# 1. Conditionals & Vectorized Logic
cat("[Step 1] Classifying Quality Scores with ifelse()...\n")
phred_scores <- c(38, 22, 14, 40, 28, 35, 18, 33)

quality_tier <- ifelse(phred_scores >= 30, "Q30+ (High Accuracy, Error <= 0.1%)",
                ifelse(phred_scores >= 20, "Q20+ (Acceptable, Error <= 1.0%)",
                                           "Low Quality (Needs Trimming)"))

for (i in seq_along(phred_scores)) {
  cat(sprintf("  Base %d [Score: Q%02d] -> %s\n", i, phred_scores[i], quality_tier[i]))
}

# 2. Functional Programming: The Apply Family in Genomics
cat("\n[Step 2] Using apply() and lapply() for Matrix Operations...\n")
set.seed(123)
expr_mat <- matrix(rpois(20, lambda = 50), nrow = 4, ncol = 5,
                   dimnames = list(paste0("Gene_", 1:4), paste0("Sample_", 1:5)))

cat("Expression Matrix:\n")
print(expr_mat)

# Calculate Coefficient of Variation (CV = SD / Mean) per gene using apply (1 = rows)
gene_cv <- apply(expr_mat, 1, function(row) {
  sd(row) / mean(row)
})
cat("\nPer-Gene Coefficient of Variation (Row-wise apply):\n")
print(round(gene_cv, 3))

# Calculate Quantile thresholds per sample using apply (2 = columns)
sample_medians <- apply(expr_mat, 2, median)
cat("\nPer-Sample Median Counts (Column-wise apply):\n")
print(sample_medians)

# 3. Defining Custom Bioinformatics Functions: Genetic Code Translation
cat("\n[Step 3] Custom Function: Standard Genetic Code Translator...\n")

translate_dna <- function(dna_sequence, to_stop = TRUE) {
  # Universal standard genetic code dictionary
  codon_map <- c(
    "TTT"="F", "TTC"="F", "TTA"="L", "TTG"="L",
    "TCT"="S", "TCC"="S", "TCA"="S", "TCG"="S",
    "TAT"="Y", "TAC"="Y", "TAA"="*", "TAG"="*",
    "TGT"="C", "TGC"="C", "TGA"="*", "TGG"="W",
    "CTT"="L", "CTC"="L", "CTA"="L", "CTG"="L",
    "CCT"="P", "CCC"="P", "CCA"="P", "CCG"="P",
    "CAT"="H", "CAC"="H", "CAA"="Q", "CAG"="Q",
    "CGT"="R", "CGC"="R", "CGA"="R", "CGG"="R",
    "ATT"="I", "ATC"="I", "ATA"="I", "ATG"="M",
    "ACT"="T", "ACC"="T", "ACA"="T", "ACG"="T",
    "AAT"="N", "AAC"="N", "AAA"="K", "AAG"="K",
    "AGT"="S", "AGC"="S", "AGA"="R", "AGG"="R",
    "GTT"="V", "GTC"="V", "GTA"="V", "GTG"="V",
    "GCT"="A", "GCC"="A", "GCA"="A", "GCG"="A",
    "GAT"="D", "GAC"="D", "GAA"="E", "GAG"="E",
    "GGT"="G", "GGC"="G", "GGA"="G", "GGG"="G"
  )
  
  seq_clean <- toupper(gsub("[^ACGT]", "", dna_sequence))
  n <- nchar(seq_clean)
  if (n < 3) return("")
  
  # Extract codons
  num_codons <- floor(n / 3)
  codons <- substring(seq_clean, seq(1, by = 3, length.out = num_codons),
                                 seq(3, by = 3, length.out = num_codons))
  
  amino_acids <- character(num_codons)
  for (i in seq_along(codons)) {
    aa <- codon_map[codons[i]]
    if (is.na(aa)) aa <- "X"
    if (to_stop && aa == "*") break
    amino_acids[i] <- aa
  }
  
  return(paste(amino_acids[amino_acids != ""], collapse = ""))
}

test_seq <- "ATGGCCATTGTAATGGGCCGCTGAAAGGGTGCCCGATAG"
peptide <- translate_dna(test_seq, to_stop = TRUE)
cat(sprintf("Input DNA (Coding Strand): %s\n", test_seq))
cat(sprintf("Translated Peptide       : %s\n", peptide))

# 4. Custom Function: Transition/Transversion (Ti/Tv) Ratio
cat("\n[Step 4] Custom Function: Transition/Transversion (Ti/Tv) Calculator...\n")

calc_titv <- function(ref_alleles, alt_alleles) {
  stopifnot(length(ref_alleles) == length(alt_alleles))
  
  transitions <- 0
  transversions <- 0
  
  # Purines: A, G; Pyrimidines: C, T
  purines <- c("A", "G")
  pyrimidines <- c("C", "T")
  
  for (i in seq_along(ref_alleles)) {
    r <- toupper(ref_alleles[i])
    a <- toupper(alt_alleles[i])
    
    if (r == a) next
    
    is_ti <- (r %in% purines && a %in% purines) || 
             (r %in% pyrimidines && a %in% pyrimidines)
    
    if (is_ti) {
      transitions <- transitions + 1
    } else {
      transversions <- transversions + 1
    }
  }
  
  ratio <- ifelse(transversions > 0, transitions / transversions, NA)
  return(list(Ti = transitions, Tv = transversions, Ratio = round(ratio, 2)))
}

# Example SNV callset
ref_vec <- c("A", "C", "G", "T", "A", "C", "A", "G")
alt_vec <- c("G", "T", "A", "C", "T", "G", "G", "A") # Transitions: A>G, C>T, G>A, T>C, A>G, G>A. Transversions: A>T, C>G
titv_result <- calc_titv(ref_vec, alt_vec)

cat(sprintf("Transitions (Ti)  : %d\n", titv_result$Ti))
cat(sprintf("Transversions (Tv): %d\n", titv_result$Tv))
cat(sprintf("Ti/Tv Ratio       : %.2f (Expected human WGS: ~2.0 - 2.1)\n", titv_result$Ratio))

cat("\n[SUCCESS] Chapter 3 control flow and functions script executed successfully.\n")