#!/usr/bin/env python3
"""Stage 2 — the "brain" of the Russian Scrabble word-list pipeline.

It reads the Stage-1 base word list (built once by ruwords so the heavy PDF is not
re-parsed) together with the grammatical notes and the singular/variant structure, runs
the whole noun-selection logic in memory, and writes a minimal result:

    sources/scrabble_ru/scrabble.txt   — the working dictionary (common nouns, nom. sing.)
    sources/scrabble_ru/undefined.txt  — the ambiguous tail, left for manual review

(sources/scrabble_ru/all.txt is the Stage-1 base.) Every other bucket — adjectives, verbs,
the merged note-nouns, singulars, variants — stays in memory. Pass --dump to also write
them; pass --trace WORD to ask how a single word did or did not reach the dictionary.

Note: all.txt is a plain word list, so the grammatical notes, "ед." singulars and "и"
variants are read from the pdftotext output (slov.txt) and the Stage-1 side files; the
expensive PDF parse itself runs only once.

Sources, most authoritative first: OpenCorpora (mawo-pymorphy3), libmorph (libmorph_check),
and the orthographic dictionary's own notes. See tools/README.md.

Run:  ru-venv/bin/python tools/ru_stage2.py [--dump] [--trace WORD]
"""
import argparse
import os
import re
import subprocess

HERE = os.path.dirname(os.path.abspath(__file__))
# The curated Russian word lists live in sources/scrabble_ru/ (this tool sits in tools/);
# the uncommitted pipeline intermediates (orfo/all/debug) are regenerated alongside them.
OUT_DIR = os.path.join(HERE, "..", "sources", "scrabble_ru")
SLOV = os.path.join(OUT_DIR, "orfo_dict_2025.txt")  # committed pdftotext output (source of truth)
WL_FROM, WL_TO = 452, 168808  # 1-based inclusive bounds of the column word-list section
OC_CACHE = "/tmp/oc_nouns.txt"
LIBMORPH_BIN = os.path.join(HERE, "libmorph_check")

ALPHABET = "абвгдеёжзийклмнопрстуфхцчшщъыьэюя"
ORDER = {c: i for i, c in enumerate(ALPHABET)}
PROPER = {"Name", "Surn", "Patr", "Geox", "Orgn", "Trad"}
ABBR_NOUN = {"Abbr", "Fixd"}      # indeclinable abbreviation marker (НДС, КПД, сельпо, …)
FUNCTION_POS = {"CONJ", "PRCL"}   # conjunction / particle: a function word, never a Scrabble noun
LIBMORPH_NOUN_CODES = set(range(7, 22)) | {24}  # 7..21 plus 24 (pluralia tantum)
ADJ_END = {"ая", "яя", "ое", "ее", "ье", "ья", "ьи"}
VERB3 = ("ет", "ёт", "ит", "ют", "ут", "ает", "яет", "ует", "уют", "нет", "жет", "чет")
GENPL = ("ов", "ёв", "ев", "ей")


def key(w):
    return [ORDER.get(c, 99) for c in w]


def destress(s):
    return "".join(c for c in s if ord(c) not in (0x0300, 0x0301)).lower()


def cyr_ok(w):
    return 2 <= len(w) <= 15 and all(("а" <= c <= "я") or c == "ё" for c in w)


def load(p):
    return [l.strip() for l in open(p, encoding="utf-8") if l.strip()] if os.path.exists(p) else []


def write(path, words):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    open(path, "w", encoding="utf-8").write("\n".join(sorted(set(words), key=key)) + "\n")


import mawo_pymorphy3  # noqa: E402

M = mawo_pymorphy3.MorphAnalyzer()
D = M._dawg_dict


def oc_noun_lemmas():
    """Every common-noun lemma (nom. sing. / pluralia tantum) in OpenCorpora's words.dawg,
    excluding indeclinable abbreviations (Abbr+Fixd, e.g. НДС, КПД). Such words are admitted
    only when the orthographic dictionary attests them as nouns (see oc_abbr_fixd_only and the
    Stage-2 loop); abbreviations carried by OpenCorpora alone (ст, ср, кпд, …) are dropped."""
    gp, pt = D.get_paradigm, D.parse_tag_string
    para0, tagc = {}, {}

    def g0(pid):
        r = para0.get(pid)
        if r is None:
            suf0, tag0, pre0 = gp(pid, 0)
            _, gr = pt(tag0)
            r = (pre0, suf0, gr)
            para0[pid] = r
        return r

    def gt(pid, idx):
        k = (pid, idx)
        r = tagc.get(k)
        if r is None:
            suf, tag, pre = gp(pid, idx)
            pos, gr = pt(tag)
            r = (suf, pre, pos, gr)
            tagc[k] = r
        return r

    out = set()
    for word, rec in D.words_dawg.iteritems():
        pid, idx = rec
        suf, pre, pos, gr = gt(pid, idx)
        if pos != "NOUN":
            continue
        pre0, suf0, gr0 = g0(pid)
        if (PROPER & gr) or (PROPER & gr0):
            continue
        if ABBR_NOUN <= gr0:        # indeclinable abbreviation: not seeded; the loop decides via the note
            continue
        stem = word[len(pre):len(word) - len(suf)] if suf else word[len(pre):]
        out.add(pre0 + stem + suf0)
    return {w for w in out if cyr_ok(w)}


