"""morphology_pt_full.py — production-grade Brazilian-Portuguese morphological generator.

NOT a toy. Backed by two real, broad, Wiktionary-lineage lexicons:

  VERBS  — UniMorph Portuguese  (github.com/unimorph/por, CC-BY-SA 3.0)
    4,001 verb lemmas × full paradigm (283,991 finite/non-finite forms +
    20,005 participle forms). Every mood/tense pt actually inflects:
      indicative present / preterite (PST;PFV) / imperfect (PST;IPFV) /
                 pluperfect-simple (PST;PRF) / future,
      conditional (futuro do pretérito),
      subjunctive present / imperfect / FUTURE  (PT-specific live tense),
      affirmative + negative imperative,
      PERSONAL infinitive (V;{p};{n};NFIN — a PT-specific finite-ish form),
      past participle (4 gender/number forms) + gerúndio (V.PTCP;PRS).

  NOUNS + ADJECTIVES — kaikki.org Portuguese (Wiktionary extract, same lineage)
    81,138 noun lemmas WITH inherent gender + real (often irregular) plural —
      so -ão→-ões / -ãos / -ães / -õos is resolved PER LEMMA by Wiktionary,
      never guessed (mão→mãos, pão→pães, coração→corações).
    40,252 adjective lemmas with real feminine + masc/fem plural forms.

Fallbacks (degrade, never crash, on out-of-vocabulary input):
  verbs : rule generator for regular -ar/-er/-ir paradigms
  nouns : gender heuristic (endings) + rule pluralization (with -ão FLAGGED)
  adjs  : -o/-a gender rule + rule pluralization

Confidence flag on every form:
  "lexicon"  straight from UniMorph/kaikki          (trust: high)
  "rule"     deterministic rule                     (trust: medium)
  "fallback" could not inflect; returned lemma      (trust: low  -> FLAG)

Public API (used by realizer_pt.py):
  conjugate(lemma, mood, tense, person, number) -> (form, conf)
  personal_infinitive(lemma, person, number) -> (form, conf)
  participle(lemma, gender="m", number="singular") -> (form, conf)
  gerund(lemma) -> (form, conf)
  noun_gender(lemma) -> "m"|"f"
  inflect_noun(lemma, number, gender=None) -> (form, conf)
  inflect_adj(lemma, gender, number) -> (form, conf)
  lexicon_stats() -> dict
"""
import json
import os
import pickle

_HERE = os.path.dirname(os.path.abspath(__file__))
_UNIMORPH = os.path.join(_HERE, "data", "por.unimorph")
_KAIKKI = os.path.join(_HERE, "data", "kaikki_pt.jsonl")
_CACHE = os.path.join(_HERE, "data", "pt_morph_cache.pkl")

# ── mood/tense pair -> UniMorph feature triple (a in tag; b in tag; c in tag) ────
_VERB_KEYMAP = {
    ("ind", "present"):     ("IND", "PRS", None),
    ("ind", "preterite"):   ("IND", "PST", "PFV"),
    ("ind", "imperfect"):   ("IND", "PST", "IPFV"),
    ("ind", "pluperfect"):  ("IND", "PST", "PRF"),   # simple mais-que-perfeito
    ("ind", "future"):      ("IND", "FUT", None),
    ("ind", "conditional"): ("COND", None, None),
    ("sbjv", "present"):    ("SBJV", "PRS", None),
    ("sbjv", "imperfect"):  ("SBJV", "PST", "IPFV"),
    ("sbjv", "future"):     ("SBJV", "FUT", None),    # PT-specific
    ("imp", "affirmative"): ("IMP", "POS", None),
    ("imp", "negative"):    ("IMP", "NEG", None),
}
_PERSON = {"first": "1", "second": "2", "third": "3"}
_NUMBER = {"singular": "SG", "plural": "PL"}


def _feat_set(tag):
    return set(tag.split(";"))


