"""
Self-adaptive learning: the humanizer learns from every run and uses what it learned in the next pass and the next
document. Nothing here can loosen a safety check: every rewrite still has to pass the guard, the clarity gate and
Groq; learning only decides HOW to generate (which settings, which sources, which words to avoid) so that more of
what is generated survives the checks and reads as human.

What it learns
  1. Bad words (cross-document). Every rejected rewrite is examined: the words it brought in that the original
     did not have are counted against the rejection reason (construct renamed, group swapped, invented name,
     attribution, off-topic ...). A word that keeps causing rejections ("tribes", "allegiance", "muscle") and is
     rarely part of an accepted rewrite is banned: Parrot is told never to generate it (bad_words_ids) and the
     guard rejects it at once.
  2. Which source wins (cross-document). Parrot samples, detector-guided (adversarial) versions, plain-wording
     (de-AI) versions and vocabulary-replaced versions compete; how often each source's surviving versions are
     the one chosen is tracked (Beta posterior). Sources that keep winning get more candidates, sources that
     never win are skipped.
  3. How many candidates to sample (cross-document). If few candidates survive the checks, more are sampled.
  4. Temperature (cross-document, Smart Humanize). Each pass picks a temperature offset by Thompson sampling;
     the reward is the AI-score drop of that pass minus a penalty for rejected rewrites, so the setting that
     produces the best accepted output on real documents is used more often.
  5. Per sentence, within a document (Smart Humanize passes). A sentence whose AI score did not fall gets a
     bolder setting next pass (higher temperature, detector-guided decoding); a sentence whose rewrites keep
     being rejected gets a calmer setting, and after ADAPT_MAX_FAILS failures it is left in the author's wording
     (saves time, avoids damage).

Storage: one JSON file (ADAPT_PATH; default /data/cs_adaptive.json on Spaces with persistent storage, else the temp
folder). Statistics decay slowly (ADAPT_DECAY) so the system keeps adapting when models or detectors change.

Space variables
  ADAPT            1 (default) | 0 = off (fixed settings, as before)
  ADAPT_PATH       file for the learned state
  ADAPT_BAN_MIN    rejections before a word can be banned (default 4)
  ADAPT_BAN_RATIO  share of rejections among its uses needed to ban it (default 0.8)
  ADAPT_MAX_FAILS  passes with only rejected rewrites before a sentence is left as written (default 3)
  ADAPT_DECAY      per-update decay of the source / temperature statistics (default 0.995)
"""
import contextvars
import json
import os
import random
import re
import tempfile
import threading
import time

ENABLED = os.getenv("ADAPT", "1") != "0"
BAN_MIN = float(os.getenv("ADAPT_BAN_MIN", "4"))
BAN_RATIO = float(os.getenv("ADAPT_BAN_RATIO", "0.8"))
MAX_FAILS = int(os.getenv("ADAPT_MAX_FAILS", "3"))
DECAY = float(os.getenv("ADAPT_DECAY", "0.995"))
FLAG = float(os.getenv("ADAPT_FLAG", "0.5"))            # a sentence at/above this AI probability still needs work
TEMP_ARMS = [-0.2, 0.0, 0.15, 0.3]
# words that are never right in an academic rewrite of someone else's study (learned list starts from these)
SEED_BANNED = {"tribes", "tribe", "muscle", "muscles", "être"}
# rejection reasons that point at the WORDS a rewrite introduced (not at its length or punctuation)
_WORD_REASONS = ("construct renamed", "group or degree swapped", "name", "attribution", "off-topic", "foreign",
                 "invented", "non-word", "informal", "lost", "abbreviation", "direction reversed", "causal")


def _default_path():
    p = os.getenv("ADAPT_PATH", "").strip()
    if p:
        return p
    if os.path.isdir("/data") and os.access("/data", os.W_OK):
        return "/data/cs_adaptive.json"
    return os.path.join(tempfile.gettempdir(), "cs_adaptive.json")


PATH = _default_path()
_lock = threading.RLock()
_state = {"words": {}, "sources": {}, "accept": [1.0, 1.0], "temps": {}, "runs": 0, "banned_extra": []}
_dirty = [False, 0.0]


def _load():
    try:
        with open(PATH, encoding="utf-8") as f:
            data = json.load(f)
        if isinstance(data, dict):
            _state.update({k: data[k] for k in _state if k in data})
    except FileNotFoundError:
        pass
    except Exception as err:
        print(f"[adaptive] state not loaded ({err}); starting fresh", flush=True)


_load()