def oc_status(word):
    """(is_common_noun, in_dictionary) for word, from OpenCorpora only."""
    parses = D.get_word_parses(word)
    if not parses:
        return False, False
    gp, pt = D.get_paradigm, D.parse_tag_string
    for pid, idx in parses:
        suf, tag, pre = gp(pid, idx)
        pos, gr = pt(tag)
        if pos == "NOUN":
            _, tag0, _ = gp(pid, 0)
            _, gr0 = pt(tag0)
            if not (PROPER & gr or PROPER & gr0):
                return True, True
    return False, True


def oc_abbr_fixd_only(word):
    """True when OpenCorpora's only common-noun reading of word is an indeclinable
    abbreviation (Abbr+Fixd) — e.g. сельпо, ндс, ан. Such a word is not admitted as a noun on
    OpenCorpora's say-so; the Stage-2 loop hands it to the orthographic note (and the function-
    word check) instead, so lexicalised nouns (сельпо, под, ска) survive while bare letter-
    abbreviations (ндс, кпд) and proper/служебные homographs (ан, зато) do not."""
    parses = D.get_word_parses(word)
    if not parses:
        return False
    gp, pt = D.get_paradigm, D.parse_tag_string
    has_noun = has_plain = False
    for pid, idx in parses:
        suf, tag, pre = gp(pid, idx)
        pos, gr = pt(tag)
        if pos != "NOUN":
            continue
        _, tag0, _ = gp(pid, 0)
        _, gr0 = pt(tag0)
        if PROPER & gr or PROPER & gr0:
            continue
        has_noun = True
        if not (ABBR_NOUN <= gr0):
            has_plain = True
    return has_noun and not has_plain


def oc_function_word(word):
    """True when OpenCorpora reads word as a conjunction or particle — a function word that
    shares a spelling with an abbreviation/proper noun (ан «ан нет», зато). It is never a
    Scrabble noun even if the orthographic note looks noun-like (ан → «самолёт Антонова»)."""
    gp, pt = D.get_paradigm, D.parse_tag_string
    for pid, idx in (D.get_word_parses(word) or []):
        if pt(gp(pid, idx)[1])[0] in FUNCTION_POS:
            return True
    return False


def oc_abbr_only(word):
    """True when every common-noun reading OpenCorpora gives word is an abbreviation (carries
    the Abbr grammeme on its paradigm). Generalises oc_abbr_fixd_only past the Fixd
    (indeclinable) requirement, so it also catches a *declinable* acronym that OpenCorpora
    inflects through a full Abbr case paradigm without Fixd — огвз, заз — as well as the
    indeclinable роно, тв, фио. A word with any non-Abbr noun reading (the homographs ваза
    «ВАЗ», рис «РИС», под «ПОД»/под печи) is not abbreviation-only and is left alone."""
    parses = D.get_word_parses(word)
    if not parses:
        return False
    gp, pt = D.get_paradigm, D.parse_tag_string
    has_noun = has_plain = False
    for pid, idx in parses:
        suf, tag, pre = gp(pid, idx)
        pos, gr = pt(tag)
        if pos != "NOUN":
            continue
        _, tag0, _ = gp(pid, 0)
        _, gr0 = pt(tag0)
        if PROPER & gr or PROPER & gr0:
            continue
        has_noun = True
        if "Abbr" not in gr0:
            has_plain = True
    return has_noun and not has_plain


def is_abbrev_note(note):
    """True when the orthographic note marks an abbreviation («сокр.», e.g. вуз, нэп, колхоз)."""
    return bool(note) and "сокр" in note


