"""Сравнение сегментации: правила vs обученный сегментатор (predict_release).

Метрика — пары (тип, автор) на уровне релиза против меток датасета, только по
страницам, которых сегментатор не видел при обучении (is_test_release).
Текст kinozal — из data/raw/kinozal/*.json (как у сервиса), см. segmenter.page_text.
На трекер берётся не больше E2E_MAX страниц (по умолчанию 400, детерминированно);
train-страницы считаются только с флагом --train-pages. Основная модель — E2E_MODEL
(по умолчанию models/best).
Запуск: .venv/Scripts/python tests/e2e_segmenter.py [каталог_сегментатора] [--train-pages]
"""

from __future__ import annotations

import hashlib
import json
import os
import sys
from collections import defaultdict
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from src.predict import DubPredictor  # noqa: E402
from src.release import predict_release  # noqa: E402
from src.segmenter import is_test_release, page_text  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
MAX_PER_SOURCE = int(os.environ.get("E2E_MAX", "400"))
MODEL_DIR = os.environ.get("E2E_MODEL", "models/best")


def run(predictor, rows, mode: str, test_only: bool) -> dict:
    by_rel = defaultdict(list)
    for r in rows:
        by_rel[r["release_id"]].append(r)
    per_src = defaultdict(list)
    pages = [p for p in (ROOT / "data" / "pages").glob("*.txt")
             if is_test_release(p.stem) == test_only and p.stem in by_rel]
    pages.sort(key=lambda p: hashlib.md5(p.stem.encode()).hexdigest())
    for page in pages:
        rid = page.stem
        src = by_rel[rid][0]["source"]
        if len(per_src[src]) >= MAX_PER_SOURCE:
            continue
        gold = {(r["translation_type"], r["studio"] or None) for r in by_rel[rid]}
        # как aggregate: «тип без студии» не отдельная дорожка, если у типа есть студия
        # (kinozal: «Перевод: многоголосый» + студии из заголовка)
        named = {t for t, s in gold if s}
        gold = {(t, s) for t, s in gold if s or t not in named}
        res = predict_release(predictor, text=page_text(page), segmenter=mode)
        pred = {(t["translation_type"], t["studio"]) for t in res["tracks"]}
        tp = len(gold & pred)
        p = tp / len(pred) if pred else 0.0
        r = tp / len(gold) if gold else 0.0
        per_src[src].append((p, r))
    out = {}
    for s, v in per_src.items():
        out[s] = (len(v), sum(x[0] for x in v) / len(v), sum(x[1] for x in v) / len(v))
    return out


def main() -> None:
    rows = json.loads((ROOT / "data" / "processed" / "dataset.json").read_text(encoding="utf-8"))
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    seg_dir = args[0] if args else "models/segmenter"
    predictor = DubPredictor(MODEL_DIR, segmenter_dir=seg_dir)
    print(f"модель: {MODEL_DIR}; сегментатор: {seg_dir}")
    if predictor.segmenter is None:
        sys.exit("нет models/segmenter: python -m src.segmenter --train")
    blocks = [(True, "тестовые страницы (не виделись сегментатору)")]
    if "--train-pages" in sys.argv:
        blocks.append((False, "train-страницы (для справки, оптимистично)"))
    for test_only, name in blocks:
        print(f"== {name} ==")
        for mode in ("rules", "model"):
            res = run(predictor, rows, mode, test_only)
            print(f"  {mode:5s}: " + "; ".join(
                f"{s} n={n} P={p:.3f} R={r:.3f}" for s, (n, p, r) in res.items()))


if __name__ == "__main__":
    main()