Blame for test_training.py

"""Small checks for mistakes that would invalidate training or evaluation."""

import json
import tempfile
import unittest
from pathlib import Path

import torch
from PIL import Image

from data import (
    FolderDataset,
    GridDataset,
    SupervisedGrids,
    prepare,
    prepare_jsonl,
    prepare_folders,
    split_records,
    target_name,
    transform,
)
from models import GridModel, select_target, transfer_backbone
from train import (
    auxiliary_records,
    choose_threshold,
    grid_logits,
    metrics,
    partial_loss,
    thresholds_for,
)


class TrainingTests(unittest.TestCase):
    def test_jsonl_dataset_preserves_splits_and_loads_images(self):
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary) / "dataset"
            for grid in ("3x3", "4x4"):
                for split in ("train", "validation", "test", "review", "excluded"):
                    folder = root / "data" / grid / split
                    (folder / "images").mkdir(parents=True)
                    Image.new("RGB", (30, 30), "red").save(folder / "images" / "grid.jpg")
                    row = {"id": f"{grid}-{split}", "grid": grid, "target": "buses",
                           "file_name": "images/grid.jpg", "correct_answers": [0],
                           "uncertain_cells": [1], "machine_answers": [2]}
                    (folder / "metadata.jsonl").write_text(json.dumps(row) + "\n")
            manifest, summary = prepare_jsonl(root, Path(temporary) / "output")
            self.assertEqual(len(manifest["records"]), 8)
            self.assertEqual(summary["splits"]["excluded"], 2)
            record = next(r for r in manifest["records"] if r["id"] == "3x3-validation")
            self.assertEqual((record["split"], record["target"], record["uncertain"]),
                             ("val", "bus", [1]))
            sample = GridDataset(manifest["root"], [record], ["bus"], 16)[0]
            self.assertEqual(tuple(sample[0].shape), (9, 3, 16, 16))
            review = next(r for r in manifest["records"] if r["id"] == "4x4-review")
            self.assertEqual(review["labeler_answers"], [2])

    def test_enlarged_views_apply_only_to_requested_targets(self):
        class MeanModel(torch.nn.Module):
            tiles = 9
            crop_target_ids = [0]

            def forward(self, images):
                mean = images.mean((1, 2, 3))
                return torch.stack([mean, mean], dim=1)

        images = torch.zeros(2, 1, 3, 8, 8)
        images[:, :, :, :4, :4] = 1
        output = grid_logits(MeanModel(), images, torch.tensor([0, 1]))
        self.assertTrue(torch.equal(output, torch.tensor([[1.0], [0.25]])))

    def test_auxiliary_training_and_validation_stay_separate(self):
        records = [
            {"id": str(i), "tiles": 16, "target": "motorcycle", "split": split}
            for i, split in enumerate(["train", "val", "test"])
        ]
        self.assertEqual(
            [r["id"] for r in auxiliary_records({"records": records}, ["motorcycle"], "train")],
            ["0"],
        )
        self.assertEqual(auxiliary_records({"records": records}, ["bus"], "train"), [])

    def test_auxiliary_crops_and_full_image_labels(self):
        with tempfile.TemporaryDirectory() as root:
            directory = Path(root) / "example"
            directory.mkdir()
            for name in ["full_images.png"] + [f"tile_{i}.png" for i in range(16)]:
                Image.new("RGB", (16, 16), "blue").save(directory / name)
            records = [{"id": "example", "tiles": 16, "target": "motorcycle", "answers": [0, 15]}]
            training = GridDataset(
                root, records, ["motorcycle"], 16, tiles=True, include_full_images=True
            )
            self.assertEqual(len(training), 17)
            self.assertEqual([training[i][2].item() for i in [0, 1, 15, 16]], [1, 0, 1, 1])
            validation = GridDataset(root, records, ["motorcycle"], 16, tile_grids=True)
            self.assertEqual(tuple(validation[0][0].shape), (16, 3, 16, 16))
            self.assertEqual(validation[0][2].sum().item(), 2)

    def test_per_target_cutoffs_and_legacy_checkpoint(self):
        checkpoint = {
            "classes": ["bus", "motorcycle"],
            "threshold": 0.5,
            "thresholds": {"motorcycle": 0.3},
        }
        cutoffs = thresholds_for(checkpoint, torch.tensor([0, 1]))
        result = metrics(torch.tensor([[0.4], [0.4]]), torch.tensor([[False], [True]]), cutoffs)
        self.assertEqual(result["exact_grid_accuracy"], 1.0)
        checkpoint.pop("thresholds")
        self.assertTrue(
            torch.equal(
                thresholds_for(checkpoint, torch.tensor([0, 1])), torch.tensor([[0.5], [0.5]])
            )
        )

    def test_conflicting_duplicates_are_excluded_without_editing_labels(self):
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary) / "dataset"
            for i in range(3):
                directory = root / str(i)
                directory.mkdir(parents=True)
                info = {"keyword": "buses", "correct_answers": [i]}
                (directory / "info.json").write_text(json.dumps(info))
                Image.new("RGB", (30, 30), "red" if i < 2 else "blue").save(
                    directory / "full_images.png"
                )
                for tile in range(9):
                    Image.new("RGB", (10, 10), (i * 70, tile * 20, 0)).save(
                        directory / f"tile_{tile}.png"
                    )
            manifest, _ = prepare(root, Path(temporary) / "output")
            self.assertEqual([r["id"] for r in manifest["records"]], ["2"])
            self.assertEqual(len(manifest["excluded_conflicting_labels"]), 2)
            self.assertEqual(
                json.loads((root / "0" / "info.json").read_text())["correct_answers"], [0]
            )
            (root / "2" / "info.json").write_text(
                json.dumps({"keyword": "bus", "correct_answers": [9]})
            )
            with self.assertRaisesRegex(ValueError, "Invalid answers"):
                prepare(root, Path(temporary) / "output")

    def test_transitive_duplicate_groups_never_leak(self):
        records = [
            {"id": str(i), "tiles": 9, "target": "bus", "hashes": [str(i)]} for i in range(30)
        ]
        records[0]["hashes"].append("shared-a")
        records[1]["hashes"].extend(["shared-a", "shared-b"])
        records[2]["hashes"].append("shared-b")
        split_records(records, 42)
        self.assertEqual(len({r["split"] for r in records[:3]}), 1)
        self.assertEqual({r["split"] for r in records}, {"train", "val", "test"})
        before = [r["split"] for r in records]
        split_records(records, 42)
        self.assertEqual(before, [r["split"] for r in records])

    def test_singleton_is_training_only(self):
        records = [{"tiles": 16, "target": "tractor", "hashes": ["x"]}]
        self.assertEqual(split_records(records, 42)[0]["split"], "train")

    def test_aliases(self):
        self.assertEqual(target_name("buses"), target_name("bus"))
        self.assertEqual(target_name("a fire hydrant"), target_name("fire hydrants"))

    def test_unknown_targets_have_no_gradient(self):
        logits = torch.zeros(2, 3, requires_grad=True)
        loss = torch.nn.functional.binary_cross_entropy_with_logits(
            select_target(logits, torch.tensor([0, 2])), torch.tensor([1.0, 0.0])
        )
        loss.backward()
        self.assertTrue(
            torch.equal(
                logits.grad != 0, torch.tensor([[True, False, False], [False, False, True]])
            )
        )

    def test_exact_grid_requires_every_cell_correct(self):
        probabilities = torch.tensor([[0.9, 0.1], [0.9, 0.8]])
        truth = torch.tensor([[True, False], [True, False]])
        result = metrics(probabilities, truth, 0.5)
        self.assertEqual(result["exact_grid_accuracy"], 0.5)
        self.assertEqual(result["cell_accuracy"], 0.75)
        self.assertGreater(choose_threshold(probabilities, truth), 0.8)

    def test_row_major_target_selection(self):
        logits = torch.arange(2 * 3 * 4 * 4).reshape(2, 3, 4, 4)
        selected = select_target(logits.flatten(2), torch.tensor([1, 2]))
        self.assertEqual(selected[0].tolist(), list(range(16, 32)))
        self.assertEqual(selected[1].tolist(), list(range(80, 96)))

    def test_models_have_expected_shapes(self):
        torch.set_num_threads(2)
        with torch.no_grad():
            for tiles, size in [(9, 224), (16, 384)]:
                model = GridModel(tiles, 3, pretrained=False).eval()
                shape = (1, 9, 3, size, size) if tiles == 9 else (1, 3, size, size)
                self.assertEqual(
                    tuple(grid_logits(model, torch.zeros(shape), torch.tensor([1])).shape),
                    (1, tiles),
                )
                if tiles == 9:
                    self.assertEqual(
                        tuple(
                            grid_logits(
                                model, torch.zeros(1, 16, 3, size, size), torch.tensor([1])
                            ).shape
                        ),
                        (1, 16),
                    )


