package tutorial

trait Sequence[T] {
  def size: Int
  def slice(from: Int, to: Int): T
  def description: String = "Sequence"
}

trait HasCodons {
  def codons: Iterator[String]
}

trait NTSequence[T] extends Sequence[T] {
  def bp(position: Int): Char = data(position)
  def data: String
  def size = data.size
  override def toString = data
  override def description = "NTSequence"
}

case class RNA(data: String) extends NTSequence[RNA] with HasCodons {
  override def description = "RNA"
  def translate = PP(this)
  def slice(from: Int, to: Int) = RNA(data.slice(from, to))
  def codons = data.grouped(3)
}

case class DNA(data: String) extends NTSequence[DNA] {
  override def description = "DNA"
  def transcribe: RNA = RNA(complement.data.replace("T", "U"))
  def slice(from: Int, to: Int) = DNA(data.slice(from, to))

  def complement = DNA(data.map(_ match {
    case 'A' => 'T'
    case 'C' => 'G'
    case 'G' => 'C'
    case 'T' => 'A'
    case _   => throw new Exception("Unexpected character in DNA")
  }))
}

case class PP(rna: RNA) extends Sequence[PP] {
  def size = rna.size / 3
  override def description = "Polypeptide/protein"
  def slice(from: Int, to: Int) = PP(rna.slice(from * 3, to * 3))
  override def toString = Helper.toAAString(rna)
}

object Helper {

  def describe(s: Sequence[_]): String = {
    s match {
      case RNA("UAA")     => "Stop codon RNA"
      case RNA(_)         => "RNA"
      case DNA(_)         => "DNA"
      case PP(_) => "Polypeptide/protein"
      case _              => "Unknown sequence"
    }
  }

  val codonAminoMap = Map("GCA" -> 'A', "AGA" -> 'R', "GAC" -> 'D',
    "AAC" -> 'N', "UGC" -> 'C')
  val aminoCodonMap = codonAminoMap.map(_.swap)

  def toAA(codon: String) = codonAminoMap(codon)
  def toCodon(aa: Char) = aminoCodonMap(aa)

  def toAAString(hc: HasCodons) = hc.codons.map(toAA(_)).mkString

}