Mass spectrometry

21 problems from Rosalind — Bioinformatics Textbook Track. Press Run on any block to execute it in your browser.

# Rosalind: BA4C — Generate the Theoretical Spectrum of a Cyclic Peptide
# https://rosalind.info/problems/ba4c/
#
# Given: An amino acid string Peptide.
# Return: Cyclospectrum(Peptide).

let peptide = "LEQN"

# A mass spectrometer breaks many copies of a peptide at every position and
# weighs the pieces. For a cyclic peptide the pieces include those that wrap past
# the end, so LEQN contributes QNL and NLE as well as the obvious subpeptides —
# and each wrapping piece is exactly the complement of a non-wrapping one.
let spectrum = cyclic_spectrum(peptide)

println("Result:   " + (spectrum |> map(|m| str(m)) |> join(" ")))
println("Expected: 0 113 114 128 129 227 242 242 257 355 356 370 371 484")

fn test_ba4c_cyclic_spectrum() {
    assert (spectrum |> map(|m| str(m)) |> join(" "))
        == "0 113 114 128 129 227 242 242 257 355 356 370 371 484",
        "BA4C: got " + str(spectrum)
    # A cyclic peptide of length n has n(n-1) proper subpeptides, plus 0 and the
    # whole peptide.
    let n = len(peptide)
    assert len(spectrum) == n * (n - 1) + 2,
        "BA4C: expected " + str(n * (n - 1) + 2) + " masses, got " + str(len(spectrum))
    assert spectrum[0] == 0, "BA4C: the empty piece weighs nothing"
    assert spectrum[len(spectrum) - 1] == peptide_mass(peptide),
        "BA4C: the heaviest piece is the whole peptide"
    # 242 twice is not a mistake: LE and QN both weigh 242, and a spectrum
    # records both.
    assert (spectrum |> count_if(|m| m == 242)) == 2, "BA4C: LE and QN both weigh 242"
}

# Rosalind: BA4D — Compute the Number of Peptides of Given Total Mass
# https://rosalind.info/problems/ba4d/
#
# Given: An integer m.
# Return: The number of linear peptides of integer mass m.

let target = 1024

# Every peptide of mass m is some peptide of mass m - r followed by a residue of
# mass r, so the counts build up from below and each is used many times. Counting
# by enumeration instead would mean listing 14.7 billion peptides to find out how
# many there are.
#
# Eighteen residue masses, not twenty: I/L and K/Q collide, and a peptide is
# counted by its masses.
let residues = amino_acid_masses()

let ways = [1]
for mass in range(1, target + 1) {
    let total = residues
        |> filter(|r| r <= mass)
        |> map(|r| ways[mass - r])
        |> sum()
    ways = push(ways, total)
}

println("Result:   " + str(ways[target]))
println("Expected: 14712706211")

fn test_ba4d_counting_peptides() {
    assert ways[target] == 14712706211, "BA4D: got " + str(ways[target])
    assert len(residues) == 18, "BA4D: 18 distinct masses, since I/L and K/Q collide"
    # The empty peptide is the one way to weigh nothing, and nothing weighs less
    # than the lightest residue.
    assert ways[0] == 1, "BA4D: one empty peptide"
    assert ways[56] == 0, "BA4D: nothing is lighter than glycine's 57"
    assert ways[57] == 1, "BA4D: glycine alone"
    assert ways[114] == 2, "BA4D: GG and N both weigh 114"
}

# Rosalind: BA4E — Find a Cyclic Peptide with Theoretical Spectrum Matching an
# Ideal Spectrum
# https://rosalind.info/problems/ba4e/
#
# Given: An ideal experimental spectrum.
# Return: Every peptide whose cyclospectrum equals it, written as masses.

let spectrum = [0, 113, 128, 186, 241, 299, 314, 427]

# Branch and bound. Grow every candidate by one residue at a time, and throw away
# immediately any whose linear spectrum contains a mass the target does not —
# because adding residues can only add masses, so such a candidate can never
# recover. That pruning is the whole algorithm: without it this is 18^n.
let parent_mass = max(spectrum)
let residues = amino_acid_masses()

fn is_consistent(peptide, target) {
    let pieces = linear_spectrum(peptide)
    let remaining = target
    let ok = true
    for piece in pieces {
        if contains(remaining, piece) {
            remaining = remove_first(remaining, piece)
        } else {
            ok = false
        }
    }
    ok
}

# Drop one copy, not every copy: the spectrum is a multiset, and a candidate
# explaining a repeated mass once must not be credited with explaining it twice.
fn remove_first(items, wanted) {
    let index = (range(0, len(items)) |> filter(|i| items[i] == wanted))[0]
    range(0, len(items)) |> filter(|i| i != index) |> map(|i| items[i])
}

let candidates = [[]]
let matches = []
while len(candidates) > 0 {
    let grown = candidates |> flat_map(|peptide| residues |> map(|r| push(peptide, r)))
    candidates = []
    for peptide in grown {
        if sum(peptide) == parent_mass {
            if cyclic_spectrum(peptide) == sort(spectrum) {
                matches = push(matches, peptide)
            }
        } else {
            if is_consistent(peptide, spectrum) {
                candidates = push(candidates, peptide)
            }
        }
    }
}

let written = matches |> map(|p| p |> map(|m| str(m)) |> join("-")) |> sort()

println("Result:   " + join(written, " "))
println("Expected: 113-128-186 113-186-128 128-113-186 128-186-113 186-113-128 186-128-113")
println("(the same cycle written from each starting point and in both directions)")

fn test_ba4e_cyclopeptide_sequencing() {
    assert len(matches) == 6,
        "BA4E: expected 6 rotations and reflections, got " + str(len(matches))
    assert join(written, " ")
        == "113-128-186 113-186-128 128-113-186 128-186-113 186-113-128 186-128-113",
        "BA4E: got " + join(written, " ")
    # Every answer must reproduce the spectrum exactly, and weigh what it should.
    for peptide in matches {
        assert cyclic_spectrum(peptide) == sort(spectrum),
            "BA4E: " + str(peptide) + " does not reproduce the spectrum"
        assert sum(peptide) == parent_mass, "BA4E: wrong total mass"
    }
}

# Rosalind: BA4F — Compute the Score of a Cyclic Peptide Against a Spectrum
# https://rosalind.info/problems/ba4f/
#
# Given: An amino acid string Peptide and a collection of integers Spectrum.
# Return: Score(Peptide, Spectrum).

let peptide = "NQEL"
let observed = [0, 99, 113, 114, 128, 227, 257, 299, 355, 356, 370, 371, 484]

# Real spectra are missing masses the peptide should produce and contain masses
# it should not — noise and incomplete fragmentation — so an exact match is not
# available and the question becomes how many masses agree.
#
# Multiplicity counts. A mass appearing twice in both spectra scores two; treating
# the spectra as sets would score a peptide that explains a repeated mass once as
# generously as one that explains it fully.
let score = spectrum_score(cyclic_spectrum(peptide), observed)