# ── build the compact lexicon from UniMorph (verbs) + kaikki (nouns/adjs) ────────
def _build_verbs():
    verbs = {}     # (lemma, "mood|tense|person|number") -> form
    pinf = {}      # (lemma, "person|number") -> personal-infinitive form
    part = {}      # lemma -> {("m","SG"): form, ...}  past participle
    ger = {}       # lemma -> gerúndio
    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, tag = parts
            f = _feat_set(tag)
            head = tag.split(";")[0]

            if head == "V.PTCP":
                if "PST" in f:  # past participle: falado/falada/falados/faladas
                    g = "m" if "MASC" in f else ("f" if "FEM" in f else "m")
                    num = "SG" if "SG" in f else ("PL" if "PL" in f else "SG")
                    part.setdefault(lemma, {})[(g, num)] = form
                elif "PRS" in f:  # gerúndio: falando
                    ger.setdefault(lemma, form)
                continue

            if head != "V":
                continue

            # personal / impersonal infinitive
            if "NFIN" in f:
                person = next((p for p in ("1", "2", "3") if p in f), None)
                number = "SG" if "SG" in f else ("PL" if "PL" in f else None)
                if person and number:
                    pinf[(lemma, f"{person}|{number}")] = form
                continue

            # finite forms
            mt = None
            for (mood, tense), (a, b, c) in _VERB_KEYMAP.items():
                if a not in f:
                    continue
                if b is not None and b not in f:
                    continue
                if c is not None and c not in f:
                    continue
                # IND;PST needs exactly PFV|IPFV|PRF — reject if the required one absent
                mt = (mood, tense)
                break
            if mt is None:
                continue
            person = next((p for p in ("1", "2", "3") if p in f), None)
            number = "SG" if "SG" in f else ("PL" if "PL" in f else None)
            if person is None or number is None:
                continue
            verbs.setdefault((lemma, f"{mt[0]}|{mt[1]}|{person}|{number}"), form)
    return verbs, pinf, part, ger


def _kaikki_gender(arg):
    if not arg:
        return None
    a = arg.lower()
    if a.startswith("f"):
        return "f"
    if a.startswith("m"):
        return "m"
    return None


def _build_nouns_adjs():
    nouns = {}   # lemma -> {"g","SG","PL"}
    adjs = {}    # lemma -> {("m","SG"),("f","SG"),("m","PL"),("f","PL")}
    with open(_KAIKKI, encoding="utf-8") as fh:
        for line in fh:
            try:
                d = json.loads(line)
            except Exception:
                continue
            pos = d.get("pos")
            word = d.get("word", "")
            if not word or " " in word:  # skip multiword entries
                continue
            forms = d.get("forms", []) or []

            if pos == "noun":
                ht = d.get("head_templates") or []
                g = None
                if ht:
                    g = _kaikki_gender((ht[0].get("args") or {}).get("1"))
                if g is None:
                    tags = d.get("tags") or []
                    if "feminine" in tags:
                        g = "f"
                    elif "masculine" in tags:
                        g = "m"
                pl = None
                for x in forms:
                    t = x.get("tags") or []
                    if "plural" in t and "alternative" not in t and "obsolete" not in t:
                        pl = x.get("form")
                        break
                # first entry wins; but a later entry with a plural fills a gap
                if word not in nouns:
                    nouns[word] = {"g": g, "SG": word, "PL": pl}
                else:
                    cur = nouns[word]
                    if cur.get("g") is None and g:
                        cur["g"] = g
                    if not cur.get("PL") and pl:
                        cur["PL"] = pl

            elif pos == "adj":
                d0 = adjs.setdefault(word, {})
                d0.setdefault(("m", "SG"), word)
                for x in forms:
                    t = set(x.get("tags") or [])
                    fm = x.get("form")
                    if not fm or ("alternative" in t) or ("obsolete" in t):
                        continue
                    if "comparative" in t or "superlative" in t or \
                       "diminutive" in t or "augmentative" in t:
                        continue
                    if "feminine" in t and "plural" in t:
                        d0[("f", "PL")] = fm
                    elif "masculine" in t and "plural" in t:
                        d0[("m", "PL")] = fm
                    elif "feminine" in t:
                        d0[("f", "SG")] = fm
                    elif "plural" in t:  # invariant-gender adj (feliz -> felizes)
                        d0[("m", "PL")] = d0.get(("m", "PL")) or fm
                        d0[("f", "PL")] = d0.get(("f", "PL")) or fm
    return nouns, adjs


