from __future__ import annotations
import pathlib
import torch
from torch.utils.data import Dataset
import typing as t
import numpy as np
from torchvision import datasets
from PIL import Image
# from torchvision.transforms import transforms


class NumpyDataset(Dataset):
  def __init__(
    self,
    root: pathlib.Path,
    x_name: str | pathlib.Path,
    y_name: str | pathlib.Path,
    transform: t.Optional[t.Callable] = None,
    target_transform: t.Optional[t.Callable] = None,
  ):
    self.x_name = x_name
    self.y_name = y_name
    self.x = np.load(root / x_name)
    self.y = np.load(root / y_name)
    self.transform = transform
    self.target_transform = target_transform

  def __getitem__(self, index: int) -> tuple[Image.Image, np.ndarray]:
    x = self.x[index]
    y = self.y[index]

    x = Image.fromarray(x)

    if self.transform:
      x = self.transform(x)
    if self.target_transform:
      y = self.target_transform(y)
    return x, y
