# -*- coding: utf-8 -*-
"""morphology_de_full.py — production German morphological generator.

Real data, no toy tables:

  PRIMARY — UniMorph German (github.com/unimorph/deu, CC-BY-SA 3.0).
    ~219k noun forms, ~199k verb forms. Supplies:
      nouns : gender (MASC/FEM/NEUT) + case×number paradigm
              (N;NOM/ACC/DAT/GEN; MASC/FEM/NEUT; SG/PL) — the genitive -(e)s,
              dative-plural -n and the five plural classes are REAL forms, not
              guessed.
      verbs : full finite paradigm IND;{SG,PL};{1,2,3};{PRS,PST}, the past
              participle (V.PTCP;PST, incl. reattached separable prefix
              'zugefügt'), and — crucially for V2 — the SEPARATED finite form
              UniMorph records directly ('füge zu', 'steht auf').
      adjs  : comparative / superlative (ADJ;CMPR, ADJ;SPRL).

  SECONDARY — kaikki.org German (Wiktionary, CC-BY-SA/GFDL). Gap-fills noun
    gender + plural where UniMorph is thin. Never overrides UniMorph.

Rule fallbacks (flagged 'rule'/'fallback') for lemmas absent from both lexicons:
  present  : -e/-st/-t/-en/-t/-en with e-epenthesis after -t/-d/-chn stems
  plural   : gender heuristic (fem -> -(e)n, else -e / umlaut left to lexicon)
  ppart    : weak ge-…-t
Adjective ENDINGS are rule-computed by the realizer (regular closed table);
this module only supplies the comparative/superlative STEM.

Perfect auxiliary (haben vs sein): sein for a curated set of intransitive
motion / change-of-state verbs (real German lexical property), else haben.

Public API:
  noun_gender(lemma) -> 'm'|'f'|'n'
  decline_noun(lemma, case, number) -> (form, conf)
  pluralize(lemma) -> (form, conf)
  finite(lemma, tense, person, number) -> (form, conf)      # may contain ' prefix'
  nonfinite(lemma, req) -> (form, conf)   # req: 'inf'|'ppart'
  past_participle(lemma) -> (form, conf)
  separable_prefix(lemma) -> str|None
  perfect_aux(lemma) -> 'haben'|'sein'
  comparative(lemma)/superlative(lemma) -> (stem, conf)
  lexicon_stats() -> dict
"""
import json
import os
import pickle

_HERE = os.path.dirname(os.path.abspath(__file__))
_UNIMORPH = os.path.join(_HERE, "data", "deu.unimorph")
_KAIKKI = os.path.join(_HERE, "data", "kaikki_de.jsonl")
_CACHE = os.path.join(_HERE, "data", "de_morph_cache.pkl")

_GENDER = {"MASC": "m", "FEM": "f", "NEUT": "n"}

# intransitive motion / change-of-state verbs that take SEIN in the perfect
_SEIN = {"gehen", "kommen", "fahren", "laufen", "rennen", "reisen", "fallen",
         "steigen", "sinken", "wachsen", "sterben", "geschehen", "passieren",
         "werden", "bleiben", "sein", "aufstehen", "einschlafen", "aufwachen",
         "ankommen", "abfahren", "aufsteigen", "erscheinen", "verschwinden",
         "fliegen", "schwimmen", "springen", "begegnen", "folgen", "gelingen",
         "wandern", "ziehen", "flüchten", "eintreten", "einsteigen", "aussteigen"}