println("Result:   " + str(score))
println("Expected: 11")

fn test_ba4f_cyclopeptide_scoring() {
    assert score == 11, "BA4F: got " + str(score)
    # 99 is in the observed spectrum but not in the peptide's — that is the noise
    # the score has to tolerate.
    assert contains(observed, 99), "BA4F: the sample contains a mass NQEL cannot make"
    assert contains(cyclic_spectrum(peptide), 99) == false, "BA4F: NQEL makes no 99"
    # A peptide always scores fully against its own spectrum.
    let perfect = cyclic_spectrum(peptide)
    assert spectrum_score(perfect, perfect) == len(perfect),
        "BA4F: a spectrum should match itself completely"
}

# Rosalind: BA4G — Implement LeaderboardCyclopeptideSequencing
# https://rosalind.info/problems/ba4g/
#
# Given: An integer N and a collection of integers Spectrum.
# Return: A peptide of maximum score against Spectrum, as masses.

let keep = 10
let observed = [0, 71, 113, 129, 147, 200, 218, 260, 313, 331, 347, 389, 460]

# BA4E needed a perfect spectrum. Real ones are noisy, so nothing can be pruned
# for being inconsistent — a correct peptide will contain masses the spectrum
# lacks. Instead every candidate survives on rank: grow them all, keep the best N
# by linear score, repeat. The leaderboard is what replaces the exact pruning.
let parent_mass = max(observed)
let residues = amino_acid_masses()

fn trim_to(board, spectrum, limit) {
    if len(board) <= limit { return board }
    let ranked = board
        |> map(|p| { peptide: p, score: spectrum_score(linear_spectrum(p), spectrum) })
        |> sort_by(|entry| 0 - entry.score)
    # Ties at the cutoff are all kept — dropping some arbitrarily can discard the
    # right answer while keeping an equally-scoring rival.
    let cutoff = ranked[limit - 1].score
    ranked |> filter(|entry| entry.score >= cutoff) |> map(|entry| entry.peptide)
}

let board = [[]]
let leader = []
let leader_score = 0
while len(board) > 0 {
    board = board |> flat_map(|peptide| residues |> map(|r| push(peptide, r)))
                  |> filter(|peptide| sum(peptide) <= parent_mass)
    for peptide in board {
        if sum(peptide) == parent_mass {
            # Scored cyclically here: a peptide of full mass *is* a cycle.
            let score = spectrum_score(cyclic_spectrum(peptide), observed)
            if score > leader_score {
                leader_score = score
                leader = peptide
            }
        }
    }
    board = trim_to(board, observed, keep)
}

let written = leader |> map(|m| str(m)) |> join("-")

println("Result:   " + written + "   score " + str(leader_score))
println("Expected: 113-147-71-129 (any rotation or reflection scores the same)")

fn test_ba4g_leaderboard_sequencing() {
    let published = [113, 147, 71, 129]
    assert leader_score == spectrum_score(cyclic_spectrum(published), observed),
        "BA4G: scored " + str(leader_score) + ", published scores "
            + str(spectrum_score(cyclic_spectrum(published), observed))
    assert sum(leader) == parent_mass, "BA4G: the answer must weigh the parent mass"
    # A cyclic peptide has no distinguished starting point, so the answer is
    # correct up to rotation and reflection — compare the spectra, not the lists.
    assert cyclic_spectrum(leader) == cyclic_spectrum(published),
        "BA4G: " + written + " is not a rotation of the published answer"
}

# Rosalind: BA4H — Generate the Convolution of a Spectrum
# https://rosalind.info/problems/ba4h/
#
# Given: A collection of integers Spectrum.
# Return: The convolution, in decreasing order of multiplicity, each element
# repeated as many times as it occurs.

let spectrum = [0, 137, 186, 323]

# The difference between two fragment masses is the mass of whatever lies between
# them — which for fragments differing by one residue is that residue's mass. So
# the commonest differences are the peptide's own residues, recoverable without
# assuming the standard twenty. That is what makes BA4I able to sequence peptides
# containing modified residues.
let differences = spectrum_convolution(spectrum)

let by_multiplicity = differences
    |> unique()
    |> map(|value| { value: value, count: differences |> count_if(|d| d == value) })
    |> sort_by(|entry| 0 - entry.count)

let listed = by_multiplicity |> flat_map(|entry| range(0, entry.count) |> map(|_| str(entry.value)))

println("Result:   " + join(listed, " "))
println("Expected: 137 137 186 186 323 49")

fn test_ba4h_spectral_convolution() {
    # Any order among equal multiplicities is accepted, so the check is the
    # multiset and the ordering by count, not the exact string.
    assert sort(listed) == sort(["137", "137", "186", "186", "323", "49"]),
        "BA4H: got " + join(listed, " ")
    # 0 differences are excluded, and every element is positive.
    assert (differences |> count_if(|d| d <= 0)) == 0, "BA4H: differences must be positive"
    # 137 and 186 each arise twice: 137-0 and 323-186, 186-0 and 323-137.
    assert (differences |> count_if(|d| d == 137)) == 2, "BA4H: 137 occurs twice"
    assert (differences |> count_if(|d| d == 49)) == 1, "BA4H: 186 - 137 = 49, once"
    # Ordered by multiplicity, so nothing rarer precedes something commoner.
    for i in range(1, len(by_multiplicity)) {
        assert by_multiplicity[i].count <= by_multiplicity[i - 1].count,
            "BA4H: multiplicities must not increase"
    }
}

# Rosalind: BA4I — Implement ConvolutionCyclopeptideSequencing
# https://rosalind.info/problems/ba4i/
#
# Given: Integers M and N, and a collection of integers Spectrum.
# Return: A cyclic peptide of maximum score, drawn from the M most frequent
# elements of the convolution between 57 and 200.

let take_masses = 20
let keep = 60
let observed = [57, 57, 71, 99, 129, 137, 170, 186, 194, 208, 228, 265,
                285, 299, 307, 323, 356, 364, 394, 422, 493]

# The published answer contains a residue of mass 72, which is not one of the
# twenty amino acids. That is the point of this problem: rather than assuming the
# standard alphabet, read the alphabet off the data. The commonest differences
# between fragment masses are the residues the peptide is actually built from,
# whatever they are — so a modified or non-standard residue is found rather than
# being impossible to represent.
#
# The 57-to-200 window is the plausible range for a single residue: glycine is
# the lightest at 57 and tryptophan the heaviest at 186.
let differences = spectrum_convolution(observed) |> filter(|d| d >= 57 and d <= 200)

let tallied = differences
    |> unique()
    |> map(|value| { value: value, count: differences |> count_if(|d| d == value) })
    |> sort_by(|entry| 0 - entry.count)

