"""
Vocabulary replacement: new, correct wording without rebuilding the sentence.

When a rewrite is rejected (broken grammar, changed meaning), the humanizer used to fall back to the author's
original sentence, which keeps its AI score. This module instead takes the CORRECT sentence and replaces selected
everyday words with suitable synonyms that fit the context, so the sentence structure stays correct while the
wording changes:

  * a word is replaced only by a true synonym (WordNet) that a language model also finds natural in that exact
    position (masked-language-model fill-in), so "investigated" -> "examined", never "increase" -> "decrease";
  * technical terms, protected terms, names, acronyms, numbers, citations and stopwords are never touched;
  * the most predictable words are replaced first (predictable wording is what AI detectors pick up);
  * only common words are used (no rare or AI-typical vocabulary), with the same inflection (-ed, -s, -ing).

Every result still goes through the clarity gate and the publication rules before it is used.

Space variables
  LEXSUB          1 (default) | 0 = off
  LEXSUB_MODEL    default distilroberta-base
  LEXSUB_RATIO    default 0.35   share of eligible words replaced per sentence
  LEXSUB_MIN_P    default 0.01   minimum fit probability of a replacement in context
"""
import os
import re
import threading

import runtime as RT
import textparse as TP

ENABLED = os.environ.get("LEXSUB", "1") != "0"
MODEL_ID = os.environ.get("LEXSUB_MODEL", "distilroberta-base")
RATIO = float(os.environ.get("LEXSUB_RATIO", "0.35"))
MIN_P = float(os.environ.get("LEXSUB_MIN_P", "0.01"))
TOPK = 60

_lock = threading.Lock()
_state = {"tok": None, "model": None, "failed": False, "wn": None}

try:
    from wordfreq import zipf_frequency as _zipf
except Exception:
    _zipf = None

_SKIP = TP._STOP | {"study", "studies", "result", "results", "data", "model", "models", "method", "methods",
                    "table", "figure", "section", "chapter", "paper", "author", "authors", "et", "al"}
_AI = None


def _ai_words():
    global _AI
    if _AI is None:
        try:
            import pubrules as PR
            _AI = {w.lower() for w in PR.AI_WORDS}
        except Exception:
            _AI = set()
    return _AI


# ------------------------------------------------------------------------------------------------ resources
def _wordnet():
    if _state["wn"] is not None:
        return _state["wn"] or None
    try:
        import nltk
        from nltk.corpus import wordnet as wn
        try:
            wn.synsets("test")
        except LookupError:
            d = os.environ.get("NLTK_DATA") or os.path.join(os.path.expanduser("~"), "nltk_data")
            nltk.download("wordnet", download_dir=d, quiet=True)
            nltk.download("omw-1.4", download_dir=d, quiet=True)
            if d not in nltk.data.path:
                nltk.data.path.append(d)
            wn.synsets("test")
        _state["wn"] = wn
    except Exception as err:
        print(f"[lexsub] synonym dictionary not available, vocabulary replacement off: {err}", flush=True)
        _state["wn"] = False
    return _state["wn"] or None


def _load():
    if _state["model"] is not None or _state["failed"] or not ENABLED:
        return _state["model"] is not None
    with _lock:
        if _state["model"] is not None or _state["failed"]:
            return _state["model"] is not None
        try:
            from transformers import AutoModelForMaskedLM, AutoTokenizer
            _state["tok"] = AutoTokenizer.from_pretrained(MODEL_ID)
            _state["model"] = RT.to_gpu(AutoModelForMaskedLM.from_pretrained(MODEL_ID).eval())
            print(f"[lexsub] loaded {MODEL_ID}", flush=True)
        except Exception as err:
            _state["failed"] = True
            print(f"[lexsub] vocabulary replacement off, model not available: {err}", flush=True)
    return _state["model"] is not None


def available():
    return ENABLED and _wordnet() is not None and _load()


def warmup():
    if ENABLED:
        threading.Thread(target=lambda: (_wordnet(), _load()), daemon=True).start()


