"""Self-contained linguistic feature extraction for the hybrid detector.

This reproduces *exactly* the preprocessing + feature pipeline used to train and
evaluate the released ``hybrid_model_best.pt`` checkpoint, so that a raw piece of
text can be scored the same way it was during testing:

    raw text
      -> normalize_text()            (lowercase, strip HTML/markdown, collapse ws)
      -> sliding_window_chunk()      (450-word windows, 350-word stride)
      -> extract_raw_features()      (24 spaCy features + 1 GPT-2 perplexity = 25)
      -> StandardScaler.transform()  (the fitted training scaler, ling_scaler.pkl)
      -> model(...)                  (per chunk)
      -> mean chunk probability      (document-level score)

The 25 features, in order, match ``model.LING_FEATURE_NAMES``:

    msttr, avg_word_len, hapax_ratio, function_ratio, punct_density, char_entropy,
    burstiness, repetition_ratio, avg_sent_len, sent_len_std, noun_ratio,
    verb_ratio, adj_ratio, adv_ratio, pron_ratio, pos_diversity, avg_tree_depth,
    max_tree_depth, sub_clause_ratio, dm_density, sent_len_cv, fp_ratio,
    num_sentences, words_per_sent, perplexity

Requirements:
    pip install torch transformers spacy scikit-learn
    python -m spacy download en_core_web_lg
"""
import math
import re
from collections import Counter

import numpy as np
import torch

# --------------------------------------------------------------------------- #
# Feature order (must match model.LING_FEATURE_NAMES).
# --------------------------------------------------------------------------- #
FEATURE_NAMES = [
    "msttr", "avg_word_len", "hapax_ratio", "function_ratio", "punct_density",
    "char_entropy", "burstiness", "repetition_ratio",
    "avg_sent_len", "sent_len_std", "noun_ratio", "verb_ratio", "adj_ratio",
    "adv_ratio", "pron_ratio", "pos_diversity", "avg_tree_depth",
    "max_tree_depth", "sub_clause_ratio",
    "dm_density", "sent_len_cv", "fp_ratio", "num_sentences", "words_per_sent",
    "perplexity",
]

# --------------------------------------------------------------------------- #
# Preprocessing (mirrors 01_data_preprocessing_v2.py).
# --------------------------------------------------------------------------- #
WINDOW_SIZE = 450
STRIDE = 350
MIN_WINDOW_TOKENS = 50


def normalize_text(text: str) -> str:
    """Light normalization matching the training preprocessing.

    Lowercase, strip HTML/markdown artifacts, normalize quotes, collapse
    whitespace. Applied to every document before chunking / feature extraction.
    """
    if not text or not isinstance(text, str):
        return ""

    text = text.lower()
    text = re.sub(r"<[^>]+>", " ", text)
    text = re.sub(r"#{1,6}\s+", " ", text)
    text = re.sub(r"\[([^\]]+)\]\([^)]+\)", r"\1", text)
    text = re.sub(r"!\[[^\]]+\]\([^)]+\)", " ", text)
    text = text.replace('"', '"').replace('"', '"')
    text = text.replace("'", "'").replace("'", "'")
    text = re.sub(r"\s+", " ", text).strip()
    return text


def sliding_window_chunk(text, window_size=WINDOW_SIZE, stride=STRIDE):
    """Split text into overlapping word windows (same as training)."""
    if not text:
        return []

    words = text.split()
    total_words = len(words)

    if total_words <= window_size:
        return [text]

    chunks = []
    start = 0
    while start < total_words:
        end = min(start + window_size, total_words)
        chunk_words = words[start:end]
        if len(chunk_words) >= MIN_WINDOW_TOKENS:
            chunks.append(" ".join(chunk_words))
        start += stride
        if end == total_words:
            break
    return chunks