def _build_cache():
    verbs, pinf, part, ger = _build_verbs()
    nouns, adjs = _build_nouns_adjs()
    data = {"verbs": verbs, "pinf": pinf, "part": part, "ger": ger,
            "nouns": nouns, "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):
        newest_src = max(os.path.getmtime(_UNIMORPH),
                         os.path.getmtime(_KAIKKI) if os.path.exists(_KAIKKI) else 0)
        if os.path.getmtime(_CACHE) >= newest_src:
            try:
                with open(_CACHE, "rb") as fh:
                    return pickle.load(fh)
            except Exception:
                pass
    return _build_cache()


_LEX = _load()
_VERBS, _PINF, _PART, _GER, _NOUNS, _ADJS = (
    _LEX["verbs"], _LEX["pinf"], _LEX["part"], _LEX["ger"],
    _LEX["nouns"], _LEX["adjs"])


# ── regular-ending rule fallback (deterministic, last resort) ────────────────────
def _vclass(lemma):
    return lemma[-2:] if lemma[-2:] in ("ar", "er", "ir") else None


def _stem(lemma):
    return lemma[:-2]


# endings indexed [1sg,2sg,3sg,1pl,2pl,3pl]
_REG = {
    ("ind", "present", "ar"): ["o", "as", "a", "amos", "ais", "am"],
    ("ind", "present", "er"): ["o", "es", "e", "emos", "eis", "em"],
    ("ind", "present", "ir"): ["o", "es", "e", "imos", "is", "em"],
    ("ind", "preterite", "ar"): ["ei", "aste", "ou", "amos", "astes", "aram"],
    ("ind", "preterite", "er"): ["i", "este", "eu", "emos", "estes", "eram"],
    ("ind", "preterite", "ir"): ["i", "iste", "iu", "imos", "istes", "iram"],
    ("ind", "imperfect", "ar"): ["ava", "avas", "ava", "ávamos", "áveis", "avam"],
    ("ind", "imperfect", "er"): ["ia", "ias", "ia", "íamos", "íeis", "iam"],
    ("ind", "imperfect", "ir"): ["ia", "ias", "ia", "íamos", "íeis", "iam"],
    ("sbjv", "present", "ar"): ["e", "es", "e", "emos", "eis", "em"],
    ("sbjv", "present", "er"): ["a", "as", "a", "amos", "ais", "am"],
    ("sbjv", "present", "ir"): ["a", "as", "a", "amos", "ais", "am"],
    ("sbjv", "imperfect", "ar"): ["asse", "asses", "asse", "ássemos", "ásseis", "assem"],
    ("sbjv", "imperfect", "er"): ["esse", "esses", "esse", "êssemos", "êsseis", "essem"],
    ("sbjv", "imperfect", "ir"): ["isse", "isses", "isse", "íssemos", "ísseis", "issem"],
    ("sbjv", "future", "ar"): ["ar", "ares", "ar", "armos", "ardes", "arem"],
    ("sbjv", "future", "er"): ["er", "eres", "er", "ermos", "erdes", "erem"],
    ("sbjv", "future", "ir"): ["ir", "ires", "ir", "irmos", "irdes", "irem"],
}
# future & conditional attach to the FULL infinitive
_FUT = ["ei", "ás", "á", "emos", "eis", "ão"]
_COND = ["ia", "ias", "ia", "íamos", "íeis", "iam"]


def _slot_idx(person, number):
    base = {"first": 0, "second": 1, "third": 2}[person]
    return base + (0 if number == "singular" else 3)


def _rule_conjugate(lemma, mood, tense, person, number):
    vc = _vclass(lemma)
    if vc is None:
        return None
    st, i = _stem(lemma), _slot_idx(person, number)
    if mood == "ind" and tense == "future":
        return lemma + _FUT[i]
    if mood == "ind" and tense == "conditional":
        return lemma + _COND[i]
    if mood == "imp":  # affirmative tú/vocês imperative ~ subjunctive present
        table = _REG.get(("sbjv", "present", vc))
        if table and tense == "negative":
            return st + table[i]
        # affirmative 2sg = 3sg present indicative; others = subjunctive
        pres = _REG.get(("ind", "present", vc))
        if person == "second" and number == "singular":
            return st + pres[2]
        return st + table[i] if table else None
    table = _REG.get((mood, tense, vc))
    if table:
        return st + table[i]
    return None


# verified corrections to UniMorph data errors (each audited individually, not
# guessed). The three 1PL-present entries are glued-allomorph errors surfaced by a
# full-lexicon scan for a non-final "mos" in V;1;PL;IND;PRS forms (the ONLY three).
_VERB_FIX = {
    ("estar", "ind", "imperfect", "third", "plural"): "estavam",   # was "estávam"
    ("estar", "ind", "present", "first", "plural"): "estamos",     # was "estamosestámos"
    ("haver", "ind", "present", "first", "plural"): "havemos",     # was "havemoshemos"
    ("ir", "ind", "present", "first", "plural"): "vamos",          # was "vamosimos"
}


# ── PUBLIC: verb conjugation ─────────────────────────────────────────────────────
def conjugate(lemma, mood, tense, person, number):
    """Return (surface, confidence). mood in ind|sbjv|imp; tense per _VERB_KEYMAP."""
    lemma = lemma.strip().lower()
    fix = _VERB_FIX.get((lemma, mood, tense, person, number))
    if fix:
        return fix, "lexicon"
    p, n = _PERSON.get(person), _NUMBER.get(number)
    if p and n:
        form = _VERBS.get((lemma, f"{mood}|{tense}|{p}|{n}"))
        if form:
            # pt-BR normalization: UniMorph `por` carries the EUROPEAN spelling of
            # the -ar 1pl PRETERITE (-ámos). Brazilian PT drops the accent
            # (falámos->falamos, chegámos->chegamos) — 3,334/4,001 verbs affected.
            if (mood == "ind" and tense == "preterite" and person == "first"
                    and number == "plural" and form.endswith("ámos")):
                form = form[:-4] + "amos"
            return form, "lexicon"
    r = _rule_conjugate(lemma, mood, tense, person, number)
    if r:
        return r, "rule"
    return lemma, "fallback"


def personal_infinitive(lemma, person, number):
    """PT personal (inflected) infinitive: para falarmos, ao chegarem."""
    lemma = lemma.strip().lower()
    p, n = _PERSON.get(person), _NUMBER.get(number)
    if p and n:
        form = _PINF.get((lemma, f"{p}|{n}"))
        if form:
            return form, "lexicon"
    # rule: infinitive + personal endings (-, -es, -, -mos, -des, -em)
    end = {("first", "singular"): "", ("second", "singular"): "es",
           ("third", "singular"): "", ("first", "plural"): "mos",
           ("second", "plural"): "des", ("third", "plural"): "em"}.get((person, number), "")
    return lemma + end, "rule"


# ── PUBLIC: participle + gerund ───────────────────────────────────────────────────
def participle(lemma, gender="m", number="singular"):
    lemma = lemma.strip().lower()
    g = "f" if gender == "f" else "m"
    num = "SG" if number == "singular" else "PL"
    d = _PART.get(lemma)
    if d:
        form = d.get((g, num)) or d.get(("m", "SG"))
        if form:
            return form, "lexicon"
    if lemma.endswith("ar"):
        base = lemma[:-2] + "ad"
    elif lemma[-2:] in ("er", "ir"):
        base = lemma[:-2] + "id"
    else:
        return lemma, "fallback"
    suf = {"m|SG": "o", "f|SG": "a", "m|PL": "os", "f|PL": "as"}[f"{g}|{num}"]
    return base + suf, "rule"


def gerund(lemma):
    lemma = lemma.strip().lower()
    if lemma in _GER:
        return _GER[lemma], "lexicon"
    if lemma.endswith("ar"):
        return lemma[:-2] + "ando", "rule"
    if lemma.endswith("er"):
        return lemma[:-2] + "endo", "rule"
    if lemma.endswith("ir"):
        return lemma[:-2] + "indo", "rule"
    return lemma, "fallback"


# ── PUBLIC: noun gender + number ─────────────────────────────────────────────────
_FEM_SUF = ("ção", "são", "ção", "dade", "tade", "agem", "igem", "ugem", "gem",
            "ez", "eza", "ice", "ície", "tude", "ude", "âncbefore")
_FEM_SUF = ("ção", "são", "dade", "tade", "agem", "gem", "eza", "ez", "ice",
            "tude", "ude", "ância", "ência", "ínia")
_MASC_SUF = ("ema", "oma", "ama", "grama", "eta", "ão")  # Greek -ma etc. (mostly m)


def _gender_heuristic(noun):
    for suf in _FEM_SUF:
        if noun.endswith(suf):
            return "f"
    if noun.endswith(("ema", "oma", "ama")):   # problema, idioma, programa
        return "m"
    if noun.endswith("a") or noun.endswith("ã"):
        return "f"
    if noun.endswith("o") or noun.endswith(("l", "r", "z", "m", "u", "i")):
        return "m"
    return "m"


def noun_gender(lemma):
    lemma = lemma.strip().lower()
    d = _NOUNS.get(lemma)
    if d and d.get("g"):
        return d["g"]
    return _gender_heuristic(lemma)


_INVARIANT_PL_SUF = ("s",)  # paroxytones ending -s are invariant (o lápis / os lápis)


def _rule_plural(noun):
    """Deterministic PT pluralization. Returns (form, ok) where ok=False flags an
    ambiguous -ão that should lower confidence (the lexicon normally resolves it)."""
    if not noun:
        return noun, True
    if noun.endswith("ão"):
        return noun[:-2] + "ões", False        # majority rule, but AMBIGUOUS -> flag
    if noun.endswith("m"):
        return noun[:-1] + "ns", True           # homem->homens, jardim->jardins
    if noun.endswith("al"):
        return noun[:-2] + "ais", True
    if noun.endswith("el"):
        return noun[:-2] + "éis", True
    if noun.endswith("ol"):
        return noun[:-2] + "óis", True
    if noun.endswith("ul"):
        return noun[:-2] + "uis", True
    if noun.endswith("il"):
        return noun[:-2] + "is", True           # stressed (funil->funis); unstressed rarer
    if noun.endswith(("r", "z")):
        return noun + "es", True                # flor->flores, luz->luzes
    if noun.endswith("s"):
        # paroxytone -s (lápis, ônibus) invariant; oxytone -s (país) -> -es
        return noun, True
    if noun.endswith(("a", "e", "i", "o", "u", "á", "é", "í", "ó", "ú", "ã")):
        return noun + "s", True
    return noun + "s", True


def inflect_noun(lemma, number, gender=None):
    lemma = lemma.strip().lower()
    d = _NOUNS.get(lemma)
    if number == "singular":
        return (d["SG"] if d and d.get("SG") else lemma), ("lexicon" if d else "rule")
    if d and d.get("PL"):
        return d["PL"], "lexicon"
    form, ok = _rule_plural(lemma)
    return form, ("rule" if ok else "fallback")


# ── PUBLIC: adjective agreement ──────────────────────────────────────────────────
def inflect_adj(lemma, gender, number):
    lemma = lemma.strip().lower()
    g = "f" if gender == "f" else "m"
    num = "SG" if number == "singular" else "PL"
    d = _ADJS.get(lemma)
    if d:
        form = d.get((g, num))
        if form:
            return form, "lexicon"
        # build a missing plural from this gender's singular
        sg = d.get((g, "SG")) or d.get(("m", "SG")) or lemma
        if num == "PL":
            pl, ok = _rule_plural(sg)
            return pl, ("rule" if ok else "fallback")
        return sg, "lexicon"
    # rule fallback: -o/-a gender, then pluralize
    a = lemma
    if g == "f":
        if a.endswith("o"):
            a = a[:-1] + "a"
        elif a.endswith(("ês", "or")) and not a.endswith("ior"):
            a = a + "a"                        # português->portuguesa, trabalhador->..a
    if num == "PL":
        a, ok = _rule_plural(a)
        return a, ("rule" if ok else "fallback")
    return a, "rule"


def lexicon_stats():
    return {
        "verb_source": "UniMorph Portuguese (github.com/unimorph/por)",
        "noun_adj_source": "kaikki.org Portuguese (Wiktionary extract)",
        "license": "CC-BY-SA (Wiktionary-derived)",
        "verb_forms": len(_VERBS),
        "verb_lemmas": len({k[0] for k in _VERBS}),
        "personal_infinitive_forms": len(_PINF),
        "participle_lemmas": len(_PART),
        "gerund_lemmas": len(_GER),
        "noun_lemmas": len(_NOUNS),
        "adj_lemmas": len(_ADJS),
    }


if __name__ == "__main__":
    print(json.dumps(lexicon_stats(), indent=2, ensure_ascii=False))
    tests = [
        ("falar", "ind", "present", "first", "singular", "falo"),
        ("comer", "ind", "present", "third", "plural", "comem"),
        ("partir", "ind", "present", "first", "plural", "partimos"),
        ("ser", "ind", "present", "third", "singular", "é"),
        ("ir", "ind", "preterite", "first", "singular", "fui"),
        ("ter", "ind", "future", "first", "singular", "terei"),
        ("fazer", "sbjv", "present", "first", "singular", "faça"),
        ("dormir", "ind", "present", "first", "singular", "durmo"),
        ("dar", "ind", "preterite", "third", "singular", "deu"),
        ("poder", "ind", "conditional", "first", "singular", "poderia"),
        ("fazer", "sbjv", "future", "third", "singular", "fizer"),
        ("estar", "ind", "present", "third", "singular", "está"),
    ]
    ok = 0
    for lemma, mood, tense, per, num, exp in tests:
        got, conf = conjugate(lemma, mood, tense, per, num)
        flag = "OK " if got == exp else "XX "
        ok += got == exp
        print(f"  {flag}{lemma:8} {mood}/{tense} {per[:3]}.{num[:2]} -> {got:14} ({conf}) exp={exp}")
    print(f"verb tests {ok}/{len(tests)}")
    print("  gender: casa=", noun_gender("casa"), "problema=", noun_gender("problema"),
          "mão=", noun_gender("mão"), "coração=", noun_gender("coração"),
          "flor=", noun_gender("flor"))
    print("  plural: mão->", inflect_noun("mão", "plural"),
          "| pão->", inflect_noun("pão", "plural"),
          "| animal->", inflect_noun("animal", "plural"),
          "| coração->", inflect_noun("coração", "plural"))
    print("  adj: bonito/f/sg->", inflect_adj("bonito", "f", "singular"),
          "| feliz/m/pl->", inflect_adj("feliz", "m", "plural"),
          "| português/f/sg->", inflect_adj("português", "f", "singular"))
    print("  part: fazer/m/sg->", participle("fazer"), "| ger falar->", gerund("falar"))
    print("  pinf falar 1pl->", personal_infinitive("falar", "first", "plural"))