# Ties at the cutoff are kept, same reasoning as trimming the leaderboard.
let cutoff = tallied[take_masses - 1].count
let residues = tallied |> filter(|entry| entry.count >= cutoff) |> map(|entry| entry.value)

let parent_mass = max(observed)

fn trim_board(board, spectrum, limit) {
    if len(board) <= limit { return board }
    let ranked = board
        |> map(|p| { peptide: p, score: spectrum_score(linear_spectrum(p), spectrum) })
        |> sort_by(|entry| 0 - entry.score)
    let edge = ranked[limit - 1].score
    ranked |> filter(|entry| entry.score >= edge) |> map(|entry| entry.peptide)
}

let board = [[]]
let leader = []
let leader_score = 0
while len(board) > 0 {
    board = board |> flat_map(|peptide| residues |> map(|r| push(peptide, r)))
                  |> filter(|peptide| sum(peptide) <= parent_mass)
    for peptide in board {
        if sum(peptide) == parent_mass {
            let score = spectrum_score(cyclic_spectrum(peptide), observed)
            if score > leader_score {
                leader_score = score
                leader = peptide
            }
        }
    }
    board = trim_board(board, observed, keep)
}

let written = leader |> map(|m| str(m)) |> join("-")

println("Residues read off the data: " + (sort(residues) |> map(|m| str(m)) |> join(" ")))
println("Result:   " + written + "   score " + str(leader_score))
println("Expected: 99-71-137-57-72-57, which also scores 21 — a different peptide")
println("          of equal score, which is all a noisy spectrum can distinguish")

fn test_ba4i_convolution_sequencing() {
    let published = [99, 71, 137, 57, 72, 57]
    assert leader_score >= spectrum_score(cyclic_spectrum(published), observed),
        "BA4I: scored " + str(leader_score) + ", published scores "
            + str(spectrum_score(cyclic_spectrum(published), observed))
    assert sum(leader) == parent_mass, "BA4I: the answer must weigh the parent mass"
    # The alphabet has to include 72, which is not an amino acid mass — if it did
    # not, the published answer would be unreachable.
    assert contains(residues, 72), "BA4I: 72 should be read off the convolution"
    assert contains(amino_acid_masses(), 72) == false,
        "BA4I: and 72 is not a standard residue mass"
}

# Rosalind: BA4J — Generate the Theoretical Spectrum of a Linear Peptide
# https://rosalind.info/problems/ba4j/
#
# Given: An amino acid string Peptide.
# Return: LinearSpectrum(Peptide).

let peptide = "NQEL"

# The same idea as BA4C with the ends not joined, so nothing wraps. That makes
# the linear spectrum a strict subset of the cyclic one — which matters for
# sequencing, because a linear score can be computed for a partial peptide that
# is not yet a full cycle.
let spectrum = linear_spectrum(peptide)

println("Result:   " + (spectrum |> map(|m| str(m)) |> join(" ")))
println("Expected: 0 113 114 128 129 242 242 257 370 371 484")

fn test_ba4j_linear_spectrum() {
    assert (spectrum |> map(|m| str(m)) |> join(" "))
        == "0 113 114 128 129 242 242 257 370 371 484",
        "BA4J: got " + str(spectrum)
    # A linear peptide of length n has n(n+1)/2 subpeptides, plus the empty one.
    let n = len(peptide)
    assert len(spectrum) == n * (n + 1) / 2 + 1,
        "BA4J: expected " + str(n * (n + 1) / 2 + 1) + " masses"
    # And every one of them also appears in the cyclic spectrum.
    let cyclic = cyclic_spectrum(peptide)
    assert spectrum_score(spectrum, cyclic) == len(spectrum),
        "BA4J: every linear fragment should also be a cyclic one"
    assert len(cyclic) > len(spectrum), "BA4J: the wrapping pieces are extra"
}

# Rosalind: BA4K — Compute the Score of a Linear Peptide Against a Spectrum
# https://rosalind.info/problems/ba4k/
#
# Given: An amino acid string Peptide and a collection of integers Spectrum.
# Return: LinearScore(Peptide, Spectrum).

let peptide = "NQEL"
let observed = [0, 99, 113, 114, 128, 227, 257, 299, 355, 356, 370, 371, 484]

# The same comparison as BA4F against the linear spectrum, which scores lower
# because it has fewer masses to agree with — 8 against 11 on identical input.
#
# The reason to have both: sequencing grows a peptide one residue at a time, and
# a partial peptide is not yet a cycle. Scoring it cyclically would credit it with
# wrap-around fragments it does not have, and rank a bad prefix above a good one.
let score = spectrum_score(linear_spectrum(peptide), observed)

println("Result:   " + str(score))
println("Expected: 8")

fn test_ba4k_linear_scoring() {
    assert score == 8, "BA4K: got " + str(score)
    # Strictly lower than the cyclic score on the same input, because the linear
    # spectrum is a subset of the cyclic one.
    let cyclic = spectrum_score(cyclic_spectrum(peptide), observed)
    assert cyclic == 11, "BA4K: the cyclic score of the same peptide is 11"
    assert score < cyclic, "BA4K: linear scoring cannot exceed cyclic"
}

# Rosalind: BA4L — Trim a Peptide Leaderboard
# https://rosalind.info/problems/ba4l/
#
# Given: A leaderboard of linear peptides, a spectrum, and an integer N.
# Return: The top N peptides, scored with LinearScore — plus every peptide tied
# with the Nth.

let leaderboard = ["LAST", "ALST", "TLLT", "TQAS"]
let observed = [0, 71, 87, 101, 113, 158, 184, 188, 259, 271, 372]
let keep = 2

# "Top N" with ties kept, not exactly N. Cutting a tie arbitrarily would make the
# search depend on the order peptides happened to be generated in, and can throw
# away the correct answer while retaining an equally-scoring rival. On this
# sample nothing actually ties at the cutoff, so the answer is the plain top two.
let scored = leaderboard |> map(|p| { peptide: p, score: spectrum_score(linear_spectrum(p), observed) })
let ranked = scored |> sort_by(|entry| 0 - entry.score)
let cutoff = ranked[keep - 1].score
let trimmed = ranked |> filter(|entry| entry.score >= cutoff) |> map(|entry| entry.peptide)

println("Scores:")
for entry in ranked { println("  " + entry.peptide + "  " + str(entry.score)) }
println("Result:   " + join(trimmed, " "))
println("Expected: LAST ALST")

