///|
/// Genetic-code presets supported by the translation helpers.
pub enum CodonTable {
  Standard
  VertebrateMitochondrial
} derive(Debug, Eq)

///|
pub fn standard_codon_table() -> CodonTable {
  CodonTable::Standard
}

///|
pub fn vertebrate_mitochondrial_table() -> CodonTable {
  CodonTable::VertebrateMitochondrial
}

///|
pub struct CodonUsage {
  codon : String
  amino : String
  count : Int
} derive(Debug, Eq)

///|
pub struct OpenReadingFrame {
  start : Int
  end : Int
  frame : Int
  protein : String
  nucleotide_length : Int
} derive(Debug, Eq)

///|
pub fn codon_amino(table : CodonTable, codon : String) -> String {
  let normalized = codon.to_upper().replace_all(old="U", new="T")
  match table {
    CodonTable::Standard => codon_to_amino(normalized)
    CodonTable::VertebrateMitochondrial => mitochondrial_amino(normalized)
  }
}

///|
fn mitochondrial_amino(codon : String) -> String {
  match codon {
    "AGA" | "AGG" => "*"
    "ATA" => "M"
    "TGA" => "W"
    _ => codon_to_amino(codon)
  }
}

///|
/// Translate a one-based reading frame, where frame 1 starts at position 1.
pub fn translate_frame(
  sequence : String,
  frame : Int,
  table : CodonTable,
) -> String {
  let offset = if frame <= 1 { 0 } else { Int::min(frame - 1, 2) }
  let upper = sequence.to_upper().replace_all(old="U", new="T")
  let output : Array[String] = []
  if offset >= upper.length() {
    return ""
  }
  let mut index = offset
  while index + 2 < upper.length() {
    output.push(
      codon_amino(table, upper.unsafe_substring(start=index, end=index + 3)),
    )
    index += 3
  }
  output.join("")
}

///|
/// Count complete codons in first-seen order.
pub fn codon_usage(sequence : String, table : CodonTable) -> Array[CodonUsage] {
  let result : Array[CodonUsage] = []
  let upper = sequence.to_upper().replace_all(old="U", new="T")
  let mut index = 0
  while index + 2 < upper.length() {
    let codon = upper.unsafe_substring(start=index, end=index + 3)
    let amino = codon_amino(table, codon)
    match result.search_by(fn(item) { item.codon == codon }) {
      Some(found) => {
        let item = result[found]
        result[found] = {
          codon: item.codon,
          amino: item.amino,
          count: item.count + 1,
        }
      }
      None => result.push({ codon, amino, count: 1 })
    }
    index += 3
  }
  result
}

///|
pub fn codon_usage_total(items : Array[CodonUsage]) -> Int {
  items.fold(init=0, fn(total, item) { total + item.count })
}

///|
/// Find start-to-stop ORFs in all three forward reading frames.
pub fn find_orfs(
  sequence : String,
  minimum_codons : Int,
  maximum_bases : Int,
  table : CodonTable,
) -> Array[OpenReadingFrame] {
  let result : Array[OpenReadingFrame] = []
  let upper = sequence.to_upper().replace_all(old="U", new="T")
  let minimum = Int::max(1, minimum_codons)
  let maximum = if maximum_bases <= 0 { upper.length() } else { maximum_bases }
  for frame in 0..<3 {
    let mut start = -1
    let mut index = frame
    while index + 2 < upper.length() {
      let codon = upper.unsafe_substring(start=index, end=index + 3)
      if start < 0 && codon == "ATG" {
        start = index
      } else if start >= 0 && codon_amino(table, codon) == "*" {
        let length = index + 3 - start
        if length <= maximum && length / 3 >= minimum {
          let dna = upper.unsafe_substring(start~, end=index + 3)
          result.push({
            start: start + 1,
            end: index + 3,
            frame: frame + 1,
            protein: translate_frame(dna, 1, table),
            nucleotide_length: length,
          })
        }
        start = -1
      }
      index += 3
    }
  }
  result.sort_by(fn(left, right) { left.start - right.start })
  result
}

///|
pub fn CodonUsage::to_json_string(self : CodonUsage) -> String {
  "{\"codon\":" +
  json_string(self.codon) +
  ",\"amino\":" +
  json_string(self.amino) +
  ",\"count\":" +
  self.count.to_string() +
  "}"
}

///|
pub fn OpenReadingFrame::to_json_string(self : OpenReadingFrame) -> String {
  "{\"start\":" +
  self.start.to_string() +
  ",\"end\":" +
  self.end.to_string() +
  ",\"frame\":" +
  self.frame.to_string() +
  ",\"protein\":" +
  json_string(self.protein) +
  ",\"nucleotide_length\":" +
  self.nucleotide_length.to_string() +
  "}"
}