"""Lematizatorul: formă din text → lemă, folosind cele 2,27 milioane de forme
flexionare din dexonline.

Româna e o limbă în care cuvântul-titlu aproape nu apare: „casă" se scrie în
text „casa", „casei", „case", „casele", „caselor". Fără lematizare, orice
numărătoare de frecvență măsoară altceva decât crede.

Normalizarea rezolvă și ortografiile succesive ale limbii: sedila în locul
virgulei (şi → și), î din a în locul lui â (romîn → român), diacriticele
secolului al XIX-lea (românĭ → români).
"""
import os, re, unicodedata
from collections import defaultdict

HERE = os.path.dirname(os.path.abspath(__file__))
RAW = os.path.join(HERE, "..", "raw")

# Ortografii vechi și greșite, aduse la forma de azi.
_CHAR_MAP = {
    "ş": "ș", "Ş": "Ș", "ţ": "ț", "Ţ": "Ț",   # sedila din ISO 8859-2
    "ĕ": "e", "ě": "e", "é": "e", "è": "e", "ê": "e",
    "ĭ": "i", "ï": "i", "í": "i", "î": "î",
    "ŭ": "u", "ù": "u", "ú": "u", "û": "â",
    "ó": "o", "ò": "o", "ô": "o",
    "á": "a", "à": "a", "ā": "a",
    "ḑ": "z", "đ": "d", "ǵ": "g",
    "’": "'", "‘": "'", "ʼ": "'",
}
_TRANS = str.maketrans(_CHAR_MAP)

TOKEN_RE = re.compile(r"[a-zăâîșțáàâéèêíïîóòôúùûŭĭĕěḑ]+(?:['-][a-zăâîșț]+)*", re.I)


def normalize(word):
    """Cheia de potrivire: minuscule, diacritice unificate, â și î confundate.

    â și î notează același sunet în română și au fost scrise alternativ de trei
    reforme ortografice (1904, 1953, 1993). „Sînt" și „sunt", „romîn" și
    „român" sunt același cuvânt scris de două comisii diferite.
    """
    w = unicodedata.normalize("NFC", word.lower()).translate(_TRANS)
    w = w.replace("â", "î")               # â/î → o singură literă
    w = w.strip("'-")
    return w


def tokenize(text):
    return [normalize(m.group(0)) for m in TOKEN_RE.finditer(text)]


_FLAT = str.maketrans({"ă": "a", "î": "i", "â": "a", "ș": "s", "ț": "t"})


def flatten(word):
    """Varianta fără diacritice. Multe transcrieri — printre ele constituțiile
    din 1952 și 1965 de pe Wikisource — sunt tastate „nationale", „dupa",
    „cetatenilor". Fără o a doua cheie, un cuvânt din cinci s-ar pierde."""
    return word.translate(_FLAT)


class Lemmatizer:
    """formă normalizată → lemele candidate."""

    def __init__(self):
        self.form2lemmas = {}
        self.flat2lemmas = {}
        self.lemma_pos = {}
        self._load()

    def _load(self):
        lex = {}
        with open(os.path.join(RAW, "dex_lexemes.tsv"), encoding="utf8") as f:
            for line in f:
                p = line.rstrip("\n").split("\t")
                if len(p) < 5:
                    continue
                lid, form, mtype = p[0], p[1], p[2]
                if not form or " " in form:
                    continue
                # Lemele cu majusculă sunt nume proprii („Dorul", „Orașul") și ar
                # fura ocurențele cuvântului comun scris la început de frază.
                if form[:1] != form[:1].lower():
                    continue
                lex[lid] = (form, mtype)
                self.lemma_pos.setdefault(form, mtype)

        acc = defaultdict(set)
        with open(os.path.join(RAW, "dex_forms.tsv"), encoding="utf8") as f:
            for line in f:
                p = line.rstrip("\n").split("\t")
                if len(p) < 2:
                    continue
                form, lid = p[0], p[1]
                l = lex.get(lid)
                if not l or not form or " " in form:
                    continue
                acc[normalize(form)].add(l[0])
        # lema însăși e mereu o formă a ei
        for form, _ in lex.values():
            acc[normalize(form)].add(form)
        self.form2lemmas = {k: tuple(sorted(v)) for k, v in acc.items()}
        flat = defaultdict(set)
        for k, v in acc.items():
            flat[flatten(k)] |= v
        self.flat2lemmas = {k: tuple(sorted(v)) for k, v in flat.items()}

    def __contains__(self, form):
        return form in self.form2lemmas or form in self.flat2lemmas

    def candidates(self, form):
        """Toate lemele care pot explica forma, inclusiv cele care se scriu cu
        diacritice acolo unde textul nu le are.

        „si" e nota muzicală, dar în text e aproape întotdeauna „și" scris fără
        diacritice. Amândouă intră în cursă, iar împărțirea de mai jos decide,
        după cât cântărește fiecare lemă în formele neambigue.
        """
        a = self.form2lemmas.get(form, ())
        b = self.flat2lemmas.get(flatten(form), ())
        if not a:
            return b
        if not b:
            return a
        return tuple(sorted(set(a) | set(b)))

    def resolve(self, counts):
        """Counter de forme → Counter de leme.

        Formele care duc la o singură lemă se atribuie direct. Formele ambigue
        („care" e și pronume, și forma verbului „căra") se împart între lemele
        candidate proporțional cu ce au strâns lemele din formele neambigue.
        Ce nu se poate atribui se raportează separat.
        """
        from collections import Counter
        solid, ambiguous, unknown = Counter(), [], 0
        for form, n in counts.items():
            c = self.candidates(form)
            if not c:
                unknown += n
            elif len(c) == 1:
                solid[c[0]] += n
            else:
                ambiguous.append((c, n))
        out = Counter(solid)
        for cands, n in ambiguous:
            weights = [solid.get(c, 0) for c in cands]
            tot = sum(weights)
            if tot:
                for c, w in zip(cands, weights):
                    out[c] += n * w / tot
            else:
                share = n / len(cands)
                for c in cands:
                    out[c] += share
        return out, unknown
