"""
キャプション付き画像データセット。/dataset/raw 以下の .png + .txt ペアを返す。
DMD2 では「real image を見る」のは fake-score 更新時 (の v-pred MSE) のみで
あって、generator/real-score 側は noise から直接サンプル → image は不要。
従ってここでは pixels + caption を返すシンプルな実装で十分。
"""
from __future__ import annotations
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import Dataset
from PIL import Image


def _resize_short_side(img: Image.Image, target: int) -> Image.Image:
    """短辺を target に合わせる aspect-preserving resize"""
    w, h = img.size
    if min(w, h) == target:
        return img
    if w < h:
        new_w = target
        new_h = int(round(h * target / w))
    else:
        new_h = target
        new_w = int(round(w * target / h))
    return img.resize((new_w, new_h), Image.BICUBIC)


def _center_crop(img: Image.Image, size: int) -> Image.Image:
    w, h = img.size
    left = (w - size) // 2
    top = (h - size) // 2
    return img.crop((left, top, left + size, top + size))


def _to_tensor_normalize(img: Image.Image) -> torch.Tensor:
    """PIL RGB -> torch (3, H, W) in [-1, 1]"""
    arr = np.asarray(img, dtype=np.float32) / 127.5 - 1.0  # (H, W, 3) in [-1, 1]
    return torch.from_numpy(arr).permute(2, 0, 1).contiguous()


class AnimaImageCaptionDataset(Dataset):
    """
    Args:
        root: 画像/キャプションのルート。再帰的に *.png + 同名.txt を拾う。
        resolution: 単一解像度に統一する出力サイズ(短辺合わせ → center crop)。
    """

    def __init__(
        self,
        root: str | Path,
        resolution: int = 1024,
        exts: tuple[str, ...] = (".png", ".jpg", ".jpeg", ".webp"),
    ):
        self.root = Path(root)
        self.resolution = resolution
        self.items: list[tuple[Path, Path]] = []
        for img in sorted(self.root.rglob("*")):
            if not img.is_file() or img.suffix.lower() not in exts:
                continue
            cap = img.with_suffix(".txt")
            if cap.exists():
                self.items.append((img, cap))

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

    def __getitem__(self, idx: int) -> dict:
        img_path, cap_path = self.items[idx]
        img = Image.open(img_path).convert("RGB")
        img = _resize_short_side(img, self.resolution)
        img = _center_crop(img, self.resolution)
        pixels = _to_tensor_normalize(img)  # (3, H, W) in [-1, 1]
        caption = cap_path.read_text(encoding="utf-8").strip()
        return {"pixels": pixels, "caption": caption, "path": str(img_path)}


def collate_fn(batch: list[dict]) -> dict:
    """default collate と同じだが captions は list で保持。"""
    pixels = torch.stack([b["pixels"] for b in batch])
    captions = [b["caption"] for b in batch]
    paths = [b["path"] for b in batch]
    return {"pixels": pixels, "captions": captions, "paths": paths}