"""Tokenisation, learned abbreviation mining, and dictionary-free abbreviation matching."""
import re
from collections import Counter, defaultdict
import numpy as np

_TOK = re.compile(r"\w+", re.UNICODE)
_NUM = re.compile(r"\d+")


def tokens(s):
    return _TOK.findall(s.lower()) if s else []


def numbers(s):
    return set(_NUM.findall(s)) if s else set()


def is_abbrev(short, long):
    """Generic (language-independent) abbreviation test: prefix, or ordered subsequence
    sharing the first letter. e.g. corp->corporation, pvt->private, rd->road."""
    if not short or not long or len(short) >= len(long) or short[0] != long[0]:
        return False
    if long.startswith(short):
        return len(short) >= 2
    it = iter(long)
    return len(short) >= 2 and all(ch in it for ch in short)


# ---------------------------------------------------------------------------
# Learned abbreviation map (mined from TRAIN-fold positive pairs only)
# ---------------------------------------------------------------------------
def _leftovers(a, b):
    sa, sb = set(a), set(b)
    return [t for t in a if t not in sb], [t for t in b if t not in sa]


def mine_abbreviations(pos_pairs, neg_pairs, min_support=5, min_precision=0.8):
    """pos_pairs / neg_pairs: iterables of (text_a, text_b).
    Returns dict short_token -> long_token (or multi-token initialism expansion)."""
    pos_cnt = Counter()
    short_cnt = Counter()
    n_pos = 0
    for a, b in pos_pairs:
        ta, tb = tokens(a), tokens(b)
        la, lb = _leftovers(ta, tb)
        n_pos += 1
        for x, ys, yfull in ((la, lb, tb), (lb, la, ta)):
            for s in x:
                short_cnt[s] += 1
                for l in ys:
                    if is_abbrev(s, l):
                        pos_cnt[(s, l)] += 1
                # initialism: s == first letters of consecutive tokens on the other side
                if 2 <= len(s) <= 6:
                    for i in range(len(yfull) - len(s) + 1):
                        seg = yfull[i:i + len(s)]
                        if "".join(t[0] for t in seg) == s:
                            pos_cnt[(s, " ".join(seg))] += 1
                            break

    cand = {k: v for k, v in pos_cnt.items() if v >= min_support}
    neg_cnt = Counter()
    n_neg = 0
    if cand:
        by_short = defaultdict(list)
        for (s, l) in cand:
            by_short[s].append(l)
        for a, b in neg_pairs:
            n_neg += 1
            ta, tb = tokens(a), tokens(b)
            sa, sb = set(ta), set(tb)
            ja, jb = " ".join(ta), " ".join(tb)
            for s in sa & by_short.keys():
                for l in by_short[s]:
                    if (" " in l and l in jb) or l in sb:
                        neg_cnt[(s, l)] += 1
            for s in sb & by_short.keys():
                for l in by_short[s]:
                    if (" " in l and l in ja) or l in sa:
                        neg_cnt[(s, l)] += 1

    scale = n_pos / max(n_neg, 1)
    best = {}
    for (s, l), v in cand.items():
        prec = v / (v + neg_cnt[(s, l)] * scale)
        if prec < min_precision:
            continue
        if s not in best or v > best[s][1]:
            best[s] = (l, v)
    return {s: l for s, (l, v) in best.items()}


def apply_abbrev(text, amap):
    if not amap or not text:
        return text
    return " ".join(amap.get(t, t) for t in tokens(text))


# ---------------------------------------------------------------------------
# Dictionary-free abbreviation-aware similarity (works for unseen countries)
# ---------------------------------------------------------------------------
def abbrev_aware_match(a, b):
    ta, tb = tokens(a), tokens(b)
    if not ta or not tb:
        return 0.0
    sb, sa = set(tb), set(ta)

    def matched(t, other, other_set):
        return t in other_set or any(is_abbrev(t, o) or is_abbrev(o, t) for o in other)

    ma = sum(matched(t, tb, sb) for t in ta)
    mb = sum(matched(t, ta, sa) for t in tb)
    return (ma + mb) / (len(ta) + len(tb))


def token_jaccard(a, b):
    sa, sb = set(tokens(a)), set(tokens(b))
    if not sa and not sb:
        return np.nan
    return len(sa & sb) / len(sa | sb)