"""
Tile library
"""
import numpy as np
from PIL import Image
from torchvision import datasets, transforms

def build_cifar10_tile_library(root="./cifar_data", max_per_class=500):
    """
    Download/load CIFAR-10 training set as tile library.
    For speed control, defaults to max_per_class tiles per class (10 classes, total <= 10*max_per_class).
    Returns:
      tiles:   List[PIL.Image], original 32x32
      means:   np.ndarray [N,3], RGB average color of each tile (0..255)
      labels:  np.ndarray [N], class labels (0..9)
    """
    ds = datasets.CIFAR10(root=root, train=True, download=True,
                          transform=transforms.ToTensor())
    
    counts = {c : 0 for c in range(10)}
    tiles, means, labels = [], [], []
    for img_tensor, lab in ds:
        if counts[lab] >= max_per_class:
            continue
        arr= (img_tensor.numpy().transpose(1,2,0) * 255).astype(np.uint8)
        pil = Image.fromarray(arr, mode="RGB")
        tiles.append(pil)
        means.append(arr.reshape(-1,3).mean(axis=0))
        labels.append(lab)
        counts[lab]+=1
    
    means = np.asarray(means, dtype=np.float32)
    labels = np.asarray(labels, dtype=np.int64)
    print(f"[INFO] CIFAR10 tiles: {len(tiles)} (each 32x32). Per-class cap={max_per_class}")
    return tiles, means, labels

def build_cifar100_tile_library(root="./cifar_data", max_per_class=500):
    """
    Download/load CIFAR-100 training set as tile library.
    For speed control, defaults to max_per_class tiles per class (100 classes, total <= 100*max_per_class).
    Returns:
      tiles:   List[PIL.Image], original 32x32
      means:   np.ndarray [N,3], RGB average color of each tile (0..255)
      labels:  np.ndarray [N], class labels (0..99)
    """
    ds = datasets.CIFAR100(root=root, train=True, download=True,
                          transform=transforms.ToTensor())
    
    counts = {c : 0 for c in range(100)}
    tiles, means, labels = [], [], []
    for img_tensor, lab in ds:
        if counts[lab] >= max_per_class:
            continue
        arr= (img_tensor.numpy().transpose(1,2,0) * 255).astype(np.uint8)
        pil = Image.fromarray(arr, mode="RGB")
        tiles.append(pil)
        means.append(arr.reshape(-1,3).mean(axis=0))
        labels.append(lab)
        counts[lab]+=1
    
    means = np.asarray(means, dtype=np.float32)
    labels = np.asarray(labels, dtype=np.int64)
    print(f"[INFO] CIFAR10 tiles: {len(tiles)} (each 32x32). Per-class cap={max_per_class}")
    return tiles, means, labels