fn test_ba4l_trim_leaderboard() {
    assert join(trimmed, " ") == "LAST ALST", "BA4L: got " + join(trimmed, " ")
    # LAST and ALST are anagrams and still score differently — 11 against 9 —
    # because linear subpeptides are *contiguous*. Rearranging the residues keeps
    # the total mass and the single-residue masses but changes every fragment in
    # between, which is exactly why a spectrum says something about order.
    assert sort(chars("LAST")) == sort(chars("ALST")), "BA4L: the two are anagrams"
    assert peptide_mass("LAST") == peptide_mass("ALST"), "BA4L: so they weigh the same"
    assert linear_spectrum("LAST") != linear_spectrum("ALST"),
        "BA4L: but their contiguous fragments differ"
    assert spectrum_score(linear_spectrum("LAST"), observed) == 11, "BA4L: LAST scores 11"
    assert spectrum_score(linear_spectrum("ALST"), observed) == 9, "BA4L: ALST scores 9"
    # Nothing below the cutoff survives.
    for entry in scored {
        if contains(trimmed, entry.peptide) == false {
            assert entry.score < cutoff, "BA4L: " + entry.peptide + " was dropped despite tying"
        }
    }
}

# Rosalind: BA4M — Solve the Turnpike Problem
# https://rosalind.info/problems/ba4m/
#
# Given: All pairwise differences between points on a line.
# Return: A set of points A with those differences.

let differences = [-10, -8, -7, -6, -5, -4, -3, -3, -2, -2, 0, 0, 0, 0, 0,
                   2, 2, 3, 3, 4, 5, 6, 7, 8, 10]

# The same shape of problem as reading a peptide from its fragment masses, with
# positions on a line instead of residues — which is why it sits in this chapter.
#
# Backtracking on the largest unplaced distance. That distance must be between
# some point and one of the two ends, so there are only two choices at each step,
# and each is checked immediately against the remaining multiset. Placing points
# in arbitrary order instead would be a search over all subsets.
let positive = differences |> filter(|d| d > 0) |> sort()
let width = max(differences)

fn remove_all(pool, wanted) {
    let remaining = pool
    let ok = true
    for value in wanted {
        let found = range(0, len(remaining)) |> filter(|i| remaining[i] == value)
        if len(found) == 0 {
            ok = false
        } else {
            let at = found[0]
            remaining = range(0, len(remaining)) |> filter(|i| i != at) |> map(|i| remaining[i])
        }
    }
    { ok: ok, rest: remaining }
}

# Distances from a candidate point to everything already placed.
fn spans(point, placed) { placed |> map(|p| abs(point - p)) }

let solution = []
let stack = [{ placed: [0, width], pool: (remove_all(positive, [width])).rest }]
while len(stack) > 0 and len(solution) == 0 {
    let state = stack[len(stack) - 1]
    stack = slice(stack, 0, len(stack) - 1)

    if len(state.pool) == 0 {
        solution = sort(state.placed)
    } else {
        let largest = max(state.pool)
        # The largest remaining distance reaches either from the left end or to
        # the right end; nothing else could produce it.
        for candidate in [largest, width - largest] {
            if contains(state.placed, candidate) == false {
                let attempt = remove_all(state.pool, spans(candidate, state.placed))
                if attempt.ok {
                    stack = push(stack, {
                        placed: push(state.placed, candidate),
                        pool: attempt.rest,
                    })
                }
            }
        }
    }
}

println("Result:   " + (solution |> map(|p| str(p)) |> join(" ")))
println("Expected: 0 2 4 7 10")

fn test_ba4m_turnpike() {
    assert (solution |> map(|p| str(p)) |> join(" ")) == "0 2 4 7 10",
        "BA4M: got " + str(solution)
    # The real requirement: the answer's own pairwise differences are the input.
    let rebuilt = solution |> flat_map(|a| solution |> map(|b| a - b)) |> sort()
    assert rebuilt == sort(differences),
        "BA4M: the reconstructed differences do not match the input"
    assert len(rebuilt) == len(solution) * len(solution),
        "BA4M: n points give n^2 differences, including the n zeros"
}

# Rosalind: BA11A — Construct the Graph of a Spectrum
# https://rosalind.info/problems/ba11a/
#
# Given: A spectrum of masses.
# Return: The spectrum graph — an edge between two masses whenever their
# difference is the mass of an amino acid.

let spectrum = [57, 71, 154, 185, 301, 332, 415, 429, 486]

# Reading a peptide off a spectrum, rather than guessing peptides and scoring
# them as BA4 did. Every path from 0 to the largest mass spells a candidate,
# because consecutive prefix masses differ by exactly one residue — so sequencing
# becomes a path problem instead of a search over 20^n peptides.
#
# The graph is a DAG (masses only increase), which is what makes BA11B able to
# read every candidate off it cheaply.
let mass_table = [
    { letter: "G", mass: 57 },  { letter: "A", mass: 71 },  { letter: "S", mass: 87 },
    { letter: "P", mass: 97 },  { letter: "V", mass: 99 },  { letter: "T", mass: 101 },
    { letter: "C", mass: 103 }, { letter: "I", mass: 113 }, { letter: "L", mass: 113 },
    { letter: "N", mass: 114 }, { letter: "D", mass: 115 }, { letter: "K", mass: 128 },
    { letter: "Q", mass: 128 }, { letter: "E", mass: 129 }, { letter: "M", mass: 131 },
    { letter: "H", mass: 137 }, { letter: "F", mass: 147 }, { letter: "R", mass: 156 },
    { letter: "Y", mass: 163 }, { letter: "W", mass: 186 },
]

# Rosalind's answer names one letter per mass, taking the first of each
# colliding pair — I before L, K before Q.
let canonical = {}
for entry in mass_table {
    if contains(keys(canonical), str(entry.mass)) == false {
        canonical[str(entry.mass)] = entry.letter
    }
}

let with_zero = [0] + spectrum
let arcs = with_zero |> flat_map(|from_mass|
    with_zero
        |> filter(|to_mass| to_mass > from_mass
                  and contains(keys(canonical), str(to_mass - from_mass)))
        |> map(|to_mass| str(from_mass) + "->" + str(to_mass) + ":"
                         + canonical[str(to_mass - from_mass)]))

println("Result:")
for line in arcs { println("  " + line) }
println("Expected: 0->57:G 0->71:A 57->154:P 57->185:K 71->185:N 154->301:F")
println("          185->332:F 301->415:N 301->429:K 332->429:P 415->486:A 429->486:G")

fn test_ba11a_spectrum_graph() {
    let expected = ["0->57:G", "0->71:A", "57->154:P", "57->185:K", "71->185:N",
                    "154->301:F", "185->332:F", "301->415:N", "301->429:K",
                    "332->429:P", "415->486:A", "429->486:G"]
    assert sort(arcs) == sort(expected), "BA11A: got " + join(sort(arcs), " ")
    # Every edge's label really weighs the difference it spans.
    for line in arcs {
        let parts = split(line, "->")
        let ends = split(parts[1], ":")
        let gap = int(ends[0]) - int(parts[0])
        assert canonical[str(gap)] == ends[1], "BA11A: " + line + " is mislabelled"
    }
    # Masses only increase, so the graph is acyclic — which is what lets BA11B
    # enumerate its paths.
    for line in arcs {
        let parts = split(line, "->")
        assert int(split(parts[1], ":")[0]) > int(parts[0]), "BA11A: an edge goes backwards"
    }
}

