proofread.py

"""
Rigorous proofread of a finished (humanized) Word document. Correct grammar comes first; a moderate AI-score rise
is accepted for that, style-only edits must keep the AI score.

  0. Local clean-up (no model): chatbot leftovers, markdown, doubled words / repeated phrases, AI filler words.
  1. Grammar audit: the AI's grammar-only check marks every sentence that contains an error (plus sentences with
     doubled terms, leftovers or broken English): these are REQUIRED fixes.
  2. Full proofread (clarity repair against the original upload, grammar, academic wording, flow). For a required
     sentence the full correction is used if its AI rise is at most PROOF_REQ_RISE, otherwise the minimal grammar
     fix; a sentence WITH an error is never left uncorrected. Style-only edits must stay within PROOF_MAX_RISE.
  3. Verification rounds (PROOF_ROUNDS): every corrected sentence is grammar-checked again until no error is found.
  4. LanguageTool: last mechanical check.
  5. Document check: if the AI score rose by more than PROOF_MAX_DOC_RISE points, style-only edits are undone
     (riskiest first). Grammar fixes are never undone.
Every change keeps terms / numbers / citations and the meaning. Only plain text runs are edited.

Unclear / garbled sentences the full proofread REPAIRS are judged against the author's original sentence (not
against the garbled text, which made good repairs fail the meaning check) and are never undone by the AI-score
steps: clarity wins.

Space variables: PROOF_REQ_RISE (0.35), PROOF_MAX_RISE (0.15), PROOF_MAX_DOC_RISE (2 points), PROOF_ROUNDS (1),
PROOF_KEEP_RISE (0.05), PROOF_ALT_RISE (0.10), PROOF_REPAIR_SIM (0.60), PROOF_REASONING (medium).
"""
import difflib
import os
import time
import re

import detector as D
import docx_report as DX
import groqai as G
import humanizer as H
import meaning as MEANING
import fluency as FL
import smart as S
import textparse as TP

MAX_RISE = float(os.getenv("PROOF_MAX_RISE", "0.15"))
FLAG = float(os.getenv("PROOF_FLAG", "0.50"))
MAX_DOC_RISE = float(os.getenv("PROOF_MAX_DOC_RISE", "2.0"))
# a sentence with a real defect (doubled word / repeated phrase / chatbot leftover / broken English) may rise more
DEFECT_RISE = float(os.getenv("PROOF_DEFECT_RISE", "0.25"))
# grammar comes first: a sentence with an error is always corrected; its full correction is used if the AI rise
# stays within PROOF_REQ_RISE (otherwise the minimal grammar fix)
REQ_RISE = float(os.getenv("PROOF_REQ_RISE", "0.35"))
ROUNDS = int(os.getenv("PROOF_ROUNDS", "1"))
BATCH = int(os.getenv("PROOF_BATCH", "24"))             # sentences per request (fewer, larger requests = faster)
# keep the AI score: a corrected sentence may be at most this much more AI-like than before; above it, correct
# natural-sounding alternatives are tried (PROOF_AI_ROUNDS rounds, PROOF_VARIANTS versions each, plus Parrot)
KEEP_RISE = float(os.getenv("PROOF_KEEP_RISE", "0.05"))
# a corrected sentence is re-worded for the AI score only if it rose more than this (each rewording costs tokens and
# moves the wording away from the checked correction)
ALT_RISE = float(os.getenv("PROOF_ALT_RISE", "0.10"))
AI_ROUNDS = int(os.getenv("PROOF_AI_ROUNDS", "1"))
VARIANTS = int(os.getenv("PROOF_VARIANTS", "3"))
# a repair of a garbled sentence must keep the author's meaning (MEANING_SOFT against the original sentence) and
# stay at least this close to the current sentence; without any original it needs PROOF_REPAIR_SIM_NOREF
REPAIR_SIM = float(os.getenv("PROOF_REPAIR_SIM", "0.60"))
REPAIR_SIM_NOREF = float(os.getenv("PROOF_REPAIR_SIM_NOREF", "0.70"))
# reasoning effort of the full proofread on reasoning models (the grammar audit stays on GROQ_REASONING)
REASONING = os.getenv("PROOF_REASONING", "medium").lower()
# single (default): the whole document is checked and corrected in ONE request (grammar, clarity, academic wording,
# flow and document-wide consistency), and the reply lists problems that need the author. Every correction still
# goes through the local checks. If the provider cannot take one request of that size, or it fails, the batched
# proofread (several smaller requests) runs instead. batched = always the batched proofread.
MODE = os.getenv("PROOF_MODE", "single").lower()
PAR_TOL = int(os.getenv("PROOF_PAR_TOL", "0"))       # flagged words a paragraph may gain after balancing
PARROT = os.getenv("PROOF_PARROT", "0") == "1"
# which AI does the proofreading: Groq by default (PROOF_PROVIDER=openai for ChatGPT). If the chosen one is out of
# quota, the other one is used automatically when its key is set.
PROVIDER = os.getenv("PROOF_PROVIDER", "groq").lower()
MAX_WAIT = float(os.getenv("PROOF_MAX_WAIT", "180"))     # seconds the proofread may wait for rate limits in total
PARALLEL = int(os.getenv("PROOF_PARALLEL", "3"))          # requests at once (Groq: low, so limits are not hit)
# no language model reachable (usage limits): do the basic proofread (rules + LanguageTool) and say so, not an error
BASIC_FALLBACK = os.getenv("PROOF_BASIC_FALLBACK", "1") != "0"
_NAME = {"groq": "Groq", "openai": "ChatGPT"}