def predict(masked):
    """List of texts with one <mask> each -> (list of [(word, prob)], list of {word: prob} for lookups)."""
    import torch
    tok, model = _state["tok"], _state["model"]
    out = []
    for i in range(0, len(masked), 32):
        chunk = [t.replace("<mask>", tok.mask_token) for t in masked[i:i + 32]]
        enc = tok(chunk, return_tensors="pt", padding=True, truncation=True, max_length=160).to(RT.device_of(model))
        with torch.no_grad():
            logits = model(**enc).logits
        pos = (enc["input_ids"] == tok.mask_token_id).nonzero()
        probs = {}
        for b, p in pos.tolist():
            if b not in probs:
                probs[b] = torch.softmax(logits[b, p].float(), -1)
        for b in range(len(chunk)):
            if b not in probs:
                out.append([])
                continue
            pr, ids = probs[b].topk(TOPK)
            words = []
            for p_, id_ in zip(pr.tolist(), ids.tolist()):
                t = tok.convert_ids_to_tokens(id_)
                if t.startswith("Ġ") and t[1:].isalpha():
                    words.append((t[1:], p_))
            out.append(words)
    return out


# ------------------------------------------------------------------------------------------------ rules
def _same_form(orig, cand, wn=None):
    """Same inflection: -ed with -ed (or an irregular past such as 'rose'), -s with -s, -ing with -ing, ..."""
    o, c = orig.lower(), cand.lower()
    if o.endswith("ed") and not c.endswith("ed"):
        return bool(wn) and (wn.morphy(c, "v") or c) != c and not c.endswith(("s", "ing"))
    if wn and not o.endswith(("ed", "s", "ing")) and set(wn._morphy(o, "v")) - {o}:
        # irregular past form ('found', 'grew', 'made'): the replacement must be a past form too
        return c.endswith("ed") or bool(set(wn._morphy(c, "v")) - {c}) and not c.endswith(("s", "ing"))
    for suf in ("ing", "ed", "ly", "s", "est"):
        if o.endswith(suf) != c.endswith(suf):
            return False
    return True


def _synonym(wn, orig, cand):
    """True if cand can replace orig: a WordNet synonym, a direct parent/child sense (nouns, verbs: path
    similarity >= 0.5, e.g. investigate -> examine) or a similar adjective; never an antonym."""
    o, c = orig.lower(), cand.lower()
    if o == c or c in o or o in c:
        return False
    for pos in ("v", "n", "a", "r"):
        lo, lc = wn.morphy(o, pos), wn.morphy(c, pos)
        if not lo or not lc or lo == lc:
            continue
        so, sc = wn.synsets(lo, pos)[:4], wn.synsets(lc, pos)[:4]
        if not so or not sc:
            continue
        lemmas = {l.name().lower() for x in so for l in x.lemmas()}
        antonyms = {a.name().lower() for x in so for l in x.lemmas() for a in l.antonyms()}
        if pos in ("a", "r"):
            for x in so:
                for y in x.similar_tos():
                    antonyms |= {a.name().lower() for l in y.lemmas() for a in l.antonyms()}
        if lc in antonyms:
            return False
        if lc in lemmas:
            return True
        if pos in ("n", "v") and max((x.path_similarity(y) or 0) for x in so for y in sc) >= 0.5:
            return True
        if pos == "a" and lc in {l.name().lower() for x in so for y in x.similar_tos() for l in y.lemmas()}:
            return True
    return False


def _eligible(word, protected):
    w = word.lower()
    if len(w) < 4 or not word.isalpha() or w in _SKIP or w in protected:
        return False
    if word[0].isupper() and not word.isupper():
        return False                                   # names, sentence-initial words (kept for safety)
    if word.isupper():
        return False                                   # acronyms
    if _zipf is not None and _zipf(w, "en") < 2.5:
        return False                                   # rare / technical words
    return True


