# Rosalind: GAFF — Global Alignment with Scoring Matrix and Affine Gap Penalty
# https://rosalind.info/problems/gaff/
#
# Given: Two protein strings.
# Return: The maximum global alignment score using BLOSUM62, a gap opening
# penalty of 11 and a gap extension penalty of 1.

let s = "PRTEINS"
let t = "PRTWPSEIN"
let gap_open = -11.0
let gap_extend = -1.0

let blosum = score_matrix("blosum62")
let residues = blosum.row_names
let width = blosum.ncol

fn residue_index(names, ch) {
    (range(0, len(names)) |> filter(|i| names[i] == ch))[0]
}

fn substitution(mat, names, cols, a, b) {
    mat.data[residue_index(names, a) * cols + residue_index(names, b)]
}

# Gotoh's three matrices, held one row at a time:
#   M — s[i] aligned to t[j]
#   X — a gap in t (consuming s)
#   Y — a gap in s (consuming t)
# Opening a gap costs gap_open, each further residue costs gap_extend.
fn affine_score(a, b, mat, names, cols, open_penalty, extend_penalty) {
    let n = len(b)
    let very_low = -1000000.0

    let prev_m = concat([0.0], range(1, n + 1) |> map(|_| very_low))
    let prev_x = concat([very_low], range(1, n + 1) |> map(|_| very_low))
    let prev_y = concat([very_low], range(1, n + 1) |> map(|j| open_penalty + float(j - 1) * extend_penalty))

    let i = 1
    while i <= len(a) {
        let ai = substr(a, i - 1, 1)
        let row_m = [very_low]
        let row_x = [open_penalty + float(i - 1) * extend_penalty]
        let row_y = [very_low]

        let j = 1
        while j <= n {
            let sub = substitution(mat, names, cols, ai, substr(b, j - 1, 1))
            let best_prev = max([prev_m[j - 1], prev_x[j - 1], prev_y[j - 1]])
            row_m = push(row_m, best_prev + sub)
            row_x = push(row_x, max([prev_m[j] + open_penalty, prev_x[j] + extend_penalty]))
            row_y = push(row_y, max([row_m[j - 1] + open_penalty, row_y[j - 1] + extend_penalty]))
            j = j + 1
        }

        prev_m = row_m
        prev_x = row_x
        prev_y = row_y
        i = i + 1
    }

    max([prev_m[n], prev_x[n], prev_y[n]])
}

let result = affine_score(s, t, blosum, residues, width, gap_open, gap_extend)

println("Result:   " + str(int(result)))
println("Expected: 8")

fn test_gaff_affine_global_alignment() {
    assert int(result) == 8, "GAFF: got " + str(result)
}