"""
Fiction-oriented single-string text metrics → feature vector

Libraries: textdescriptives (+ pos_proportions), lexicalrichness,
           textstat, NRCLex

Performance (long ~50k-char answers):
  - spaCy: tagger + sentencizer only (no parser/NER/lemmatizer)
  - quality n-gram metrics: pure-Python (spaCy quality pipe was ~1s+/doc)
  - extract_vectors: multi-thread full extracts with thread-local spaCy.
    Worker count = f(os.cpu_count()), capped at n_texts (e2 scale up/down
    is automatic). Override with METRIC_WORKERS. Process pools not used
    (model reload cost dominates on Windows; little gain for 5-sample batches).
"""

from __future__ import annotations

import os
import re
import threading
import warnings
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from typing import Iterable

warnings.filterwarnings("ignore")

import numpy as np
import spacy
import textdescriptives as td
import textstat
from lexicalrichness import LexicalRichness
from nrclex import NRCLex

# ---------------------------------------------------------------------------
# Thread-local spaCy / NRC (Language is not thread-safe; each worker gets one)
# ---------------------------------------------------------------------------
_thread_local = threading.local()
# Persistent pool so workers keep their loaded models across trials.
_METRIC_POOL: ThreadPoolExecutor | None = None
_METRIC_POOL_SIZE: int = 0

# Minimal English stoplist for quality n_stop_words (avoid loading spaCy stops).
_STOP = frozenset(
    """
    a an the and or but if in on at to for of as is was are were be been being
    it its this that these those i you he she we they me him her us them my your
    his their our not no nor so than then too very can could will would should
    may might must do does did have has had with from by into about over after
    before under again further once here there when where why how all each few
    more most other some such only own same both through during above below
    """.split()
)

_WORD_RE = re.compile(r"[A-Za-z0-9']+")
_ELLIPSIS_RE = re.compile(r"\.\.\.|…")


def _build_nlp():
    """Tagger + sentencizer only — no parser/NER/lemmatizer/quality pipe."""
    nlp = spacy.load(
        "en_core_web_sm",
        disable=["parser", "ner", "lemmatizer"],
    )
    if "senter" not in nlp.pipe_names and "sentencizer" not in nlp.pipe_names:
        nlp.add_pipe("sentencizer")
    for pipe in [
        "textdescriptives/descriptive_stats",
        "textdescriptives/readability",
        "textdescriptives/pos_proportions",
        "textdescriptives/information_theory",
    ]:
        if pipe not in nlp.pipe_names:
            nlp.add_pipe(pipe)
    nlp.max_length = max(nlp.max_length, 2_000_000)
    return nlp


def _get_nlp():
    """Per-thread spaCy pipeline (safe concurrent extract_vector)."""
    nlp = getattr(_thread_local, "nlp", None)
    if nlp is None:
        nlp = _build_nlp()
        _thread_local.nlp = nlp
    return nlp


def _get_nrc() -> NRCLex:
    """Per-thread NRCLex (mutates instance state on load_token_list)."""
    nrc = getattr(_thread_local, "nrc", None)
    if nrc is None:
        nrc = NRCLex()
        _thread_local.nrc = nrc
    return nrc


def _total_ram_gb() -> float | None:
    """Best-effort host RAM (GiB). Used to avoid multi-spaCy OOM on small e2."""
    # Linux (typical e2)
    try:
        with open("/proc/meminfo", encoding="utf-8") as f:
            for line in f:
                if line.startswith("MemTotal:"):
                    kb = float(line.split()[1])
                    return kb / (1024.0 * 1024.0)
    except Exception:
        pass
    # Windows
    try:
        import ctypes

        class _MemStatus(ctypes.Structure):
            _fields_ = [
                ("dwLength", ctypes.c_ulong),
                ("dwMemoryLoad", ctypes.c_ulong),
                ("ullTotalPhys", ctypes.c_ulonglong),
                ("ullAvailPhys", ctypes.c_ulonglong),
                ("ullTotalPageFile", ctypes.c_ulonglong),
                ("ullAvailPageFile", ctypes.c_ulonglong),
                ("ullTotalVirtual", ctypes.c_ulonglong),
                ("ullAvailVirtual", ctypes.c_ulonglong),
                ("sullAvailExtendedVirtual", ctypes.c_ulonglong),
            ]

        m = _MemStatus()
        m.dwLength = ctypes.sizeof(_MemStatus)
        if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(m)):
            return float(m.ullTotalPhys) / (1024.0 ** 3)
    except Exception:
        pass
    return None