# --------------------------------------------------------------------------- #
# Lexicons (identical to 02_feature_extraction.py).
# --------------------------------------------------------------------------- #
DISCOURSE_MARKERS = {
    "however", "therefore", "moreover", "furthermore", "nevertheless",
    "consequently", "meanwhile", "additionally", "similarly", "likewise",
    "thus", "hence", "accordingly", "otherwise", "instead",
    "first", "second", "third", "finally", "next", "then",
    "in conclusion", "in summary", "to summarize", "overall",
    "for example", "for instance", "specifically", "in particular",
    "on the other hand", "in contrast", "conversely", "although",
    "because", "since", "while", "whereas", "unless", "if",
}

FUNCTION_WORDS = {
    "the", "a", "an", "and", "or", "but", "if", "then", "else",
    "when", "where", "how", "what", "who", "which", "that", "this",
    "is", "are", "was", "were", "be", "been", "being",
    "have", "has", "had", "do", "does", "did",
    "will", "would", "could", "should", "may", "might", "must",
    "to", "of", "in", "for", "on", "with", "at", "by", "from",
    "as", "into", "through", "during", "before", "after",
    "above", "below", "between", "under", "over",
    "i", "you", "he", "she", "it", "we", "they", "me", "him", "her", "us", "them",
    "my", "your", "his", "her", "its", "our", "their",
    "not", "no", "yes", "so", "very", "just", "also", "only",
}

SUB_MARKERS = {
    "that", "which", "who", "whom", "whose", "where", "when", "while",
    "because", "although", "if", "unless",
}

FIRST_PERSON = {
    "i", "me", "my", "mine", "myself", "we", "us", "our", "ours", "ourselves",
}


# --------------------------------------------------------------------------- #
# Feature helpers (identical to 02_feature_extraction.py).
# --------------------------------------------------------------------------- #
def get_tree_depth(token):
    depth = 0
    current = token
    while current.head != current:
        depth += 1
        current = current.head
        if depth > 100:
            break
    return depth


def calculate_entropy(text):
    if not text:
        return 0.0
    counts = Counter(text)
    total = len(text)
    probs = [c / total for c in counts.values()]
    return -sum(p * math.log2(p) for p in probs)


def calculate_burstiness(words):
    if not words:
        return 0.0
    word_counts = list(Counter(words).values())
    if not word_counts:
        return 0.0
    return np.std(word_counts) / np.mean(word_counts) if np.mean(word_counts) > 0 else 0.0


def calculate_repetition_ratio(words):
    if not words:
        return 0.0
    counts = Counter(words)
    repeated = sum(c for w, c in counts.items() if c > 1)
    return repeated / len(words)


def calculate_msttr(words, window_size=50):
    if not words:
        return 0.0
    if len(words) < window_size:
        return len(set(words)) / len(words)
    ttrs = []
    for i in range(0, len(words), window_size):
        segment = words[i:i + window_size]
        if len(segment) == window_size:
            ttrs.append(len(set(segment)) / len(segment))
    return np.mean(ttrs) if ttrs else 0.0


def calculate_pos_diversity(doc):
    pos_counts = Counter([token.pos_ for token in doc])
    total = len(doc)
    if total == 0:
        return 0.0
    probs = [count / total for count in pos_counts.values()]
    return -sum(p * math.log2(p) for p in probs)


def calculate_perplexity(text, model, tokenizer, device):
    """GPT-2 perplexity (single truncated window, matching training)."""
    max_length = model.config.n_positions
    encodings = tokenizer(text, return_tensors="pt", truncation=True, max_length=max_length)
    seq_len = encodings.input_ids.size(1)
    if seq_len < 2:
        return 0.0
    input_ids = encodings.input_ids.to(device)
    with torch.no_grad():
        outputs = model(input_ids, labels=input_ids)
        neg_log_likelihood = outputs.loss
    if neg_log_likelihood is None:
        return 0.0
    return torch.exp(neg_log_likelihood).item()


