"""
T-Coffee multiple sequence alignment for phonetic SEGMENT sequences.

A "word" here is a tuple of segments, e.g. ("kʰ", "e", "r") -- not a string of
characters. IPA segments are frequently multi-character (kʰ, aː, t̪ʰ) and
splitting them would be wrong.

A faithful Python port of the Julia implementation used in Jäger (2026),
`worldtree_msa/code/alignment_functions_phmm.jl`. Small enough to read, and it
runs in the browser under Pyodide.

The only substantive difference: the original scores pairwise alignments with a
pair-hidden-Markov model (Viterbi); here we use LingPy's Needleman-Wunsch, which
is the aligner the course already teaches in Lab 1. The consistency
transformation -- the part that makes T-Coffee T-Coffee -- is identical.
"""

import numpy as np
from lingpy.align.pairwise import nw_align
from lingpy.sequence.sound_classes import token2class
from lingpy.data import Model

_SCA = Model("sca")


def _cls(seg):
    """Sound class of a segment. h, ɦ, ç are different segments but the same
    *kind* of sound; exact identity would call them unrelated."""
    try:
        return token2class(seg, _SCA)
    except Exception:
        return seg


# ---------------------------------------------------------------- base library

def _align_matrix(a, b):
    """Binary matrix A: A[i, j] = 1 iff position i of word 1 is aligned to
    position j of word 2. Also returns the alignment's similarity."""
    n = sum(1 for x in a if x != "-")
    m = sum(1 for x in b if x != "-")
    A = np.zeros((n, m), dtype=np.float32)
    i = j = 0
    matches = 0
    for x, y in zip(a, b):
        if x != "-" and y != "-":
            A[i, j] = 1.0
            if _cls(x) == _cls(y):      # sound-class identity, not segment identity
                matches += 1
            i += 1
            j += 1
        elif x != "-":
            i += 1
        elif y != "-":
            j += 1
    sim = matches / max(len(a), 1)          # 1 - normalised Hamming distance
    return A, sim


def base_library(words):
    """Every pairwise alignment, as a matrix + similarity."""
    lib = {}
    for i, w1 in enumerate(words):
        for j in range(i, len(words)):
            w2 = words[j]
            if w1 == w2:
                A = np.eye(len(w1), dtype=np.float32)
                sim = 1.0
            else:
                a, b, _ = nw_align([_cls(s) for s in w1], [_cls(s) for s in w2])
                # re-map the sound-class alignment back onto the real segments
                A, sim = _align_matrix(a, b)
            lib[(w1, w2)] = (A, sim)
            lib[(w2, w1)] = (A.T, sim)
    return lib


# ------------------------------------------------- the consistency transformation

def extended_library(words, base):
    """THE T-COFFEE STEP.

    The alignment of w1 and w2 is re-scored by how consistently every *third*
    word w3 aligns to both of them:

        ext[w1,w2][i,j] = SUM over w3 of  min(sim(w1,w3), sim(w3,w2))
                                          * (A(w1,w3) @ A(w3,w2))[i,j]

    A pair of positions that survives being routed through the rest of the
    wordlist is trustworthy. One that does not, is not.
    """
    uw = list(dict.fromkeys(words))
    ext = {}
    for i, w1 in enumerate(uw):
        for j in range(i, len(uw)):
            w2 = uw[j]
            S = np.zeros((len(w1), len(w2)), dtype=np.float32)
            for w3 in uw:
                A1, s1 = base[(w1, w3)]
                A2, s2 = base[(w3, w2)]
                S += min(s1, s2) * (A1 @ A2)
            ext[(w1, w2)] = S
            ext[(w2, w1)] = S.T
    return ext


# ------------------------------------------------------ progressive alignment