def save(force=False):
    """Write the learned state (at most every 30 s unless forced)."""
    if not ENABLED:
        return
    with _lock:
        if not _dirty[0] or (not force and time.time() - _dirty[1] < 30):
            return
        try:
            tmp = PATH + ".tmp"
            with open(tmp, "w", encoding="utf-8") as f:
                json.dump(_state, f)
            os.replace(tmp, PATH)
            _dirty[0], _dirty[1] = False, time.time()
        except Exception as err:
            print(f"[adaptive] state not saved: {err}", flush=True)


def _touch():
    _dirty[0] = True
    save()


# ------------------------------------------------------------------------------------------- 1. bad words
def _introduced(original, cand):
    o = set(re.findall(r"[a-zà-ÿ]{4,}", (original or "").lower()))
    return {w for w in re.findall(r"[a-zà-ÿ]{4,}", (cand or "").lower()) if w not in o}


def record_candidate(original, cand, source, accepted, reason=None):
    """One generated rewrite and whether it survived the checks."""
    if not ENABLED or not cand or not original or cand.strip() == original.strip():
        return
    with _lock:
        a = _state["accept"]
        a[0 if accepted else 1] += 1
        new = _introduced(original, cand)
        words = _state["words"]
        word_reason = (not accepted) and reason and any(r in reason for r in _WORD_REASONS)
        for w in new:
            st = words.setdefault(w, [0.0, 0.0])          # [times in a rejected rewrite (word reason), accepted]
            if accepted:
                st[1] += 1
            elif word_reason:
                st[0] += 1
        if len(words) > 50000:                            # forget rare words
            for w in [w for w, (b, g) in words.items() if b + g <= 1][:20000]:
                words.pop(w, None)
        s = _state["sources"].setdefault(source, [1.0, 1.0, 0.0])   # [wins+1, losses+1, survivors]
        s[2] += 1 if accepted else 0
    _touch()


def banned():
    """Words never to generate: the seed list + words that keep causing rejections."""
    if not ENABLED:
        return set()
    with _lock:
        out = set(SEED_BANNED) | set(_state.get("banned_extra") or [])
        for w, (bad, good) in _state["words"].items():
            if bad >= BAN_MIN and bad / max(1.0, bad + good) >= BAN_RATIO and not _common(w):
                out.add(w)
        return out


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


def _common(w):
    """Very common words ("shows", "between") are never banned automatically: they only travel along with the
    word that really caused a rejection."""
    return _zipf is not None and _zipf(w, "en") >= 5.0


def ban_word(word):
    """Admin: ban a word by hand (e.g. after reviewing a client's output)."""
    with _lock:
        lst = set(_state.get("banned_extra") or [])
        lst.add(word.strip().lower())
        _state["banned_extra"] = sorted(lst)
    _touch()
    save(force=True)


def banned_in(original, cand):
    """A banned word the rewrite introduces, or None."""
    b = banned()
    if not b:
        return None
    hit = _introduced(original, cand) & b
    return sorted(hit)[0] if hit else None


_bw_cache = {}


def bad_words_ids(tokenizer):
    """Token sequences of the banned words, for model.generate(bad_words_ids=...)."""
    b = banned()
    key = (id(tokenizer), frozenset(b))
    if key in _bw_cache:
        return _bw_cache[key]
    ids = []
    for w in sorted(b):
        for form in (w, " " + w, w.capitalize(), " " + w.capitalize()):
            try:
                t = tokenizer(form, add_special_tokens=False).input_ids
            except Exception:
                continue
            if t and t not in ids and len(t) <= 4:
                ids.append(t)
    _bw_cache.clear()
    _bw_cache[key] = ids or None
    return _bw_cache[key]


def suppress_ids(tokenizer):
    """Banned words that are ONE token, for model.generate(suppress_tokens=...): blocked in a single vectorised step.
    (bad_words_ids with the learned list - 5,500 sequences on the Space - made every Parrot call take 20-30 s on the
    GPU: it is checked in Python at every generated token. Banned words of several tokens are not blocked while
    generating; banned_in() still rejects any version that contains one.)"""
    return sorted({t[0] for t in (bad_words_ids(tokenizer) or ()) if len(t) == 1}) or None


# ------------------------------------------------------------------------------------------- 2-3. sources, k
def record_win(source, won):
    """Among the versions that survived the checks for one sentence: did this source's version get chosen?"""
    if not ENABLED:
        return
    with _lock:
        s = _state["sources"].setdefault(source, [1.0, 1.0, 0.0])
        s[0], s[1] = s[0] * DECAY, s[1] * DECAY            # slow forgetting
        s[0 if won else 1] += 1
    _touch()


def source_weight(source):
    """Thompson sample of how likely this source's version is to be the chosen one (0..1)."""
    if not ENABLED:
        return 1.0
    with _lock:
        a, b, _ = _state["sources"].get(source, [1.0, 1.0, 0.0])
    return random.betavariate(max(a, 0.1), max(b, 0.1))


