"""MDL-guided morpheme segmentation for tiny vocabularies, pure python.

Greedy Morfessor-flavored search: start with whole words as morphs, repeatedly
apply the single split that most reduces description length
L(lexicon) + L(corpus | lexicon). Vocabularies here are tiny (10-100 word
types), so an O(V * maxlen) sweep per iteration is instant.

Alignment conditioning: tokens known (from align.py) to share a gloss get a
bonus for splits that expose their shared substring — this is the
"segmentation conditioned on alignment" step from the plan, and is what keeps
MDL from over-segmenting on 20-word corpora.
"""

from __future__ import annotations

import math
from collections import Counter
from typing import Dict, Iterable, List, Optional, Sequence, Set, Tuple

_MIN_MORPH = 1


def _lex_cost(morphs: Iterable[str]) -> float:
    # ~1 char = a few bits; +1 per morph for the boundary/index overhead
    return sum(len(m) + 1 for m in set(morphs)) * 4.0


def _corpus_cost(usage: Counter) -> float:
    total = sum(usage.values())
    if total == 0:
        return 0.0
    return -sum(c * math.log2(c / total) for c in usage.values())


class Segmenter:
    def __init__(self, share_bonus: float = 8.0):
        self.share_bonus = share_bonus
        self.seg: Dict[str, List[str]] = {}

    def fit(
        self,
        words: Sequence[str],
        counts: Optional[Counter] = None,
        share_groups: Optional[List[Set[str]]] = None,
        max_iters: int = 200,
    ) -> "Segmenter":
        """words: vocabulary (task-language word types).
        counts: token frequencies (defaults to 1 each).
        share_groups: sets of words believed to share a morpheme (same gloss
        alignment); splits exposing a shared prefix/suffix get a bonus."""
        counts = counts or Counter({w: 1 for w in words})
        self.seg = {w: [w] for w in dict.fromkeys(words) if w}
        shared_subs = self._shared_substrings(share_groups or [])

        for _ in range(max_iters):
            best = self._best_split(counts, shared_subs)
            if best is None:
                break
            word, mi, cut = best
            m = self.seg[word][mi]
            self.seg[word][mi : mi + 1] = [m[:cut], m[cut:]]
        return self

    def _shared_substrings(self, groups: List[Set[str]]) -> Set[str]:
        subs: Set[str] = set()
        for g in groups:
            g = [w for w in g if w]
            if len(g) < 2:
                continue
            # longest common prefix and suffix over the group
            pre = g[0]
            suf = g[0]
            for w in g[1:]:
                while pre and not w.startswith(pre):
                    pre = pre[:-1]
                while suf and not w.endswith(suf):
                    suf = suf[1:]
            if len(pre) >= 2:
                subs.add(pre)
            if len(suf) >= 2:
                subs.add(suf)
        return subs

    def _cost(self, counts: Counter, shared_subs: Set[str]) -> float:
        usage: Counter = Counter()
        for w, morphs in self.seg.items():
            for m in morphs:
                usage[m] += counts[w]
        cost = _lex_cost(usage.keys()) + _corpus_cost(usage)
        cost -= self.share_bonus * sum(1 for m in usage if m in shared_subs)
        return cost

    def _best_split(self, counts: Counter, shared_subs: Set[str]):
        base = self._cost(counts, shared_subs)
        best_gain, best = 1e-6, None
        for w, morphs in self.seg.items():
            for mi, m in enumerate(morphs):
                if len(m) < 2 * _MIN_MORPH:
                    continue
                for cut in range(_MIN_MORPH, len(m) - _MIN_MORPH + 1):
                    morphs[mi : mi + 1] = [m[:cut], m[cut:]]
                    gain = base - self._cost(counts, shared_subs)
                    morphs[mi : mi + 2] = [m]
                    if gain > best_gain:
                        best_gain, best = gain, (w, mi, cut)
        return best

    def segment(self, word: str) -> List[str]:
        """Segment a word; unseen words are matched greedily against the
        learned morph inventory (longest-match, both ends first)."""
        if word in self.seg:
            return list(self.seg[word])
        morphs = {m for parts in self.seg.values() for m in parts}
        return _greedy_decompose(word, morphs)

    @property
    def morphs(self) -> Set[str]:
        return {m for parts in self.seg.values() for m in parts}


def _greedy_decompose(word: str, morphs: Set[str]) -> List[str]:
    """Best-effort decomposition of an unseen word over a morph set: dynamic
    programming for fewest chunks, unknown spans kept as single chunks."""
    n = len(word)
    INF = float("inf")
    # cost[i] = (num chunks, num unknown chars) to segment word[:i]
    cost = [(INF, INF)] * (n + 1)
    back: List[Optional[Tuple[int, str]]] = [None] * (n + 1)
    cost[0] = (0, 0)
    for i in range(n):
        if cost[i][0] == INF:
            continue
        for j in range(i + 1, n + 1):
            piece = word[i:j]
            known = piece in morphs
            c = (cost[i][0] + 1, cost[i][1] + (0 if known else len(piece)))
            # prefer fewer unknown chars, then fewer chunks
            key = (c[1], c[0])
            if key < (cost[j][1], cost[j][0]):
                cost[j] = c
                back[j] = (i, piece)
    out: List[str] = []
    i = n
    while i > 0 and back[i]:
        prev, piece = back[i]
        out.append(piece)
        i = prev
    out.reverse()
    return out or [word]