# Rosalind: BA11B — Implement DecodingIdealSpectrum
# https://rosalind.info/problems/ba11b/
#
# Given: An ideal spectrum.
# Return: A peptide whose ideal spectrum it is.

let spectrum = [57, 71, 154, 185, 301, 332, 415, 429, 486]

let mass_table = [
    { letter: "G", mass: 57 },  { letter: "A", mass: 71 },  { letter: "S", mass: 87 },
    { letter: "P", mass: 97 },  { letter: "V", mass: 99 },  { letter: "T", mass: 101 },
    { letter: "C", mass: 103 }, { letter: "I", mass: 113 }, { letter: "L", mass: 113 },
    { letter: "N", mass: 114 }, { letter: "D", mass: 115 }, { letter: "K", mass: 128 },
    { letter: "Q", mass: 128 }, { letter: "E", mass: 129 }, { letter: "M", mass: 131 },
    { letter: "H", mass: 137 }, { letter: "F", mass: 147 }, { letter: "R", mass: 156 },
    { letter: "Y", mass: 163 }, { letter: "W", mass: 186 },
]
let canonical = {}
for entry in mass_table {
    if contains(keys(canonical), str(entry.mass)) == false {
        canonical[str(entry.mass)] = entry.letter
    }
}

# Every path through BA11A's graph from 0 to the heaviest mass spells a
# candidate, but not every candidate is right: the spectrum holds prefix *and*
# suffix masses mixed together, and a path only accounts for the prefixes. So
# each candidate is generated and then checked by rebuilding its full ideal
# spectrum — generate-and-test, but over a handful of paths rather than 20^n
# peptides.
let masses = [0] + spectrum
let heaviest = max(spectrum)

fn ideal_spectrum_of(peptide, letters) {
    let running = 0
    let prefixes = []
    for residue in chars(peptide) {
        running = running + letters[residue]
        prefixes = push(prefixes, running)
    }
    let total = running
    # Prefixes and suffixes together. The empty piece is dropped, and the whole
    # peptide appears once rather than twice — it is both the last prefix and the
    # last suffix, and the spectrum lists it a single time.
    let proper_prefixes = prefixes |> filter(|m| m != total)
    let proper_suffixes = prefixes |> map(|m| total - m) |> filter(|m| m != 0)
    sort(proper_prefixes + proper_suffixes + [total])
}

let letter_mass = {}
for entry in mass_table { letter_mass[entry.letter] = entry.mass }

# Depth-first over the graph, collecting complete paths.
let candidates = []
let stack = [{ at: 0, spelled: "" }]
while len(stack) > 0 {
    let state = stack[len(stack) - 1]
    stack = slice(stack, 0, len(stack) - 1)
    if state.at == heaviest {
        candidates = push(candidates, state.spelled)
    } else {
        for onward in masses {
            if onward > state.at and contains(keys(canonical), str(onward - state.at)) {
                stack = push(stack, {
                    at: onward,
                    spelled: state.spelled + canonical[str(onward - state.at)],
                })
            }
        }
    }
}

let answers = candidates |> filter(|p| ideal_spectrum_of(p, letter_mass) == sort(spectrum))

println("Candidates from the graph: " + join(candidates, " "))
println("Consistent with the spectrum: " + join(answers, " "))
println("Result:   " + answers[0])
println("Expected: GPFNA — and ANFPG is equally correct, being its reverse:")
println("          an ideal spectrum holds prefixes and suffixes together, and")
println("          reversing a peptide simply swaps which is which.")

fn test_ba11b_decoding_ideal_spectrum() {
    assert contains(answers, "GPFNA"), "BA11B: GPFNA should be among " + join(answers, " ")
    assert ideal_spectrum_of("GPFNA", letter_mass) == sort(spectrum),
        "BA11B: GPFNA's ideal spectrum must be the input"
    # The filtering step is doing real work — the graph offers paths that are not
    # answers, which is why generate-and-test is needed rather than any path.
    assert len(candidates) > len(answers),
        "BA11B: some graph paths should fail the spectrum check"
    # A peptide and its reverse have the same ideal spectrum, so both survive —
    # this is a real ambiguity in the data, not a bug in the search.
    assert contains(answers, "ANFPG"), "BA11B: the reverse should also be consistent"
    assert reverse("GPFNA") == "ANFPG", "BA11B: and it is the reverse"
    assert ideal_spectrum_of("ANFPG", letter_mass) == ideal_spectrum_of("GPFNA", letter_mass),
        "BA11B: a peptide and its reverse share an ideal spectrum"
}

# Rosalind: BA11C — Convert a Peptide into a Peptide Vector
# https://rosalind.info/problems/ba11c/
#
# Given: A peptide.
# Return: Its peptide vector — a 1 at every prefix mass, 0 elsewhere.

# The toy alphabet Rosalind uses for this chapter: X weighs 4 and Z weighs 5.
# Small masses keep the vectors readable; the real 18-mass table works the same
# way and produces vectors thousands of entries long.
let peptide = "XZZXX"
let toy_masses = { "X": 4, "Z": 5 }

# A peptide vector turns a peptide into something the same shape as a spectrum,
# so the two can be compared by a dot product. That is the whole idea behind the
# chapter: scoring a peptide against a spectrum becomes one multiplication per
# position, and finding the best peptide becomes a path problem over the vector.
let prefix_masses = []
let running = 0
for residue in chars(peptide) {
    running = running + toy_masses[residue]
    prefix_masses = push(prefix_masses, running)
}

let total = running
let vector = range(1, total + 1) |> map(|mass| if contains(prefix_masses, mass) then 1 else 0)

println("Result:   " + (vector |> map(|v| str(v)) |> join(" ")))
println("Expected: 0 0 0 1 0 0 0 0 1 0 0 0 0 1 0 0 0 1 0 0 0 1")

fn test_ba11c_peptide_to_vector() {
    assert (vector |> map(|v| str(v)) |> join(" "))
        == "0 0 0 1 0 0 0 0 1 0 0 0 0 1 0 0 0 1 0 0 0 1",
        "BA11C: got " + join(vector |> map(|v| str(v)), " ")
    # As many 1s as residues, and the last entry is always 1 — the peptide's own
    # mass is its final prefix.
    assert sum(vector) == len(peptide), "BA11C: one 1 per residue"
    assert vector[len(vector) - 1] == 1, "BA11C: the total mass is the last prefix"
    assert len(vector) == total, "BA11C: the vector is as long as the peptide is heavy"
    assert prefix_masses == [4, 9, 14, 18, 22], "BA11C: prefix masses"
}

