grabber-ai/tools/diversity.py

"""Разнообразная выборка раздач: годы, качество, фильмы/сериалы.

Список раздела идёт «от новых к старым», поэтому чтение подряд даёт только
свежие релизы. Здесь:

* страницы раздела берутся вразброс по всей глубине (последовательность с
  золотым сечением, уже посещённые пропускаются);
* на странице выбираются не все темы, а те, чья «корзина» (годовая группа ×
  качество × фильм/сериал) в датасете представлена меньше всего.

Всё работает по заголовкам тем из списка, заходить в раздачу для выбора не нужно.
"""

from __future__ import annotations

import csv
import random
import re
from collections import Counter
from pathlib import Path

RE_YEAR = re.compile(r"(?<!\d)(19[3-9]\d|20[0-3]\d)(?!\d)")
RE_SERIES = re.compile(r"(?i)сезон|сери[ия]|\bS\d{1,2}\b|\bE\d{1,3}\b|эпизод|\bTV\b|выпуск")
GOLDEN = 0.6180339887498949

QUALITY_RULES = [
    ("4K", re.compile(r"(?i)2160|\b4K\b|\bUHD\b")),
    ("1080", re.compile(r"(?i)1080|\bFHD\b")),
    ("720", re.compile(r"(?i)\b720|\bHD(?:TV)?\b|\bHDRip\b")),
    ("SD", re.compile(r"(?i)DVD|SAT|TVRip|VHS|\b480|\b576|CAMRip|TS\b")),
]


def year_of(title: str) -> int | None:
    """Год выпуска из заголовка: сначала из скобок «[2020, …]» / «(2020)», иначе первый."""
    for m in re.finditer(r"[\[(]\s*(19[3-9]\d|20[0-3]\d)(?!\d)", title or ""):
        return int(m.group(1))
    m = RE_YEAR.search(title or "")
    return int(m.group(1)) if m else None


def year_band(year: int | None) -> str:
    if year is None:
        return "?"
    if year < 2010:
        return f"{year // 10 * 10}-е"
    return str(year)


def quality_of(title: str) -> str:
    for name, rx in QUALITY_RULES:
        if rx.search(title or ""):
            return name
    return "прочее"


def bucket(title: str) -> tuple[str, str, str]:
    return (year_band(year_of(title)), quality_of(title),
            "сериал" if RE_SERIES.search(title or "") else "фильм")


def counts_from_meta(path: Path, site: str | None = None) -> Counter:
    """Сколько релизов каждой корзины уже в датасете (по data/meta.csv)."""
    out: Counter = Counter()
    try:
        with path.open(encoding="utf-8-sig", newline="") as f:
            for r in csv.DictReader(f):
                if site and r.get("source") != site:
                    continue
                title = r.get("title") or ""
                year = int(r["year"]) if str(r.get("year", "")).isdigit() else year_of(title)
                q = quality_of(f"{title} {r.get('quality', '')} {r.get('resolution', '')}")
                is_s = str(r.get("is_series", "")).lower() in ("true", "1")
                out[(year_band(year), q, "сериал" if is_s else "фильм")] += 1
    except OSError:
        pass
    return out


def pick(items: list[dict], n: int, counts: Counter, rng: random.Random) -> list[dict]:
    """n тем из страницы: сначала из самых малочисленных корзин (counts обновляется)."""
    pool = list(items)
    rng.shuffle(pool)                                 # равные корзины -> случайный порядок
    chosen: list[dict] = []
    while pool and len(chosen) < n:
        best = min(pool, key=lambda it: counts[bucket(it.get("text", ""))])
        pool.remove(best)
        counts[bucket(best.get("text", ""))] += 1
        chosen.append(best)
    return chosen


def next_pages(depth: int, visited: set[int], k: int, rng: random.Random) -> list[int]:
    """До k ещё не посещённых страниц, разбросанных по глубине раздела."""
    if depth <= 0:
        return []
    left = depth - len(visited)
    out: list[int] = []
    x = rng.random()
    tries = 0
    while len(out) < min(k, left) and tries < depth * 8 + 50:
        tries += 1
        x = (x + GOLDEN) % 1.0
        p = int(x * depth)
        if p not in visited and p not in out:
            out.append(p)
    if len(out) < min(k, left):                       # хвост добираем подряд
        out += [p for p in range(depth) if p not in visited and p not in out][:min(k, left) - len(out)]
    return out