src/datasets/__init__.py"""Датасеты курса."""

src/datasets/crop_dataset.py"""Dataset поверх train/val/test.jsonl (по split.json) и функции сборки батча.

DataLoader не требует наследования от torch.utils.data.Dataset: достаточно __len__ и __getitem__.
Поэтому сам Dataset и сборка батча на NumPy не зависят от PyTorch; в тензоры батч превращает
тонкая обёртка to_torch().
"""
import json
from pathlib import Path

import numpy as np
from PIL import Image

from alphabet import encode_text
from image_to_tensor import preprocess_array
from split_groups import read_jsonl

VALUE_CLASS = 0
EMPTY_CLASS = 1   # P(empty) в уроках 7–9 — это вероятность класса 1


class CropDataset:
    """Один сплит. augment: функция (PIL-образ, rng) -> массив [1, 48, 160] или None."""

    def __init__(self, data_dir, split: str, augment=None, seed: int = 0) -> None:
        data_dir = Path(data_dir)
        manifest = json.loads((data_dir / "manifest.json").read_text(encoding="utf-8"))
        self.task = manifest["task"]
        self.image_root = Path(manifest["image_root"])
        self.records = read_jsonl(data_dir / f"{split}.jsonl")
        self.augment = augment
        self.seed = seed
        self.epoch = 0

    def __len__(self) -> int:
        return len(self.records)

    def set_epoch(self, epoch: int) -> None:
        """Аугментация меняется от эпохи к эпохе, но воспроизводима."""
        self.epoch = epoch

    def __getitem__(self, i: int) -> dict:
        rec = self.records[i]
        path = self.image_root / rec["image"]
        if self.augment is None:
            image = preprocess_array(path)
        else:
            # у каждого (seed, эпоха, индекс) свой генератор: результат не зависит от воркеров DataLoader
            rng = np.random.default_rng([self.seed, self.epoch, i])
            img = Image.open(path)
            img.load()
            image = self.augment(img, rng)
        return {
            "id": rec["id"],
            "image": image,
            "text": rec["text"],
            "target": np.array(encode_text(rec["text"], self.task), dtype=np.int64),
            "label": EMPTY_CLASS if rec["is_empty"] else VALUE_CLASS,
            "is_empty": bool(rec["is_empty"]),
        }


def collate_ctc_np(batch: list[dict]) -> dict:
    """Батч для CTC: цели склеены в один вектор + длины целей (пустая клетка = длина 0)."""
    return {
        "images": np.stack([b["image"] for b in batch]).astype(np.float32),
        "targets": np.concatenate([b["target"] for b in batch]).astype(np.int64),
        "target_lengths": np.array([len(b["target"]) for b in batch], dtype=np.int64),
        "texts": [b["text"] for b in batch],
        "ids": [b["id"] for b in batch],
    }


def collate_classify_np(batch: list[dict]) -> dict:
    """Батч для бинарной классификации empty / value."""
    return {
        "images": np.stack([b["image"] for b in batch]).astype(np.float32),
        "labels": np.array([b["label"] for b in batch], dtype=np.int64),
        "ids": [b["id"] for b in batch],
    }


def to_torch(batch: dict) -> dict:
    import torch  # импорт здесь: всё остальное в файле не требует torch
    return {k: torch.from_numpy(v) if isinstance(v, np.ndarray) else v for k, v in batch.items()}


def collate_ctc(batch: list[dict]) -> dict:
    return to_torch(collate_ctc_np(batch))


def collate_classify(batch: list[dict]) -> dict:
    return to_torch(collate_classify_np(batch))

src/augment.py"""Аугментации, СОХРАНЯЮЩИЕ МЕТКУ.

Правила (их проверяют тесты):
  * без зеркал и переворотов; поворот не больше ±5° (иначе 6 превращается в 9);
  * геометрия не режет цифры: если после преобразования чернила вышли бы за край кропа,
    преобразование отклоняется и пробуется заново;
  * паддинг холста остаётся ровно 0 (как при проверке на val и в проде).

Аугментация идёт в два этапа:
  1) геометрия на исходном кропе (поворот, масштаб, наклон, сдвиг);
  2) фотометрия на готовом холсте 48×160 (бледность чернил, размытие, JPEG, шум):
     шум добавляется ПОСЛЕ растяжения контраста, поэтому защита из урока 2 не нужна.
"""
import io
import math
from dataclasses import dataclass, replace