class _LLM:
    """Calls Groq / ChatGPT for the proofread so that EVERY sentence gets an answer: failed batches are retried
    (only the missing sentences), rate limits are waited out (up to PROOF_MAX_WAIT in total), and if the provider
    is out of quota the other provider is used when its key is set."""

    def __init__(self, rep):
        self.rep = rep
        self.order = [PROVIDER] + [p for p in ("groq", "openai") if p != PROVIDER]
        self.order = [p for p in self.order if (G.OPENAI_KEY if p == "openai" else G.KEYS)]
        self.k = 0
        self.waited = 0.0
        self.used = set()
        self.dead = False

    @property
    def name(self):
        return _NAME.get(self.order[self.k], self.order[self.k]) if self.order else "AI"

    def _ready(self):
        return G.available()

    def _switch(self):
        if self.k + 1 < len(self.order):
            self.k += 1
            print(f"[proofread] switching to {self.name}", flush=True)
            return True
        return False

    def _blocked_for(self):
        """Seconds until the current provider can be used again (large = daily quota / billing / bad key)."""
        now = time.time()
        if self.order[self.k] == "openai":
            return max(0.0, G._oa["until"] - now) if G.OPENAI_KEY else 1e9
        return max(G._down_until - now, G._next_free_in(), 0.0)

    def _wait(self, why):
        blocked = self._blocked_for()
        if blocked > 90:                               # quota for the day / billing: waiting will not help
            print(f"[proofread] {self.name} blocked for {int(blocked)}s ({why}); "
                  f"groq: {G.STATS.get('last_error', '')[:120]} | openai: {G.OA_STATS.get('last_error', '')[:120]}",
                  flush=True)
            return False
        t = min(30.0, max(3.0, blocked or 5.0))
        if self.waited + t > MAX_WAIT:
            return False
        self.rep(None, f"Waiting {int(t)}s for the {self.name} rate limit ({why})")
        time.sleep(t)
        self.waited += t
        return True

    def run(self, fn, items, *args, missing=lambda r: r is None, **kw):
        """fn(list, *args, **kw) -> list of results; returns results for all items (None only if impossible)."""
        out = [None] * len(items)
        todo = list(range(len(items)))
        while todo and self.order:
            tok_p = G.PROVIDER.set(self.order[self.k])
            tok_l = G.PAR_LIMIT.set(PARALLEL if self.order[self.k] == "groq" else 0)
            try:
                if not self._ready():
                    raise G.Unavailable("paused")
                res = fn([items[i] for i in todo], *args, **kw)
                self.used.add(self.order[self.k])
                for i, r in zip(todo, res):
                    if not missing(r):
                        out[i] = r
                left = [i for i in todo if missing(out[i])]
                if len(left) == len(todo):              # nothing came back: treat as unavailable
                    raise G.Unavailable("no answers")
                todo = left
                if todo and not self._wait(f"{len(todo)} sentence(s) to retry"):
                    break
            except G.Unavailable as err:
                daily = "quota" in str(err).lower() or "billing" in str(err).lower() or \
                    (self.order[self.k] == "groq" and G._next_free_in() > 300)
                if (daily or self._blocked_for() > 90) and self._switch():
                    continue
                if not self._wait(str(err)[:60]) and not self._switch():
                    self.dead = True
                    break
            finally:
                G.PROVIDER.reset(tok_p)
                G.PAR_LIMIT.reset(tok_l)
        return out