class FolderDataTests(unittest.TestCase):
    def test_folder_names_map_to_targets(self):
        self.assertEqual(target_name("Hydrant"), "fire hydrant")
        self.assertEqual(target_name("Traffic Light"), "traffic light")
        self.assertEqual(target_name("Stair"), "stairs")
        self.assertEqual(target_name("Palm"), target_name("palm trees"))

    def test_folders_merge_copies_group_near_copies_and_drop_old_heldout(self):
        with tempfile.TemporaryDirectory() as temporary:
            temporary = Path(temporary)
            old = temporary / "dataset" / "grid"
            old.mkdir(parents=True)
            held_out = Image.linear_gradient("L").convert("RGB").resize((64, 64))
            for name in ["full_images.png"] + [f"tile_{i}.png" for i in range(9)]:
                held_out.save(old / name)
            (temporary / "split.json").write_text(
                json.dumps(
                    {
                        "root": str(temporary / "dataset"),
                        "records": [{"id": "grid", "tiles": 9, "split": "test"}],
                    }
                )
            )
            root = temporary / "folders"
            for split in ("train", "val"):
                for name in ("Bus", "Car", "Other"):
                    (root / split / name).mkdir(parents=True)
            Image.new("RGB", (100, 100), "red").save(root / "train" / "Bus" / "a.png")
            Image.new("RGB", (100, 100), "red").save(root / "val" / "Bus" / "copy.png")
            Image.new("RGB", (100, 100), "red").save(root / "train" / "Car" / "also.png")
            held_out.resize((120, 120)).rotate(180).rotate(180).save(
                root / "train" / "Car" / "old.png"
            )
            for i in range(20):
                image = Image.new("RGB", (100, 100), (0, 0, 0))
                image.paste((255, 255, 255), (0, 0, 5 * (i + 1), 100))
                image.save(root / "train" / "Car" / f"{i}.png")
                image.resize((101, 101)).save(root / "val" / "Car" / f"near{i}.png")
            Image.new("RGB", (450, 450), "blue").save(root / "train" / "Other" / "full.png")
            manifest, summary = prepare_folders(
                root, temporary / "split.json", temporary / "out", workers=2
            )
            self.assertEqual(summary["exact_duplicates_merged"], 2)
            self.assertEqual(manifest["dropped_matching_old_heldout"], ["train/Car/old.png"])
            merged = [r for r in manifest["records"] if len(r["targets"]) > 1]
            self.assertEqual([r["targets"] for r in merged], [["bus", "car"]])
            splits = {}
            for record in manifest["records"]:
                splits.setdefault(record["hashes"][1], set()).add(record["split"])
            self.assertTrue(all(len(v) == 1 for v in splits.values()))
            self.assertEqual({r["split"] for r in manifest["records"]}, {"train", "val"})
            self.assertEqual(
                [r["kind"] for r in manifest["records"] if r["targets"] == ["other"]], ["full"]
            )

    def test_three_folder_layouts_and_glued_yolo_lines(self):
        with tempfile.TemporaryDirectory() as temporary:
            temporary = Path(temporary)
            (temporary / "split.json").write_text(
                json.dumps({"root": str(temporary), "records": []})
            )
            images = {
                "a/train/Bus/1.png": (10, 0, 0),
                "a/val/TLight/2.png": (0, 10, 0),
                "b/Stair/3.png": (0, 0, 10),
                "b/Bus/1.png": (10, 0, 0),
                "c/images/Chimney/4.png": (50, 0, 0),
                "c/images/Stair/5.png": (0, 50, 0),
                "c/images/Stair/6.png": (0, 0, 50),
            }
            for name, color in images.items():
                (temporary / name).parent.mkdir(parents=True, exist_ok=True)
                Image.new("RGB", (100, 100), color).save(temporary / name)
            labels = temporary / "c" / "labels"
            (labels / "Chimney").mkdir(parents=True)
            (labels / "Stair").mkdir(parents=True)
            (labels / "Chimney" / "4.txt").write_text(
                "1 0.500000 0.500000 0.100000 0.1000002 0.300000 0.300000 0.200000 0.200000\n"
            )
            (labels / "Stair" / "5.txt").write_text("2 0.500000 0.500000 0.100000 0.100000\n")
            (labels / "Stair" / "6.txt").write_text("2 0.400000 0.500000 0.100000 0.100000\n")
            manifest, summary = prepare_folders(
                [temporary / "a", temporary / "b", temporary / "c"],
                temporary / "split.json",
                temporary / "out",
                workers=2,
            )
            targets = {r["path"]: r["targets"] for r in manifest["records"]}
            self.assertEqual(targets["a/val/TLight/2.png"], ["traffic light"])
            self.assertEqual(targets["c/images/Chimney/4.png"], ["chimney", "stairs"])
            self.assertEqual(targets["c/images/Stair/6.png"], ["stairs"])
            self.assertEqual(summary["yolo_class_ids"], {"1": "chimney", "2": "stairs"})
            self.assertEqual(summary["exact_duplicates_merged"], 1)
            self.assertEqual(summary["unique_images_by_first_source"], {"a": 2, "b": 1, "c": 3})
            self.assertTrue(all((temporary / path).exists() for path in targets))

    def test_folder_labels_are_weak_negatives_except_background(self):
        with tempfile.TemporaryDirectory() as root:
            Image.new("RGB", (8, 8)).save(Path(root) / "x.png")
            records = [
                {"path": "x.png", "targets": ["bus"]},
                {"path": "x.png", "targets": ["other"]},
            ]
            dataset = FolderDataset(root, records, ["bus", "car"], 8, cells=16, negative_weight=0.2)
            _, _, cell_weights, labels, weights = dataset[0]
            self.assertEqual(labels.tolist(), [1, 0])
            self.assertEqual([round(w, 3) for w in weights.tolist()], [1, 0.2])
            self.assertEqual(cell_weights.sum().item(), 0)
            self.assertEqual(tuple(cell_weights.shape), (2, 16))
            _, _, _, labels, weights = dataset[1]
            self.assertEqual((labels.tolist(), weights.tolist()), ([0, 0], [1, 1]))

    def test_answer_key_items_supervise_only_the_queried_target(self):
        with tempfile.TemporaryDirectory() as root:
            directory = Path(root) / "example"
            directory.mkdir()
            for name in ["full_images.png"] + [f"tile_{i}.png" for i in range(16)]:
                Image.new("RGB", (16, 16), "blue").save(directory / name)
            records = [{"id": "example", "tiles": 16, "target": "car", "answers": [0, 15]}]
            grids = SupervisedGrids(GridDataset(root, records, ["bus", "car"], 16), 16)
            _, labels, weights, _, any_weights = grids[0]
            self.assertEqual(weights.sum(1).tolist(), [0, 16])
            self.assertEqual(labels[1].nonzero().flatten().tolist(), [0, 15])
            self.assertEqual(any_weights.sum().item(), 0)
            tiles = SupervisedGrids(GridDataset(root, records, ["bus", "car"], 16, tiles=True), 1)
            self.assertEqual(
                (tiles[15][1].tolist(), tiles[15][2].tolist()), ([[0], [1]], [[0], [1]])
            )

    def test_partial_loss_matches_old_loss_and_ignores_unknowns(self):
        logits = torch.randn(2, 3, 16, requires_grad=True)
        labels, weights = torch.zeros(2, 3, 16), torch.zeros(2, 3, 16)
        labels[0, 1, :4], weights[0, 1], weights[1, 2] = 1, 1, 1
        zero = torch.zeros(2, 3)
        loss, count = partial_loss(logits, labels, weights, zero, zero)
        expected = torch.nn.functional.binary_cross_entropy_with_logits(
            select_target(logits, torch.tensor([1, 2])), labels[[0, 1], [1, 2]]
        )
        self.assertTrue(torch.allclose(loss, expected))
        self.assertEqual(count.item(), 32)
        loss.backward()
        self.assertTrue(
            torch.equal(
                (logits.grad != 0).any(-1),
                torch.tensor([[False, True, False], [False, False, True]]),
            )
        )

    def test_image_level_label_trains_only_the_strongest_cell(self):
        logits = torch.zeros(1, 2, 16)
        logits[0, 0, 5] = 2
        logits.requires_grad_()
        empty = torch.zeros(1, 2, 16)
        loss, _ = partial_loss(
            logits, empty, empty, torch.tensor([[1.0, 0.0]]), torch.tensor([[1.0, 0.0]])
        )
        loss.backward()
        self.assertEqual(logits.grad.nonzero().tolist(), [[0, 0, 5]])

    def test_tile_backbone_initializes_grid_model(self):
        tile_model, grid_model = (
            GridModel(9, 3, pretrained=False),
            GridModel(16, 5, pretrained=False),
        )
        transfer_backbone(grid_model, tile_model.state_dict())
        for name, value in grid_model.backbone.state_dict().items():
            self.assertTrue(torch.equal(value, tile_model.backbone.state_dict()[name]), name)

    def test_torch_free_preprocessing_matches_training(self):
        try:
            from predict_onnx import preprocess
        except ImportError:
            self.skipTest("onnxruntime is not installed (uv sync --group export)")
        image = Image.effect_mandelbrot((97, 97), (-2, -1.5, 1, 1.5), 50).convert("RGB")
        expected = transform(64)(image).numpy()
        self.assertLess(abs(preprocess(image, 64) - expected).max(), 1e-5)


