"""按目录组织的分类数据集：`root/<类别名>/*.png`，子目录名即标签。

这是 scripts/export_clusters.py 的产物结构——聚类 + 人工归类之后得到的数据集。
没有预先划分 train/val，这里按类别分层、用固定随机种子切分，保证两次运行拿到同一份划分。
"""

from __future__ import annotations

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

SPLITS = ("train", "val", "all")


@DATASETS.register("folder_classification")
class FolderClassificationDataset(Dataset):
    """Args:
    root: 数据集根目录，其下每个子目录是一个类别。
    split: "train" / "val" / "all"。
    val_ratio: 验证集比例，**按类别分别切分**——小类只有几十张，随机整体切分
        很容易让某个类在验证集里一张都没有。
    seed: 划分用的随机种子；同一个 seed 两次运行得到完全相同的划分。
    classes: 显式指定类别顺序；缺省按目录名排序。训练与评估必须一致，
        否则标签下标对不上。
    transforms: transform 配置列表，见 `datasets.transforms.build_transforms`。
    exclude: 要跳过的子目录名。
    """

    def __init__(
        self,
        root: str | Path,
        split: str = "train",
        val_ratio: float = 0.2,
        seed: int = 42,
        classes: Sequence[str] | None = None,
        transforms: Sequence[Mapping[str, Any]] = (),
        gray: bool = False,
        exclude: Sequence[str] = (),
    ) -> None:
        if split not in SPLITS:
            raise ValueError(f"未知 split: {split}，可选 {SPLITS}")
        self.root = Path(root)
        if not self.root.is_dir():
            raise FileNotFoundError(f"数据集目录不存在: {self.root}")

        self.split = split
        self.gray = gray
        self.transforms = build_transforms(transforms)

        excluded = set(exclude)
        found = sorted(
            p.name for p in self.root.iterdir() if p.is_dir() and p.name not in excluded
        )
        self.classes = list(classes) if classes else found
        if not self.classes:
            raise FileNotFoundError(f"{self.root} 下没有找到类别子目录")
        missing = [c for c in self.classes if not (self.root / c).is_dir()]
        if missing:
            raise FileNotFoundError(f"以下类别在 {self.root} 下不存在: {missing}")
        self.class_to_idx = {name: i for i, name in enumerate(self.classes)}

        self.samples = self._collect(val_ratio, seed)
        if not self.samples:
            raise ValueError(f"{self.root} 的 {split} 划分里没有样本")

    def _collect(self, val_ratio: float, seed: int) -> list[tuple[Path, int]]:
        samples: list[tuple[Path, int]] = []
        for label, name in enumerate(self.classes):
            files = sorted(
                p
                for p in (self.root / name).rglob("*")
                if p.is_file() and p.suffix.lower() in IMAGE_SUFFIXES
            )
            if not files:
                continue
            # 逐类切分：小类只有几十张，整体随机切很容易让它在验证集里挂零
            generator = torch.Generator().manual_seed(seed + label)
            order = torch.randperm(len(files), generator=generator).tolist()
            n_val = max(1, round(len(files) * val_ratio)) if len(files) > 1 else 0
            val_idx = set(order[:n_val])
            for index, path in enumerate(files):
                in_val = index in val_idx
                if self.split == "all" or (self.split == "val") == in_val:
                    samples.append((path, label))
        return samples

    def class_counts(self) -> list[int]:
        counts = [0] * len(self.classes)
        for _, label in self.samples:
            counts[label] += 1
        return counts

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

    def __getitem__(self, index: int) -> dict:
        path, label = self.samples[index]
        with Image.open(path) as raw:
            image = raw.convert("L" if self.gray else "RGB")
        sample = self.transforms({"image": image})
        return {
            "image": sample["image"],
            "label": label,
            "meta": {"path": str(path), "index": index, "class": self.classes[label]},
        }


def classification_collate(batch: list[dict]) -> dict:
    return {
        "images": torch.stack([item["image"] for item in batch]),
        "labels": torch.tensor([item["label"] for item in batch], dtype=torch.long),
        "meta": [item["meta"] for item in batch],
    }