"""
dataset.py — PyTorch datasets wrapping the BDAPPV HuggingFace dataset.

Two dataset classes:
  - SegmentationDataset : positives only (has_mask=True), returns (image, mask)
  - ClassificationDataset : all examples, returns (image, label)
"""

import torch
from torch.utils.data import Dataset
from torchvision import transforms as T

IMG_MEAN = [0.485, 0.456, 0.406]
IMG_STD  = [0.229, 0.224, 0.225]


class SegmentationDataset(Dataset):
    """
    Wraps a HF split pre-filtered to positives (has_mask=True).
    Returns (image_tensor [3, H, W], mask_tensor [1, H, W]) as float32.
    """

    def __init__(self, hf_split, img_size: int = 400, augment: bool = False):
        self.data    = hf_split.filter(lambda x: x["has_mask"])
        self.augment = augment
        self.img_size = img_size

        self.img_tf = T.Compose([
            T.Resize((img_size, img_size)),
            T.ToTensor(),
            T.Normalize(mean=IMG_MEAN, std=IMG_STD),
        ])
        self.mask_tf = T.Compose([
            T.Resize((img_size, img_size), interpolation=T.InterpolationMode.NEAREST),
            T.ToTensor(),
        ])

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

    def __getitem__(self, idx):
        ex    = self.data[idx]
        image = ex["image"].convert("RGB")
        mask  = ex["mask"].convert("L")

        if self.augment:
            if torch.rand(1) > 0.5:
                image = T.functional.hflip(image)
                mask  = T.functional.hflip(mask)
            if torch.rand(1) > 0.5:
                image = T.functional.vflip(image)
                mask  = T.functional.vflip(mask)

        image = self.img_tf(image)
        mask  = self.mask_tf(mask)
        mask  = (mask > 0.5).float()   # binarize → {0, 1}

        return image, mask


class ClassificationDataset(Dataset):
    """
    All examples (positives + negatives).
    Returns (image_tensor [3, H, W], label [float scalar]).
    label = 1.0 if has_mask else 0.0
    """

    def __init__(self, hf_split, img_size: int = 299, augment: bool = False):
        self.data    = hf_split
        self.augment = augment

        self.img_tf = T.Compose([
            T.Resize((img_size, img_size)),
            T.ToTensor(),
            T.Normalize(mean=IMG_MEAN, std=IMG_STD),
        ])

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

    def __getitem__(self, idx):
        ex    = self.data[idx]
        image = ex["image"].convert("RGB")
        label = float(ex["has_mask"])

        if self.augment:
            if torch.rand(1) > 0.5:
                image = T.functional.hflip(image)
            if torch.rand(1) > 0.5:
                image = T.functional.vflip(image)

        image = self.img_tf(image)
        return image, torch.tensor(label)