"""无标注图像数据集，用于无监督聚类。

递归扫描一个目录下的所有图片，不需要任何标注。

样本还会带出形状特征（宽、高、宽高比、面积的对数），但**默认不参与聚类**：
这些尺寸是上游 ROI 裁剪的产物，不是缺陷本身的属性，拿它做特征只会让聚类
去区分采集方式。保留它只是为了写进 assignments.csv 供人工核对、以及做对照实验。
"""

from __future__ import annotations

import math
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Any

import torch
from PIL import Image
from torch.utils.data import Dataset

from datasets.registry import DATASETS
from datasets.transforms import build_transforms
from datasets.yolo import IMAGE_SUFFIXES

SHAPE_FEATURE_DIM = 4


def shape_features(width: int, height: int) -> torch.Tensor:
    """[log 宽, log 高, log 宽高比, log 面积]，取对数是因为尺寸跨了两个数量级。"""
    return torch.tensor(
        [
            math.log(width),
            math.log(height),
            math.log(width / height),
            math.log(width * height),
        ],
        dtype=torch.float32,
    )


@DATASETS.register("unlabeled")
class UnlabeledImageDataset(Dataset):
    """Args:
    root: 图片根目录，递归扫描。
    transforms: transform 配置列表；聚类默认用 Letterbox + ToTensor + Normalize。
    gray: 是否转单通道（RGBA / 调色板图会先统一到 RGB 或 L）。
    views: 每个样本返回几个增强视图。1 = 推理取特征；2 = SimCLR 对比学习。
    view_transforms: `views > 1` 时每个视图各自过一遍的随机增强。
    exclude: 需要跳过的目录名（默认跳过已划分好的 images/labels 与临时目录）。
    """

    def __init__(
        self,
        root: str | Path,
        transforms: Sequence[Mapping[str, Any]] = (),
        gray: bool = False,
        views: int = 1,
        view_transforms: Sequence[Mapping[str, Any]] = (),
        exclude: Sequence[str] = (),
    ) -> None:
        self.root = Path(root)
        if not self.root.is_dir():
            raise FileNotFoundError(f"目录不存在: {self.root}")
        self.gray = gray
        self.views = int(views)
        self.transforms = build_transforms(transforms)
        self.view_transforms = build_transforms(view_transforms) if view_transforms else None
        self.exclude = set(exclude)

        self.images = sorted(
            path
            for path in self.root.rglob("*")
            if path.suffix.lower() in IMAGE_SUFFIXES
            and not path.name.startswith(".")
            and not self.exclude & set(path.relative_to(self.root).parts)
        )
        if not self.images:
            raise FileNotFoundError(f"{self.root} 下没有找到图片")

    def source_dir(self, path: Path) -> str:
        """图片相对 root 的父目录，用来和聚类结果做交叉表（人工已分好的部分即是弱标签）。"""
        relative = path.relative_to(self.root).parent
        return str(relative) if str(relative) != "." else "(root)"

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

    def __getitem__(self, index: int) -> dict:
        path = self.images[index]
        with Image.open(path) as raw:
            image = raw.convert("L" if self.gray else "RGB")  # RGBA / P 一并统一掉
        width, height = image.size

        if self.views == 1:
            views = [self.transforms({"image": image.copy()})["image"]]
        else:
            views = []
            for _ in range(self.views):
                sample = {"image": image.copy()}
                if self.view_transforms is not None:
                    sample = self.view_transforms(sample)
                views.append(self.transforms(sample)["image"])

        return {
            "views": views,
            "shape": shape_features(width, height),
            "meta": {
                "path": str(path),
                "index": index,
                "source_dir": self.source_dir(path),
                "size": (width, height),
            },
        }


def unlabeled_collate(batch: list[dict]) -> dict:
    """把每个视图各自堆成一个 batch：views[i] 形状 (B, C, H, W)。"""
    num_views = len(batch[0]["views"])
    return {
        "views": [torch.stack([item["views"][i] for item in batch]) for i in range(num_views)],
        "shape": torch.stack([item["shape"] for item in batch]),
        "meta": [item["meta"] for item in batch],
    }