# Rosalind: BA11D — Convert a Peptide Vector into a Peptide
# https://rosalind.info/problems/ba11d/
#
# Given: A peptide vector.
# Return: A peptide with that vector.

let vector = [0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1]
let toy_masses = { "X": 4, "Z": 5 }

# The inverse of BA11C. Each 1 is a prefix mass, so the gaps between consecutive
# 1s are the residue masses — reading them off in order spells the peptide, and
# nothing has to be searched for.
let prefix_positions = range(0, len(vector)) |> filter(|i| vector[i] == 1) |> map(|i| i + 1)
let gaps = range(0, len(prefix_positions)) |> map(|i| if i == 0 then prefix_positions[0] else prefix_positions[i] - prefix_positions[i - 1])

let by_mass = {}
for letter in keys(toy_masses) { by_mass[str(toy_masses[letter])] = letter }

let peptide = gaps |> map(|gap| by_mass[str(gap)]) |> join("")

println("Result:   " + peptide)
println("Expected: XZZXX")

fn test_ba11d_vector_to_peptide() {
    assert peptide == "XZZXX", "BA11D: got " + peptide
    # It really inverts BA11C: rebuilding the vector returns the input.
    let running = 0
    let rebuilt_prefixes = []
    for residue in chars(peptide) {
        running = running + toy_masses[residue]
        rebuilt_prefixes = push(rebuilt_prefixes, running)
    }
    let rebuilt = range(1, running + 1)
        |> map(|mass| if contains(rebuilt_prefixes, mass) then 1 else 0)
    assert rebuilt == vector, "BA11D: the round trip does not return the vector"
    assert len(peptide) == sum(vector), "BA11D: one residue per 1 in the vector"
}

# Rosalind: BA11E — Sequence a Peptide
# https://rosalind.info/problems/ba11e/
#
# Given: A spectral vector S.
# Return: A peptide whose peptide vector scores highest against S.

let spectral = [0, 0, 0, 4, -2, -3, -1, -7, 6, 5, 3, 2, 1, 9, 3, -8, 0, 3, 1, 2, 1, 0]
let toy_masses = { "X": 4, "Z": 5 }

# BA11C made peptides and spectra the same shape so they could be compared by a
# dot product. This is why: a peptide's score is the sum of the spectral entries
# at its prefix masses, so the best peptide is the heaviest path through a graph
# whose nodes are positions and whose edges are residues.
#
# Negative entries matter — a spectral vector is real measurement, not a count,
# so a path is penalised for claiming a prefix the data argues against. That is
# what stops the answer being simply the longest path.
let masses = [0] + spectral
let sink = len(spectral)

let best = range(0, sink + 1) |> map(|i| if i == 0 then 0 else 0 - 1000000)
let came_from = {}

for position in range(1, sink + 1) {
    for letter in keys(toy_masses) {
        let previous = position - toy_masses[letter]
        if previous >= 0 and best[previous] > 0 - 1000000 {
            let candidate = best[previous] + spectral[position - 1]
            if candidate > best[position] {
                best[position] = candidate
                came_from[str(position)] = { at: previous, letter: letter }
            }
        }
    }
}

let peptide = ""
let at = sink
while at > 0 {
    let step = came_from[str(at)]
    peptide = step.letter + peptide
    at = step.at
}

println("Result:   " + peptide + "   score " + str(best[sink]))
println("Expected: XZZXX")

fn test_ba11e_peptide_sequencing() {
    assert peptide == "XZZXX", "BA11E: got " + peptide
    # The reported score must be what the peptide actually scores against S.
    let running = 0
    let prefixes = []
    for residue in chars(peptide) {
        running = running + toy_masses[residue]
        prefixes = push(prefixes, running)
    }
    let scored = prefixes |> map(|m| spectral[m - 1]) |> sum()
    assert scored == best[sink],
        "BA11E: the path claims " + str(best[sink]) + " but the peptide scores " + str(scored)
    # The peptide's mass has to be the vector's length — a shorter one is not a
    # candidate at all, however well it scores.
    assert running == len(spectral), "BA11E: the peptide must weigh the whole vector"
    # And it beats an alternative of the same mass.
    let rival = "ZXXZX"
    let rival_running = 0
    let rival_prefixes = []
    for residue in chars(rival) {
        rival_running = rival_running + toy_masses[residue]
        rival_prefixes = push(rival_prefixes, rival_running)
    }
    assert (rival_prefixes |> map(|m| spectral[m - 1]) |> sum()) < scored,
        "BA11E: " + rival + " should score lower"
}

# Rosalind: BA11F — Find a Highest-Scoring Peptide in a Proteome against a Spectrum
# https://rosalind.info/problems/ba11f/
#
# Given: A spectral vector S and a proteome.
# Return: The substring of the proteome scoring highest against S.

let spectral = [0, 0, 0, 4, -2, -3, -1, -7, 6, 5, 3, 2, 1, 9, 3, -8, 0, 3, 1, 2, 1, 8]
let proteome = "XZZXZXXXZXZZXZXXZ"
let toy_masses = { "X": 4, "Z": 5 }

# The realistic version of BA11E. There, any peptide the masses allowed was a
# candidate; here only substrings of a known proteome are, which is how
# proteomics actually works — the genome is sequenced first, and the spectrum is
# matched against what it could produce. Constraining the search that way is also
# what makes it tractable.
let total = len(spectral)

fn mass_of(piece, letters) { chars(piece) |> map(|c| letters[c]) |> sum() }

fn score_of(piece, letters, vector) {
    let running = 0
    let prefixes = []
    for residue in chars(piece) {
        running = running + letters[residue]
        prefixes = push(prefixes, running)
    }
    prefixes |> map(|m| vector[m - 1]) |> sum()
}

# Only substrings weighing exactly the vector's length can be compared at all.
let candidates = range(0, len(proteome)) |> flat_map(|start|
    range(start + 1, len(proteome) + 1)
        |> map(|stop| substr(proteome, start, stop - start))
        |> filter(|piece| mass_of(piece, toy_masses) == total))

let ranked = candidates
    |> map(|piece| { piece: piece, score: score_of(piece, toy_masses, spectral) })
    |> sort_by(|entry| 0 - entry.score)

println("Candidates of the right mass: " + join(candidates, " "))
println("Result:   " + ranked[0].piece + "   score " + str(ranked[0].score))
println("Expected: ZXZXX")