# hardcoded high-frequency irregular / auxiliary / modal paradigms (closed class,
# verified) — consulted before the lexicon so aux+modal chains are always correct.
_CORE = {
    "sein": {"prs": {("first", "singular"): "bin", ("second", "singular"): "bist",
                     ("third", "singular"): "ist", ("first", "plural"): "sind",
                     ("second", "plural"): "seid", ("third", "plural"): "sind"},
             "pst": {("first", "singular"): "war", ("second", "singular"): "warst",
                     ("third", "singular"): "war", ("first", "plural"): "waren",
                     ("second", "plural"): "wart", ("third", "plural"): "waren"},
             "ppart": "gewesen"},
    "haben": {"prs": {("first", "singular"): "habe", ("second", "singular"): "hast",
                      ("third", "singular"): "hat", ("first", "plural"): "haben",
                      ("second", "plural"): "habt", ("third", "plural"): "haben"},
              "pst": {("first", "singular"): "hatte", ("second", "singular"): "hattest",
                      ("third", "singular"): "hatte", ("first", "plural"): "hatten",
                      ("second", "plural"): "hattet", ("third", "plural"): "hatten"},
              "ppart": "gehabt"},
    "werden": {"prs": {("first", "singular"): "werde", ("second", "singular"): "wirst",
                       ("third", "singular"): "wird", ("first", "plural"): "werden",
                       ("second", "plural"): "werdet", ("third", "plural"): "werden"},
               "pst": {("first", "singular"): "wurde", ("second", "singular"): "wurdest",
                       ("third", "singular"): "wurde", ("first", "plural"): "wurden",
                       ("second", "plural"): "wurdet", ("third", "plural"): "wurden"},
               "ppart": "geworden"},
}
_MODAL_PRS = {
    "können": ("kann", "kannst", "kann", "können", "könnt", "können"),
    "müssen": ("muss", "musst", "muss", "müssen", "müsst", "müssen"),
    "wollen": ("will", "willst", "will", "wollen", "wollt", "wollen"),
    "sollen": ("soll", "sollst", "soll", "sollen", "sollt", "sollen"),
    "dürfen": ("darf", "darfst", "darf", "dürfen", "dürft", "dürfen"),
    "mögen": ("mag", "magst", "mag", "mögen", "mögt", "mögen"),
}
_MODAL_PST = {
    "können": ("konnte", "konntest", "konnte", "konnten", "konntet", "konnten"),
    "müssen": ("musste", "musstest", "musste", "mussten", "musstet", "mussten"),
    "wollen": ("wollte", "wolltest", "wollte", "wollten", "wolltet", "wollten"),
    "sollen": ("sollte", "solltest", "sollte", "sollten", "solltet", "sollten"),
    "dürfen": ("durfte", "durftest", "durfte", "durften", "durftet", "durften"),
    "mögen": ("mochte", "mochtest", "mochte", "mochten", "mochtet", "mochten"),
}
_PN_ORDER = [("first", "singular"), ("second", "singular"), ("third", "singular"),
             ("first", "plural"), ("second", "plural"), ("third", "plural")]
_MODAL_PPART = {"können": "gekonnt", "müssen": "gemusst", "wollen": "gewollt",
                "sollen": "gesollt", "dürfen": "gedurft", "mögen": "gemocht"}
for _m, _forms in _MODAL_PRS.items():
    _CORE[_m] = {"prs": dict(zip(_PN_ORDER, _forms)),
                 "pst": dict(zip(_PN_ORDER, _MODAL_PST[_m])),
                 "ppart": _MODAL_PPART[_m]}


def _person_num(tags):
    p = n = None
    for t in tags:
        if t in ("1", "2", "3"):
            p = {"1": "first", "2": "second", "3": "third"}[t]
        elif t == "SG":
            n = "singular"
        elif t == "PL":
            n = "plural"
    return p, n