import numpy as np
from PIL import Image, ImageFilter

from image_to_tensor import CANVAS_H, CANVAS_W, load_image, preprocess_array, to_gray_on_white

ROT_CAP_DEG = 5.0       # больше нельзя: 6 и 9 начинают путаться
MIN_INK_CONTRAST = 25   # слабее — в кропе нет заметных чернил
INK_LEVEL = 0.35        # чернилами считаем пиксели темнее бумаги на 35 % контраста
LINE_FRAC = 0.85        # строка/столбец, занятые чернилами на 85 %+, — линия-рамка, а не цифры
GEOMETRY_TRIES = 8


@dataclass(frozen=True)
class AugmentConfig:
    rot_deg: float = 3.0
    scale: tuple = (0.93, 1.07)
    aspect: tuple = (0.96, 1.04)
    shear: float = 0.05
    shift_x: float = 0.03        # доли ширины кропа
    shift_y: float = 0.04        # доли высоты кропа
    gain: tuple = (0.85, 1.0)    # множитель яркости чернил (меньше = бледнее)
    blur_prob: float = 0.3
    blur_radius: tuple = (0.4, 0.9)
    jpeg_prob: float = 0.0
    jpeg_quality: tuple = (60, 95)
    noise: float = 0.02          # максимум σ гауссова шума (в долях 0..1)


DEFAULT = AugmentConfig()
PRINTED = AugmentConfig(rot_deg=5.0, scale=(0.85, 1.15), aspect=(0.90, 1.10), shear=0.10,
                        shift_x=0.05, shift_y=0.06, gain=(0.55, 1.0), blur_prob=0.7,
                        blur_radius=(0.5, 1.4), jpeg_prob=0.5, jpeg_quality=(30, 85), noise=0.06)


def validate_config(cfg: AugmentConfig) -> None:
    if abs(cfg.rot_deg) > ROT_CAP_DEG:
        raise ValueError(f"rot_deg={cfg.rot_deg}: поворот больше ±{ROT_CAP_DEG}° не сохраняет метку")


def content_box(w0: int, h0: int, h: int = CANVAS_H, w: int = CANVAS_W) -> tuple[int, int, int, int]:
    """Где на холсте лежит кроп размера w0×h0 после letterbox: (x0, y0, x1, y1)."""
    scale = min(w / w0, h / h0)
    new_w = max(1, round(w0 * scale))
    new_h = max(1, round(h0 * scale))
    top = (h - new_h) // 2
    return 0, top, new_w, top + new_h


def ink_mask(gray: np.ndarray) -> np.ndarray | None:
    """Маска «цифр» на исходном кропе (без линий-рамок); None, если чернил нет."""
    g = gray.astype(np.float32)
    paper = float(np.median(g))
    contrast = paper - float(g.min())
    if contrast < MIN_INK_CONTRAST:
        return None
    mask = (paper - g) > INK_LEVEL * contrast
    rows_full = mask.mean(axis=1) > LINE_FRAC
    cols_full = mask.mean(axis=0) > LINE_FRAC
    mask[rows_full, :] = False
    mask[:, cols_full] = False
    return mask if mask.any() else None


def ink_bbox(gray: np.ndarray):
    """(x0, y0, x1, y1) рамка чернил (x1, y1 не включены) или None."""
    mask = ink_mask(gray)
    if mask is None:
        return None
    ys, xs = np.where(mask)
    return int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1


def _translation(tx: float, ty: float) -> np.ndarray:
    return np.array([[1, 0, tx], [0, 1, ty], [0, 0, 1]], dtype=np.float64)


