from pathlib import Path
import random

import torch
from torch.utils.data import Dataset


class DeepChoiceDataset(Dataset):
    def __init__(self, pt_files, shuffle=False):
        self.pt_files = [str(Path(path)) for path in pt_files]
        if shuffle:
            random.shuffle(self.pt_files)

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

    def __getitem__(self, index):
        path = self.pt_files[index]
        sample = torch.load(path, weights_only=False)
        sample["source_path"] = path
        sample["tile_name"] = Path(path).stem.split("_batch_")[0]
        return sample