"""学習・評価用の Dataset と入力の正規化（SPEC_02「データ拡張」「前処理」）.

Dataset はレターボックス済みの uint8 画像を返し，正規化（ImageNet の mean / std）はデバイス上で
`normalize_batch` が行う（CPU のワーカーで float 化するより転送量が小さいため）．

データ拡張は学習時のみ・幾何変換のみ（左右反転・回転・拡大縮小・平行移動．はみ出しは黒）．
舌の色そのものがラベルであるため，**色を変える拡張は行わない**．
"""

from __future__ import annotations

import numpy as np
import torch
from torch.utils.data import Dataset
from torchvision.transforms import v2

IMAGENET_MEAN: tuple[float, float, float] = (0.485, 0.456, 0.406)
IMAGENET_STD: tuple[float, float, float] = (0.229, 0.224, 0.225)

# 幾何拡張の強さ（SPEC_02「学習の設定」の箇条書き）．
_FLIP_PROB = 0.5
_ROTATION_DEGREES = 10.0
_SCALE_RANGE = (0.9, 1.1)
_TRANSLATE_RATIO = (0.05, 0.05)


def make_train_augmentation() -> v2.Compose:
    """学習時の幾何拡張（左右反転・回転 ±10°・拡大縮小 0.9〜1.1・平行移動 ±5%）を作る.

    はみ出した部分は黒 (0) で埋める．色を変える変換は含めない．

    Returns:
        `(3, H, W)` の uint8 テンソルを受け取り同じ形状で返す変換．
    """
    return v2.Compose(
        [
            v2.RandomHorizontalFlip(p=_FLIP_PROB),
            v2.RandomAffine(
                degrees=_ROTATION_DEGREES,
                translate=_TRANSLATE_RATIO,
                scale=_SCALE_RANGE,
                interpolation=v2.InterpolationMode.BILINEAR,
                fill=0,
            ),
        ]
    )


class TongueImageDataset(Dataset):
    """レターボックス済み画像とクラス ID の Dataset.

    Args:
        images: `(N, H, W, 3)` の uint8 配列（RGB）．
        labels: `(N,)` の整数配列（クラス ID．`classes.py` の固定 ID）．
        augment: True なら学習時の幾何拡張をかける．
    """

    def __init__(self, images: np.ndarray, labels: np.ndarray, *, augment: bool) -> None:
        if len(images) != len(labels):
            raise RuntimeError("画像とラベルの件数が一致しない")
        self._images = images
        self._labels = labels.astype(np.int64)
        self._augment = make_train_augmentation() if augment else None

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

    def __getitem__(self, index: int) -> tuple[torch.Tensor, int]:
        """1 件を返す.

        Returns:
            (`(3, H, W)` の uint8 テンソル, クラス ID)．
        """
        image = torch.from_numpy(np.ascontiguousarray(self._images[index])).permute(2, 0, 1)
        if self._augment is not None:
            image = self._augment(image)
        return image, int(self._labels[index])


def normalize_batch(images: torch.Tensor) -> torch.Tensor:
    """uint8 画像のバッチを float にして ImageNet の mean / std で正規化する.

    Args:
        images: `(B, 3, H, W)` の uint8 テンソル（任意のデバイス）．

    Returns:
        `(B, 3, H, W)` の float32 テンソル（同じデバイス）．
    """
    mean = torch.tensor(IMAGENET_MEAN, device=images.device).view(1, 3, 1, 1)
    std = torch.tensor(IMAGENET_STD, device=images.device).view(1, 3, 1, 1)
    return (images.float() / 255.0 - mean) / std