def _build_from_unimorph():
    nouns, verbs, adjs = {}, {}, {}
    if not os.path.exists(_UNIMORPH):
        return nouns, verbs, adjs
    with open(_UNIMORPH, encoding="utf-8") as fh:
        for line in fh:
            line = line.rstrip("\n")
            if not line or "\t" not in line:
                continue
            parts = line.split("\t")
            if len(parts) != 3:
                continue
            lemma, form, tagstr = parts
            tags = tagstr.split(";")
            head = tags[0]
            tset = set(tags)
            if head == "N":
                rec = nouns.setdefault(lemma, {"g": None, "cases": {}, "pl": None})
                g = next((_GENDER[t] for t in tags if t in _GENDER), None)
                if g and not rec["g"]:
                    rec["g"] = g
                case = next((t for t in tags if t in ("NOM", "ACC", "DAT", "GEN")), None)
                num = "plural" if "PL" in tset else ("singular" if "SG" in tset else None)
                if case and num:
                    rec["cases"].setdefault((case, num), form)
                    if case == "NOM" and num == "plural" and not rec["pl"]:
                        rec["pl"] = form
            elif head.startswith("V"):
                rec = verbs.setdefault(lemma, {"prs": {}, "pst": {}, "ppart": None})
                if "PTCP" in head and "PST" in tset:
                    rec["ppart"] = rec["ppart"] or form
                elif "IND" in tset and ("PRS" in tset or "PST" in tset):
                    p, n = _person_num(tags)
                    if p and n:
                        slot = "prs" if "PRS" in tset else "pst"
                        rec[slot].setdefault((p, n), form)
            elif head == "ADJ":
                rec = adjs.setdefault(lemma, {})
                if "CMPR" in tset:
                    rec.setdefault("cmpr", form.replace("am ", "").strip())
                elif "SPRL" in tset:
                    rec.setdefault("sprl", form.replace("am ", "").replace("sten", "st")
                                   if form.endswith("sten") else form.replace("am ", ""))
    return nouns, verbs, adjs


def _build_from_kaikki(nouns):
    """Gap-fill noun gender + plural from kaikki German."""
    if not os.path.exists(_KAIKKI):
        return
    _g = {"masculine": "m", "feminine": "f", "neuter": "n", "m": "m", "f": "f", "n": "n"}
    with open(_KAIKKI, encoding="utf-8") as fh:
        for line in fh:
            try:
                d = json.loads(line)
            except Exception:
                continue
            if d.get("pos") != "noun":
                continue
            w = d.get("word", "")
            if not w or not w[0].isalpha() or " " in w:
                continue
            rec = nouns.setdefault(w, {"g": None, "cases": {}, "pl": None})
            # GENDER: Wiktionary gender is hand-curated and OVERRIDES UniMorph's
            # auto-tagged gender, which has known errors (e.g. UniMorph deu mis-
            # records Zeit=MASC, Wagen=NEUT; Wiktionary has f, m correctly).
            for h in d.get("head_templates", []) or []:
                a = h.get("args", {}) or {}
                raw = a.get("1") or a.get("g") or ""
                code = str(raw).split(",")[0].strip().lower()
                if code in _g:
                    rec["g"] = _g[code]
                    break
            if not rec["pl"]:
                for f in d.get("forms", []) or []:
                    t = set(f.get("tags", []) or [])
                    if "plural" in t and f.get("form") and "genitive" not in t:
                        rec["pl"] = f["form"]
                        break


def _build_cache():
    nouns, verbs, adjs = _build_from_unimorph()
    _build_from_kaikki(nouns)
    data = {"nouns": nouns, "verbs": verbs, "adjs": adjs}
    try:
        with open(_CACHE, "wb") as fh:
            pickle.dump(data, fh, protocol=pickle.HIGHEST_PROTOCOL)
    except OSError:
        pass
    return data


def _load():
    if os.path.exists(_CACHE):
        srcs = [p for p in (_UNIMORPH, _KAIKKI) if os.path.exists(p)]
        newest = max((os.path.getmtime(p) for p in srcs), default=0)
        if os.path.getmtime(_CACHE) >= newest:
            try:
                with open(_CACHE, "rb") as fh:
                    return pickle.load(fh)
            except Exception:
                pass
    return _build_cache()


_LEX = _load()
_NOUNS, _VERBS, _ADJS = _LEX["nouns"], _LEX["verbs"], _LEX["adjs"]


# ── nouns ────────────────────────────────────────────────────────────────────────
def noun_gender(lemma):
    rec = _NOUNS.get(lemma) or _NOUNS.get(lemma.capitalize())
    if rec and rec.get("g"):
        return rec["g"]
    # last-resort rule: -ung/-heit/-keit/-schaft/-tät/-ion -> f ; -chen/-lein -> n
    low = lemma.lower()
    if low.endswith(("ung", "heit", "keit", "schaft", "tät", "ion", "ik", "ei")):
        return "f"
    if low.endswith(("chen", "lein", "ment", "um")):
        return "n"
    return "m"