def libmorph_analyze(words):
    """Map each word to (known, noun_lemma, codes) per libmorph; noun_lemma is None when it
    is not a common noun there. Empty result if the helper binary is not built."""
    words = list(words)
    if not words or not os.path.exists(LIBMORPH_BIN):
        return {}
    proc = subprocess.run([LIBMORPH_BIN], input="\n".join(words), capture_output=True, text=True)
    out = {}
    for w, line in zip(words, proc.stdout.split("\n")):
        fields = line.split("\t")
        known = fields[:1] == ["1"]
        codes, noun_lemmas = set(), []
        for field in fields[1:]:
            code, _, lex = field.partition(":")
            if code.isdigit():
                codes.add(int(code))
                if int(code) in LIBMORPH_NOUN_CODES:
                    noun_lemmas.append(lex)
        lemma = (w if w in noun_lemmas else noun_lemmas[0]) if noun_lemmas else None
        out[w] = (known, lemma, codes)
    return out


def build_notes():
    """Map each headword (destressed, lowercased) to its grammatical note."""
    def is_hw(ch):
        o = ord(ch)
        return (0x0430 <= o <= 0x044F) or (0x0410 <= o <= 0x042F) or o in (0x0401, 0x0451, 0x0300, 0x0301)

    hmap = {}
    lines = open(SLOV, encoding="utf-8").read().split("\n")
    for l in lines[WL_FROM - 1:WL_TO]:
        s = l.lstrip()
        e = 0
        for ch in s:
            if is_hw(ch):
                e += 1
            else:
                break
        hw = destress(s[:e])
        if hw and hw not in hmap:
            hmap[hw] = destress(s[e:]).strip()
    return hmap


def classify(w, note):
    """Coarse part of speech of an out-of-dictionary word from its PDF note."""
    if note is None:
        return "amb"
    n = re.sub(r"\([^)]*\)", "", note).strip()  # drop domain/etymology parentheticals
    if "кр. ф" in n or "кр.ф" in n or "прич." in n or "прил." in n:
        return "adj"
    ends = re.findall(r"-([а-яё]+)", n)
    if any(e in ADJ_END for e in ends):
        return "adj"
    if "сов." in n or "несов." in n or "безл." in n:
        return "verb"
    if w.endswith("ся"):  # reflexive: no Russian noun ends in -ся
        return "verb"
    if any(e.endswith(VERB3) for e in ends) and not any(m in n for m in ("ед.", "тв.", "род.", "м.", "ж.", "с.")):
        return "verb"
    if n == "" and w.endswith(("ый", "ий", "ой", "ая", "ое", "ые", "ие", "яя", "ее")):
        return "adj"
    if "нескл" in n:
        return "noun" if any(g in n for g in ("м.", "ж.", "с.", "мн.")) else "amb"
    if ends:
        return "noun"
    if n == "" and w.endswith(("ать", "ять", "еть", "ить", "оть", "уть", "ыть", "ти", "чь")):
        return "verb"
    return "amb"


def singular(w, note):
    """Nominative singular of a noun headword from the PDF note (authoritative) or, for a
    plural headword without an explicit singular, the mawo lemma; pluralia tantum kept."""
    n = note or ""
    full = re.search(r"ед\.\s+([а-яё]+)", n)
    if full:
        return full.group(1)
    suf = re.search(r"ед\.\s+-([а-яё]+)", n)
    if suf:
        s = suf.group(1)
        i = w.rfind(s[0])
        return w[:i] + s if i > 0 else w
    ends = re.findall(r"-([а-яё]+)", re.sub(r"\([^)]*\)", "", n))
    if ends and ends[0].endswith(GENPL):
        for p in M.parse(w):
            if str(p.tag.POS) == "NOUN":
                return p.normal_form
        return w
    return w