def _metric_workers(n_texts: int) -> int:
    """
    How many concurrent full extracts.

    Each worker holds its own spaCy pipeline (~0.3–0.8 GiB resident).
    Floor is 2 when n_texts >= 2 (unless METRIC_WORKERS overrides).
    Hard ceiling is always n_texts.

    RAM-aware ceiling:
      RAM < 12 GiB → up to 2
      else         → min(n_texts, max(2, cpu-1))

    Override: METRIC_WORKERS=N.
    """
    n_texts = max(1, int(n_texts))
    if n_texts <= 1:
        return 1

    env = (os.getenv("METRIC_WORKERS") or "").strip()
    if env:
        try:
            return max(1, min(n_texts, int(env)))
        except ValueError:
            pass

    ram = _total_ram_gb()
    cpu = os.cpu_count() or 2
    # Never below 2 when we have at least 2 samples to score.
    floor = min(2, n_texts)
    if ram is not None and ram < 12.0:
        return max(floor, min(n_texts, 2))
    if cpu <= 2:
        want = max(floor, cpu)
    else:
        want = max(floor, cpu - 1)
    return max(floor, min(n_texts, want))


def _get_metric_pool(n_workers: int) -> ThreadPoolExecutor:
    global _METRIC_POOL, _METRIC_POOL_SIZE
    if _METRIC_POOL is not None and _METRIC_POOL_SIZE == n_workers:
        return _METRIC_POOL
    if _METRIC_POOL is not None:
        _METRIC_POOL.shutdown(wait=False, cancel_futures=False)
        _METRIC_POOL = None
    _METRIC_POOL = ThreadPoolExecutor(
        max_workers=n_workers, thread_name_prefix="obj1-metrics"
    )
    _METRIC_POOL_SIZE = n_workers
    return _METRIC_POOL


def _doc_to_td_dict(doc) -> dict:
    out = td.extract_dict(doc)
    if isinstance(out, list):
        if not out:
            return {}
        return out[0] if isinstance(out[0], dict) else {}
    return out if isinstance(out, dict) else {}


def _td(STRING: str) -> dict:
    return _doc_to_td_dict(_get_nlp()(STRING))


def _safe(v, default=0.0) -> float:
    try:
        return float(v) if v is not None else default
    except Exception:
        return default


def _log10_perplexity(v, default=0.0) -> float:
    """log10 of textdescriptives perplexity; raw exp-scale values are unusable."""
    x = _safe(v, default=float("nan"))
    if not np.isfinite(x) or x <= 0.0:
        return default
    y = float(np.log10(x))
    if not np.isfinite(y):
        return default
    return float(np.clip(y, 0.0, 300.0))


# ---------------------------------------------------------------------------
# Fast pure-Python quality metrics (replaces textdescriptives/quality pipe)
# ---------------------------------------------------------------------------

def _chr_fraction_of_dup_spans(text: str, spans: list[str]) -> float:
    """Fraction of characters belonging to items that appear more than once."""
    if not text or not spans:
        return 0.0
    counts = Counter(spans)
    dup_chars = sum(len(s) * c for s, c in counts.items() if c > 1 and s)
    return float(min(1.0, dup_chars / max(len(text), 1)))


def _ngram_chr_len(gram: tuple[str, ...]) -> int:
    """Character length of words joined by single spaces."""
    if not gram:
        return 0
    return sum(len(w) for w in gram) + (len(gram) - 1)