def pluralize(lemma):
    rec = _NOUNS.get(lemma) or _NOUNS.get(lemma.capitalize())
    if rec and rec.get("pl"):
        return rec["pl"], "lexicon"
    g = noun_gender(lemma)
    if g == "f":
        return (lemma + "en" if not lemma.endswith("e") else lemma + "n"), "rule"
    return (lemma if lemma.endswith(("er", "en", "el")) else lemma + "e"), "rule"


def decline_noun(lemma, case, number):
    """case in NOM/ACC/DAT/GEN, number in singular/plural."""
    rec = _NOUNS.get(lemma) or _NOUNS.get(lemma.capitalize())
    if case == "DAT" and number == "singular":
        # modern German drops the archaic dative -e ('dem Kinde' -> 'dem Kind');
        # the article carries the case. Keep bare nominative form.
        base = (rec or {}).get("cases", {}).get(("NOM", "singular")) or lemma
        return base, ("lexicon" if rec else "rule")
    if rec and rec.get("cases", {}).get((case, number)):
        return rec["cases"][(case, number)], "lexicon"
    if number == "plural":
        pl, c = pluralize(lemma)
        if case == "DAT" and not pl.endswith("n") and not pl.endswith("s"):
            return pl + "n", c            # dative plural -n
        return pl, c
    # singular
    g = noun_gender(lemma)
    if case == "GEN" and g in ("m", "n"):
        return (lemma + "es" if lemma.endswith(("s", "ß", "z", "x")) else lemma + "s"), "rule"
    return lemma, "lexicon" if rec else "rule"


# ── verbs ──────────────────────────────────────────────────────────────────────--
_PRS_ENDINGS = {("first", "singular"): "e", ("second", "singular"): "st",
                ("third", "singular"): "t", ("first", "plural"): "en",
                ("second", "plural"): "t", ("third", "plural"): "en"}


def _stem(lemma):
    if lemma.endswith("en"):
        return lemma[:-2]
    if lemma.endswith("n"):
        return lemma[:-1]
    return lemma


def separable_prefix(lemma):
    """Return the separable prefix if the lemma is a separable-prefix verb."""
    rec = _VERBS.get(lemma)
    if rec:
        for (_p, _n), form in rec.get("prs", {}).items():
            if " " in form:
                return form.rsplit(" ", 1)[1]
    _SEP = ("auf", "aus", "ab", "an", "ein", "mit", "nach", "vor", "zu", "zurück",
            "weg", "hin", "her", "los", "bei", "fest", "fort", "um", "zusammen")
    _INSEP = ("be", "ge", "er", "ver", "zer", "ent", "emp", "miss")
    for p in sorted(_SEP, key=len, reverse=True):
        if lemma.startswith(p) and len(lemma) > len(p) + 2 \
                and not lemma.startswith(_INSEP):
            return p
    return None


def finite(lemma, tense, person, number):
    """Present/past finite. For separable verbs the returned string is the
    UniMorph SEPARATED form 'stem prefix' (realizer places prefix per V2)."""
    slot = "prs" if tense == "present" else "pst"
    if lemma in _CORE and _CORE[lemma].get(slot, {}).get((person, number)):
        return _CORE[lemma][slot][(person, number)], "lexicon"
    rec = _VERBS.get(lemma)
    if rec and rec.get(slot, {}).get((person, number)):
        return rec[slot][(person, number)], "lexicon"
    # rule fallback (present only reliable; past weak -te)
    stem = _stem(lemma)
    pref = separable_prefix(lemma)
    if pref:
        stem = _stem(lemma[len(pref):])
    if tense == "present":
        end = _PRS_ENDINGS[(person, number)]
        if stem.endswith(("t", "d", "chn", "ffn", "gn")) and end in ("st", "t"):
            end = "e" + end
        form = stem + end
    else:
        form = stem + ("ete" if stem.endswith(("t", "d")) else "te")
        if (person, number) == ("second", "singular"):
            form += "st"
        elif number == "plural" and person != "second":
            form += "n"
        elif (person, number) == ("second", "plural"):
            form += "t"
    if pref:
        return f"{form} {pref}", "rule"
    return form, "rule"