def build():
    """Run the whole pipeline in memory. Returns the result sets plus a `fate` map giving
    every word's outcome, so a word's path can be traced or the buckets dumped."""
    oc = set(load(OC_CACHE)) or oc_noun_lemmas()
    if not os.path.exists(OC_CACHE):
        write(OC_CACHE, oc)
    hmap = build_notes()
    all_words = load(os.path.join(OUT_DIR, "all.txt"))
    ed_nouns = set(load("/tmp/ru_singulars.txt"))
    pairs = [tuple(p) for l in load("/tmp/ru_variants.txt") if len(p := l.split("\t")) == 2]
    pdf = [w for w in all_words if cyr_ok(w)]
    lm = libmorph_analyze(pdf)

    def to_singular(w):
        s = singular(w, hmap.get(w))
        return s if cyr_ok(s) else w

    fate = {}
    scrabble = set(oc)
    adj, verb, amb = [], [], []
    for w in pdf:
        if oc_abbr_fixd_only(w):
            # Indeclinable abbreviation in OpenCorpora: do not trust OC's noun verdict. Drop it
            # if OC also reads it as a function word (ан, зато — служебное/пропер); otherwise let
            # the orthographic note decide, keeping lexicalised nouns (сельпо, под, ска) and
            # dropping the rest (кпд, чп, гор, …).
            if oc_function_word(w):
                fate[w] = "отброшено: служебное слово (CONJ/PRCL в OpenCorpora)"
            elif classify(w, hmap.get(w)) == "noun":
                s = to_singular(w)
                scrabble.add(s)
                fate[w] = "scrabble: несклон. аббрев.-сущ. по помете орфословаря" + ("" if s == w else f" → {s}")
            else:
                fate[w] = "отброшено: аббревиатура без подтверждающей пометы орфословаря"
            continue
        oc_noun, oc_known = oc_status(w)
        if oc_noun:
            # A common-noun reading already put w in the seed (scrabble = set(oc)) — unless its
            # lemma is a different word, i.e. w is an inflected form (ан = род. мн. от «ана»);
            # such a form is not a headword and stays out.
            fate[w] = ("scrabble: сущ. по OpenCorpora" if w in scrabble
                       else "отброшено: словоформа OpenCorpora (лемма — другое слово)")
            continue
        lm_known, lm_lemma, _ = lm.get(w, (False, None, frozenset()))
        if lm_lemma is not None:
            s = lm_lemma if cyr_ok(lm_lemma) else to_singular(w)
            scrabble.add(s)
            fate[w] = "scrabble: сущ. по libmorph" + ("" if s == w else f" → {s}")
            continue
        if oc_known or lm_known:
            fate[w] = "отброшено: словарь знает как не-существительное"
            continue
        if w in ed_nouns:
            scrabble.add(w)
            fate[w] = "scrabble: ед.ч. по помете «ед.»"
            continue
        c = classify(w, hmap.get(w))
        if c == "noun":
            s = to_singular(w)
            scrabble.add(s)
            fate[w] = "scrabble: сущ. по помете орфословаря" + ("" if s == w else f" → {s}")
        elif c == "adj":
            adj.append(w)
            fate[w] = "отброшено: прилагательное (помета орфословаря)"
        elif c == "verb":
            verb.append(w)
            fate[w] = "отброшено: глагол (помета орфословаря)"
        else:
            amb.append(w)
            fate[w] = "undefined: неоднозначное (нет в словарях, помета не определяет)"

    # Manual confirmations: nouns the maintainer approved from the undefined tail.
    for w in load(os.path.join(OUT_DIR, "manual_confirm.txt")):
        if cyr_ok(w):
            scrabble.add(w)
            fate[w] = "scrabble: подтверждено вручную (manual_confirm.txt)"

    # Variant rescue: a word joined by "и" to a confirmed noun is itself a noun.
    pending = set(amb) - scrabble
    changed = True
    while changed:
        changed = False
        for a, b in pairs:
            for x, y in ((a, b), (b, a)):
                if x in scrabble and y in pending:
                    scrabble.add(y)
                    pending.discard(y)
                    fate[y] = f"scrabble: вариант от «{x}» (через «и»)"
                    changed = True

    # Abbreviations are not Scrabble words. Two independent signals catch them: the orthographic
    # note marks «(сокр.: …)» (вуз, нэп, колхоз, гост, …); and OpenCorpora tags acronyms Abbr —
    # both the indeclinable (роно, тв, фио) and the declinable ones the Abbr+Fixd seed filter
    # misses (огвз, заз, гулаг). Drop every such word, except the lexicalised ones the maintainer
    # rescues in manual_confirm.txt (колхоз, совхоз, рация, спецназ, … and под — a real noun
    # OpenCorpora knows only as the abbreviation ПОД). The homographs that are real words with a
    # mere abbreviation reading (ваза, рис) survive — oc_abbr_only requires *all* noun readings
    # to be Abbr. Subtracted before the manual-reject veto.
    confirm = {w for w in load(os.path.join(OUT_DIR, "manual_confirm.txt")) if cyr_ok(w)}
    abbrev = {w for w in scrabble if w not in confirm and (is_abbrev_note(hmap.get(w)) or oc_abbr_only(w))}
    scrabble -= abbrev
    for w in abbrev:
        fate[w] = "отброшено: аббревиатура (помета «сокр.» орфословаря или только Abbr-чтения в OpenCorpora)"

    # Manual rejections: words the maintainer vetoed — substantivized-adjective false nouns
    # a dictionary misreads as nouns (нёбный, акцизный, …). Subtracted last, so it overrides
    # every admission path above (OC seed / libmorph / note / manual_confirm / variant).
    reject = {w for w in load(os.path.join(OUT_DIR, "manual_reject.txt")) if cyr_ok(w)}
    scrabble -= reject
    for w in reject:
        fate[w] = "отброшено: ручной запрет (manual_reject.txt)"

    undefined = [w for w in amb if w not in scrabble and w not in reject]
    return {
        "oc": oc, "scrabble": scrabble, "undefined": undefined,
        "adjectives": adj, "verbs": verb, "singulars": ed_nouns,
        "fate": fate, "all": set(all_words), "hmap": hmap,
    }