def _word_ngram_chr_fraction(words: list[str], n: int) -> float:
    """
    Character-mass of word n-grams that repeat at least once, over total
    character mass of all word n-grams (textdescriptives-style duplicate_*gram).
    Uses tuple keys (no join) for speed on long docs.
    """
    if len(words) < n:
        return 0.0
    counts: Counter[tuple[str, ...]] = Counter()
    total = 0
    for i in range(len(words) - n + 1):
        g = tuple(words[i : i + n])
        counts[g] += 1
        total += _ngram_chr_len(g)
    if total <= 0:
        return 0.0
    dup = sum(_ngram_chr_len(g) * c for g, c in counts.items() if c > 1)
    return float(min(1.0, dup / total))


def _top_ngram_chr_fraction(words: list[str], n: int) -> float:
    """Character mass of the single most frequent word n-gram / all n-gram mass."""
    if len(words) < n:
        return 0.0
    counts: Counter[tuple[str, ...]] = Counter()
    total = 0
    for i in range(len(words) - n + 1):
        g = tuple(words[i : i + n])
        counts[g] += 1
        total += _ngram_chr_len(g)
    if total <= 0 or not counts:
        return 0.0
    top_g, top_c = counts.most_common(1)[0]
    return float(min(1.0, (_ngram_chr_len(top_g) * top_c) / total))


def _quality_vals(STRING: str) -> list[float]:
    """17 quality features — same names/order as former textdescriptives quality."""
    text = STRING or ""
    n_chars = max(len(text), 1)
    words = _WORD_RE.findall(text)
    n_words = max(len(words), 1)
    lower_words = [w.lower() for w in words]

    alpha_chars = sum(1 for c in text if c.isalpha())
    alpha_ratio = alpha_chars / n_chars
    mean_word_length = float(np.mean([len(w) for w in words])) if words else 0.0
    n_stop_words = float(sum(1 for w in lower_words if w in _STOP))
    proportion_ellipsis = float(len(_ELLIPSIS_RE.findall(text))) / n_words
    # lines starting with bullet-like markers
    lines = text.splitlines()
    if lines:
        bullets = sum(
            1
            for ln in lines
            if ln.lstrip().startswith(("-", "*", "•", "·"))
            or (ln.lstrip()[:2].isdigit() and ln.lstrip()[1:2] in ".)")
        )
        proportion_bullet_points = bullets / len(lines)
    else:
        proportion_bullet_points = 0.0

    # non-empty lines / paragraphs
    nonempty_lines = [ln for ln in lines if ln.strip()]
    paragraphs = [p for p in re.split(r"\n\s*\n", text) if p.strip()]
    duplicate_line_chr_fraction = _chr_fraction_of_dup_spans(text, nonempty_lines)
    duplicate_paragraph_chr_fraction = _chr_fraction_of_dup_spans(text, paragraphs)

    dup_ngrams = [_word_ngram_chr_fraction(lower_words, n) for n in range(5, 11)]
    top_ngrams = [_top_ngram_chr_fraction(lower_words, n) for n in range(2, 5)]

    # symbol '#' to word ratio (textdescriptives symbol_to_word_ratio_#)
    symbol_hash_to_word_ratio = text.count("#") / n_words

    return [
        _safe(alpha_ratio),
        _safe(mean_word_length),
        _safe(n_stop_words),
        _safe(proportion_ellipsis),
        _safe(proportion_bullet_points),
        _safe(duplicate_line_chr_fraction),
        _safe(duplicate_paragraph_chr_fraction),
        *[_safe(x) for x in dup_ngrams],
        *[_safe(x) for x in top_ngrams],
        _safe(symbol_hash_to_word_ratio),
    ]


