import os
import torch
from openslide import OpenSlide
from utils.preprocessor import MacenkoNormalizer, preprocessor
from torch.utils.data import Dataset


class WSIPatchDataset(Dataset):
    def __init__(
        self,
        coords,
        wsi_path,
        pretrained=False,
        patch_size=256,
        patch_level=0,
        macenko=True,
        return_coord=False,
    ):
        self.pretrained = pretrained
        self.wsi = OpenSlide(wsi_path)
        self.patch_size = patch_size
        self.patch_level = patch_level
        self.return_coord = return_coord

        if macenko:
            normalizer = MacenkoNormalizer(
                target_path=os.path.join(
                    os.path.dirname(os.path.dirname(os.path.join(__file__))),
                    "macenko_target",
                    "macenko_param.pt",
                )
            )
        else:
            normalizer = None

        self.roi_transforms = preprocessor(pretrained=pretrained, normalizer=normalizer)
        self.coords = coords
        self.length = len(self.coords)

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        coord = self.coords[idx]
        img = self.wsi.read_region(
            coord, self.patch_level, (self.patch_size, self.patch_size)
        ).convert("RGB")
        img = self.roi_transforms(img)
        if self.return_coord:
            return img, torch.tensor(coord)
        else:
            return img