"""
PyTorch Dataset Definition Module.

This module contains the custom Dataset classes for MedMNIST, handling
the conversion from NumPy arrays to PyTorch tensors and applying
image transformations for training and inference.

It implements selective RAM loading to balance I/O speed with memory
efficiency and ensures deterministic subsampling for reproducible research.

Key Components:
    VisionDataset: Full-featured dataset with transforms, subsampling, and PIL conversion.
    LazyNPZDataset: Memory-mapped dataset for lightweight health checks.
"""

from pathlib import Path
from typing import Final, Tuple

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


# DATASET CLASS
class VisionDataset(Dataset[Tuple[torch.Tensor, torch.Tensor]]):
    """
    Enhanced PyTorch Dataset for MedMNIST data.

    Features:
        - In-memory caching of specific splits to eliminate disk I/O bottlenecks.
        - Seed-aware deterministic subsampling for rapid smoke testing.
        - Automatic dimensionality standardization (N, H, W, C).
    """

    def __init__(
        self,
        path: Path,
        split: str = "train",
        transform: transforms.Compose | None = None,
        max_samples: int | None = None,
        seed: int = 42,
    ):
        """
        Initializes the dataset by loading the specified .npz split into RAM.

        Args:
            path (Path): Path to the MedMNIST .npz archive.
            split (str): Dataset split to load ('train', 'val', or 'test').
            transform (transforms.Compose | None): Pipeline of Torchvision transforms.
            max_samples (int | None): If set, limits the number of samples (subsampling).
            seed (int): Random seed for deterministic subsampling.
        """
        if not path.exists():
            raise FileNotFoundError(f"Dataset file not found at: {path}")

        self.path: Final[Path] = path
        self.transform: Final[transforms.Compose | None] = transform
        self.split: Final[str] = split

        # Open NPZ once and load target arrays into system memory
        with np.load(path) as data:
            raw_images = data[f"{split}_images"]
            raw_labels = data[f"{split}_labels"].ravel().astype(np.int64)

            total_available = len(raw_labels)

            # Deterministic subsampling logic
            if max_samples and max_samples < total_available:
                rng = np.random.default_rng(seed)
                chosen_indices = rng.choice(total_available, size=max_samples, replace=False)
                self.images = raw_images[chosen_indices]
                self.labels = raw_labels[chosen_indices]
            else:
                self.images = np.array(raw_images)
                self.labels = raw_labels

            # This ensures consistent PIL conversion regardless of source format
            if self.images.ndim == 3:  # (N, H, W) -> (N, H, W, 1)
                self.images = np.expand_dims(self.images, axis=-1)

    def __len__(self) -> int:
        """Returns the total number of samples currently in the dataset."""
        return len(self.labels)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Retrieves a standardized sample-label pair.

        The image is converted to a PIL object to ensure compatibility with
        Torchvision V2 transforms before being returned as a PyTorch Tensor.
        """
        img = self.images[idx]
        label = self.labels[idx]

        pil_img = Image.fromarray(img.squeeze() if img.shape[-1] == 1 else img)

        if self.transform:
            img = self.transform(pil_img)
        else:
            img = transforms.functional.to_tensor(pil_img)

        return img, torch.tensor(label, dtype=torch.long)


class LazyNPZDataset(Dataset):
    """Memory-mapped PyTorch Dataset for lazy loading from ``.npz`` files.

    Uses NumPy memory-mapping (``mmap_mode="r"``) to avoid loading the entire
    archive into RAM, making it suitable for large-scale health checks and
    quick data integrity verification.

    Attributes:
        npz_path (Path): Filesystem path to the ``.npz`` archive.
        images (np.ndarray): Memory-mapped view of ``train_images``.
        labels (np.ndarray): Memory-mapped view of ``train_labels``.
    """

    def __init__(self, npz_path: Path):
        """Initializes the dataset with a memory-mapped ``.npz`` archive.

        Args:
            npz_path: Path to the ``.npz`` file containing
                ``train_images`` and ``train_labels`` arrays.
        """
        self.npz_path = npz_path
        self.data = np.load(npz_path, mmap_mode="r")
        self.images = self.data["train_images"]
        self.labels = self.data["train_labels"]

    def __len__(self) -> int:
        """Returns the total number of samples in the dataset."""
        return len(self.labels)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
        """Retrieves a normalized (C, H, W) tensor and its integer label.

        Args:
            idx: Sample index.

        Returns:
            A tuple of (image_tensor, label) where the image is scaled to [0, 1].

        Raises:
            ValueError: If the image has an unexpected number of dimensions.
        """
        img = self.images[idx]
        # Ensure channel dimension (C,H,W)
        if img.ndim == 2:  # (H,W) grayscale
            img = np.expand_dims(img, axis=0)
        elif img.ndim == 3:  # (H,W,C)
            img = np.transpose(img, (2, 0, 1))
        else:
            raise ValueError(f"Unexpected image shape: {img.shape}")

        img = torch.from_numpy(img).float() / 255.0
        label = int(self.labels[idx][0])
        return img, label
