"""Runnable end-to-end example: score raw text with the hybrid detector.

Unlike a placeholder demo, this extracts the 25 linguistic features *exactly* the
way they were produced for training/testing, normalizes them with the fitted
training scaler (``ling_scaler.pkl``), and aggregates chunk probabilities into a
document-level score (mean chunk probability), which is how the thesis reports
per-document predictions.

Pipeline (see features.py for details):
    raw text -> normalize -> sliding-window chunks -> per-chunk 25 features
             -> StandardScaler.transform -> model -> mean chunk P(machine)

Requirements:
    pip install torch transformers spacy scikit-learn
    python -m spacy download en_core_web_lg

Run:
    python example_usage.py
    python example_usage.py --text "Some text to classify..."
    python example_usage.py --file path/to/document.txt
"""
import argparse
import pickle

import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    GPT2LMHeadModel,
    GPT2TokenizerFast,
    RobertaTokenizer,
)

from model import load_model, CONFIG, LING_FEATURE_NAMES
from features import extract_raw_features, prepare_document, FEATURE_NAMES

# Sanity check: features.py and model.py must agree on feature order.
assert FEATURE_NAMES == LING_FEATURE_NAMES, "Feature order mismatch."

DEFAULT_TEXT = (
    "This is an example transcript to classify as human- or machine-written. "
    "It is deliberately short, so it forms a single chunk. Provide your own "
    "longer document via --text or --file to see multi-chunk aggregation."
)


def load_components(device):
    """Load every model/resource needed to reproduce the test-time pipeline."""
    print("Loading hybrid model ...")
    model = load_model("hybrid_model_best.pt", device=device)

    print("Loading RoBERTa tokenizer ...")
    tokenizer = RobertaTokenizer.from_pretrained(CONFIG["roberta_model"])

    print("Loading spaCy (en_core_web_lg) ...")
    import spacy
    try:
        nlp = spacy.load("en_core_web_lg", disable=["ner"])
    except OSError:
        print("  en_core_web_lg not found -> downloading ...")
        spacy.cli.download("en_core_web_lg")
        nlp = spacy.load("en_core_web_lg", disable=["ner"])

    print("Loading GPT-2 (perplexity) ...")
    gpt2_tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
    gpt2_model = GPT2LMHeadModel.from_pretrained("gpt2").to(device)
    gpt2_model.eval()

    print("Loading training feature scaler (ling_scaler.pkl) ...")
    with open("ling_scaler.pkl", "rb") as f:
        scaler = pickle.load(f)

    return model, tokenizer, nlp, gpt2_model, gpt2_tokenizer, scaler


def score_text(text, model, tokenizer, nlp, gpt2_model, gpt2_tokenizer, scaler, device):
    """Return (doc_prob, per_chunk_probs) for a raw document."""
    chunks = prepare_document(text)
    if not chunks:
        raise ValueError("Text is empty after normalization; nothing to score.")

    chunk_probs = []
    for chunk in chunks:
        # 1. Extract the 25 raw features exactly as during training.
        raw = extract_raw_features(chunk, nlp, gpt2_model, gpt2_tokenizer, device)
        # 2. Normalize with the fitted training scaler.
        norm = scaler.transform(raw.reshape(1, -1)).astype(np.float32)
        ling = torch.from_numpy(norm).to(device)
        # 3. Tokenize text for RoBERTa (same 512-token truncation as training).
        enc = tokenizer(chunk, max_length=512, truncation=True, return_tensors="pt")
        input_ids = enc["input_ids"].to(device)
        attention_mask = enc["attention_mask"].to(device)
        # 4. Forward pass -> P(machine) for this chunk.
        with torch.no_grad():
            logits = model(input_ids, attention_mask, ling)
            p_llm = F.softmax(logits, dim=1)[0, 1].item()
        chunk_probs.append(p_llm)

    doc_prob = float(np.mean(chunk_probs))
    return doc_prob, chunk_probs


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    src = parser.add_mutually_exclusive_group()
    src.add_argument("--text", type=str, help="Raw text to classify.")
    src.add_argument("--file", type=str, help="Path to a UTF-8 text file to classify.")
    args = parser.parse_args()

    if args.file:
        with open(args.file, "r", encoding="utf-8") as f:
            text = f.read()
    elif args.text:
        text = args.text
    else:
        text = DEFAULT_TEXT

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Device: {device}\n")

    components = load_components(device)
    doc_prob, chunk_probs = score_text(text, *components, device=device)

    print("\n" + "=" * 60)
    print(f"Chunks scored: {len(chunk_probs)}")
    for i, p in enumerate(chunk_probs):
        print(f"  chunk {i:>2}: P(machine) = {p:.4f}")
    print("-" * 60)
    print(f"Document P(machine-generated) = {doc_prob:.4f}")
    print(f"Prediction: {'machine-generated' if doc_prob >= 0.5 else 'human-written'}")
    print("=" * 60)


if __name__ == "__main__":
    main()