# Rosalind: QRTD — Quartet Distance
# https://rosalind.info/problems/qrtd/
#
# Given: n taxa and two unrooted binary trees.
# Return: The quartet distance between them.

let taxa = ["A", "B", "C", "D", "E"]
# (A,C,((B,D),E));   and   (C,(B,D),(A,E));
# Each internal edge is a split; these are the non-trivial sides.
let first_clades = [["B", "D"], ["B", "D", "E"]]
let second_clades = [["B", "D"], ["A", "E"]]

# Two trees can share every taxon and still disagree about almost everything, so
# comparing them needs a measure. Quartets give one: each set of four taxa is
# separated into two pairs by exactly one edge, and two trees either agree about
# that pairing or they do not.
#
# It beats counting shared splits because it degrades gracefully — moving one
# taxon changes a handful of quartets, where it can destroy every split at once.
fn pairing_in(four, clades) {
    # The clade cutting these four 2-and-2 names the quartet.
    let cutting = clades |> filter(|c| (four |> count_if(|t| contains(c, t))) == 2)
    if len(cutting) == 0 { return "" }
    let side = four |> filter(|t| contains(cutting[0], t))
    let other = four |> filter(|t| contains(cutting[0], t) == false)
    let one = join(sort(side), ",")
    let two = join(sort(other), ",")
    if one < two then one + "|" + two else two + "|" + one
}

fn subsets_of_four(items) {
    range(0, len(items)) |> flat_map(|a|
        range(a + 1, len(items)) |> flat_map(|b|
            range(b + 1, len(items)) |> flat_map(|c|
                range(c + 1, len(items)) |> map(|d|
                    [items[a], items[b], items[c], items[d]]))))
}

let quartets = subsets_of_four(taxa)
let shared = quartets |> count_if(|four| {
    let one = pairing_in(four, first_clades)
    let two = pairing_in(four, second_clades)
    one != "" and one == two
})

# Both trees are fully resolved, so each induces exactly one quartet per subset.
let first_count = quartets |> count_if(|four| pairing_in(four, first_clades) != "")
let second_count = quartets |> count_if(|four| pairing_in(four, second_clades) != "")
let distance = first_count + second_count - 2 * shared

println("Result:   " + str(distance))
println("Expected: 4")
println("(" + str(len(quartets)) + " quartets each, " + str(shared) + " agreeing)")

fn test_qrtd_quartet_distance() {
    assert distance == 4, "QRTD: got " + str(distance)
    # A resolved tree resolves every four-taxon subset, which CNTQ established.
    assert first_count == choose(len(taxa), 4), "QRTD: the first tree resolves all 5"
    assert second_count == choose(len(taxa), 4), "QRTD: and so does the second"
    assert shared == 3, "QRTD: 3 of the 5 quartets agree"

    # A tree against itself is distance zero — the identity any distance must
    # satisfy, and worth checking rather than assuming.
    let self_shared = quartets |> count_if(|four| pairing_in(four, first_clades) != "")
    assert first_count + first_count - 2 * self_shared == 0,
        "QRTD: a tree must be distance 0 from itself"
    # The two disagreeing quartets are the ones involving both B,D and A,E.
    let disagreeing = quartets |> filter(|four|
        pairing_in(four, first_clades) != pairing_in(four, second_clades))
    assert len(disagreeing) == 2, "QRTD: exactly two quartets differ"
}