class LabeledLayoutTests(unittest.TestCase):
    def test_invalid_metadata_and_sam_files_are_skipped(self):
        from data import scan

        with tempfile.TemporaryDirectory() as root:
            self.make(root, "valid", "3x3", [0])
            self.make(root, "bad-target", "3x3", [0])
            self.make(root, "bad-metadata", "3x3", [0])
            self.make(root, "bad-sam", "4x4", [0])
            (Path(root) / "bad-target" / "metadata.json").write_text(
                json.dumps({"target": None, "grid": "3x3", "correct_answers": [0]})
            )
            (Path(root) / "bad-metadata" / "metadata.json").write_text("[]")
            (Path(root) / "bad-sam" / "sam.json").write_text("{}")
            records, skipped = scan(root, workers=2, sam=True)
            self.assertEqual([record["id"] for record in records], ["valid"])
            self.assertEqual(
                skipped,
                {
                    "invalid target": 1,
                    "invalid metadata.json": 1,
                    "invalid sam.json": 1,
                },
            )

    def test_folder_preparation_excludes_tiles_from_labeled_holdouts(self):
        from data import grid_boxes, read_rgb, scan

        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            data = root / "data"
            data.mkdir()
            self.make(data, "holdout", "3x3", [0])
            records, _ = scan(data, workers=2)
            records[0]["split"] = "test"
            manifest = root / "split.json"
            manifest.write_text(json.dumps({"root": str(data), "records": records}))
            folder = root / "folders" / "bus"
            folder.mkdir(parents=True)
            image = read_rgb(data / "holdout" / "image.jpg")
            image.crop(grid_boxes(*image.size, 9)[0]).save(folder / "reused.png")
            Image.effect_mandelbrot((100, 100), (-2, -1.5, 1, 1.5), 50).save(folder / "new.png")
            prepared, summary = prepare_folders(
                root / "folders",
                manifest,
                root / "out",
                workers=2,
            )
            self.assertEqual(summary["dropped_matching_old_heldout"], 1)
            self.assertEqual([record["path"] for record in prepared["records"]], ["bus/new.png"])

    def make(self, root, name, grid, answers=None, uncertain=(), colors=None, size=300):
        folder = Path(root) / name
        folder.mkdir()
        n = 3 if grid == "3x3" else 4
        image = Image.new("RGB", (size, size))
        for i, color in enumerate(colors or [(i * 20, 0, 0) for i in range(n * n)]):
            r, c = divmod(i, n)
            image.paste(
                color, (c * size // n, r * size // n, (c + 1) * size // n, (r + 1) * size // n)
            )
        image.save(folder / "image.jpg", quality=100)
        info = {"target": "buses", "grid": grid}
        if answers is not None:
            info.update(correct_answers=answers, uncertain_cells=list(uncertain))
        (folder / "metadata.json").write_text(json.dumps(info))

    def test_labeled_items_are_scanned_and_unlabeled_skipped(self):
        from data import scan

        with tempfile.TemporaryDirectory() as root:
            self.make(root, "a", "3x3", [1, 4], [4])
            self.make(root, "b", "4x4", [5], size=450)
            self.make(root, "c", "3x3")
            records, skipped = scan(root, workers=2)
            self.assertEqual(
                [(r["id"], r["tiles"], r["target"], r["answers"], r["uncertain"]) for r in records],
                [("a", 9, "bus", [1, 4], [4]), ("b", 16, "bus", [5], [])],
            )
            self.assertEqual(skipped, {"not labeled": 1})
            self.assertEqual(len(records[0]["hashes"]), 10)

    def test_tiles_come_from_the_single_image_and_uncertain_cells_are_ignored(self):
        from data import scan

        with tempfile.TemporaryDirectory() as root:
            colors = [(250, 0, 0), (0, 250, 0), (0, 0, 250)] + [(0, 0, 0)] * 6
            self.make(root, "a", "3x3", [1], [2], colors=colors)
            self.make(root, "b", "4x4", [5, 6], [6], size=450)
            records, _ = scan(root, workers=2)
            tiles = GridDataset(root, [records[0]], ["bus"], 8, tiles=True)
            self.assertEqual(len(tiles), 9)
            self.assertGreater(
                tiles[1][0][1].mean().item(), tiles[1][0][0].mean().item()
            )  # green tile
            items = SupervisedGrids(tiles, 1)
            self.assertEqual(items[1][1].tolist(), [[1.0]])
            self.assertEqual(items[2][2].sum().item(), 0)  # uncertain tile: no weight
            self.assertEqual(items[0][2].sum().item(), 1)
            grid = SupervisedGrids(GridDataset(root, [records[1]], ["bus"], 32), 16)
            _, labels, weights, _, _ = grid[0]
            self.assertEqual(labels[0].nonzero().flatten().tolist(), [5, 6])
            self.assertEqual(weights[0].tolist(), [1.0] * 6 + [0.0] + [1.0] * 9)
            kept = SupervisedGrids(
                GridDataset(root, [records[1]], ["bus"], 32), 16, keep_uncertain=True
            )
            self.assertEqual(kept[0][2].sum().item(), 16)

    def test_sam_labels_replace_4x4_answers_only(self):
        from data import scan

        with tempfile.TemporaryDirectory() as root:
            self.make(root, "a", "3x3", [1, 4], [4])
            self.make(root, "b", "4x4", [5], [6], size=450)
            self.make(root, "c", "4x4", size=450)  # Qwen never labeled it; SAM did
            self.make(root, "d", "4x4", [2], size=450)  # no sam.json
            for name, cells in (("b", [5, 6, 9]), ("c", [0])):
                (Path(root) / name / "sam.json").write_text(
                    json.dumps({"model": "facebook/sam3", "correct_answers": cells})
                )
            records, skipped = scan(root, workers=2, sam=True)
            self.assertEqual(
                [(r["id"], r["answers"], r["uncertain"], r["label_source"]) for r in records],
                [
                    ("a", [1, 4], [4], None),
                    ("b", [5, 6, 9], [], "facebook/sam3:sam"),
                    ("c", [0], [], "facebook/sam3:sam"),
                ],
            )
            self.assertEqual(skipped, {"no sam.json": 1})
            records, _ = scan(root, workers=2)
            self.assertEqual([r["answers"] for r in records], [[1, 4], [5], [2]])

    def test_reviewed_grids_are_held_out_with_their_reviewed_answers(self):
        from data import prepare

        with tempfile.TemporaryDirectory() as root, tempfile.TemporaryDirectory() as out:
            for i in range(20):
                # g01 is a copy of the reviewed g00, so it must not reach training.
                j = 0 if i == 1 else i
                colors = [((j * 37 + k * 11) % 256, (k * 23) % 256, j % 256) for k in range(9)]
                self.make(root, f"g{i:02d}", "3x3", [0], colors=colors)
            key = Path(out) / "key.json"
            key.write_text(json.dumps({"items": {"g00": {"answers": [0, 8]}}}))
            manifest, counts = prepare(root, out, workers=2, review_keys=[key])
            by_id = {r["id"]: r for r in manifest["records"]}
            self.assertEqual(
                (by_id["g00"]["split"], by_id["g00"]["answers"], by_id["g00"]["labeler_answers"]),
                ("review", [0, 8], [0]),
            )
            self.assertNotIn("g01", by_id)
            self.assertEqual(counts["dropped_sharing_images_with_review"], 1)
            self.assertNotIn(
                "review", {r["split"] for r in manifest["records"] if r["id"] != "g00"}
            )

    def test_flip_mirrors_4x4_cells_and_uncertain_cells(self):
        import random

        from data import scan

        with tempfile.TemporaryDirectory() as root:
            colors = [(250, 0, 0)] + [(0, 0, 0)] * 15
            self.make(root, "b", "4x4", [0, 5], [1], colors=colors, size=448)
            records, _ = scan(root, workers=2)
            grids = GridDataset(root, records, ["bus"], 32)
            plain = SupervisedGrids(grids, 16)[0]
            seed = next(
                s for s in range(50) if random.Random(s).random() < 0.5
            )  # a seed whose first draw flips
            random.seed(seed)
            flipped = SupervisedGrids(grids, 16, flip=True)[0]
            _, labels, weights, _, _ = flipped
            self.assertEqual(labels[0].nonzero().flatten().tolist(), [3, 6])
            self.assertEqual(
                weights[0].tolist(), [1.0] * 2 + [0.0] + [1.0] * 13
            )  # uncertain 1 -> 2
            self.assertTrue(torch.equal(flipped[0], plain[0].flip(-1)))

    def test_reused_tiles_keep_grids_in_one_split(self):
        from data import scan

        with tempfile.TemporaryDirectory() as root:
            for i in range(40):
                colors = [((i * 37 + k * 11) % 256, (k * 23) % 256, i % 256) for k in range(9)]
                if i in (0, 1):
                    colors[4] = (255, 255, 255)  # the same photo reused in two grids
                self.make(root, f"g{i:02d}", "3x3", [0], colors=colors)
            records, _ = scan(root, workers=2)
            split_records(records, 42)
            self.assertEqual(records[0]["split"], records[1]["split"])


class LayoutTests(unittest.TestCase):
    def test_folder_metadata_takes_priority_over_image_size(self):
        from train import layout

        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            Image.new("RGB", (300, 300)).save(root / "image.jpg")
            (root / "metadata.json").write_text(json.dumps({"grid": "4x4"}))
            self.assertEqual(layout(root), 16)

    def test_layout_from_size_or_seams(self):
        from train import layout

        with tempfile.TemporaryDirectory() as root:
            root = Path(root)
            Image.linear_gradient("L").convert("RGB").resize((300, 300)).save(root / "a.png")
            Image.linear_gradient("L").convert("RGB").resize((450, 450)).save(root / "b.png")
            tiles = Image.new("RGB", (600, 600))
            for i in range(9):
                tiles.paste(
                    (i * 28, 255 - i * 28, (i * 90) % 256),
                    ((i % 3) * 200, (i // 3) * 200, (i % 3 + 1) * 200, (i // 3 + 1) * 200),
                )
            tiles.save(root / "c.png")
            Image.linear_gradient("L").convert("RGB").resize((600, 600)).save(root / "d.png")
            self.assertEqual([layout(root / f"{n}.png") for n in "abcd"], [9, 16, 9, 16])


if __name__ == "__main__":
    unittest.main()