fn test_ba11f_peptide_identification() {
    assert ranked[0].piece == "ZXZXX", "BA11F: got " + ranked[0].piece
    # It really is a substring of the proteome, and weighs the whole vector.
    assert contains(proteome, ranked[0].piece), "BA11F: the answer must occur in the proteome"
    assert mass_of(ranked[0].piece, toy_masses) == total,
        "BA11F: the peptide must weigh the vector's length"
    # Nothing else scores higher, and there was more than one candidate — so the
    # mass filter alone did not decide it.
    for entry in ranked {
        assert entry.score <= ranked[0].score, "BA11F: something outscores the answer"
    }
    assert len(candidates) > 1, "BA11F: several substrings have the right mass"
}

# Rosalind: BA11G — Implement PSMSearch
# https://rosalind.info/problems/ba11g/
#
# Given: A set of spectral vectors, a proteome, and a threshold.
# Return: Every peptide-spectrum match scoring at least the threshold.

let spectra = [
    [-1, 5, -4, 5, 3, -1, -4, 5, -1, 0, 0, 4, -1, 0, 1, 4, 4, 4],
    [-4, 2, -2, -4, 4, -5, -1, 4, -1, 2, 5, -3, -1, 3, 2, -3],
]
let proteome = "XXXZXZXXZXZXXXZXXZX"
let threshold = 5
let toy_masses = { "X": 4, "Z": 5 }

# BA11F finds the best peptide for one spectrum whether or not it is any good.
# A real experiment produces thousands of spectra, most of which match nothing —
# they are noise, or peptides absent from the proteome. The threshold is what
# separates the two, and without it every spectrum would be assigned a peptide
# and most of those assignments would be wrong.
fn mass_of(piece, letters) { chars(piece) |> map(|c| letters[c]) |> sum() }

fn score_of(piece, letters, vector) {
    let running = 0
    let prefixes = []
    for residue in chars(piece) {
        running = running + letters[residue]
        prefixes = push(prefixes, running)
    }
    prefixes |> map(|m| vector[m - 1]) |> sum()
}

fn best_match(vector, source, letters) {
    let total = len(vector)
    let candidates = range(0, len(source)) |> flat_map(|start|
        range(start + 1, len(source) + 1)
            |> map(|stop| substr(source, start, stop - start))
            |> filter(|piece| mass_of(piece, letters) == total))
    if len(candidates) == 0 { return { piece: "", score: 0 - 1000000 } }
    (candidates
        |> map(|piece| { piece: piece, score: score_of(piece, letters, vector) })
        |> sort_by(|entry| 0 - entry.score))[0]
}

let matches = spectra
    |> map(|vector| best_match(vector, proteome, toy_masses))
    |> filter(|entry| entry.score >= threshold)
    |> map(|entry| entry.piece)
    |> unique()

println("Result:   " + join(matches, " "))
println("Expected: XZXZ")

fn test_ba11g_psm_search() {
    assert join(matches, " ") == "XZXZ", "BA11G: got " + join(matches, " ")
    # Exactly one of the two spectra clears the threshold — the other's best
    # match scores below it and is correctly left unassigned, which is the whole
    # purpose of the threshold.
    let scored = spectra |> map(|vector| best_match(vector, proteome, toy_masses))
    assert scored[0].score >= threshold, "BA11G: the first spectrum should match"
    assert scored[1].score < threshold,
        "BA11G: the second scores " + str(scored[1].score) + ", below the threshold"
    assert contains(proteome, matches[0]), "BA11G: the match must occur in the proteome"
}

# Rosalind: BA11H — Compute the Size of a Spectral Dictionary
# https://rosalind.info/problems/ba11h/
#
# Given: A spectral vector, a threshold, and a maximum score.
# Return: How many peptides score within [threshold, max_score].

let spectral = [4, -3, -2, 3, 3, -4, 5, -3, -1, -1, 3, 4, 1, 3]
let threshold = 1
let max_score = 8
let toy_masses = { "X": 4, "Z": 5 }

# BA11G assigns a peptide to a spectrum, but how confident should anyone be? If
# thousands of peptides would have scored just as well, a high score means
# nothing. The spectral dictionary is that count — the number of peptides
# reaching a given score — and it is what turns a score into a statistical
# statement rather than a number.
#
# Counted rather than enumerated, for the same reason as BA4D: the dictionary can
# be astronomically large even when the count is quick to compute.
let residues = keys(toy_masses) |> map(|letter| toy_masses[letter])

# ways[mass][score] = how many peptides of that mass reach exactly that score.
# Scores can go negative, so they are shifted to keep indices non-negative.
let shift = 1000
# Built as a fresh record per mass, since nested index assignment is not
# available — only `ways[i] = record`.
let start_row = {}
start_row[str(shift)] = 1
let ways = [start_row]
for _ in range(1, len(spectral) + 1) { ways = push(ways, {}) }

for mass in range(1, len(spectral) + 1) {
    let here = {}
    for residue in residues {
        let previous = mass - residue
        if previous >= 0 {
            for key in keys(ways[previous]) {
                let new_score = int(key) + spectral[mass - 1]
                let existing = if contains(keys(here), str(new_score)) then here[str(new_score)] else 0
                here[str(new_score)] = existing + ways[previous][key]
            }
        }
    }
    ways[mass] = here
}

let final_scores = ways[len(spectral)]
let size = keys(final_scores)
    |> filter(|key| int(key) - shift >= threshold and int(key) - shift <= max_score)
    |> map(|key| final_scores[key])
    |> sum()

println("Result:   " + str(size))
println("Expected: 3")

fn test_ba11h_spectral_dictionary_size() {
    assert size == 3, "BA11H: got " + str(size)
    # Checked by enumeration, which is possible only because this vector is tiny
    # — the point of the counting version is that it stays possible when the
    # dictionary does not fit in memory.
    fn peptides_of_mass(target, letters) {
        let growing = [""]
        let complete = []
        for _ in range(0, target) {
            let extended = growing
                |> flat_map(|p| keys(letters) |> map(|c| p + c))
                |> filter(|p| (chars(p) |> map(|c| letters[c]) |> sum()) <= target)
            # Finished peptides are set aside rather than extended further —
            # carrying them on would push every one past the target and leave
            # nothing behind.
            complete = complete + (extended
                |> filter(|p| (chars(p) |> map(|c| letters[c]) |> sum()) == target))
            growing = extended
                |> filter(|p| (chars(p) |> map(|c| letters[c]) |> sum()) < target)
        }
        unique(complete)
    }
    let all_peptides = peptides_of_mass(len(spectral), toy_masses)
    let scored = all_peptides |> map(|p| {
        let running = 0
        let prefixes = []
        for residue in chars(p) {
            running = running + toy_masses[residue]
            prefixes = push(prefixes, running)
        }
        prefixes |> map(|m| spectral[m - 1]) |> sum()
    })
    let within = scored |> count_if(|s| s >= threshold and s <= max_score)
    assert within == size,
        "BA11H: counting says " + str(size) + " but enumeration finds " + str(within)
}