# ---------------------------------------------------------------------------
# Fixed feature names (order must match the vector)
# ---------------------------------------------------------------------------
FEATURE_NAMES: list[str] = [
    # TextDescriptives – descriptive
    "token_length_mean",
    "token_length_median",
    "token_length_std",
    "sentence_length_mean",
    "sentence_length_median",
    "sentence_length_std",
    "syllables_per_token_mean",
    "syllables_per_token_median",
    "syllables_per_token_std",
    "proportion_unique_tokens",
    "n_tokens",
    "n_unique_tokens",
    "n_characters",
    "n_sentences",
    # TextDescriptives – readability
    "flesch_reading_ease",
    "flesch_kincaid_grade",
    "smog",
    "gunning_fog",
    "automated_readability_index",
    "coleman_liau_index",
    "lix",
    "rix",
    # Quality (fast pure-Python; same names as former textdescriptives quality)
    "alpha_ratio",
    "mean_word_length",
    "n_stop_words",
    "proportion_ellipsis",
    "proportion_bullet_points",
    "duplicate_line_chr_fraction",
    "duplicate_paragraph_chr_fraction",
    "duplicate_5gram_chr_fraction",
    "duplicate_6gram_chr_fraction",
    "duplicate_7gram_chr_fraction",
    "duplicate_8gram_chr_fraction",
    "duplicate_9gram_chr_fraction",
    "duplicate_10gram_chr_fraction",
    "top_2gram_chr_fraction",
    "top_3gram_chr_fraction",
    "top_4gram_chr_fraction",
    "symbol_hash_to_word_ratio",
    # TextDescriptives – information theory (perplexity → log10)
    "entropy",
    "log10_perplexity",
    # TextDescriptives – POS proportions
    "pos_prop_NOUN",
    "pos_prop_VERB",
    "pos_prop_ADJ",
    "pos_prop_ADV",
    "pos_prop_PRON",
    "pos_prop_DET",
    "pos_prop_ADP",
    "pos_prop_AUX",
    "pos_prop_CCONJ",
    "pos_prop_SCONJ",
    "pos_prop_PART",
    "pos_prop_INTJ",
    "pos_prop_PROPN",
    "pos_prop_NUM",
    "pos_prop_PUNCT",
    # lexicalrichness
    "mattr",
    "mtld",
    "hdd",
    "maas",
    "herdan",
    "summer",
    "dugast",
    "yulek",
    "yulei",
    "herdanvm",
    "simpsond",
    "vocd",
    # textstat
    "dale_chall_readability_score",
    "linsear_write_formula",
    "spache_readability",
    "mcalpine_eflaw",
    "reading_time",
    "syllable_count",
    "lexicon_count",
    "sentence_count",
    "char_count",
    "letter_count",
    "polysyllabcount",
    "monosyllabcount",
    "difficult_words",
    "text_standard_float",
    # NRCLex
    "nrc_fear",
    "nrc_anger",
    "nrc_anticipation",
    "nrc_trust",
    "nrc_surprise",
    "nrc_sadness",
    "nrc_disgust",
    "nrc_joy",
    "nrc_positive",
    "nrc_negative",
]

# spaCy/TD dict slice: descriptive(14) + readability(8) — quality is pure Python
_N_DESC_READ = 22
_N_QUALITY = 17
_N_INFO_POS = 17  # entropy+perp(2) + POS(15)
_N_SPACY_CORE = _N_DESC_READ + _N_INFO_POS  # values taken from spaCy dict
_N_TOTAL = len(FEATURE_NAMES)


def _desc_read_vals(d: dict) -> list[float]:
    return [
        _safe(d.get("token_length_mean")),
        _safe(d.get("token_length_median")),
        _safe(d.get("token_length_std")),
        _safe(d.get("sentence_length_mean")),
        _safe(d.get("sentence_length_median")),
        _safe(d.get("sentence_length_std")),
        _safe(d.get("syllables_per_token_mean")),
        _safe(d.get("syllables_per_token_median")),
        _safe(d.get("syllables_per_token_std")),
        _safe(d.get("proportion_unique_tokens")),
        _safe(d.get("n_tokens")),
        _safe(d.get("n_unique_tokens")),
        _safe(d.get("n_characters")),
        _safe(d.get("n_sentences")),
        _safe(d.get("flesch_reading_ease")),
        _safe(d.get("flesch_kincaid_grade")),
        _safe(d.get("smog")),
        _safe(d.get("gunning_fog")),
        _safe(d.get("automated_readability_index")),
        _safe(d.get("coleman_liau_index")),
        _safe(d.get("lix")),
        _safe(d.get("rix")),
    ]


