import pytest
from PIL import Image

from datasets.tile_classification import YoloTileClassificationDataset, tile_label
from datasets.yolo import write_data_yaml


def test_label_uses_human_class_ids_and_empty_background():
    assert tile_label("", 9, 4) == 4
    assert tile_label("1 0.5 0.5 0.2 0.2", 9, 4) == 1
    assert tile_label("1 0.5 0.5 0.2 0.2\n3 0.5 0.5 0.4 0.4", 9, 4) == 3
    with pytest.raises(ValueError):
        tile_label("1 0.5 0.5 nan 0.2", 9, 4)


def test_existing_split_is_preserved_and_no_resize(tmp_path):
    for split in ("train", "val"):
        (tmp_path / "images" / split).mkdir(parents=True)
        (tmp_path / "labels" / split).mkdir(parents=True)
        Image.new("RGB", (16, 16)).save(tmp_path / "images" / split / "sample.png")
        (tmp_path / "labels" / split / "sample.txt").write_text(
            "1 0.5 0.5 0.2 0.2" if split == "train" else "")
    yaml = write_data_yaml(tmp_path / "data.yaml", ["normal", "blackdot"], root=tmp_path)
    kwargs = {"yaml_path": yaml, "image_size": [16, 16], "transforms": [{"name": "ToTensor"}]}
    train = YoloTileClassificationDataset(split="train", **kwargs)
    val = YoloTileClassificationDataset(split="val", **kwargs)
    assert train.class_counts() == [0, 1]
    assert val.class_counts() == [1, 0]
    assert train[0]["image"].shape == (3, 16, 16)
    Image.new("RGB", (12, 16)).save(train.samples[0][0])
    with pytest.raises(ValueError, match="resize is forbidden"):
        train[0]


def test_missing_labels_are_not_silent_background(tmp_path):
    (tmp_path / "images" / "train").mkdir(parents=True)
    Image.new("RGB", (16, 16)).save(tmp_path / "images" / "train" / "sample.png")
    yaml = write_data_yaml(tmp_path / "data.yaml", ["normal", "blackdot"], root=tmp_path)
    with pytest.raises(FileNotFoundError, match="Missing label"):
        YoloTileClassificationDataset(yaml, "train")