"""
Developed by: Daniel Crovo
Dataset class definition

"""
from torch.utils.data import Dataset
import os
from PIL import Image 
from pydicom import dcmread
from dicom_utils import apply_windowing
import numpy as np

class HSDataset(Dataset): 
    def __init__(self, image_dir, mask_dir, transform = None, w_level=35, w_width=350, normalized =False) -> None:
        """Initialises the Heart Segmentation Dataset class. The dataset asumes all the dicom files are in one single folder aswell as the masks
            both the images (dicom files) and masks should have the same name for each slice
        Args:
            image_dir (string): Path to the dicom directory
            mask_dir (string): Path to the masks directory 
            transform (_type_, optional):Aplied transformations Defaults to None.
            w_level (int, optional): The window level (center). Defaults to 35.
            w_width (int, optional): The width of the window Defaults to 400.
            normalized (bool, optional): wheter to normalized the windowed array after applying windoing
        """
        super().__init__()
        
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.transform = transform
        self.w_level = w_level
        self.w_width = w_width
        self.normalized = normalized
        self.images = os.listdir(self.image_dir)
        self.masks = os.listdir(self.mask_dir)

    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        img_path = os.path.join(self.image_dir, self.images[idx])
        mask_path = os.path.join(self.mask_dir, self.masks[idx])
        ds = dcmread(img_path)
        dicom_img = apply_windowing(ds = ds,window_center = self.w_level, 
                                    window_width =self.w_width, normalized = self.normalized)
        image = Image.fromarray(dicom_img).convert('RGB')
        image = np.array(image)
        mask = Image.open(mask_path).convert('L')
        mask = np.array(mask)
        mask[mask == 255.0] = 1.0


        if self.transform is not None:
            transformations = self.transform(image = image, mask = mask)
            image = transformations['image']
            mask = transformations['mask']
            
        return image, mask