def _jobs(doc, orig_path):
    blocks = list(D.docx_blocks(doc))
    orig = None
    if orig_path:
        try:
            ob = list(D.docx_blocks(DX.open_clean(orig_path)))
            orig = ob if len(ob) == len(blocks) else None
        except Exception as err:
            print(f"[proofread] original not used: {err}", flush=True)
    jobs = []
    for bi, b in enumerate(blocks):
        if b.kind != "para" or b.el is None or not b.raw:
            continue
        raws = D.split_sentences(b.raw)
        o_raw = orig[bi].raw if orig is not None and orig[bi].kind == "para" else None
        for sp, (cp, cs) in zip(DX.sentence_char_spans(b.raw, raws), H.originals_for(raws, o_raw)):
            if not sp:
                continue
            sub = S.safe_subspan(b.el, *sp)
            if sub and TP.artifact_only(b.raw[sub[0]:sub[1]]):
                jobs.append(dict(b=b, span=sub, text=b.raw[sub[0]:sub[1]], par=cp, sent=None, drop=True))
            elif sub and len(b.raw[sub[0]:sub[1]].split()) >= 4:
                full = tuple(sub) == tuple(sp)
                jobs.append(dict(b=b, span=sub, text=b.raw[sub[0]:sub[1]], par=cp, sent=cs if full else None))
    return jobs


accept = TP.strict_accept


def _checked(jobs, cands, glossary, minimal=False, repaired=()):
    """Indices whose candidate passes the strict acceptance, the meaning check and the fluency check. A full
    proofread is compared with the author's original sentence (when known) and with the current sentence; a
    minimal correction with the current sentence only. A REPAIR of a garbled sentence (`repaired`) is judged
    against the author's original: it must say what the original says, and only loosely resemble the garbled
    current wording (that wording is what was wrong)."""
    def cur_of(i):
        return jobs[i].get("base", jobs[i]["text"])
    idx = [i for i, c in enumerate(cands) if c and c.strip() != cur_of(i).strip()]
    ok = {}
    for i in idx:
        a = accept(cur_of(i), cands[i], glossary, jobs[i]["sent"])
        sent = jobs[i]["sent"]
        if not a and not minimal and i in repaired and sent:
            # a repair is measured against the author's sentence (back to the original wording is allowed)
            a = TP.basic_grammar(cands[i].strip()) if cands[i].strip() == sent.strip() else \
                accept(sent, cands[i], glossary)
        if a:
            ok[i] = a
    keys = list(ok)
    if not keys:
        return ok
    if minimal:
        mean = MEANING.keeps_meaning([(cur_of(i), ok[i]) for i in keys], min_sim=MEANING.min_sim())
    else:
        soft = MEANING.soft_sim()
        mean = []
        pairs_o = [(jobs[i]["sent"] or jobs[i]["text"], ok[i]) for i in keys]
        pairs_c = [(cur_of(i), ok[i]) for i in keys]
        m_o = MEANING.keeps_meaning(pairs_o, min_sim=soft)
        sim_c = MEANING.similarities(pairs_c)
        par_keys = [i for i in keys if i in repaired and not jobs[i]["sent"] and jobs[i]["par"]]
        par_sim = {}
        if par_keys:
            ps = MEANING.similarities([(jobs[i]["par"], ok[i]) for i in par_keys]
                                      + [(jobs[i]["par"], cur_of(i)) for i in par_keys])
            par_sim = {i: (ps[n], ps[len(par_keys) + n]) for n, i in enumerate(par_keys)}
        for i, mo, sc_ in zip(keys, m_o, sim_c):
            same = ok[i].strip() == cur_of(i).strip()
            if i not in repaired:                         # light edit: close to the original AND the current text
                mean.append(mo and (same or sc_ >= soft))
            elif jobs[i]["sent"]:                         # repair, original sentence known
                mean.append(mo and (same or sc_ >= REPAIR_SIM))
            elif i in par_sim:                            # repair, only the original paragraph known
                new_s, old_s = par_sim[i]
                mean.append(new_s >= old_s - 0.02 and sc_ >= REPAIR_SIM)
            else:                                         # repair without any original (standalone proofread)
                mean.append(same or sc_ >= REPAIR_SIM_NOREF)
    flu = FL.acceptable([(jobs[i]["text"], ok[i]) for i in keys])
    for i, m, f in zip(keys, mean, flu):
        if not (m and f):
            ok.pop(i)
    return ok


