import os
from os.path import isdir
import tarfile

from pathlib import Path
from PIL import Image

import torch
from torch import tensor
from torchvision.datasets import ImageFolder
from torchvision import transforms
from torch.utils.data import DataLoader
from .data_transform import transform_default



class Dataset:
    def __init__(self, train_dir: str, test_dir:str,  size : int = 224):
        self.train_dir = train_dir
        self.test_dir = test_dir
        self.size = size
        self.train_ds = TrainDatasetFolder(train_dir,size)
        self.test_ds = TestDatasetFolder(test_dir,size)


    def get_datasets(self):
        return self.train_ds, self.test_ds

    def get_dataloaders(self):
        return DataLoader(self.train_ds), DataLoader(self.test_ds)

class TrainDatasetFolder(ImageFolder):
    def __init__(self,train_dir:str, size : int, transform = transform_default, second_transform = transform_default):
        super().__init__(
            root= train_dir,
            transform=transform
        )
        self.size = size
        self.transform = transform
        self.second_transform = second_transform
    
    def __getitem__(self, index):
        path, _ = self.samples[index]
        image_pil = Image.open(path).convert("RGB")

        second_image_trans = self.second_transform(image_pil)
        sample = self.loader(path)
        if self.transform is not None:
            sample = self.transform(sample)
        return sample, {"second_image_trans": second_image_trans, "path": path, "size": self.size}

class TestDatasetFolder(ImageFolder):
    def __init__(self,test_dir:str, size : int):
        super().__init__(
            root= test_dir,
            transform=transform_default,
        )
        self.size = size
            
    def __getitem__(self, index):
        path, _ = self.samples[index]
        sample = self.loader(path)
        
        if "good" in path:
            #target = Image.new('L', (self.size, self.size))
            sample_class = 0
        else:
            # target_path = path.replace("test", "ground_truth")
            # target_path = target_path.replace(".png", "_mask.png")
            # target = self.loader(target_path)
            sample_class = 1

        if self.transform is not None:
            sample = self.transform(sample)
        # if self.target_transform is not None:
        #     target = self.target_transform(target)

        return sample, sample_class

class TrainDatasetListPath(torch.utils.data.Dataset):
    def __init__(self, image_paths, size=224, transform=transform_default):
        self.image_paths = image_paths
        self.size = size
        self.transform = transform
        self.labels = os.path.dirname(image_paths[0]) if len(image_paths) >0 else 0
    
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        image_path = self.image_paths[idx]
        image = Image.open(image_path).convert("RGB")

        image_trans = self.transform(image)
        return image_trans, self.labels