// The Benchmarks Game's `k-nucleotide`: reads DNA in FASTA form from standard
// input and reports how often each short fragment occurs in one section.
//
//     cat dna.fasta | nikaia run
//     A 33.333
//     T 30.000
//     …
//     3	GGT
//     3	GGTA
//     0	GGTATT
//     …
//
// The input is FASTA text: a header line starting with `>`, then lines of DNA.
// The section read is the one whose header starts with `THREE` (the
// benchmark's choice); `cat dna.fasta | nikaia run -- ONE` picks another.
//
// Shows: reading standard input, a `grammar` for a line-based format, map
// keys that are views into a string rather than copies, and `?.` with `??`.
//
// A package: `nikaia run` in this directory runs it, and `nikaia test` checks
// what it prints against `tests/` (ADR-269 D13).

use std::io
use std::cli
use std::collections

// ---------------------------------------------------------------------------
// 1. FASTA
// ---------------------------------------------------------------------------

grammar Fasta {

    // In FASTA a newline is the structure, not a blank to skip. Defining `WS`
    // as empty turns off the whitespace skipping of lowercase rules, which
    // would otherwise eat the line endings this grammar is made of.
    rule WS = "" -> { }

    entry rule file -> Vec[Section] =
        sections:section* -> { sections }

    // A header, then every line that does not start a new one. `not(">")`
    // consumes nothing: it only tests, and it is what ends the body.
    rule section -> Section =
        ">" header:LINE body:body_line* -> { Section { header, lines: body } }

    rule body_line -> String = not(">") l:LINE -> { l }

    // `until` stops before the terminator, so the line ending is consumed
    // explicitly. Input whose last line has no line ending is refused as not
    // FASTA rather than losing that line.
    rule LINE -> String = s:until(line_ending) line_ending -> { s }
}

pub struct Section {
    header: String,
    lines: Vec[String],
}

// ---------------------------------------------------------------------------
// 2. Counting
// ---------------------------------------------------------------------------

// A count is a struct with a method rather than a bare number:
// `entry(..).and_modify` hands the lambda the value already there, and the
// lambda calls a method on it. `sync` says a function cannot pause, and the
// compiler checks it.
pub struct Tally {
    n: i64,
}

impl Tally {
    pub fn(first: i64) -> Tally sync {
        return Tally { n: first }
    }

    fn bump(ref mut self) sync {
        self.n += 1
    }
}

// Every fragment of length `k`, and how often each occurs.
//
// The keys are views into `dna`: `ref dna[i..<i + k]` is a slice, not a copy,
// so a sequence of a million characters builds a table of positions and
// lengths rather than a million small strings.
fn count(dna: String, k: i64) -> collections::HashMap[String, Tally] {
    let mut counts: collections::HashMap[String, Tally] = collections::HashMap()

    if dna.len() < k {
        return counts
    }

    for i in 0..<dna.len() - k + 1 {
        // `i` is a counter, not a position the text handed out: the bytes
        // are cut, and made text again (Part I 2.6). A cut through a scalar
        // is no fragment of DNA.
        let fragment = dna.bytes[i..<i + k].text() catch {
            continue
        }
        counts
            .entry(fragment)
            .and_modify fn (tally) { tally.bump() }
            .or_insert_with fn { Tally(1) }
    }

    return counts
}

// The fragments of one length, as percentages, most frequent first.
//
// Two stable sorts give the order the benchmark asks for: by fragment, then by
// count descending. The second sort keeps the first's order among equal counts,
// so ties come out alphabetically.
fn report_frequencies(dna: String, k: i64) {
    let counts = count(dna, k)

    let mut rows = Vec()
    let mut total = 0
    for (fragment, tally) in counts {
        total += tally.n
        rows.push((fragment, tally.n))
    }

    rows.sort_by_key fn (row) { row.0 }
    rows.sort_by_key fn (row) { -row.1 }

    for row in rows {
        let percent = (row.1 as f64) * 100.0 / (total as f64)
        println(f"{row.0} {percent:.3}")
    }
    println("")
}

// One named fragment, and how often it occurs. `k` is the fragment's own
// length, so the table is built for exactly the width being asked about.
fn report_fragment(dna: String, fragment: String) {
    let counts = count(dna, fragment.len())
    // `get` hands back a `Tally?`, because the key may not be there: `?.`
    // reaches into it if it is, and `??` gives the value to use if it is not.
    let n = counts.get(fragment)?.n ?? 0
    println(f"{n}\t{fragment}")
}

// ---------------------------------------------------------------------------
// 3. The program
// ---------------------------------------------------------------------------

fn main() throws {
    // The whole of standard input, read into a string the program owns: a pipe
    // cannot be mapped like a file. There is no `await`, yet while the pipe is
    // empty other tasks run. With `user-parallelism = "yes"` in nikaia.toml
    // `[build]` the read may resume on another thread; the output is the same.
    let text = io::read_to_string() catch {
        eprintln(f"cannot read standard input: {error}")
        return
    }

    let sections = Fasta::file(text) catch {
        eprintln(f"standard input is not FASTA: {error}")
        return
    }

    // The section to count, by the start of its header. Lines are uppercased,
    // since FASTA may write DNA in either case.
    let wanted = cli::args().nth(1) ?? "THREE"

    let mut dna = String()
    for section in sections {
        if section.header.starts_with(ref wanted) {
            for line in section.lines {
                dna.push_str(ref line.to_uppercase())
            }
        }
    }

    if dna.len() == 0 {
        eprintln(f"no section whose header starts with {wanted}")
        return
    }

    report_frequencies(dna, 1)
    report_frequencies(dna, 2)

    report_fragment(dna, "GGT")
    report_fragment(dna, "GGTA")
    report_fragment(dna, "GGTATT")
    report_fragment(dna, "GGTATTTTAATT")
    report_fragment(dna, "GGTATTTTAATTTATAGT")
}