"""MS1M dataset: identity folders of 112x112 ArcFace-aligned faces."""
import os
import glob

import cv2
import torch
from torch.utils.data import Dataset


def load_rgb01(path):
    img = cv2.imread(path)              # BGR uint8, 112x112
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    t = torch.from_numpy(img).float().permute(2, 0, 1) / 255.0
    return t


class MS1MDataset(Dataset):
    """Aligned MS1M faces. Uses `num_ids` identity folders (sorted
    numerically) starting at `start` so different scripts see the same
    deterministic split."""

    def __init__(self, root, num_ids, start=0):
        ids = sorted(os.listdir(root), key=int)[start:start + num_ids]
        self.samples = []
        for label, ident in enumerate(ids):
            for f in sorted(glob.glob(f"{root}/{ident}/*.jpg")):
                self.samples.append((f, label))
        self.num_classes = len(ids)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, i):
        path, label = self.samples[i]
        return load_rgb01(path), label