# Rosalind: BA11I — Compute the Probability of a Spectral Dictionary
# https://rosalind.info/problems/ba11i/
#
# Given: A spectral vector, a threshold, and a maximum score.
# Return: The probability of the spectral dictionary.

let spectral = [4, -3, -2, 3, 3, -4, 5, -3, -1, -1, 3, 4, 1, 3]
let threshold = 1
let max_score = 8
let toy_masses = { "X": 4, "Z": 5 }

# BA11H counts the peptides reaching a score; this weights them. Every residue is
# taken as equally likely, so a peptide of length n has probability
# (1/|alphabet|)^n — and summing that over the dictionary gives the chance a
# random peptide would score as well as the observed match.
#
# That number is what turns a match into evidence. A score of 8 means nothing on
# its own; a score only 0.375 of random peptides reach means something, and a
# score one in a billion reach means a great deal more.
let residues = keys(toy_masses)
let share = 1.0 / len(residues)

# Same recurrence as BA11H, carrying probability instead of a count.
let shift = 1000
let start_row = {}
start_row[str(shift)] = 1.0
let ways = [start_row]
for _ in range(1, len(spectral) + 1) { ways = push(ways, {}) }

for mass in range(1, len(spectral) + 1) {
    let here = {}
    for letter in residues {
        let previous = mass - toy_masses[letter]
        if previous >= 0 {
            for key in keys(ways[previous]) {
                let new_score = int(key) + spectral[mass - 1]
                let existing = if contains(keys(here), str(new_score)) then here[str(new_score)] else 0.0
                here[str(new_score)] = existing + ways[previous][key] * share
            }
        }
    }
    ways[mass] = here
}

let final_scores = ways[len(spectral)]
let probability = keys(final_scores)
    |> filter(|key| int(key) - shift >= threshold and int(key) - shift <= max_score)
    |> map(|key| final_scores[key])
    |> sum()

println("Result:   " + str(probability))
println("Expected: 0.375")

fn test_ba11i_spectral_dictionary_probability() {
    assert abs(probability - 0.375) < 1e-9, "BA11I: got " + str(probability)
    # BA11H found 3 peptides in the dictionary, all of length 3, and each has
    # probability (1/2)^3 = 0.125 — so 3 * 0.125 = 0.375. The two problems have
    # to agree that way, which is worth checking rather than assuming.
    assert abs(3.0 * pow(share, 3) - probability) < 1e-9,
        "BA11I: three length-3 peptides at (1/2)^3 each should give the answer"
    # A probability, so between 0 and 1.
    assert probability >= 0.0 and probability <= 1.0, "BA11I: not a probability"
    # Widening the score window can only include more.
    let wider = keys(final_scores)
        |> filter(|key| int(key) - shift >= threshold - 5 and int(key) - shift <= max_score + 5)
        |> map(|key| final_scores[key])
        |> sum()
    assert wider >= probability, "BA11I: a wider window cannot be less likely"
}

# Rosalind: BA11J — Find a Highest-Scoring Modified Peptide against a Spectrum
# https://rosalind.info/problems/ba11j/
#
# Given: A peptide, a spectral vector, and an integer k.
# Return: A variant of the peptide, with at most k modified residues, scoring
# highest against the vector.

let peptide = "XXZ"
let spectral = [4, -3, -2, 3, 3, -4, 5, -3, -1, -1, 3, 4, 1, 3]
let allowed = 2
let toy_masses = { "X": 4, "Z": 5 }

# Proteins are chemically modified after they are made — phosphorylated,
# methylated, acetylated — and a modified residue weighs something other than the
# table says. A spectrum of a modified peptide therefore matches nothing under
# exact search, which is why identification has to allow for shifts.
#
# Spectral alignment: each residue may take a mass offset, at most k of them
# non-zero. The state is (which residue, what total mass so far, how many
# modifications used), and the answer is the best-scoring path through it. The
# peptide's own mass here is 13 against a vector of length 14, so at least one
# modification is forced.
let residue_masses = chars(peptide) |> map(|c| toy_masses[c])
let total = len(spectral)

# best[i][m][k] via a flat record keyed by the three indices.
fn key(i, m, used) { str(i) + "," + str(m) + "," + str(used) }

let best = {}
let came_from = {}
best[key(0, 0, 0)] = 0

for i in range(0, len(residue_masses)) {
    for m in range(0, total + 1) {
        for used in range(0, allowed + 1) {
            let from_key = key(i, m, used)
            if contains(keys(best), from_key) {
                # Every reachable mass for the next prefix. An unmodified step
                # adds the residue's own mass; anything else costs a modification.
                for next_mass in range(m + 1, total + 1) {
                    let shift = next_mass - m - residue_masses[i]
                    let cost = if shift == 0 then 0 else 1
                    if used + cost <= allowed {
                        let to_key = key(i + 1, next_mass, used + cost)
                        let candidate = best[from_key] + spectral[next_mass - 1]
                        let known = if contains(keys(best), to_key) then best[to_key] else 0 - 1000000
                        if candidate > known {
                            best[to_key] = candidate
                            came_from[to_key] = { at: from_key, shift: shift, mass: next_mass }
                        }
                    }
                }
            }
        }
    }
}

# The best complete variant: all residues placed, total mass reached.
let finals = range(0, allowed + 1)
    |> filter(|used| contains(keys(best), key(len(residue_masses), total, used)))
    |> map(|used| { used: used, score: best[key(len(residue_masses), total, used)] })
    |> sort_by(|entry| 0 - entry.score)

let at = key(len(residue_masses), total, finals[0].used)
let shifts = []
while contains(keys(came_from), at) {
    let step = came_from[at]
    shifts = [step.shift] + shifts
    at = step.at
}

let written = range(0, len(shifts)) |> map(|i| {
    let letter = substr(peptide, i, 1)
    if shifts[i] == 0 then letter
    else if shifts[i] > 0 then letter + "(+" + str(shifts[i]) + ")"
    else letter + "(" + str(shifts[i]) + ")"
}) |> join("")

println("Result:   " + written + "   score " + str(finals[0].score))
println("Expected: XX(-1)Z(+2)")

fn test_ba11j_spectral_alignment() {
    assert written == "XX(-1)Z(+2)", "BA11J: got " + written
    # At most k modifications, and here exactly two are used.
    let modified = shifts |> count_if(|s| s != 0)
    assert modified <= allowed, "BA11J: too many modifications"
    assert modified == 2, "BA11J: expected two modifications, got " + str(modified)
    # The shifts must carry the peptide's mass to the vector's length — that is
    # what forces a modification here at all.
    let unmodified_mass = sum(residue_masses)
    assert unmodified_mass == 13, "BA11J: XXZ weighs 13"
    assert unmodified_mass + sum(shifts) == total,
        "BA11J: the shifts must make up the difference to " + str(total)
    assert unmodified_mass != total, "BA11J: so at least one modification is forced"
}