def _ai_ok(cur_scores, new_texts, defect=None):
    """Per candidate: does it keep the AI probability in bounds? Returns (flags, new scores). A defective sentence
    (doubled terms, chatbot leftovers, broken English) may rise up to DEFECT_RISE, and may enter the flagged range
    only within that allowance."""
    sc = H._detector_scores(new_texts) or [0.0] * len(new_texts)
    defect = defect or [False] * len(new_texts)
    flags = []
    for c, n, d in zip(cur_scores, sc, defect):
        if d:
            flags.append(n <= c + DEFECT_RISE)
        else:
            flags.append(n <= c + MAX_RISE and not (n >= FLAG > c))
    return flags, sc


def proofread_docx(path, orig_path=None, rep=None, hi=0.80, mod=0.50, protect_ai=False):
    """protect_ai: the automatic proofread right after humanizing. Its job is to correct the humanized text, not to
    undo it, so after the style edits, grammar fixes and repairs are undone too (those that raised the AI score most
    first) until the score is back within PROOF_MAX_DOC_RISE. A repair puts back the author's original wording,
    which is the wording the detector flagged: one repair alone took a paragraph from 0% back to 100%."""
    return _proofread(path, orig_path, rep, hi, mod, protect_ai)


def _proofread(path, orig_path=None, rep=None, hi=0.80, mod=0.50, protect_ai=False):
    rep0 = rep or (lambda *a, **k: None)
    _last = [0.05]

    def rep(p, desc=""):
        if p is not None:
            _last[0] = p
        rep0(_last[0], desc=desc)
    llm = _LLM(rep)
    acct = G.start_account()                           # tokens and cost of this proofread (logged at the end)
    if not llm.order:
        raise G.Unavailable("No AI key is set for the proofread: add GROQ_API_KEY (or OPENAI_API_KEY).")
    rep(0.05, "Opening the humanized document")
    doc = DX.open_clean(path)
    before = D.summarise(S._score_doc(doc), hi, mod)
    jobs = _jobs(doc, orig_path)
    if not jobs:
        raise ValueError("No prose sentences found to proofread.")
    # 0. local clean-up first (no model, no AI risk): chatbot leftovers, markdown, doubled words / phrases,
    #    AI-typical filler words and spaced dashes
    texts = [j["text"] for j in jobs]
    local = {}
    drops = {i for i, j in enumerate(jobs) if j.get("drop")}
    for i, t in enumerate(texts):
        if i in drops:
            continue
        c = TP.basic_grammar(TP.de_ai(t))
        if c.strip() != t.strip():
            g = accept(t, c, set())
            if g:
                local[i] = g
    src = list(dict.fromkeys(j["par"] or j["sent"] or j["text"] for j in jobs))
    glossary = TP.build_glossary(src) | TP.GLOSSARY.get()
    TP.set_document([b.raw for b in D.docx_blocks(doc) if b.kind == "para" and b.raw])
    rep(0.12, f"Scoring {len(texts)} sentences")
    raw_sc = H._detector_scores(texts) or [0.0] * len(texts)
    if local:                                          # keep a local clean-up unless it raises the AI score
        keys = list(local)
        sc = H._detector_scores([local[i] for i in keys]) or [0.0] * len(keys)
        for i, s_ in zip(keys, sc):
            if s_ > raw_sc[i] + MAX_RISE and not TP.has_defect(texts[i]):
                local.pop(i)
            else:
                raw_sc[i] = s_
    base = [local.get(i, t) for i, t in enumerate(texts)]   # what Groq proofreads
    for i, j in enumerate(jobs):
        j["base"] = base[i]
    cur = raw_sc
    defect = [TP.has_defect(t) for t in texts]
    for i in FL.bad(base):
        defect[i] = True

    work = [i for i in range(len(jobs)) if i not in drops]
    ctx = [j["par"] or j["sent"] for j in jobs]
    ctx = [c if c and c.strip() != t.strip() else None for c, t in zip(ctx, texts)]
    groups = [id(j["b"]) for j in jobs]
    idx_all = list(range(len(jobs)))
    minimal, notes, single, unchecked = {}, [], False, 0
    limited = False                                    # True: no language model could be reached (basic proofread)

    # ONE request for the whole document (PROOF_MODE=single): grammar, clarity, wording, flow and consistency
    if MODE == "single":
        rep(0.2, "Final check: the whole document is read and corrected in one pass")
        constructs = sorted({re.sub(r"\s+", " ", m.group(0)).strip() for m in re.finditer(
            r"(?:[A-Za-z][\w\-]*\s+){1,4}\([A-Z][A-Za-z0-9\-]{1,10}\)",
            " ".join(b.raw for b in D.docx_blocks(doc) if b.kind == "para" and b.raw))})[:40]
        tok_r = G.REASONING.set(REASONING)
        for prov in llm.order:                         # one call; the other provider only if the first fails
            tok_p = G.PROVIDER.set(prov)
            try:
                res, notes = G.proofread_document([base[i] for i in idx_all], groups, ctx, glossary, constructs)
                llm.used.add(prov)
                single = True
                break
            except G.Unavailable as err:
                print(f"[proofread] one-call check not possible on {prov}: {err}", flush=True)
            except Exception as err:
                print(f"[proofread] one-call check failed on {prov}: {err}", flush=True)
            finally:
                G.PROVIDER.reset(tok_p)
        G.REASONING.reset(tok_r)
        if single:
            unchecked = sum(1 for r in res if r == (None, None))
            res = [r if r != (None, None) else ("ok", base[i]) for i, r in enumerate(res)]
            fixed_ids = {i for i, (st, _) in enumerate(res) if st in ("fixed", "repaired")}
            required = {i for i in work if defect[i] or i in fixed_ids}
        else:
            rep(0.2, "The document is too large for one request here: checking it in parts")

    if not single:
        # 1. grammar audit: Groq's grammar-only correction tells which sentences contain an error (and gives the
        #    minimal fix for them). Sentences with an error or a defect are REQUIRED fixes: correctness comes first.
        rep(0.15, "Grammar audit of every sentence")
        g = llm.run(G.grammar, [base[i] for i in work], glossary, rep, 0.15, 0.35, size=BATCH)
        unchecked = sum(1 for t in g if t is None)
        if work and unchecked == len(work):
            # No language model can be reached (usage limits on every key and model, no credit on the other
            # provider). The proofread does not fail: what needs no model still runs (rule-based clean-up above,
            # LanguageTool below, the AI-score checks), and the result says plainly that the full check is still
            # to be done. PROOF_BASIC_FALLBACK=0 restores the error.
            if not BASIC_FALLBACK:
                raise G.Unavailable("The proofreading engine has reached its usage limit for now. Please try again "
                                    "in a while; nothing in your document was changed.")
            limited = True
            llm.dead = True
            print("[proofread] no AI provider available: basic proofread only (rules + LanguageTool)", flush=True)
        cands = [None] * len(jobs)
        for i, t in zip(work, g):
            cands[i] = t
        minimal = _checked(jobs, cands, glossary, minimal=True)
        required = {i for i in work if defect[i] or i in minimal}

        # 2. full proofread (clarity + grammar + academic wording + flow)
        tok_r = G.REASONING.set(REASONING)                 # deeper reading for the full proofread
        try:
            res = [None] * len(idx_all) if limited else llm.run(
                lambda ids, *a: G.polish([base[i] for i in ids], [groups[i] for i in ids], glossary,
                                         rep, 0.35, 0.65, originals=[ctx[i] for i in ids], size=BATCH),
                idx_all, missing=lambda r: r is None or r[1] is None)
        finally:
            G.REASONING.reset(tok_r)
        res = [r if r is not None else (None, None) for r in res]
    asked_repair = {i for i, (st, _) in enumerate(res) if st == "repaired"}
    full = _checked(jobs, [t for _, t in res], glossary, repaired=asked_repair)
    for i in drops:
        full.pop(i, None)
    # an accepted repair of an unclear sentence is a required fix: clarity comes before the AI score
    repaired = {i for i in full if i in asked_repair}
    required |= repaired
    final, rise = {i: t for i, t in local.items()}, {}
    for i in drops:                                    # chatbot leftovers on their own: removed
        final[i] = ""
    # every correct version of a sentence (full proofread, minimal grammar fix) is scored. A repaired sentence takes
    # its repair; a sentence with an error takes the full correction if its AI rise stays within PROOF_REQ_RISE,
    # otherwise the minimal grammar fix; a correct sentence only takes a style edit that keeps the AI score
    opts = {}
    for i in work:
        for t in (full.get(i), minimal.get(i)):
            if t and t not in opts.setdefault(i, []):
                opts[i].append(t)
    flat = [(i, t) for i, ts in opts.items() for t in ts]
    sc = H._detector_scores([t for _, t in flat]) or [0.0] * len(flat)
    score_of = {(i, t): s_ for (i, t), s_ in zip(flat, sc)}
    n_minimal = 0
    for i, ts in opts.items():
        if i in repaired:
            final[i], rise[i] = full[i], score_of[(i, full[i])] - cur[i]
        elif i in required:
            f_, m_ = full.get(i), minimal.get(i)
            if f_ and (not m_ or score_of[(i, f_)] - cur[i] <= REQ_RISE):
                best = f_
            elif m_:
                best = m_
            else:
                best = min(ts, key=lambda t: (score_of[(i, t)], len(t)))
            final[i], rise[i] = best, score_of[(i, best)] - cur[i]
            n_minimal += best == m_ and best != f_
        elif full.get(i):
            s_ = score_of[(i, full[i])]
            if s_ - cur[i] <= KEEP_RISE:              # style-only edit: only if the AI score stays the same
                final[i], rise[i] = full[i], s_ - cur[i]

    # 3. verification rounds: every corrected or required sentence is checked again until Groq finds no error
    n_verify = 0
    for rnd in range(0 if single else ROUNDS):
        # a sentence whose text IS the grammar check's own output is not sent again (it was just checked)
        chk = [i for i in work if (i in final or i in required) and final.get(i, base[i]) != minimal.get(i)]
        if not chk:
            break
        rep(0.65 + 0.08 * rnd, f"Verification round {rnd + 1}: re-checking {len(chk)} sentence(s)")
        if llm.dead:
            break
        g = llm.run(G.grammar, [final.get(i, base[i]) for i in chk], glossary, rep, 0.65 + 0.08 * rnd,
                    0.72 + 0.08 * rnd, size=BATCH)
        changed = 0
        for i, t in zip(chk, g):
            now = final.get(i, base[i])
            if not t or t.strip() == now.strip():
                continue
            ok = accept(now, t, glossary, jobs[i]["sent"])
            if ok and H._small_edit(now, ok) and MEANING.keeps_meaning([(now, ok)])[0]:
                final[i] = ok
                required.add(i)                        # it had an error: keep this fix
                changed += 1
        n_verify += changed
        if not changed:
            break

    # 4. LanguageTool: last mechanical check (agreement, confused words, a/an, punctuation)
    try:
        import lt_grammar as LT
        if LT.MODE != "off":
            ids = [i for i in work if final.get(i, base[i])]
            for i, t in zip(ids, LT.correct([final.get(i, base[i]) for i in ids])):
                now = final.get(i, base[i])
                if t.strip() != now.strip():
                    ok = accept(now, t, glossary)
                    if ok and H._small_edit(now, ok):
                        final[i] = ok
                        required.add(i)
    except Exception as err:
        print(f"[proofread] LanguageTool skipped: {err}", flush=True)

    # 5. keep the AI score: correct, natural-sounding alternatives (AI) and Parrot rewrites of corrected sentences;
    #    only versions verified error-free (grammar check returns them unchanged, LanguageTool finds nothing,
    #    fluent, same meaning) are kept as options.
    verified = {}                                      # sentence -> set of verified correct versions

    def gen_alts(ids, stage):
        if not ids or llm.dead:
            return
        rep(stage, f"Keeping the AI score: natural rewording of {len(ids)} sentence(s)")
        alts = llm.run(G.variants, [final.get(i, base[i]) for i in ids], glossary,
                       [jobs[i]["sent"] or None for i in ids], VARIANTS)
        pool = []
        for i, vs in zip(ids, alts):
            for v in (vs or []):
                ok = accept(final.get(i, base[i]), v, glossary, jobs[i]["sent"])
                if ok:
                    pool.append((i, ok))
        if PARROT:
            try:
                pr, _ = H._paraphrase_local([final.get(i, base[i]) for i in ids], 0.9, None, 0.8, 0.85)
                for i, t in zip(ids, pr):
                    ok = accept(final.get(i, base[i]), t, glossary, jobs[i]["sent"])
                    if ok:
                        pool.append((i, ok))
            except Exception as err:
                print(f"[proofread] Parrot rewording skipped: {err}", flush=True)
        pool = list(dict.fromkeys(pool))
        if not pool:
            return
        mean = MEANING.keeps_meaning([(final.get(i, base[i]), t) for i, t in pool])
        flu = FL.acceptable([(final.get(i, base[i]), t) for i, t in pool])
        pool = [(i, t) for (i, t), m, f in zip(pool, mean, flu) if m and f]
        if not pool:
            return
        chk = llm.run(G.grammar, [t for _, t in pool], glossary, rep, stage, stage + 0.02, size=BATCH)
        clean = [(i, t) for (i, t), g_ in zip(pool, chk) if g_ is not None and g_.strip() == t.strip()]
        try:
            import lt_grammar as LT
            if LT.MODE != "off" and clean:
                ltc = LT.correct([t for _, t in clean])
                clean = [c for c, l_ in zip(clean, ltc) if l_.strip() == c[1].strip()]
        except Exception:
            pass
        for i, t in clean:
            verified.setdefault(i, set()).add(t)

    n_reworded = 0
    for rnd in range(0 if single else AI_ROUNDS):
        ids = [i for i in work if final.get(i)]
        if not ids:
            break
        now_sc = H._detector_scores([final[i] for i in ids]) or [0.0] * len(ids)
        hot = [i for i, s_ in zip(ids, now_sc) if s_ > cur[i] + ALT_RISE and i not in repaired]
        if not hot:
            break
        now_of = dict(zip(ids, now_sc))
        gen_alts(hot, 0.78 + 0.02 * rnd)
        cand = [(i, t) for i in hot for t in verified.get(i, ())]
        sc = H._detector_scores([t for _, t in cand]) or [1.0] * len(cand)
        best = {}
        for (i, t), s_ in zip(cand, sc):
            if s_ < now_of[i] - 0.02 and (i not in best or s_ < best[i][1]):
                best[i] = (t, s_)
        for i, (t, s_) in best.items():
            final[i], rise[i] = t, s_ - cur[i]
            required.add(i)
            n_reworded += 1
        if not best:
            break

    # 5b. paragraph balance: the document score is measured over sentence windows inside each paragraph, so
    #     single-sentence scores are not enough. Every paragraph whose AI-flagged words rose is re-balanced by
    #     choosing, sentence by sentence, among the CORRECT versions only (current, minimal fix, verified natural
    #     versions; for a sentence without an error also its unedited form) the combination with the fewest
    #     flagged words.
    by_par = {}
    for i, j in enumerate(jobs):
        by_par.setdefault(id(j["b"]), []).append(i)

    def par_text(pid, choice):
        ids = by_par[pid]
        raw = jobs[ids[0]]["b"].raw
        for i in sorted(ids, key=lambda i: -jobs[i]["span"][0]):
            t = choice.get(i, jobs[i]["text"])
            a_, e_ = jobs[i]["span"]
            if t == "":
                while e_ < len(raw) and raw[e_] == " ":
                    e_ += 1
            raw = raw[:a_] + t + raw[e_:]
        return raw

    def flagged(text):
        try:
            st = D.summarise(S._score_paras([text]), hi, mod)
            return st["high"] + st["moderate"]
        except Exception:
            return 0

    def options(i):
        o = [final.get(i, base[i])]
        if i in repaired:                          # only the repair and its verified natural versions
            return o + [t for t in sorted(verified.get(i, ())) if t not in o]
        for t in [minimal.get(i)] + sorted(verified.get(i, ())):
            if t and t not in o:
                o.append(t)
        if i not in required and jobs[i]["text"] not in o:
            o.append(jobs[i]["text"])              # no error: its unedited form is allowed too
        return o

    n_balanced = 0
    changed_pars = [pid for pid, ids in by_par.items() if any(i in final for i in ids)]
    rep(0.86, f"Balancing the AI score of {len(changed_pars)} paragraph(s)")
    for attempt in range(2):
        worse = []
        for pid in changed_pars:
            was = flagged(par_text(pid, {}))
            if flagged(par_text(pid, final)) > was + PAR_TOL:
                worse.append((pid, was))
        if not worse:
            break
        if attempt == 1 and not single:                # more correct versions for the sentences in those paragraphs
            gen_alts([i for pid, _ in worse for i in by_par[pid] if i in final and final[i]], 0.88)
        for pid, was in worse:
            choice = {i: final[i] for i in by_par[pid] if i in final}
            score = flagged(par_text(pid, choice))
            for _ in range(len(by_par[pid]) * 2):      # greedy: best single swap per step
                best = None
                for i in by_par[pid]:
                    if i in drops:
                        continue
                    for t in options(i):
                        if t == choice.get(i, jobs[i]["text"]):
                            continue
                        trial = dict(choice)
                        if t == jobs[i]["text"]:
                            trial.pop(i, None)
                        else:
                            trial[i] = t
                        f_ = flagged(par_text(pid, trial))
                        if f_ < score and (best is None or f_ < best[0]):
                            best = (f_, trial)
                if best is None:
                    break
                score, choice = best
                n_balanced += 1
                if score <= was + PAR_TOL:
                    break
            for i in by_par[pid]:
                if i in drops:
                    continue
                if i in choice:
                    final[i] = choice[i]
                else:
                    final.pop(i, None)

    ids = [i for i in work if final.get(i)]
    still = sum(1 for i, s_ in zip(ids, H._detector_scores([final[i] for i in ids]) or [0.0] * len(ids))
                if s_ > cur[i] + ALT_RISE and i not in repaired) if ids else 0

    # 6. write back, then document-level AI check (undo style-only edits if the score still rose)
    rep(0.85, "Writing corrections and re-checking the AI score")

    def write(chosen):
        d = DX.open_clean(path)
        js = _jobs(d, None)
        if len(js) != len(jobs):
            return None, 0
        n = 0
        for i in sorted(chosen, key=lambda i: (id(js[i]["b"]), -js[i]["span"][0])):
            a_, e_ = js[i]["span"]
            if chosen[i] == "":                        # deletion: take the following space with it
                raw = js[i]["b"].raw
                while e_ < len(raw) and raw[e_] == " ":
                    e_ += 1
            if S.replace_span(js[i]["b"].el, a_, e_, chosen[i]):
                n += 1
        return d, n

    out_doc, n = write(final)
    if out_doc is None:
        raise ValueError("The document layout changed while proofreading; please run it again.")
    after = D.summarise(S._score_doc(out_doc), hi, mod)
    undone = 0
    while after["pct"] > before["pct"] + MAX_DOC_RISE:
        optional = [i for i in final if i not in required and i not in drops and i not in local]
        if not optional:                               # only required (grammar) fixes left: they always stay
            break
        worst = sorted(optional, key=lambda i: rise.get(i, 0), reverse=True)[:max(1, len(optional) // 5)]
        for i in worst:
            final.pop(i)
            undone += 1
        out_doc, n = write(final)
        after = D.summarise(S._score_doc(out_doc), hi, mod)
    held = 0                                           # protect_ai: corrections left for the author to apply
    while protect_ai and after["pct"] > before["pct"] + MAX_DOC_RISE:
        left = [i for i in final if i not in drops and i not in local]
        if not left:
            break
        worst = sorted(left, key=lambda i: rise.get(i, 0), reverse=True)[:max(1, len(left) // 5)]
        for i in worst:
            final.pop(i)
            held += 1
        out_doc, n = write(final)
        after = D.summarise(S._score_doc(out_doc), hi, mod)
    out = H._save_output(out_doc, S._clean_stem(path), "_proofread")
    who = " + ".join(_NAME.get(p, p) for p in llm.order if p in llm.used) or llm.name
    note = (("⚠️ Basic proofread only — the advanced check has reached its usage limit for now. Rule-based "
             "corrections and the grammar checker were applied" if limited else
             ("🔍 Final check in one request" if single else "🔍 Rigorous proofread") + f" by {who}")
            + f": {n} of {len(jobs)} sentences corrected"
            + (f" ({len(required)} had a grammar error or defect and were always corrected"
               + (f", {len(repaired)} unclear sentence(s) repaired to the original meaning" if repaired else "")
               + (f", {n_minimal} with the minimal grammar fix" if n_minimal else "")
               + (f"; {n_verify} more fix(es) found in the verification rounds" if n_verify else "")
               + (f"; {n_reworded} reworded naturally to keep the AI score" if n_reworded else "")
               + (f"; {n_balanced} swap(s) between correct versions to balance paragraphs" if n_balanced else "")
               + ")"
               if required else "")
            + (f"; {sum(defect)} had doubled terms, chatbot leftovers or broken English, "
               f"{sum(1 for i, d in enumerate(defect) if d and i in final)} of them fixed" if any(defect) else "")
            + f". AI score {before['pct']}% → {after['pct']}%."
            + (f" {undone} style-only edit(s) were undone because they raised the AI score"
               + ("." if held else "; grammar fixes were all kept.") if undone else "")
            + (f" {held} correction(s) were not applied because they would have undone the humanizing (raised the "
               "AI score); press Final check & proofread to apply them anyway." if held else "")
            + (f" {still} corrected sentence(s) still read slightly more AI-like: no error-free natural version "
               "was found, and correct grammar was kept." if still else "")
            + (" The full check (unclear sentences, academic wording, flow, consistency) has NOT been done yet: "
               "run the proofread again later to finish it; it stays free." if limited else
               f" ⚠️ {unchecked} sentence(s) could not be checked (AI quota exhausted); run the proofread again "
               "later to finish them." if unchecked else "")
            + ("" if orig_path else " Tip: run it right after humanizing, so your original is used as the meaning "
                                    "reference for unclear sentences."))
    if notes:                                          # document-level problems that need the author
        note += " 📝 For the author: " + " · ".join(notes)
    print(f"[proofread] {note}", flush=True)
    print(f"[cost] proofread, {sum(len(j['text'].split()) for j in jobs)} words: {G.usage_line(acct)}", flush=True)
    return dict(path=out, before=before, after=after, n=n, total=len(jobs), note=note, notes=notes,
                single=single, limited=limited, held=held)