# 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"
}