"""
perplexity.py — token-predictability engine.

The core AI-vs-human signal: LLMs generate text auto-regressively by picking high-probability
next tokens, so machine-written text tends to be *more predictable* — lower
perplexity, higher fraction of tokens the model would itself have ranked at
the top. Human text is "burstier" and picks lower-probability words more often.

This module loads a small causal LM per language (GPT-2 for English,
KoGPT2 for Korean) and computes, for a piece of text:

    - mean perplexity                (exp of mean negative log-likelihood)
    - mean / std of token log-prob   (predictability + its variability)
    - top-k rank hit rate            (fraction of tokens in model's top-k)
    - per-token surprisal            (for span/paragraph highlighting)

Long inputs are processed with a sliding window so we stay within the model's
context length and CPU memory budget.
"""
from __future__ import annotations

import os
import math
from dataclasses import dataclass, field
from functools import lru_cache
from typing import List, Optional

import torch

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

# Only force offline mode when BOTH language-model dirs are present locally
# (e.g. a bundled offline install). Otherwise leave the hub reachable so the
# weights can be downloaded on first run (e.g. on Hugging Face Spaces).
_gpt2_dir = os.environ.get("AIDETECT_GPT2_DIR", "models/gpt2")
_kogpt2_dir = os.environ.get("AIDETECT_KOGPT2_DIR", "models/kogpt2")
if os.path.isdir(_gpt2_dir) and os.path.isdir(_kogpt2_dir):
    os.environ.setdefault("HF_HUB_OFFLINE", "1")
    os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

from transformers import (  # noqa: E402
    AutoModelForCausalLM,
    AutoTokenizer,
    PreTrainedTokenizerFast,
)
from transformers.utils import logging as _hf_logging  # noqa: E402

_hf_logging.set_verbosity_error()  # silence benign load-report / buffer notes

# ----------------------------------------------------------------------------
# Model locations. Point these at local dirs (assembled during setup) or HF ids.
# ----------------------------------------------------------------------------
_DEFAULT_MODEL_DIRS = {
    "en": os.environ.get("AIDETECT_GPT2_DIR", "models/gpt2"),
    "ko": os.environ.get("AIDETECT_KOGPT2_DIR", "models/kogpt2"),
}
# Fallback HF hub ids if the local dir is absent.
_HUB_IDS = {"en": "gpt2", "ko": "skt/kogpt2-base-v2"}

# KoGPT2 ships only tokenizer.json; special tokens must be supplied explicitly.
_KOGPT2_SPECIAL = dict(
    bos_token="</s>", eos_token="</s>", unk_token="<unk>",
    pad_token="<pad>", mask_token="<mask>",
)

torch.set_num_threads(max(1, os.cpu_count() or 1))


@dataclass
class PPLResult:
    """Perplexity / predictability summary for one text."""
    lang: str
    n_tokens: int
    perplexity: float          # exp(mean NLL) — lower = more AI-like
    mean_logprob: float        # mean token log-prob (nats); higher = more predictable
    std_logprob: float         # variability of surprisal ~ "burstiness" proxy
    topk_hit_rate: float       # fraction of tokens in model's top-k prediction
    median_logprob: float
    token_surprisals: List[float] = field(default_factory=list)  # per-token NLL (nats)
    token_strings: List[str] = field(default_factory=list)

    def as_features(self) -> dict:
        return {
            "ppl": self.perplexity,
            "log_ppl": math.log(self.perplexity + 1e-9),
            "mean_logprob": self.mean_logprob,
            "std_logprob": self.std_logprob,
            "median_logprob": self.median_logprob,
            "topk_hit_rate": self.topk_hit_rate,
        }


class _LM:
    """A loaded causal LM + tokenizer for one language."""

    def __init__(self, lang: str):
        self.lang = lang
        path = _DEFAULT_MODEL_DIRS.get(lang)
        src = path if (path and os.path.isdir(path)) else _HUB_IDS[lang]
        self.model = AutoModelForCausalLM.from_pretrained(src)
        self.model.eval()
        if lang == "ko":
            tok_file = os.path.join(src, "tokenizer.json")
            if os.path.isfile(tok_file):
                self.tok = PreTrainedTokenizerFast(
                    tokenizer_file=tok_file, **_KOGPT2_SPECIAL
                )
            else:  # hub fallback
                self.tok = AutoTokenizer.from_pretrained(src, **_KOGPT2_SPECIAL)
        else:
            self.tok = AutoTokenizer.from_pretrained(src)
        self.max_ctx = getattr(self.model.config, "n_positions", 1024)


