# Rosalind: BA7E — Implement the Neighbor Joining Algorithm
# https://rosalind.info/problems/ba7e/
#
# Given: An integer n and an n x n distance matrix.
# Return: The adjacency list of the tree neighbour joining builds.

let leaf_count = 4
let distances = [
    [0, 23, 27, 20],
    [23, 0, 30, 28],
    [27, 30, 0, 30],
    [20, 28, 30, 0],
]

# UPGMA joins whichever pair is closest, which is wrong whenever two lineages
# evolve at different rates — a fast-evolving leaf looks distant from its true
# relatives and gets attached elsewhere. Neighbour joining corrects each distance
# by how far each leaf is from *everything* else before comparing, so a
# uniformly-distant leaf is no longer penalised. It makes no molecular-clock
# assumption, and the tree it returns is unrooted.
fn neighbour_matrix(working, live) {
    let n = len(live)
    let totals = {}
    for i in live { totals[str(i)] = live |> map(|j| working[i][j]) |> sum() }
    let corrected = {}
    for i in live {
        for j in live {
            if i != j {
                corrected[str(i) + "," + str(j)] =
                    (n - 2) * working[i][j] - totals[str(i)] - totals[str(j)]
            }
        }
    }
    { corrected: corrected, totals: totals }
}

let working = distances |> map(|line| line |> map(|d| d * 1.0))
let live = range(0, leaf_count)
let next_id = leaf_count
let adjacency = {}

fn connect(graph, a, b, weight) {
    let entry = { to_node: str(b), weight: round(weight, 3) }
    if contains(keys(graph), str(a)) {
        graph[str(a)] = push(graph[str(a)], entry)
    } else {
        graph[str(a)] = [entry]
    }
    graph
}

while len(live) > 2 {
    let built = neighbour_matrix(working, live)
    let corrected = built.corrected
    let totals = built.totals

    let best_a = live[0]
    let best_b = live[1]
    let best = corrected[str(best_a) + "," + str(best_b)]
    for a in live {
        for b in live {
            if a != b and corrected[str(a) + "," + str(b)] < best {
                best = corrected[str(a) + "," + str(b)]
                best_a = a
                best_b = b
            }
        }
    }

    let n = len(live)
    # How lopsided the pair is — this is what lets the two limbs differ, which
    # UPGMA cannot express.
    let delta = (totals[str(best_a)] - totals[str(best_b)]) / (n - 2)
    let limb_a = (working[best_a][best_b] + delta) / 2
    let limb_b = (working[best_a][best_b] - delta) / 2

    let row = range(0, next_id + 1) |> map(|_| 0.0)
    for other in live {
        if other != best_a and other != best_b {
            let joined = working[best_a][other] + working[best_b][other]
            row[other] = (joined - working[best_a][best_b]) / 2
        }
    }
    working = push(working, row)
    for other in live {
        if other != best_a and other != best_b {
            working[other] = push(working[other], row[other])
        }
    }

    adjacency = connect(adjacency, next_id, best_a, limb_a)
    adjacency = connect(adjacency, best_a, next_id, limb_a)
    adjacency = connect(adjacency, next_id, best_b, limb_b)
    adjacency = connect(adjacency, best_b, next_id, limb_b)

    live = (live |> filter(|node| node != best_a and node != best_b)) + [next_id]
    next_id = next_id + 1
}

# Two left: join them with the distance between them.
let last_a = live[0]
let last_b = live[1]
adjacency = connect(adjacency, last_a, last_b, working[last_a][last_b])
adjacency = connect(adjacency, last_b, last_a, working[last_a][last_b])

let listed = sort(keys(adjacency)) |> flat_map(|node|
    adjacency[node] |> map(|edge| node + "->" + edge.to_node + ":" + str(round(edge.weight, 3))))

println("Result:")
for line in listed { println("  " + line) }
println("Expected: 0->4:8, 1->5:13.5, 2->5:16.5, 3->4:12, 4->5:2")

fn test_ba7e_neighbour_joining() {
    let want = ["0->4:8.0", "4->0:8.0", "1->5:13.5", "5->1:13.5", "2->5:16.5", "5->2:16.5",
                "3->4:12.0", "4->3:12.0", "4->5:2.0", "5->4:2.0"]
    assert len(listed) == len(want),
        "BA7E: expected " + str(len(want)) + " edges, got " + str(len(listed))
    for edge in want { assert contains(listed, edge), "BA7E: missing edge " + edge }

    # This matrix is *not* additive, which the four-point condition shows: for an
    # additive matrix the two largest of the three pairings must be equal, and
    # here they are 55, 53 and 50. So no tree reproduces it exactly, and neighbour
    # joining returns a best fit rather than an exact answer — worth asserting,
    # because the natural expectation is that the distances come back unchanged.
    let pairing_one = distances[0][1] + distances[2][3]
    let pairing_two = distances[0][2] + distances[1][3]
    let pairing_three = distances[0][3] + distances[1][2]
    assert pairing_one == 53 and pairing_two == 55 and pairing_three == 50,
        "BA7E: the three pairings are 53, 55, 50"
    let two_largest_agree = pairing_one == pairing_two or pairing_one == pairing_three
        or pairing_two == pairing_three
    assert two_largest_agree == false, "BA7E: so the working is not additive"

    # Some distances the tree does get exactly right, and none is off by much.
    assert (8.0 + 12.0) == distances[0][3], "BA7E: 0 to 3 is 8 + 12 = 20, exactly"
    assert (13.5 + 16.5) == distances[1][2], "BA7E: 1 to 2 is 30, exactly"
    assert abs((8.0 + 2.0 + 13.5) - distances[0][1]) <= 0.5,
        "BA7E: 0 to 1 comes out at 23.5 against 23 — the cost of non-additivity"
    assert abs((8.0 + 2.0 + 16.5) - distances[0][2]) <= 0.5, "BA7E: 0 to 2 within half"

    # Limbs differ within a pair, which is exactly what UPGMA cannot express.
    assert 13.5 != 16.5, "BA7E: the two limbs off node 5 have different lengths"
}