def _info_pos_vals(d: dict) -> list[float]:
    return [
        _safe(d.get("entropy")),
        _log10_perplexity(d.get("perplexity")),
        _safe(d.get("pos_prop_NOUN")),
        _safe(d.get("pos_prop_VERB")),
        _safe(d.get("pos_prop_ADJ")),
        _safe(d.get("pos_prop_ADV")),
        _safe(d.get("pos_prop_PRON")),
        _safe(d.get("pos_prop_DET")),
        _safe(d.get("pos_prop_ADP")),
        _safe(d.get("pos_prop_AUX")),
        _safe(d.get("pos_prop_CCONJ")),
        _safe(d.get("pos_prop_SCONJ")),
        _safe(d.get("pos_prop_PART")),
        _safe(d.get("pos_prop_INTJ")),
        _safe(d.get("pos_prop_PROPN")),
        _safe(d.get("pos_prop_NUM")),
        _safe(d.get("pos_prop_PUNCT")),
    ]


def _spacy_core_vals(d: dict) -> list[float]:
    """desc+read + info+POS (quality filled separately)."""
    return _desc_read_vals(d) + _info_pos_vals(d)


def _lex_textstat_nrc_vals(STRING: str) -> list[float]:
    """lexicalrichness + textstat + NRCLex."""
    vals: list[float] = []

    try:
        lex = LexicalRichness(STRING)
        vals.extend(
            [
                _safe(lex.mattr(window_size=50)),
                _safe(lex.mtld(threshold=0.72)),
                _safe(lex.hdd(draws=42)),
                _safe(lex.Maas),
                _safe(lex.Herdan),
                _safe(lex.Summer),
                _safe(lex.Dugast),
                _safe(lex.yulek),
                _safe(lex.yulei),
                _safe(lex.herdanvm),
                _safe(lex.simpsond),
                _safe(lex.vocd()),
            ]
        )
    except Exception:
        vals.extend([0.0] * 12)

    try:
        vals.extend(
            [
                _safe(textstat.dale_chall_readability_score(STRING)),
                _safe(textstat.linsear_write_formula(STRING)),
                _safe(textstat.spache_readability(STRING)),
                _safe(textstat.mcalpine_eflaw(STRING)),
                _safe(textstat.reading_time(STRING)),
                _safe(textstat.syllable_count(STRING)),
                _safe(textstat.lexicon_count(STRING)),
                _safe(textstat.sentence_count(STRING)),
                _safe(textstat.char_count(STRING)),
                _safe(textstat.letter_count(STRING)),
                _safe(textstat.polysyllabcount(STRING)),
                _safe(textstat.monosyllabcount(STRING)),
                _safe(textstat.difficult_words(STRING)),
                _safe(textstat.text_standard(STRING, float_output=True)),
            ]
        )
    except Exception:
        vals.extend([0.0] * 14)

    try:
        tokens = [
            w.strip(".,;:!?\"'()[]{}").lower() for w in STRING.split() if w.strip()
        ]
        emotion = _get_nrc()
        if tokens:
            emotion.load_token_list(tokens)
        else:
            emotion.load_raw_text(STRING)
        raw = getattr(emotion, "raw_emotion_scores", None) or {}
        word_count = max(len(tokens) if tokens else len(STRING.split()), 1)
        vals.extend(
            [
                _safe(raw.get("fear", 0) / word_count),
                _safe(raw.get("anger", 0) / word_count),
                _safe(raw.get("anticipation", 0) / word_count),
                _safe(raw.get("trust", 0) / word_count),
                _safe(raw.get("surprise", 0) / word_count),
                _safe(raw.get("sadness", 0) / word_count),
                _safe(raw.get("disgust", 0) / word_count),
                _safe(raw.get("joy", 0) / word_count),
                _safe(raw.get("positive", 0) / word_count),
                _safe(raw.get("negative", 0) / word_count),
            ]
        )
    except Exception:
        vals.extend([0.0] * 10)

    return vals


def _non_spacy_block(STRING: str) -> list[float]:
    """Quality + lexicalrichness + textstat + NRC (everything not needing spaCy)."""
    return _quality_vals(STRING) + _lex_textstat_nrc_vals(STRING)