def use_source(source, min_share=0.08):
    """Skip a source only when it has clearly and repeatedly lost (keeps exploring otherwise)."""
    if not ENABLED:
        return True
    with _lock:
        a, b, _ = _state["sources"].get(source, [1.0, 1.0, 0.0])
    if a + b < 30:
        return True
    return a / (a + b) >= min_share or random.random() < 0.1


def candidates(k, cap=8):
    """More samples when few survive the checks."""
    if not ENABLED:
        return k
    with _lock:
        ok, bad = _state["accept"]
    rate = ok / max(1.0, ok + bad)
    if ok + bad < 50:
        return k
    if rate < 0.25:
        return min(cap, k + 2)
    if rate < 0.4:
        return min(cap, k + 1)
    return k


# ------------------------------------------------------------------------------------------- 4. temperature
def pick_temp_offset():
    """Thompson sampling over TEMP_ARMS (Gaussian posterior of the reward)."""
    if not ENABLED:
        return 0.0
    with _lock:
        stats = dict(_state["temps"])
    best, best_v = 0.0, None
    for arm in TEMP_ARMS:
        n, mean = stats.get(str(arm), [0.0, 0.0])
        v = random.gauss(mean, 1.0 / (1.0 + n) ** 0.5) if n else random.gauss(0.0, 1.0)
        if best_v is None or v > best_v:
            best, best_v = arm, v
    return best


def record_pass(offset, ai_before, ai_after, n_jobs, n_changed):
    """Reward of one Smart Humanize pass: AI-% drop, minus a penalty when most rewrites were rejected."""
    if not ENABLED or n_jobs <= 0:
        return
    drop = (ai_before - ai_after) / 10.0                 # 10 points = reward 1
    kept = n_changed / max(1, n_jobs)
    reward = drop - 0.5 * (1.0 - kept)
    with _lock:
        n, mean = _state["temps"].get(str(offset), [0.0, 0.0])
        n = n * DECAY
        _state["temps"][str(offset)] = [n + 1, mean + (reward - mean) / (n + 1)]
        _state["runs"] += 1
    _touch()


# ------------------------------------------------------------------------------------------- 5. per sentence
SENT = contextvars.ContextVar("adapt_sentences", default=None)   # job: anchor -> dict(score, fails, temp, bold)


def start_job():
    SENT.set({})


def observe(anchor, score):
    """Pass k scored this sentence; compare with pass k-1 and adjust its settings for pass k+1."""
    d = SENT.get()
    if d is None or not ENABLED:
        return
    st = d.setdefault(anchor, {"score": None, "fails": 0, "dt": 0.0, "bold": False, "tried": 0, "changed": False})
    if st["score"] is not None and st.get("changed") and score >= FLAG:
        if st["score"] - score < 0.05:                   # rewritten but still reads as AI: be bolder next time
            st["dt"] = min(0.5, st["dt"] + 0.15)
            st["bold"] = True
    st["score"] = score


def outcome(anchor, changed):
    """This pass's rewrite of the sentence was written into the document (changed) or every version failed."""
    d = SENT.get()
    if d is None or not ENABLED:
        return
    st = d.setdefault(anchor, {"score": None, "fails": 0, "dt": 0.0, "bold": False, "tried": 0, "changed": False})
    st["tried"] += 1
    st["changed"] = bool(changed)
    if changed:
        st["fails"] = 0
    else:
        st["fails"] += 1
        st["dt"] = max(-0.3, st["dt"] - 0.15)            # its rewrites keep failing the checks: calmer


def settings(anchor):
    """(temperature delta, detector-guided?, skip?) for this sentence in the next pass."""
    d = SENT.get()
    if d is None or not ENABLED:
        return 0.0, False, False
    st = d.get(anchor)
    if not st:
        return 0.0, False, False
    return st["dt"], st["bold"], st["fails"] >= MAX_FAILS


def status():
    if not ENABLED:
        return "Adaptive learning: off (ADAPT=0)."
    with _lock:
        ok, bad = _state["accept"]
        src = ", ".join(f"{k} {a / (a + b):.0%}" for k, (a, b, _) in sorted(_state["sources"].items()))
        temps = ", ".join(f"{k}: {m:+.2f} ({n:.0f})" for k, (n, m) in sorted(_state["temps"].items()))
        nb = len(banned())
    return (f"Adaptive learning: {_state['runs']} passes learned from; {ok + bad - 2:.0f} rewrites seen, "
            f"{ok / max(1, ok + bad):.0%} pass the checks; {nb} banned words; source win rates: {src or '-'}; "
            f"temperature rewards: {temps or '-'}. State file: {PATH}")