"""Word alignment from tiny parallel corpora, pure python.

Two complementary signals:
1. Minimal-pair set difference: if two sentence pairs differ in exactly one
   token on each side, those tokens correspond. Exact and high-precision;
   these puzzles are constructed to contain such pairs.
2. Dice co-occurrence over the whole pair set: soft alignment for everything
   the minimal pairs don't cover.
"""

from __future__ import annotations

from collections import Counter, defaultdict
from itertools import combinations
from typing import Dict, List, Tuple

from .preprocess import Pair, strip_punct, tokenize


def _toks(s: str) -> List[str]:
    return [strip_punct(t).casefold() for t in tokenize(s) if strip_punct(t)]


def minimal_pair_links(pairs: List[Pair]) -> Counter:
    """Set-difference alignment: for every pair of examples whose source sides
    differ by exactly one token multiset element and likewise on target, link
    the differing tokens. Returns Counter[(src_tok, tgt_tok)] link strengths."""
    links: Counter = Counter()
    toks = [(Counter(_toks(p.src)), Counter(_toks(p.tgt))) for p in pairs]
    for (s1, t1), (s2, t2) in combinations(toks, 2):
        ds1, ds2 = s1 - s2, s2 - s1
        dt1, dt2 = t1 - t2, t2 - t1
        # exactly one differing token on each side, in both examples
        if sum(ds1.values()) == 1 and sum(ds2.values()) == 1 \
                and sum(dt1.values()) == 1 and sum(dt2.values()) == 1:
            a1, a2 = next(iter(ds1)), next(iter(ds2))
            b1, b2 = next(iter(dt1)), next(iter(dt2))
            links[(a1, b1)] += 2  # strong: attested by contrast
            links[(a2, b2)] += 2
        # shared residue: tokens present in both examples also co-align weakly
    return links


def dice_scores(pairs: List[Pair]) -> Dict[Tuple[str, str], float]:
    """Dice coefficient between source and target tokens across examples."""
    src_count: Counter = Counter()
    tgt_count: Counter = Counter()
    co: Counter = Counter()
    for p in pairs:
        st, tt = set(_toks(p.src)), set(_toks(p.tgt))
        for a in st:
            src_count[a] += 1
        for b in tt:
            tgt_count[b] += 1
        for a in st:
            for b in tt:
                co[(a, b)] += 1
    return {
        (a, b): 2 * c / (src_count[a] + tgt_count[b])
        for (a, b), c in co.items()
    }


def _morph_backoff(pairs: List[Pair], scores) -> None:
    """Substring evidence from single-word glosses: if (moko = dog) is
    attested and token `namoko` co-occurs with `dog`, boost (namoko, dog) —
    inflected forms inherit their stem's translation. Applied in place."""
    word_pairs = [
        (_toks(p.src)[0], _toks(p.tgt)[0])
        for p in pairs
        if len(_toks(p.src)) == 1 and len(_toks(p.tgt)) == 1
    ]
    for (a, b) in list(scores.keys()):
        for w, x in word_pairs:
            if x == b and len(w) >= 3 and w in a and w != a:
                scores[(a, b)] += 2.0  # inflected src contains attested stem
            if w == a and len(x) >= 3 and x in b and x != b:
                scores[(a, b)] += 2.0  # inflected tgt contains attested stem


def align(pairs: List[Pair]) -> Dict[str, List[Tuple[str, float]]]:
    """Combined alignment: src token -> ranked [(tgt token, score)].

    Minimal-pair links dominate (score offset +1.0 per link unit); Dice fills
    in the rest; single-word glosses back off into inflected forms containing
    them. Scores are comparable only within one puzzle.
    """
    links = minimal_pair_links(pairs)
    dice = dice_scores(pairs)
    scores: Dict[Tuple[str, str], float] = defaultdict(float)
    for k, v in dice.items():
        scores[k] += v
    for k, v in links.items():
        scores[k] += 1.0 * v
    _morph_backoff(pairs, scores)
    # competition ("explaining away"): a target token strongly claimed by
    # some other source is a worse candidate — demote it proportionally to
    # its best competing suitor. Breaks the pervasive co-occurrence ties of
    # 10-sentence corpora in favor of unclaimed targets.
    best_suitor: Dict[str, float] = defaultdict(float)
    second_suitor: Dict[str, float] = defaultdict(float)
    for (a, b), s in scores.items():
        if s > best_suitor[b]:
            second_suitor[b] = best_suitor[b]
            best_suitor[b] = s
        elif s > second_suitor[b]:
            second_suitor[b] = s
    out: Dict[str, List[Tuple[str, float]]] = defaultdict(list)
    for (a, b), s in scores.items():
        rival = second_suitor[b] if s >= best_suitor[b] else best_suitor[b]
        out[a].append((b, s - 0.3 * rival))
    for a in out:
        out[a].sort(key=lambda x: -x[1])
    return dict(out)


def one_to_one(pairs: List[Pair]) -> Dict[str, str]:
    """Greedy 1:1 token alignment: highest-scoring links assigned first, each
    token used once. Sharper than independent argmax when several tokens tie
    on co-occurrence (small corpora make ties common)."""
    amap = align(pairs)
    edges = [(s, a, b) for a, cands in amap.items() for b, s in cands]
    edges.sort(key=lambda e: (-e[0], e[1], e[2]))
    taken_a, taken_b, out = set(), set(), {}
    for s, a, b in edges:
        if a not in taken_a and b not in taken_b:
            out[a] = b
            taken_a.add(a)
            taken_b.add(b)
    return out


def best_translation(align_map: Dict[str, List[Tuple[str, float]]], tok: str) -> str:
    cands = align_map.get(tok.casefold(), [])
    return cands[0][0] if cands else ""