def affine_forward(w: int, h: int, rot_deg: float, sx: float, sy: float,
                   shear: float, tx: float, ty: float) -> np.ndarray:
    """Матрица 3×3: поворот, наклон и масштаб вокруг центра кропа, затем сдвиг."""
    th = math.radians(rot_deg)
    rot = np.array([[math.cos(th), -math.sin(th), 0], [math.sin(th), math.cos(th), 0], [0, 0, 1]])
    sh = np.array([[1, shear, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float64)
    sc = np.diag([sx, sy, 1.0])
    return _translation(w / 2 + tx, h / 2 + ty) @ rot @ sh @ sc @ _translation(-w / 2, -h / 2)


def sample_matrix(w: int, h: int, rng: np.random.Generator, cfg: AugmentConfig) -> np.ndarray:
    """Случайное преобразование кропа w×h. Масштабы положительные, поворот ≤ rot_deg: отражений не бывает."""
    rot = float(rng.uniform(-cfg.rot_deg, cfg.rot_deg))
    s = float(rng.uniform(*cfg.scale))
    a = float(rng.uniform(*cfg.aspect))
    shear = float(rng.uniform(-cfg.shear, cfg.shear))
    tx = float(rng.uniform(-cfg.shift_x, cfg.shift_x)) * w
    ty = float(rng.uniform(-cfg.shift_y, cfg.shift_y)) * h
    return affine_forward(w, h, rot, s * math.sqrt(a), s / math.sqrt(a), shear, tx, ty)


def geometric(img: Image.Image, rng: np.random.Generator, cfg: AugmentConfig) -> tuple[Image.Image, bool]:
    """Геометрия на исходном кропе. Возвращает (картинка, применилась ли геометрия)."""
    gray = to_gray_on_white(img)
    w, h = gray.size
    arr = np.asarray(gray)
    bbox = ink_bbox(arr)
    paper = int(np.median(arr))
    for _ in range(GEOMETRY_TRIES):
        m = sample_matrix(w, h, rng, cfg)
        if bbox is not None:
            x0, y0, x1, y1 = bbox
            corners = m @ np.array([[x0, x1, x1, x0], [y0, y0, y1, y1], [1, 1, 1, 1]], dtype=np.float64)
            if (corners[0].min() < 1 or corners[0].max() > w - 1
                    or corners[1].min() < 1 or corners[1].max() > h - 1):
                continue                      # цифры вышли бы за край: пробуем другие параметры
        inv = np.linalg.inv(m)
        coeffs = (inv[0, 0], inv[0, 1], inv[0, 2], inv[1, 0], inv[1, 1], inv[1, 2])
        out = gray.transform((w, h), Image.Transform.AFFINE, coeffs,
                             resample=Image.Resampling.BILINEAR, fillcolor=paper)
        return out, True
    return gray, False


def photometric(arr: np.ndarray, box: tuple, rng: np.random.Generator, cfg: AugmentConfig) -> np.ndarray:
    """Фотометрия на холсте [1, 48, 160]; паддинг в конце снова обнуляется."""
    a = arr[0].copy()
    x0, y0, x1, y1 = box
    a[y0:y1, x0:x1] *= float(rng.uniform(*cfg.gain))
    if rng.random() < cfg.blur_prob:
        im = Image.fromarray((a * 255).round().astype(np.uint8), mode="L")
        a = np.asarray(im.filter(ImageFilter.GaussianBlur(float(rng.uniform(*cfg.blur_radius)))),
                       dtype=np.float32) / 255.0
    if rng.random() < cfg.jpeg_prob:
        buf = io.BytesIO()
        Image.fromarray((a * 255).round().astype(np.uint8), mode="L").save(
            buf, "JPEG", quality=int(rng.integers(cfg.jpeg_quality[0], cfg.jpeg_quality[1] + 1)))
        buf.seek(0)
        a = np.asarray(Image.open(buf).convert("L"), dtype=np.float32) / 255.0
    sigma = float(rng.uniform(0.0, cfg.noise))
    if sigma > 0:
        a = a + rng.normal(0.0, sigma, size=a.shape).astype(np.float32)
    a = np.clip(a, 0.0, 1.0)
    keep = np.zeros(a.shape, dtype=bool)
    keep[y0:y1, x0:x1] = True
    a[~keep] = 0.0                       # паддинг остаётся ровно 0
    return a[None, :, :].astype(np.float32)


def augment(source, rng: np.random.Generator, cfg: AugmentConfig, validate: bool = True) -> np.ndarray:
    """Путь или PIL-образ → аугментированный массив [1, 48, 160]."""
    if validate:
        validate_config(cfg)
    img, _ = geometric(load_image(source), rng, cfg)
    arr = preprocess_array(img)
    return photometric(arr, content_box(*img.size), rng, cfg)


def augment_default(source, rng: np.random.Generator) -> np.ndarray:
    return augment(source, rng, DEFAULT)


def augment_printed(source, rng: np.random.Generator) -> np.ndarray:
    return augment(source, rng, PRINTED)


def with_overrides(cfg: AugmentConfig, **kwargs) -> AugmentConfig:
    return replace(cfg, **kwargs)

experiments/show_augment_grid.py"""Сетка 4×4: исходный кроп и 15 его аугментаций."""
import argparse
import json
import sys
from pathlib import Path

import matplotlib
import numpy as np

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from augment import DEFAULT, PRINTED, augment, validate_config, with_overrides  # noqa: E402
from image_to_tensor import load_image, preprocess_array  # noqa: E402


def parse_overrides(items: list[str]) -> dict:
    out = {}
    for item in items:
        key, _, value = item.partition("=")
        out[key] = float(value)
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--image", help="путь к кропу")
    ap.add_argument("--find-text", help="взять первый кроп из train.jsonl с такой меткой, например 6")
    ap.add_argument("--data-dir", default="data/processed")
    ap.add_argument("--style", choices=["default", "printed"], default="default")
    ap.add_argument("--override", action="append", default=[], metavar="ключ=значение",
                    help="переопределить параметр, например rot_deg=40 или noise=0.1")
    ap.add_argument("--unsafe", action="store_true", help="не проверять правило «поворот ≤ 5°»")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", default="experiments/aug_grid.png")
    args = ap.parse_args()

    if args.find_text:
        data_dir = Path(args.data_dir)
        manifest = json.loads((data_dir / "manifest.json").read_text(encoding="utf-8"))
        with open(data_dir / "train.jsonl", encoding="utf-8") as f:
            rec = next((r for r in map(json.loads, f) if r["text"] == args.find_text), None)
        if rec is None:
            sys.exit(f"В {data_dir / 'train.jsonl'} нет кропа с меткой {args.find_text!r}: попробуйте другую метку.")
        image_path = Path(manifest["image_root"]) / rec["image"]
    elif args.image:
        image_path = Path(args.image)
    else:
        sys.exit("Укажите --image или --find-text")

    cfg = with_overrides(DEFAULT if args.style == "default" else PRINTED, **parse_overrides(args.override))
    img = load_image(image_path)
    rng = np.random.default_rng(args.seed)

    try:
        validate_config(cfg)
    except ValueError as e:
        if not args.unsafe:
            sys.exit(f"Параметры отклонены: {e}.\nДля эксперимента добавьте --unsafe (в обучении так делать нельзя).")

    fig, axes = plt.subplots(4, 4, figsize=(10, 4.2))
    for k, ax in enumerate(axes.flat):
        arr = preprocess_array(img) if k == 0 else augment(img, rng, cfg, validate=False)
        ax.imshow(arr[0], cmap="gray", vmin=0, vmax=1)
        ax.set_title("исходный" if k == 0 else f"#{k}", fontsize=7)
        ax.set_xticks([])
        ax.set_yticks([])
    extra = f"  {' '.join(args.override)}" if args.override else ""
    fig.suptitle(f"{image_path.name}: стиль {args.style}{extra}", fontsize=9)
    fig.tight_layout()
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out, dpi=110)
    print(f"сохранено: {out}")


if __name__ == "__main__":
    main()

tests/test_augment.pyimport sys
from dataclasses import replace
from pathlib import Path

import numpy as np
import pytest
from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from augment import (  # noqa: E402
    DEFAULT, PRINTED, ROT_CAP_DEG, augment, content_box, geometric, ink_bbox, ink_mask, sample_matrix,
    validate_config,
)
from image_to_tensor import CANVAS_H, CANVAS_W, preprocess_array, to_gray_on_white  # noqa: E402
from make_synthetic_crops import generate  # noqa: E402

import json  # noqa: E402
import math  # noqa: E402
import tempfile  # noqa: E402

_cache: dict = {}
_tmp: list = []


def sample_crops() -> list:
    """Value-кропы двух листов (один с рамкой, другой без): (PIL-образ, метка)."""
    if "crops" not in _cache:
        d = tempfile.TemporaryDirectory()
        _tmp.append(d)
        generate(d.name, sheets=8, per_sheet=10, seed=11)
        styles = json.loads((Path(d.name) / "sheets.json").read_text(encoding="utf-8"))
        crops = []
        for line in open(Path(d.name) / "index.jsonl", encoding="utf-8"):
            r = json.loads(line)
            if r["kind"] == "value":
                img = Image.open(Path(d.name) / r["file"])
                img.load()
                crops.append((img, r["text"], styles[r["sheet"]]["border"]))
        _cache["crops"] = crops
    return _cache["crops"]


def text_like(w: int, h: int) -> Image.Image:
    img = Image.new("L", (w, h), 235)
    img.paste(40, (w // 5, h // 4, w // 2, 3 * h // 4))
    img.paste(40, (w // 2 + 4, h // 4, 7 * w // 10, 3 * h // 4))
    return img


def test_content_box_matches_letterbox():
    for w0, h0 in [(40, 50), (100, 40), (200, 50), (380, 70), (160, 48), (30, 30)]:
        dark = preprocess_array(Image.new("L", (w0, h0), 0))[0]       # однородно тёмный → единицы только в содержимом
        x0, y0, x1, y1 = content_box(w0, h0)
        expected = np.zeros((CANVAS_H, CANVAS_W), dtype=bool)
        expected[y0:y1, x0:x1] = True
        assert ((dark > 0.5) == expected).all(), (w0, h0)


def test_shape_dtype_range_for_both_styles():
    for cfg in (DEFAULT, PRINTED):
        arr = augment(text_like(120, 40), np.random.default_rng(0), cfg)
        assert arr.shape == (1, CANVAS_H, CANVAS_W)
        assert arr.dtype == np.float32
        assert arr.min() >= 0.0 and arr.max() <= 1.0


def test_deterministic_for_same_rng_and_different_for_other():
    img = text_like(120, 40)
    a = augment(img, np.random.default_rng([1, 2, 3]), PRINTED)
    b = augment(img, np.random.default_rng([1, 2, 3]), PRINTED)
    c = augment(img, np.random.default_rng([1, 2, 4]), PRINTED)
    assert (a == b).all()
    assert not (a == c).all()


def test_padding_stays_exactly_zero():
    img = text_like(40, 50)                      # узкий кроп: справа много паддинга
    x0, y0, x1, y1 = content_box(40, 50)
    for seed in range(20):
        arr = augment(img, np.random.default_rng(seed), PRINTED)[0]
        assert np.all(arr[:, x1:] == 0.0)
        assert np.all(arr[:y0, :] == 0.0) and np.all(arr[y1:, :] == 0.0)


def test_geometry_never_cuts_digits():
    # Рамку клетки после поворота нельзя отличить от чернил, поэтому границы проверяем на листах без рамки;
    # потерю чернил (резать цифры нельзя) проверяем на всех.
    checked_bbox = checked_mass = 0
    for img, text, border in sample_crops():
        gray = np.asarray(to_gray_on_white(img))
        before = ink_mask(gray)
        for seed in range(15):
            out, applied = geometric(img, np.random.default_rng(seed), PRINTED)
            if not applied:
                continue
            out_gray = np.asarray(out)
            assert ink_mask(out_gray).sum() >= 0.6 * before.sum(), text
            checked_mass += 1
            if border:
                continue
            after = ink_bbox(out_gray)
            h, w = out_gray.shape
            assert after is not None, text
            assert after[0] >= 1 and after[1] >= 1 and after[2] <= w - 1 and after[3] <= h - 1, (text, after)
            checked_bbox += 1
    assert checked_mass > 200 and checked_bbox > 80, (checked_mass, checked_bbox)


def test_borders_do_not_disable_geometry():
    # Линия-рамка клетки не должна считаться «цифрами»: иначе на листах с рамкой геометрия не применялась бы.
    applied = {True: 0, False: 0}
    total = {True: 0, False: 0}
    for img, text, border in sample_crops():
        for seed in range(5):
            _, ok = geometric(img, np.random.default_rng(seed), PRINTED)
            applied[border] += ok
            total[border] += 1
    for border in (True, False):
        assert total[border] > 0
        assert applied[border] / total[border] > 0.8, (border, applied[border] / total[border])


def test_rotation_cap():
    validate_config(DEFAULT)
    validate_config(PRINTED)
    with pytest.raises(ValueError):
        validate_config(replace(DEFAULT, rot_deg=ROT_CAP_DEG + 1))
    with pytest.raises(ValueError):
        augment(text_like(100, 40), np.random.default_rng(0), replace(DEFAULT, rot_deg=30))
    augment(text_like(100, 40), np.random.default_rng(0), replace(DEFAULT, rot_deg=30), validate=False)


def pooled(a: np.ndarray) -> np.ndarray:
    return a.reshape(CANVAS_H // 2, 2, CANVAS_W // 2, 2).mean(axis=(1, 3))


def best_corr(a: np.ndarray, b: np.ndarray) -> float:
    """Лучшая корреляция по небольшим сдвигам: устойчива к сдвигу текста, чувствительна к зеркалу."""
    a, b = pooled(a), pooled(b)
    best = -1.0
    for dy in range(-3, 4):
        for dx in range(-5, 6):
            shifted = np.roll(np.roll(b, dy, axis=0), dx, axis=1)
            best = max(best, float(np.corrcoef(a.ravel(), shifted.ravel())[0, 1]))
    return best


def test_transform_is_never_a_reflection_and_respects_rotation_cap():
    # Переворот по вертикали по корреляции не проверить (цифры почти симметричны), поэтому проверяем матрицу:
    # определитель > 0 и положительный масштаб по x исключают отражения, угол ограничен.
    for cfg in (DEFAULT, PRINTED):
        rng = np.random.default_rng(0)
        for _ in range(2000):
            m = sample_matrix(160, 48, rng, cfg)
            assert np.linalg.det(m[:2, :2]) > 0
            assert m[0, 0] > 0
            assert abs(math.degrees(math.atan2(m[1, 0], m[0, 0]))) <= min(cfg.rot_deg, ROT_CAP_DEG) + 1e-9


def test_no_horizontal_mirror():
    # Аугментированный кроп (с поправкой на сдвиг) в среднем похож на исходный сильнее, чем на его зеркало.
    # Почти симметричные кропы пропускаем: у них зеркало ≈ оригинал.
    checked = 0
    for img, text, _ in sample_crops()[:30]:
        clean = preprocess_array(img)[0]
        if clean.std() < 1e-3:
            continue
        mirrored = np.ascontiguousarray(clean[:, ::-1])
        if best_corr(clean, mirrored) > 0.5:
            continue
        augs = [augment(img, np.random.default_rng(seed), PRINTED)[0] for seed in range(5)]
        same = np.mean([best_corr(a, clean) for a in augs])
        other = np.mean([best_corr(a, mirrored) for a in augs])
        assert same > other, (text, same, other)
        checked += 1
    assert checked >= 10, checked


def test_printed_is_stronger_than_default():
    img = text_like(150, 45)
    clean = preprocess_array(img)[0]
    def mean_change(cfg):
        return np.mean([np.abs(augment(img, np.random.default_rng(s), cfg)[0] - clean).mean() for s in range(60)])
    assert mean_change(PRINTED) > 1.5 * mean_change(DEFAULT)

tests/test_dataset.pyimport subprocess
import sys
import tempfile
from pathlib import Path

import numpy as np
import pytest
from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from alphabet import decode_indices  # noqa: E402
from augment import augment_default, augment_printed, content_box  # noqa: E402
from datasets.crop_dataset import (  # noqa: E402
    EMPTY_CLASS, VALUE_CLASS, CropDataset, collate_classify_np, collate_ctc_np,
)
from image_to_tensor import preprocess_array  # noqa: E402

_cache: dict = {}
_tmp: list = []


def run(*args) -> None:
    subprocess.run([sys.executable, *map(str, args)], cwd=ROOT, check=True, capture_output=True)


def data_dir() -> Path:
    """Маленький датасет (16 листов × 10 кропов), собранный теми же командами, что в уроках 4–5."""
    if "dir" not in _cache:
        d = tempfile.TemporaryDirectory()
        _tmp.append(d)
        base = Path(d.name)
        run("src/make_synthetic_crops.py", "--out", base / "syn", "--sheets", 16, "--per-sheet", 10, "--seed", 3)
        run("src/prepare_dataset.py", "--source", "synthetic", "--src", base / "syn", "--out", base / "proc")
        run("src/split_groups.py", "--src", base / "proc", "--out", base / "proc")
        _cache["dir"] = base / "proc"
    return _cache["dir"]


def test_lengths_match_split_files():
    d = data_dir()
    for split in ("train", "val", "test"):
        n_lines = sum(1 for _ in open(d / f"{split}.jsonl", encoding="utf-8"))
        assert len(CropDataset(d, split)) == n_lines


def test_no_group_overlap_through_dataset():
    d = data_dir()
    groups = {s: {r["group_id"] for r in CropDataset(d, s).records} for s in ("train", "val", "test")}
    assert groups["train"].isdisjoint(groups["val"])
    assert groups["train"].isdisjoint(groups["test"])
    assert groups["val"].isdisjoint(groups["test"])


def test_item_fields_and_shapes():
    ds = CropDataset(data_dir(), "train")
    for i in range(min(10, len(ds))):
        item = ds[i]
        assert item["image"].shape == (1, 48, 160) and item["image"].dtype == np.float32
        assert item["target"].dtype == np.int64
        assert decode_indices(item["target"].tolist(), ds.task) == item["text"]
        assert item["label"] == (EMPTY_CLASS if item["is_empty"] else VALUE_CLASS)
        assert item["is_empty"] == (item["text"] == "")


def test_without_augment_equals_plain_preprocess():
    ds = CropDataset(data_dir(), "val")
    rec = ds.records[0]
    assert (ds[0]["image"] == preprocess_array(ds.image_root / rec["image"])).all()


def test_augmentation_is_reproducible_and_changes_with_epoch():
    d = data_dir()
    a = CropDataset(d, "train", augment=augment_printed, seed=5)
    b = CropDataset(d, "train", augment=augment_printed, seed=5)
    assert (a[3]["image"] == b[3]["image"]).all()                 # тот же seed, эпоха, индекс
    first = a[3]["image"].copy()
    a.set_epoch(1)
    assert not (a[3]["image"] == first).all()                     # другая эпоха — другая аугментация
    c = CropDataset(d, "train", augment=augment_printed, seed=6)
    assert not (c[3]["image"] == first).all()                     # другой seed


def test_augmented_padding_is_zero_and_label_unchanged():
    d = data_dir()
    ds = CropDataset(d, "train", augment=augment_default, seed=0)
    plain = CropDataset(d, "train")
    for i in range(20):
        assert ds[i]["text"] == plain[i]["text"]
        assert ds[i]["target"].tolist() == plain[i]["target"].tolist()
        with Image.open(ds.image_root / ds.records[i]["image"]) as raw:
            x0, y0, x1, y1 = content_box(*raw.size)
        image = ds[i]["image"][0]
        assert np.all(image[:, x1:] == 0.0) and np.all(image[:y0, :] == 0.0) and np.all(image[y1:, :] == 0.0)


def test_collate_ctc_np():
    ds = CropDataset(data_dir(), "train")
    batch = [ds[i] for i in range(12)]
    out = collate_ctc_np(batch)
    assert out["images"].shape == (12, 1, 48, 160)
    assert out["targets"].dtype == np.int64
    assert out["target_lengths"].sum() == len(out["targets"])
    for b, length in zip(batch, out["target_lengths"]):
        assert length == len(b["text"])
    assert any(length == 0 for length in out["target_lengths"]), "в батче нет пустых клеток"


def test_collate_classify_np():
    ds = CropDataset(data_dir(), "train")
    batch = [ds[i] for i in range(12)]
    out = collate_classify_np(batch)
    assert out["labels"].tolist() == [int(b["is_empty"]) for b in batch]
    assert out["images"].shape == (12, 1, 48, 160)


def test_torch_collate_and_dataloader():
    torch = pytest.importorskip("torch")
    from datasets.crop_dataset import collate_ctc
    from seed import make_generator, seed_worker
    ds = CropDataset(data_dir(), "train", augment=augment_default, seed=0)
    loader = torch.utils.data.DataLoader(ds, batch_size=8, shuffle=True, collate_fn=collate_ctc,
                                         generator=make_generator(0), worker_init_fn=seed_worker)
    batch = next(iter(loader))
    assert tuple(batch["images"].shape) == (8, 1, 48, 160)
    assert batch["images"].dtype == torch.float32
    assert batch["targets"].dtype == torch.int64
    assert int(batch["target_lengths"].sum()) == len(batch["targets"])