from pathlib import Path

import numpy as np
import torch
from torch.utils.data import Dataset


class RandomSubsetBatchDataset(Dataset):
    def __init__(self, pt_files, subset_size, seed=42):
        self.all_files = [str(Path(path)) for path in pt_files]
        if not self.all_files:
            raise ValueError("RandomSubsetBatchDataset requires at least one batch file")
        self.subset_size = max(1, min(int(subset_size), len(self.all_files)))
        self.seed = int(seed)
        self.epoch = 0
        self.active_files = self.all_files[: self.subset_size]
        self.set_epoch(0)

    def set_epoch(self, epoch):
        self.epoch = int(epoch)
        if self.subset_size >= len(self.all_files):
            self.active_files = list(self.all_files)
            return
        rng = np.random.default_rng(self.seed + self.epoch)
        indices = rng.choice(len(self.all_files), size=self.subset_size, replace=False)
        self.active_files = [self.all_files[int(idx)] for idx in indices]

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

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