"""无标注数据集与聚类相关的预处理。"""

import pytest

pytest.importorskip("torch")
pytest.importorskip("PIL")

import numpy as np  # noqa: E402
import torch  # noqa: E402
from PIL import Image  # noqa: E402

from datasets.build import build_dataset  # noqa: E402
from datasets.transforms import build_transforms  # noqa: E402
from datasets.unlabeled import UnlabeledImageDataset, shape_features, unlabeled_collate  # noqa: E402

TRANSFORMS = [
    {"name": "Resize", "size": [32, 32]},
    {"name": "ToTensor"},
    {"name": "PerImageStandardize"},
]


@pytest.fixture
def image_root(tmp_path):
    """尺寸各异、含 RGBA 与灰度、还有一个要跳过的子目录。"""
    (tmp_path / "a").mkdir()
    (tmp_path / "b" / "deep").mkdir(parents=True)
    (tmp_path / "images").mkdir()

    Image.fromarray(np.zeros((113, 152, 3), dtype=np.uint8)).save(tmp_path / "a" / "1.png")
    Image.fromarray(np.zeros((60, 900, 4), dtype=np.uint8), mode="RGBA").save(tmp_path / "a" / "2.png")
    Image.fromarray(np.zeros((80, 40), dtype=np.uint8)).save(tmp_path / "b" / "deep" / "3.png")
    Image.fromarray(np.zeros((10, 10, 3), dtype=np.uint8)).save(tmp_path / "images" / "skip.png")
    (tmp_path / "a" / "notes.txt").write_text("忽略", encoding="utf-8")
    return tmp_path


def test_recursive_scan_skips_excluded_dirs_and_non_images(image_root):
    dataset = UnlabeledImageDataset(image_root, transforms=TRANSFORMS, exclude=["images"])
    names = [path.name for path in dataset.images]
    assert names == ["1.png", "2.png", "3.png"]


def test_mixed_sizes_and_modes_become_one_tensor_shape(image_root):
    dataset = UnlabeledImageDataset(image_root, transforms=TRANSFORMS, exclude=["images"])
    for index in range(len(dataset)):
        assert dataset[index]["views"][0].shape == (3, 32, 32)  # RGBA / L 都统一到 3 通道


def test_gray_mode_and_metadata(image_root):
    dataset = UnlabeledImageDataset(
        image_root,
        transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
        gray=True,
        exclude=["images"],
    )
    sample = dataset[0]
    assert sample["views"][0].shape == (1, 32, 32)
    assert sample["meta"]["size"] == (152, 113)
    assert sample["meta"]["source_dir"] == "a"


def test_two_views_are_augmented_independently(image_root):
    dataset = UnlabeledImageDataset(
        image_root,
        transforms=TRANSFORMS,
        views=2,
        view_transforms=[{"name": "RandomResizedCrop", "size": [32, 32], "scale": [0.3, 0.6]}],
        exclude=["images"],
    )
    sample = dataset[0]
    assert len(sample["views"]) == 2

    batch = unlabeled_collate([dataset[0], dataset[1]])
    assert len(batch["views"]) == 2
    assert batch["views"][0].shape == (2, 3, 32, 32)
    assert batch["shape"].shape == (2, 4)
    assert len(batch["meta"]) == 2


def test_registry_build(image_root):
    dataset = build_dataset(
        {"name": "unlabeled", "root": str(image_root), "transforms": TRANSFORMS, "exclude": ["images"]}
    )
    assert len(dataset) == 3


def test_empty_dir_raises(tmp_path):
    (tmp_path / "empty").mkdir()
    with pytest.raises(FileNotFoundError, match="没有找到图片"):
        UnlabeledImageDataset(tmp_path / "empty")
    with pytest.raises(FileNotFoundError, match="目录不存在"):
        UnlabeledImageDataset(tmp_path / "nope")


def test_shape_features_are_log_scaled():
    features = shape_features(200, 100)
    assert features.shape == (4,)
    assert features[2] == pytest.approx(np.log(2.0))  # 宽高比
    assert features[3] == pytest.approx(np.log(20000))


def test_per_image_standardize_removes_brightness_and_contrast():
    """两张只差整体亮度/增益的图，标准化后应当几乎一致。"""
    transform = build_transforms([{"name": "ToTensor"}, {"name": "PerImageStandardize"}])
    base = np.random.default_rng(0).integers(40, 80, (32, 32, 3), dtype=np.uint8)

    dark = transform({"image": Image.fromarray(base)})["image"]
    bright = transform({"image": Image.fromarray(np.clip(base * 1.5 + 40, 0, 255).astype(np.uint8))})[
        "image"
    ]

    assert torch.allclose(dark, bright, atol=0.05)
    assert float(dark.mean()) == pytest.approx(0.0, abs=1e-4)
    assert float(dark.std()) == pytest.approx(1.0, abs=1e-2)


def test_letterbox_and_centercrop_handle_extreme_aspect_ratio():
    """40:1 的长条，两种预处理都得吐出固定尺寸。"""
    strip = Image.fromarray(np.zeros((113, 4464, 3), dtype=np.uint8))

    letterboxed = build_transforms([{"name": "Letterbox", "size": [128, 128]}, {"name": "ToTensor"}])(
        {"image": strip}
    )
    assert letterboxed["image"].shape == (3, 128, 128)
    assert letterboxed["letterbox"]["pad"][1] > 0  # 上下补边

    cropped = build_transforms([{"name": "CenterCrop", "size": [128, 128]}, {"name": "ToTensor"}])(
        {"image": strip}
    )
    assert cropped["image"].shape == (3, 128, 128)