def _nw_block(prof1, prof2, ext):
    """Needleman-Wunsch on two profiles (blocks of already-aligned words),
    scoring each column pair with the extended library. Gaps are free -- all
    library scores are >= 0, so the DP simply maximises consistency."""
    w1 = [tuple(c for c in row if c != "-") for row in zip(*prof1)]
    w2 = [tuple(c for c in row if c != "-") for row in zip(*prof2)]

    # for column i of the profile, which position of word k does it hold? (-1 = gap)
    def postab(prof):
        tab = []
        for col in zip(*prof):
            counts, k = [], 0
            for ch in col:
                if ch != "-":
                    counts.append(k); k += 1
                else:
                    counts.append(-1)
            tab.append(counts)
        return tab

    p1 = [[-1] * len(w1) for _ in prof1]
    for k in range(len(w1)):
        c = 0
        for i in range(len(prof1)):
            if prof1[i][k] != "-":
                p1[i][k] = c; c += 1
    p2 = [[-1] * len(w2) for _ in prof2]
    for k in range(len(w2)):
        c = 0
        for i in range(len(prof2)):
            if prof2[i][k] != "-":
                p2[i][k] = c; c += 1

    n, m = len(prof1), len(prof2)
    dp = np.zeros((n + 1, m + 1), dtype=np.float32)
    ptr = np.zeros((n + 1, m + 1), dtype=np.int8)
    ptr[0, 1:] = 3
    ptr[1:, 0] = 2

    for i in range(1, n + 1):
        for j in range(1, m + 1):
            s = 0.0
            for k in range(len(w1)):
                pi = p1[i - 1][k]
                if pi < 0:
                    continue
                for l in range(len(w2)):
                    pj = p2[j - 1][l]
                    if pj < 0:
                        continue
                    s += ext[(w1[k], w2[l])][pi, pj]
            match, ins, dele = dp[i-1, j-1] + s, dp[i-1, j], dp[i, j-1]
            best = max(match, ins, dele)
            dp[i, j] = best
            ptr[i, j] = 1 if best == match else (2 if best == ins else 3)

    out1, out2 = [], []
    i, j = n, m
    while i > 0 or j > 0:
        step = ptr[i, j]
        if step == 1:
            i -= 1; j -= 1; out1.append(prof1[i]); out2.append(prof2[j])
        elif step == 2:
            i -= 1; out1.append(prof1[i]); out2.append(["-"] * len(w2))
        else:
            j -= 1; out1.append(["-"] * len(w1)); out2.append(prof2[j])
    out1.reverse(); out2.reverse()
    return [a + b for a, b in zip(out1, out2)]


def _guide_tree(words, base):
    """UPGMA on 1 - similarity. Returns a nested tuple of indices."""
    idx = [(i,) for i in range(len(words))]
    prof = {i: [i] for i in range(len(words))}
    node = {i: i for i in range(len(words))}
    D = {}
    for i in range(len(words)):
        for j in range(i + 1, len(words)):
            D[(i, j)] = 1 - base[(words[i], words[j])][1]
    active = list(range(len(words)))
    tree = {i: i for i in active}
    nxt = len(words)
    while len(active) > 1:
        best, bi, bj = None, None, None
        for a in range(len(active)):
            for b in range(a + 1, len(active)):
                x, y = active[a], active[b]
                d = np.mean([D[(min(p, q), max(p, q))]
                             for p in prof[x] for q in prof[y] if p != q])
                if best is None or d < best:
                    best, bi, bj = d, x, y
        tree[nxt] = (tree[bi], tree[bj])
        prof[nxt] = prof[bi] + prof[bj]
        active = [a for a in active if a not in (bi, bj)] + [nxt]
        nxt += 1
    return tree[active[0]]


def tcoffee(words):
    """Full T-Coffee: base library -> consistency transformation -> progressive
    alignment along a UPGMA guide tree. Returns a list of aligned strings."""
    base = base_library(words)
    ext = extended_library(words, base)
    gt = _guide_tree(words, base)

    def build(node):
        if isinstance(node, int):
            return [[seg] for seg in words[node]], [node]
        l, r = node
        pl, il = build(l)
        pr, ir = build(r)
        return _nw_block(pl, pr, ext), il + ir

    prof, order = build(gt)
    cols = list(zip(*prof))                      # one tuple per word
    out = [None] * len(words)
    for pos, w in enumerate(order):
        out[w] = list(cols[pos])
    return out