"""Reproducible raw-vs-CPT evaluation for the Qwen3 resident base.

Measures token-weighted RU/EN perplexity on pinned human-written UD test sets,
plus deterministic RU/EN generation probes. Writes machine-readable JSON for
decision_gate.py. It does not mutate a model or dataset.
"""
from __future__ import annotations

import argparse
import json
import math
import os
from pathlib import Path

os.environ.setdefault("HF_HOME", "/mnt/D/.cache/huggingface")
os.environ.setdefault("HF_DATASETS_CACHE", "/mnt/D/.cache/huggingface/datasets")

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from corpus_common import has_mixed_script, cyrillic_ratio

RU_PROMPTS = [
    "Объясни простыми словами, почему зимой дни короче.",
    "Опиши, как спокойно подготовиться к сложному рабочему дню.",
    "Напиши короткое напоминание купить продукты после работы.",
    "Расскажи, чем резервная копия отличается от синхронизации.",
    "Продолжи естественно: Когда я вернулась домой, оказалось, что",
    "Сформулируй вежливый отказ от встречи без лишних подробностей.",
    "Объясни разницу между привычкой и разовым действием.",
    "Дай три коротких совета человеку, который плохо выспался.",
]
HERE = Path(__file__).resolve().parent
EN_PROMPTS = [
    "Explain in one sentence why winter days are shorter.",
    "Write a polite one-sentence reminder to buy groceries after work.",
    "Explain the difference between a backup and synchronization.",
    "Continue naturally: When I returned home, I discovered that",
]


def load_lines(path: Path, limit: int) -> list[str]:
    lines = [line.strip() for line in path.read_text(encoding="utf-8").splitlines()
             if line.strip()]
    return lines[:limit] if limit else lines


def load_model(model_id: str, device: str):
    tok = AutoTokenizer.from_pretrained(model_id)
    dtype = torch.bfloat16 if device == "cuda" else torch.float32
    model = AutoModelForCausalLM.from_pretrained(
        model_id, dtype=dtype, attn_implementation="eager",
        local_files_only=False,
    ).eval().to(device)
    return tok, model


@torch.inference_mode()
def generate(tok, model, prompt: str, max_new_tokens: int) -> str:
    encoded = tok(prompt, return_tensors="pt").to(model.device)
    output = model.generate(
        **encoded, max_new_tokens=max_new_tokens, do_sample=False,
        pad_token_id=tok.eos_token_id,
    )
    continuation = output[0][encoded.input_ids.shape[1]:]
    return tok.decode(continuation, skip_special_tokens=True).strip()


@torch.inference_mode()
def perplexity(tok, model, texts: list[str], max_length: int) -> tuple[float, int]:
    nll = 0.0
    predicted = 0
    for text in texts:
        encoded = tok(text, return_tensors="pt", truncation=True,
                      max_length=max_length).to(model.device)
        tokens = int(encoded.attention_mask.sum()) - 1
        if tokens <= 0:
            continue
        loss = model(**encoded, labels=encoded.input_ids).loss.float().item()
        nll += loss * tokens
        predicted += tokens
    if not predicted:
        raise ValueError("evaluation set contains no predictable tokens")
    return math.exp(nll / predicted), predicted


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--ru-text", default=str(HERE / "data/eval/ru_ud_test.txt"))
    ap.add_argument("--en-text", default=str(HERE / "data/eval/en_ud_test.txt"))
    ap.add_argument("--limit", type=int, default=512,
                    help="sentences per PPL language; 0 uses the complete files")
    ap.add_argument("--max-length", type=int, default=512)
    ap.add_argument("--max-new-tokens", type=int, default=96)
    ap.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
    ap.add_argument("--output")
    args = ap.parse_args()
    device = ("cuda" if torch.cuda.is_available() else "cpu") if args.device == "auto" else args.device

    tok, model = load_model(args.model, device)
    ru_texts = load_lines(Path(args.ru_text), args.limit)
    en_texts = load_lines(Path(args.en_text), args.limit)
    ru_ppl, ru_tokens = perplexity(tok, model, ru_texts, args.max_length)
    en_ppl, en_tokens = perplexity(tok, model, en_texts, args.max_length)
    ru_gens = [generate(tok, model, p, args.max_new_tokens) for p in RU_PROMPTS]
    en_gens = [generate(tok, model, p, args.max_new_tokens) for p in EN_PROMPTS]
    ru_valid = [not has_mixed_script(g) and cyrillic_ratio(g) >= 0.80 for g in ru_gens]
    en_valid = [len(g) >= 10 and cyrillic_ratio(g) < 0.20 for g in en_gens]

    report = {
        "schema_version": 1,
        "model": args.model,
        "device": device,
        "deterministic": True,
        "perplexity": {
            "ru": ru_ppl, "ru_predicted_tokens": ru_tokens,
            "en": en_ppl, "en_predicted_tokens": en_tokens,
            "sentence_limit": args.limit,
        },
        "generation": {
            "ru_valid_pct": 100 * sum(ru_valid) / len(ru_valid),
            "ru_mixed_script_count": sum(has_mixed_script(g) for g in ru_gens),
            "en_retained": all(en_valid),
            "ru": [{"prompt": p, "output": g, "valid": ok}
                   for p, g, ok in zip(RU_PROMPTS, ru_gens, ru_valid)],
            "en": [{"prompt": p, "output": g, "valid": ok}
                   for p, g, ok in zip(EN_PROMPTS, en_gens, en_valid)],
        },
    }
    rendered = json.dumps(report, ensure_ascii=False, indent=2) + "\n"
    if args.output:
        Path(args.output).write_text(rendered, encoding="utf-8")
    print(rendered, end="")


if __name__ == "__main__":
    main()