def trace(word, r):
    """A detailed, per-signal audit of how WORD did or did not reach the dictionary: its
    Stage-1 (all.txt) membership, every OpenCorpora reading, libmorph, the orthographic note
    with its classify verdict, and the final outcome. Used by `--trace WORD`."""
    w = destress(word)
    note = r["hmap"].get(w)
    parses = D.get_word_parses(w) or []
    gp, pt = D.get_paradigm, D.parse_tag_string
    pos = sorted({pt(gp(pid, idx)[1])[0] for pid, idx in parses})
    oc_noun, _ = oc_status(w)

    out = [f"{word}  →  {'В СЛОВАРЕ' if w in r['scrabble'] else 'НЕ в словаре'}"]
    out.append(f"  Stage-1 / all.txt (строчный заголовок РАН): {'да' if w in r['all'] else 'нет'}")
    if parses:
        kind = ("только несклон. аббрев. (Abbr+Fixd)" if oc_abbr_fixd_only(w) else
                "обычное сущ." if oc_noun else
                "не существительное")
        out.append(f"  OpenCorpora: части речи {pos}; как сущ. — {kind}; "
                   f"служебное (CONJ/PRCL): {'да' if oc_function_word(w) else 'нет'}")
    else:
        out.append("  OpenCorpora: нет в словаре")
    lm = libmorph_analyze([w]).get(w)
    if lm:
        known, lemma, _ = lm
        out.append("  libmorph: " + (f"сущ. → {lemma}" if lemma else ("известно, не сущ." if known else "нет")))
    else:
        out.append("  libmorph: helper отсутствует — в анализе не участвует")
    out.append("  Орфословарь РАН: " + (f"помета «{note}» → classify={classify(w, note)}"
                                        if note is not None else "нет заголовка"))
    reason = r["fate"].get(w)
    if reason is None:
        if w in r["scrabble"]:
            reason = "лексикон OpenCorpora (обычное сущ.)"
        elif oc_abbr_fixd_only(w) and w not in r["all"]:
            reason = "исключено из посева: несклон. аббрев. (Abbr+Fixd), нет в all.txt"
        elif not parses and w not in r["all"]:
            reason = "нет ни в OpenCorpora, ни в all.txt"
        else:
            reason = "не прошло ни один путь отбора"
    out.append(f"  ИТОГ: {reason}")
    return "\n".join(out)


def main():
    ap = argparse.ArgumentParser(description="Stage 2 brain: build the noun dictionary, trace a word, or dump buckets.")
    ap.add_argument("--dump", action="store_true", help="also write the in-memory buckets (adjectives, verbs, singulars, variants, fate)")
    ap.add_argument("--trace", metavar="WORD", help="report how WORD did or did not reach the dictionary, then exit")
    args = ap.parse_args()

    r = build()
    if args.trace:
        print(trace(args.trace, r))
        return

    write(os.path.join(OUT_DIR, "scrabble.txt"), r["scrabble"])
    print(f"=> sources/scrabble_ru/scrabble.txt   {len(r['scrabble'])}")
    print(f"   undefined kept in memory: {len(set(r['undefined']))} (use --dump to write it)")
    if args.dump:
        write(os.path.join(OUT_DIR, "undefined.txt"), r["undefined"])
        write(os.path.join(OUT_DIR, "adjectives.txt"), r["adjectives"])
        write(os.path.join(OUT_DIR, "verbs.txt"), r["verbs"])
        write(os.path.join(OUT_DIR, "singulars.txt"), r["singulars"])
        fate_path = os.path.join(OUT_DIR, "fate.tsv")
        os.makedirs(OUT_DIR, exist_ok=True)
        with open(fate_path, "w", encoding="utf-8") as f:
            for w in sorted(r["fate"], key=key):
                f.write(f"{w}\t{r['fate'][w]}\n")
        print(f"   dumped: undefined.txt ({len(set(r['undefined']))}), adjectives.txt, verbs.txt, singulars.txt, fate.tsv")


if __name__ == "__main__":
    main()