"""
Fluency check: is a sentence well-formed English?

The paraphraser sometimes produces sentences with awkward wording, missing words, doubled words or broken structure
("Their hybrid purposed to do a 25-item questionnaire with follow-up interviews together qualitative interviews
combined."). A small grammatical-acceptability model (RoBERTa trained on CoLA, the Corpus of Linguistic
Acceptability) scores every candidate rewrite, and candidates that read clearly worse than the author's original
sentence are rejected BEFORE they can be chosen, so damaged sentences never reach the document.

Loaded lazily; if the model cannot be loaded the check is skipped (every sentence passes) and a message is printed.

Space variables:
  FLUENCY_MODEL   default textattack/roberta-base-CoLA
  FLUENCY_LABEL   index of the "acceptable" label (default 1)
  FLUENCY_MIN     default 0.5   minimum acceptability of a rewrite
  FLUENCY_DROP    default 0.25  a rewrite may score at most this much below the original sentence
  FLUENCY_CHECK   default 1     (0 = off)
"""
import os
import re
import threading

import runtime as RT

MODEL_ID = os.environ.get("FLUENCY_MODEL", "textattack/roberta-base-CoLA")
LABEL = int(os.environ.get("FLUENCY_LABEL", "1"))
MIN = float(os.environ.get("FLUENCY_MIN", "0.5"))
DROP = float(os.environ.get("FLUENCY_DROP", "0.25"))
ENABLED = os.environ.get("FLUENCY_CHECK", "1") != "0"
BATCH = 32

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

_DOUBLED = re.compile(r"\b(\w+)\s+\1\b", re.I)
_ALLOWED_DOUBLE = {"that", "had", "is", "very", "so"}          # "that that", "had had" can be correct


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 AutoModelForSequenceClassification, AutoTokenizer
            _state["tok"] = AutoTokenizer.from_pretrained(MODEL_ID)
            _state["model"] = RT.to_gpu(AutoModelForSequenceClassification.from_pretrained(MODEL_ID).eval())
            print(f"[fluency] loaded {MODEL_ID}", flush=True)
        except Exception as err:
            _state["failed"] = True
            print(f"[fluency] check disabled, model not available: {err}", flush=True)
    return _state["model"] is not None


def _mask(s):
    """Citations, numbers and symbols confuse an English-grammar model: replace them by neutral tokens."""
    s = re.sub(r"\[[\d,\s–\-]+\]", "", s)
    s = re.sub(r"\(([^()]*\d{4}[a-z]?[^()]*)\)", "", s)                 # (Kumar et al., 2021)
    s = re.sub(r"\d+(?:[.,]\d+)*\s*%?", "5", s)
    s = re.sub(r"\s+([,.;:!?])", r"\1", s)
    return re.sub(r"\s+", " ", s).strip()


def doubled_word(s):
    m = _DOUBLED.search(s or "")
    return bool(m) and m.group(1).lower() not in _ALLOWED_DOUBLE and not m.group(1).isdigit()


def scores(texts):
    """Acceptability (0..1) of each text; 1.0 for every text when the model is off."""
    if not texts:
        return []
    if not _load():
        return [1.0] * len(texts)
    import torch
    todo = [t for t in dict.fromkeys(texts) if t not in _cache]
    tok, model = _state["tok"], _state["model"]
    try:
        for i in range(0, len(todo), BATCH):
            chunk = todo[i:i + BATCH]
            enc = tok([_mask(t) for t in chunk], padding=True, truncation=True, max_length=128,
                      return_tensors="pt").to(RT.device_of(model))
            with torch.no_grad():
                p = torch.softmax(model(**enc).logits, dim=-1)[:, LABEL].tolist()
            _cache.update(zip(chunk, p))
        if len(_cache) > 50000:
            _cache.clear()
    except Exception as err:
        print(f"[fluency] check skipped for this batch: {err}", flush=True)
        return [1.0] * len(texts)
    return [_cache.get(t, 1.0) for t in texts]


def acceptable(pairs):
    """[(original, rewrite), ...] -> list of bools: the rewrite is well-formed and not clearly worse than the
    original (a rewrite of an already awkward original only has to be no worse)."""
    if not pairs:
        return []
    flags = [not (doubled_word(n) and not doubled_word(o)) for o, n in pairs]
    sc = scores([t for p in pairs for t in p])
    out = []
    for k, (f, (o, n)) in enumerate(zip(flags, pairs)):
        so, sn = sc[2 * k], sc[2 * k + 1]
        out.append(f and sn >= min(MIN, so - 0.05) and sn >= so - DROP)
    return out


def bad(texts, limit=None):
    """Indices of texts that read as broken English (acceptability below `limit`, default FLUENCY_MIN)."""
    lim = MIN if limit is None else limit
    sc = scores(texts)
    return [i for i, (t, s) in enumerate(zip(texts, sc)) if s < lim or doubled_word(t)]