# Rosalind: BA10F — Construct a Profile HMM with Pseudocounts
# https://rosalind.info/problems/ba10f/
#
# Given: A threshold theta, a pseudocount sigma, an alphabet, and a multiple
# alignment.
# Return: The transition and emission probabilities of HMM(Alignment, theta, sigma).

let threshold = 0.358
let pseudocount = 0.01
let alphabet = ["A", "B", "C", "D", "E"]
let alignment = [
    "ADA",
    "ADA",
    "AAA",
    "ADC",
    "-DA",
    "D-A",
]

# Without a pseudocount, anything the alignment never happened to do is scored as
# impossible — so a new sequence differing in one position gets probability zero
# rather than a low score. Six sequences cannot rule out the rest of the family,
# and that is what the pseudocount corrects.
#
# It is added after the counts are turned into probabilities, not before: added
# to raw counts its influence would depend on how many sequences the alignment
# happens to contain, which is not something anyone means to tune.
let profile = hmm_profile(alignment, alphabet, threshold, pseudocount)

println("Transitions out of S, I0, M1, D1:")
for source in ["S", "I0", "M1", "D1"] {
    let row = ["I0", "M1", "D1", "I1", "M2", "D2"]
        |> map(|target| target + "=" + str(round(profile.transition[source][target], 3)))
        |> join("  ")
    println("  " + source + ":  " + row)
}
println("Expected S:   I0=0.01  M1=0.819  D1=0.172")
println("Expected I0:  I0=0.333  M1=0.333  D1=0.333")
println("Expected M1:  I1=0.01  M2=0.786")
println("Expected D1:  I1=0.01  M2=0.981")

fn test_ba10f_profile_hmm_with_pseudocounts() {
    assert abs(profile.transition.S.I0 - 0.01)  < 5e-4, "BA10F: S->I0"
    assert abs(profile.transition.S.M1 - 0.819) < 5e-4, "BA10F: S->M1"
    assert abs(profile.transition.S.D1 - 0.172) < 5e-4, "BA10F: S->D1"
    # A row with no counts smooths to uniform over what the topology allows —
    # three states, not all twelve.
    assert abs(profile.transition.I0.I0 - 0.333) < 5e-4, "BA10F: I0->I0"
    assert abs(profile.transition.I0.M1 - 0.333) < 5e-4, "BA10F: I0->M1"
    assert abs(profile.transition.I0.D1 - 0.333) < 5e-4, "BA10F: I0->D1"
    assert abs(profile.transition.M1.I1 - 0.01)  < 5e-4, "BA10F: M1->I1"
    assert abs(profile.transition.M1.M2 - 0.786) < 5e-4, "BA10F: M1->M2"
    assert abs(profile.transition.D1.M2 - 0.981) < 5e-4, "BA10F: D1->M2"

    assert abs(profile.emission.I0.A - 0.2)   < 5e-4, "BA10F: I0 emits A"
    assert abs(profile.emission.M1.A - 0.771) < 5e-4, "BA10F: M1 emits A"
    assert abs(profile.emission.M1.B - 0.01)  < 5e-4, "BA10F: M1 emits B"
    assert abs(profile.emission.M2.D - 0.771) < 5e-4, "BA10F: M2 emits D"

    # A pseudocount does not make a silent state emit, and does not open
    # transitions the topology forbids — smoothing those would invent paths the
    # model does not have.
    let emitted_by_d1 = alphabet |> map(|s| profile.emission.D1[s]) |> sum()
    assert emitted_by_d1 == 0, "BA10F: D1 is a deletion state and cannot emit"
    assert profile.transition.S.M2 == 0, "BA10F: S cannot skip a layer to M2"
    assert profile.transition.M2.M1 == 0, "BA10F: a profile HMM cannot go backwards"
}