def _assemble(d: dict, STRING: str, non_spacy: list[float] | None = None) -> list[float]:
    """desc/read from spaCy, quality..nrc from non_spacy block (or compute)."""
    rest = non_spacy if non_spacy is not None else _non_spacy_block(STRING)
    # rest = quality(17) + lex(12) + textstat(14) + nrc(10) = 53
    q = rest[:_N_QUALITY]
    tail = rest[_N_QUALITY:]
    vals = _desc_read_vals(d) + q + _info_pos_vals(d) + tail
    return vals


def extract_vector(STRING: str) -> np.ndarray:
    """Return a fixed-order float vector of all metrics (single text)."""
    d = _td(STRING)
    vals = _assemble(d, STRING)
    out = np.asarray(vals, dtype=np.float64)
    if out.shape != (_N_TOTAL,):
        raise RuntimeError(
            f"extract_vector length {out.shape[0]} != expected {_N_TOTAL}"
        )
    return out


def submit_extract_vector(text: str, *, n_texts_hint: int = 5):
    """
    Queue one full-document metric extract on the shared metric thread pool.

    Used as soon as a generation sample finishes so obj1 can overlap remaining
    NovelAI streams. Returns a concurrent.futures.Future[np.ndarray].
    """
    from concurrent.futures import Future  # noqa: F401 — type clarity for callers

    workers = _metric_workers(max(1, int(n_texts_hint)))
    pool = _get_metric_pool(workers)
    return pool.submit(extract_vector, text if isinstance(text, str) else str(text))


def _preload_worker_nlp(_: int = 0) -> int:
    """Pool map target: load spaCy (+nrc) into this worker thread once."""
    _get_nlp()
    _get_nrc()
    return 0


def warmup_metric_workers(n_texts: int = 5) -> int:
    """
    Preload thread-local spaCy models in the metric pool.

    Call once at process start so the first trial does not pay multi-model load
    on the critical path. Returns worker count (reflects current VM size).
    """
    workers = _metric_workers(max(1, int(n_texts)))
    pool = _get_metric_pool(workers)
    list(pool.map(_preload_worker_nlp, range(workers), chunksize=1))
    # Parent thread too (human cache / single extracts).
    _get_nlp()
    _get_nrc()
    return workers


def metric_pool_info(n_texts: int = 5) -> dict[str, int | str | float | None]:
    """Diagnostics for logs: workers vs host CPUs / RAM (useful on e2 resize)."""
    n_texts = max(1, int(n_texts))
    cpu = os.cpu_count() or 2
    ram = _total_ram_gb()
    workers = _metric_workers(n_texts)
    return {
        "cpu_count": cpu,
        "ram_gb": None if ram is None else round(ram, 1),
        "n_texts": n_texts,
        "workers": workers,
        "override": (os.getenv("METRIC_WORKERS") or "").strip() or "auto",
    }


def extract_vectors(texts: Iterable[str]) -> np.ndarray:
    """
    Batch extract for multiple texts (shape: n_texts × n_features).

    Multi-thread full `extract_vector` per text with thread-local spaCy models.
    Process-based spaCy (n_process / ProcessPool) is slower on Windows because
    each worker reloads en_core_web_sm — measured ~2–4× worse wall time.
    """
    texts_list = [t if isinstance(t, str) else str(t) for t in texts]
    n = len(texts_list)
    if n == 0:
        return np.zeros((0, _N_TOTAL), dtype=np.float64)
    if n == 1:
        return extract_vector(texts_list[0]).reshape(1, -1)

    workers = _metric_workers(n)
    pool = _get_metric_pool(workers)
    # map preserves order; each worker builds/reuses its own spaCy pipeline.
    vectors = list(pool.map(extract_vector, texts_list, chunksize=1))
    return np.stack(vectors, axis=0)


if __name__ == "__main__":
    text = (
        "The old house stood silent under the moon. "
        "She whispered, 'I am afraid,' and the shadows seemed to listen. "
        "Fear and hope fought inside her heart."
    )
    vec = extract_vector(text)
    print(f"Vector shape: {vec.shape}")
    print(f"Features: {len(FEATURE_NAMES)}")
    for name, value in zip(FEATURE_NAMES, vec):
        print(f"{name:45s} = {value:10.4f}")