def _protected_words(sentence, glossary):
    words = set()
    for g in glossary or ():
        for w in str(g).lower().split():
            words.add(w)
    for m in TP.TECH_TOKEN.finditer(sentence):
        words.add(m.group(0).lower())
    return words | {w.lower() for w in TP.FRAGILE}


# ------------------------------------------------------------------------------------------------ public API
def ratios_for(target, n=3):
    """Replacement shares aimed at a vocabulary-change target (%). Vocabulary change counts content words, so a
    target of T% needs a little over T% of the eligible words replaced (some content words, such as technical
    terms and names, are never replaced); the versions straddle that level."""
    if not target:
        import profile_settings as PS
        return [PS.get("lex_ratio")] * n
    need = min(1.0, 1.15 * float(target) / 100.0)
    steps = [0.85, 1.0, 1.2, 1.4][:max(1, n)]
    return sorted({round(min(1.0, max(0.2, need * f)), 2) for f in steps})


def variants(sentences, glossary=(), n=2, ratio=None, predict_fn=None, ratios=None):
    """Vocabulary-replaced versions of each sentence (list per sentence; [] when nothing to replace): n versions at
    `ratio`, or one version per share in `ratios` (e.g. from ratios_for(target))."""
    if ratio is None:
        import profile_settings as PS
        ratio = PS.get("lex_ratio")
    plan = list(ratios) if ratios else [ratio] * n
    out = [[] for _ in sentences]
    if not sentences or not ENABLED:
        return out
    wn = _wordnet()
    if wn is None or (predict_fn is None and not _load()):
        return out
    predict_fn = predict_fn or predict
    jobs = []                                           # (sentence index, (start, end, word), masked text)
    cite_spans, n_elig = [], {}
    for k, s in enumerate(sentences):
        prot = _protected_words(s, glossary)
        cs = [(m.start(), m.end()) for m in TP.CITE.finditer(s)]
        cite_spans.append(cs)
        for m in re.finditer(r"[A-Za-z]+", s):
            if any(a <= m.start() < b for a, b in cs) or not _eligible(m.group(0), prot):
                continue
            jobs.append((k, (m.start(), m.end(), m.group(0)), s[:m.start()] + "<mask>" + s[m.end():]))
            n_elig[k] = n_elig.get(k, 0) + 1
    if not jobs:
        return out
    try:
        preds = predict_fn([j[2] for j in jobs])
    except Exception as err:
        print(f"[lexsub] skipped: {err}", flush=True)
        return out
    ai = _ai_words()
    per = {}                                            # sentence -> [(predictability, span, [candidates])]
    for (k, (a, b, w), _), cands in zip(jobs, preds):
        d = dict((c.lower(), p) for c, p in cands)
        p_orig = d.get(w.lower(), 0.0)
        good = [(c, p) for c, p in cands
                if p >= MIN_P and c.lower() not in ai and _same_form(w, c, wn) and _synonym(wn, w, c)
                and (_zipf is None or _zipf(c.lower(), "en") >= 3.0)]
        if good:
            per.setdefault(k, []).append((p_orig, (a, b, w), good))
    for k, items in per.items():
        s = sentences[k]
        n_words = len(s.split())
        items.sort(key=lambda x: -x[0])                 # most predictable words first
        for v, r in enumerate(plan):
            limit = max(1, round(r * n_elig.get(k, 0))) if n_words >= 6 else 1
            gap = 0 if r >= 0.6 else 1                  # high targets may replace neighbouring words
            chosen = []
            for p_orig, (a, b, w), good in items:
                if len(chosen) >= limit:
                    break
                if any(a <= cb + gap and ca <= b + gap for ca, cb, _ in chosen):
                    continue
                c = good[min(v, len(good) - 1)][0]
                if w[0].isupper():
                    c = c.capitalize()
                chosen.append((a, b, c))
            if not chosen:
                continue
            t = s
            for a, b, c in sorted(chosen, reverse=True):
                t = t[:a] + c + t[b:]
            t = TP.basic_grammar(t)                     # a/an after a replaced word
            if t != s and t not in out[k]:
                out[k].append(t)
    return out