# Rosalind: BA10K — Implement Baum-Welch Learning
# https://rosalind.info/problems/ba10k/
#
# Given: A number of iterations i, a string x, its alphabet, the states of an
# HMM, and initial transition and emission matrices.
# Return: Matrices estimated after i rounds of Baum-Welch learning.

let iterations = 10
let observed = "xzyyzyzyxy"

let start = {
    states: ["A", "B"],
    symbols: ["x", "y", "z"],
    transition: {
        A: { A: 0.019, B: 0.981 },
        B: { A: 0.668, B: 0.332 },
    },
    emission: {
        A: { x: 0.175, y: 0.003, z: 0.821 },
        B: { x: 0.196, y: 0.512, z: 0.293 },
    },
}

# Baum-Welch is expectation-maximisation for an HMM. Where BA10I commits to the
# single best path and counts along it, this counts the *expected* number of
# times each transition was taken across every path at once — which
# forward-backward supplies without enumerating any of them. Keeping the paths it
# would otherwise discard is what makes it the better estimator.
let learned = hmm_baum_welch(observed, start, iterations)

println("Transition:")
println("      A       B")
for name in learned.states {
    println("  " + name + "   " + str(round(learned.transition[name].A, 3))
                     + "   " + str(round(learned.transition[name].B, 3)))
}
println("Emission:")
println("      x       y       z")
for name in learned.states {
    println("  " + name + "   " + (learned.symbols |> map(|s| str(round(learned.emission[name][s], 3))) |> join("   ")))
}
println("Expected transition: 0.0 1.0 / 0.786 0.214")
println("Expected emission:   0.242 0.0 0.758 / 0.172 0.828 0.0")

fn test_ba10k_baum_welch() {
    assert abs(learned.transition.A.B - 1.0)   < 5e-4, "BA10K: A->B is " + str(learned.transition.A.B)
    assert abs(learned.transition.B.A - 0.786) < 5e-4, "BA10K: B->A is " + str(learned.transition.B.A)
    assert abs(learned.emission.A.x - 0.242) < 5e-4, "BA10K: A emits x at " + str(learned.emission.A.x)
    assert abs(learned.emission.A.z - 0.758) < 5e-4, "BA10K: A emits z at " + str(learned.emission.A.z)
    assert abs(learned.emission.B.y - 0.828) < 5e-4, "BA10K: B emits y at " + str(learned.emission.B.y)
    # Each round can only raise Pr(x). Checked round by round rather than only
    # end to end, because a single bad update can be hidden by later good ones.
    let running = range(1, 6) |> map(|n| hmm_likelihood(observed, hmm_baum_welch(observed, start, n)))
    for i in range(1, len(running)) {
        assert running[i] >= running[i - 1] - 1e-12,
            "BA10K: round " + str(i + 1) + " lowered the likelihood"
    }
    assert hmm_likelihood(observed, learned) > hmm_likelihood(observed, start),
        "BA10K: learning should explain the observation better than the start did"
}