@lru_cache(maxsize=4)
def get_lm(lang: str) -> _LM:
    """Load (and cache) the LM for a language. 'en' or 'ko'."""
    if lang not in _HUB_IDS:
        lang = "en"
    return _LM(lang)


@torch.no_grad()
def score_text(
    text: str,
    lang: str = "en",
    topk: int = 10,
    stride: Optional[int] = None,
    want_tokens: bool = True,
) -> PPLResult:
    """
    Compute perplexity / predictability statistics for `text`.

    Uses a sliding window (`stride`) so text longer than the model context is
    scored without truncation; overlapping context tokens are masked from the
    loss so each token is scored exactly once with full left context.
    """
    lm = get_lm(lang)
    tok, model = lm.tok, lm.model
    max_len = min(lm.max_ctx, 1024)
    if stride is None:
        stride = max_len // 2

    enc = tok(text, return_tensors="pt")
    input_ids = enc.input_ids
    n_all = input_ids.shape[1]
    if n_all < 2:
        return PPLResult(lang, 0, float("inf"), 0.0, 0.0, 0.0, 0.0, [], [])

    nll_sum = 0.0
    logprobs: List[float] = []
    surprisals: List[float] = []
    topk_hits = 0
    counted = 0
    prev_end = 0

    for begin in range(0, n_all, stride):
        end = min(begin + max_len, n_all)
        ids = input_ids[:, begin:end]
        # number of *new* target tokens in this window (avoid double counting)
        trg_len = end - prev_end
        logits = model(ids).logits  # (1, L, V)
        # predict position t from t-1: align logits[:-1] with labels[1:]
        shift_logits = logits[0, :-1, :]
        shift_labels = ids[0, 1:]
        logp = torch.log_softmax(shift_logits, dim=-1)
        tok_logp = logp[torch.arange(shift_labels.shape[0]), shift_labels]  # (L-1,)
        # only keep the last (trg_len-1) newly-scored tokens in this window
        if trg_len > 1:
            keep_from = shift_labels.shape[0] - (trg_len - 1)
        else:
            keep_from = shift_labels.shape[0]
        keep_from = max(0, keep_from)
        sel_logp = tok_logp[keep_from:]
        sel_labels = shift_labels[keep_from:]
        sel_logits = shift_logits[keep_from:]
        # top-k membership
        if sel_logits.shape[0] > 0:
            topk_ids = sel_logits.topk(topk, dim=-1).indices
            hits = (topk_ids == sel_labels.unsqueeze(1)).any(dim=1)
            topk_hits += int(hits.sum())
            nll = -sel_logp
            nll_sum += float(nll.sum())
            counted += sel_logp.shape[0]
            logprobs.extend(sel_logp.tolist())
            surprisals.extend(nll.tolist())
        prev_end = end
        if end == n_all:
            break

    if counted == 0:
        return PPLResult(lang, 0, float("inf"), 0.0, 0.0, 0.0, 0.0, [], [])

    mean_nll = nll_sum / counted
    ppl = math.exp(min(mean_nll, 50))
    lp_tensor = torch.tensor(logprobs)
    result = PPLResult(
        lang=lang,
        n_tokens=counted,
        perplexity=ppl,
        mean_logprob=float(lp_tensor.mean()),
        std_logprob=float(lp_tensor.std(unbiased=False)) if counted > 1 else 0.0,
        median_logprob=float(lp_tensor.median()),
        topk_hit_rate=topk_hits / counted,
    )
    if want_tokens:
        all_toks = tok.convert_ids_to_tokens(input_ids[0].tolist())
        result.token_surprisals = surprisals
        result.token_strings = all_toks[1:]  # labels are input_ids[1:]
    return result


if __name__ == "__main__":
    import sys
    txt = sys.argv[1] if len(sys.argv) > 1 else \
        "The results demonstrate a statistically significant improvement."
    lang = sys.argv[2] if len(sys.argv) > 2 else "en"
    r = score_text(txt, lang)
    print(f"lang={r.lang} n={r.n_tokens} ppl={r.perplexity:.1f} "
          f"mean_lp={r.mean_logprob:.3f} std_lp={r.std_logprob:.3f} "
          f"topk={r.topk_hit_rate:.3f}")