def _cpu_features_from_doc(doc):
    """The 24 spaCy-based features for a single spaCy Doc."""
    text = doc.text
    words = [token.text.lower() for token in doc if token.is_alpha]
    sentences = list(doc.sents)

    features = {}
    if len(words) == 0:
        return {}

    features["msttr"] = calculate_msttr(words)
    features["avg_word_len"] = np.mean([len(w) for w in words]) if words else 0

    word_counts = Counter(words)
    hapax = sum(1 for w, c in word_counts.items() if c == 1)
    features["hapax_ratio"] = hapax / len(words) if words else 0

    function_count = sum(1 for w in words if w in FUNCTION_WORDS)
    features["function_ratio"] = function_count / len(words) if words else 0

    punct_count = sum(1 for token in doc if token.is_punct)
    features["punct_density"] = (punct_count / len(words)) * 100 if words else 0

    features["char_entropy"] = calculate_entropy(text)
    features["burstiness"] = calculate_burstiness(words)
    features["repetition_ratio"] = calculate_repetition_ratio(words)

    sent_lengths = [len([t for t in sent if t.is_alpha]) for sent in sentences]
    features["avg_sent_len"] = np.mean(sent_lengths) if sent_lengths else 0
    features["sent_len_std"] = np.std(sent_lengths) if len(sent_lengths) > 1 else 0

    pos_counts = Counter([token.pos_ for token in doc])
    total_tokens = len(doc)
    features["noun_ratio"] = pos_counts.get("NOUN", 0) / total_tokens if total_tokens else 0
    features["verb_ratio"] = pos_counts.get("VERB", 0) / total_tokens if total_tokens else 0
    features["adj_ratio"] = pos_counts.get("ADJ", 0) / total_tokens if total_tokens else 0
    features["adv_ratio"] = pos_counts.get("ADV", 0) / total_tokens if total_tokens else 0
    features["pron_ratio"] = pos_counts.get("PRON", 0) / total_tokens if total_tokens else 0

    features["pos_diversity"] = calculate_pos_diversity(doc)

    depths = [get_tree_depth(token) for token in doc if token.dep_ != "punct"]
    features["avg_tree_depth"] = np.mean(depths) if depths else 0
    features["max_tree_depth"] = max(depths) if depths else 0

    sub_count = sum(1 for token in doc if token.text.lower() in SUB_MARKERS)
    features["sub_clause_ratio"] = sub_count / len(sentences) if sentences else 0

    text_lower = text.lower()
    dm_count = sum(1 for dm in DISCOURSE_MARKERS if dm in text_lower)
    features["dm_density"] = (dm_count / len(words)) * 100 if words else 0

    features["sent_len_cv"] = (
        features["sent_len_std"] / features["avg_sent_len"]
        if features["avg_sent_len"] > 0 else 0
    )

    fp_count = sum(1 for w in words if w in FIRST_PERSON)
    features["fp_ratio"] = fp_count / len(words) if words else 0

    features["num_sentences"] = len(sentences)
    features["words_per_sent"] = len(words) / len(sentences) if sentences else 0

    return features


def extract_raw_features(text, nlp, gpt2_model, gpt2_tokenizer, device):
    """Return the 25-dim *raw* (un-normalized) feature vector for one chunk.

    ``text`` is expected to be a single normalized chunk (see normalize_text /
    sliding_window_chunk). Order matches FEATURE_NAMES.
    """
    doc = nlp(text)
    feats = _cpu_features_from_doc(doc)
    ppl = calculate_perplexity(text, gpt2_model, gpt2_tokenizer, device)

    vec = []
    for name in FEATURE_NAMES:
        if name == "perplexity":
            vec.append(ppl)
        else:
            vec.append(feats.get(name, 0.0))
    return np.array(vec, dtype=np.float32)


def prepare_document(text):
    """Normalize a raw document and split it into the chunks used at test time."""
    return sliding_window_chunk(normalize_text(text))