def _weak_t(stem):
    return stem + ("et" if stem.endswith(("t", "d", "chn", "ffn", "gn")) else "t")


def past_participle(lemma):
    if lemma in _CORE:
        return _CORE[lemma]["ppart"], "lexicon"
    rec = _VERBS.get(lemma)
    if rec and rec.get("ppart"):
        return rec["ppart"], "lexicon"
    stem = _stem(lemma)
    pref = separable_prefix(lemma)
    _INSEP = ("be", "ge", "er", "ver", "zer", "ent", "emp", "miss")
    if pref:
        inner = _stem(lemma[len(pref):])
        return pref + "ge" + _weak_t(inner), "rule"
    if lemma.startswith(_INSEP):
        return _weak_t(stem), "rule"
    return "ge" + _weak_t(stem), "rule"


def nonfinite(lemma, req):
    if req == "ppart":
        return past_participle(lemma)
    return lemma, "lexicon" if lemma in _VERBS else "rule"   # infinitive


def perfect_aux(lemma):
    return "sein" if lemma in _SEIN else "haben"


# ── adjectives ────────────────────────────────────────────────────────────────---
_ADJ_IRREG_SPRL = {"gut": "best", "groß": "größt", "hoch": "höchst",
                   "nah": "nächst", "viel": "meist", "gern": "liebst"}


def comparative(lemma):
    rec = _ADJS.get(lemma)
    if rec and rec.get("cmpr"):
        return rec["cmpr"], "lexicon"
    return lemma + "er", "rule"


def superlative(lemma):
    """Return the bare superlative STEM (realizer adds 'am ...en' or '-e' ending)."""
    if lemma in _ADJ_IRREG_SPRL:
        return _ADJ_IRREG_SPRL[lemma], "lexicon"
    # derive from the comparative so umlaut is carried (alt->älter->ältest)
    cmpr, cconf = comparative(lemma)
    base = cmpr[:-2] if cmpr.endswith("er") else lemma
    end = "est" if base.endswith(("t", "d", "s", "ß", "z", "sch")) else "st"
    return base + end, cconf


def lexicon_stats():
    return {
        "source": "UniMorph deu (primary) + kaikki.org German (gap-fill gender/plural)",
        "license": "CC-BY-SA 3.0 (UniMorph); CC-BY-SA/GFDL (Wiktionary)",
        "noun_lemmas": len(_NOUNS),
        "nouns_with_gender": sum(1 for v in _NOUNS.values() if v.get("g")),
        "nouns_with_plural": sum(1 for v in _NOUNS.values() if v.get("pl")),
        "verb_lemmas": len(_VERBS),
        "verbs_with_ppart": sum(1 for v in _VERBS.values() if v.get("ppart")),
        "adj_lemmas": len(_ADJS),
    }


if __name__ == "__main__":
    print(json.dumps(lexicon_stats(), indent=2, ensure_ascii=False))
    for w in ("Hund", "Frau", "Kind", "Mann", "Buch", "Blume"):
        print(f"  {w}: gender={noun_gender(w)} pl={pluralize(w)} "
              f"gen.sg={decline_noun(w, 'GEN', 'singular')} "
              f"dat.pl={decline_noun(w, 'DAT', 'plural')}")
    for v in ("machen", "gehen", "aufstehen", "sein", "haben", "arbeiten"):
        print(f"  {v}: 3sg.prs={finite(v, 'present', 'third', 'singular')} "
              f"3sg.pst={finite(v, 'past', 'third', 'singular')} "
              f"ppart={past_participle(v)} aux={perfect_aux(v)} sep={separable_prefix(v)}")
    for a in ("schnell", "gut", "groß", "alt"):
        print(f"  {a}: cmpr={comparative(a)} sprl={superlative(a)}")
