"""
Subset-aware dataset for the data-scaling study.

The training side reads its filename list from `subsets/subset_{25,50,100}.txt`,
which is generated once by `subsets/make_subsets.py`. The validation side reads
the full validation directory and is shared across every run.
"""
from pathlib import Path
import torch
from torch.utils.data import Dataset
from PIL import Image
import torchvision.transforms as transforms


class SubsetSolarPanelDataset(Dataset):
    def __init__(self, image_dir, mask_dir, file_list=None, image_size=128, augment=False):
        self.image_dir = Path(image_dir)
        self.mask_dir = Path(mask_dir)
        self.image_size = image_size
        self.augment = augment

        if file_list is not None:
            with open(file_list) as f:
                self.image_files = [line.strip() for line in f if line.strip()]
        else:
            self.image_files = sorted(p.name for p in self.image_dir.iterdir() if p.suffix == ".jpg")

        self.image_transform = transforms.Compose([
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
        ])
        self.mask_transform = transforms.Compose([
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
        ])
        self.augment_transform = transforms.Compose([
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomVerticalFlip(p=0.5),
            transforms.RandomRotation(15),
        ])

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

    def __getitem__(self, idx):
        img_name = self.image_files[idx]
        img_path = self.image_dir / img_name
        mask_path = self.mask_dir / img_name.replace(".jpg", "_mask.png")

        image = Image.open(img_path).convert("RGB")
        mask = Image.open(mask_path).convert("L")

        image = self.image_transform(image)
        mask = self.mask_transform(mask)

        if self.augment:
            seed = torch.randint(0, 2**32, (1,)).item()
            torch.manual_seed(seed)
            image = self.augment_transform(image)
            torch.manual_seed(seed)
            mask = self.augment_transform(mask)

        mask = (mask > 0.5).float()
        return image, mask