from torch.utils.data import Dataset
from .transform import *
import albumentations as albu
from PIL import Image
import numpy as np
import os
import torch
import albumentations as A

class BaseDataset(Dataset):
    def __init__(self, transform=None,mode="train"):
        self.mosaic_ratio = 0.25
        self.mode = mode
        self.img_size = (1024,1024)
        aug_list = []
        for k,v in transform.items():
            if v != None:
                aug_list.append(eval(k)(**v))
            else: aug_list.append(eval(k)())

        self.transform = Compose(aug_list)

        self.t1_normalize = A.Compose([
            A.Normalize()
        ])

        self.t2_normalize = A.Compose([
            A.Normalize()
        ])

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

    def __getitem__(self, index):
        imgA, imgB, mask, img_id = self.load_img_and_mask(index)
        if len(self.transform.transforms) != 0:
            imgA, imgB, mask = self.transform([imgA, imgB], mask)
        imgA, imgB, mask = np.array(imgA), np.array(imgB), np.array(mask)
        imgA = self.t1_normalize(image=imgA)['image']
        imgB = self.t2_normalize(image=imgB)['image']
        imgA, imgB = [torch.from_numpy(img).permute(2, 0, 1).float() for img in [imgA, imgB]]
        mask = torch.from_numpy(mask).long()
        return imgA, imgB, mask, img_id

    def get_path(self, data_root, imgA_dir, imgB_dir, mask_dir):
        imgA_filename_list = os.listdir(os.path.join(data_root, imgA_dir))
        imgB_filename_list = os.listdir(os.path.join(data_root, imgB_dir))
        mask_filename_list = os.listdir(os.path.join(data_root, mask_dir))
        assert len(imgA_filename_list) == len(mask_filename_list)
        img_ids = [str(id.split('.')[0]) for id in mask_filename_list]
        return img_ids
    
    def load_img_and_mask(self, index):
        img_id = self.file_paths[index]
        imgA_name = os.path.join(self.data_root, self.imgA_dir, img_id + self.img_suffix)
        imgB_name = os.path.join(self.data_root, self.imgB_dir, img_id + self.img_suffix)
        mask_name = os.path.join(self.data_root, self.mask_dir, img_id + self.mask_suffix)
        imgA = Image.open(imgA_name).convert('RGB')
        imgB = Image.open(imgB_name).convert('RGB')
        mask_rgb = Image.open(mask_name).convert('RGB')
        mask = self.rgb2label(mask_rgb)
        return imgA, imgB, mask, img_id

    # def load_mosaic_img_and_mask(self, index):
    #     indexes = [index] + [random.randint(0, len(self.file_paths) - 1) for _ in range(3)]
    #     img_a, mask_a = self.load_img_and_mask(indexes[0])
    #     img_b, mask_b = self.load_img_and_mask(indexes[1])
    #     img_c, mask_c = self.load_img_and_mask(indexes[2])
    #     img_d, mask_d = self.load_img_and_mask(indexes[3])

    #     img_a, mask_a = np.array(img_a), np.array(mask_a)
    #     img_b, mask_b = np.array(img_b), np.array(mask_b)
    #     img_c, mask_c = np.array(img_c), np.array(mask_c)
    #     img_d, mask_d = np.array(img_d), np.array(mask_d)

    #     h = self.img_size[0]
    #     w = self.img_size[1]

    #     start_x = w // 4
    #     strat_y = h // 4
    #     # The coordinates of the splice center
    #     offset_x = random.randint(start_x, (w - start_x))
    #     offset_y = random.randint(strat_y, (h - strat_y))

    #     crop_size_a = (offset_x, offset_y)
    #     crop_size_b = (w - offset_x, offset_y)
    #     crop_size_c = (offset_x, h - offset_y)
    #     crop_size_d = (w - offset_x, h - offset_y)

    #     random_crop_a = albu.RandomCrop(width=crop_size_a[0], height=crop_size_a[1])
    #     random_crop_b = albu.RandomCrop(width=crop_size_b[0], height=crop_size_b[1])
    #     random_crop_c = albu.RandomCrop(width=crop_size_c[0], height=crop_size_c[1])
    #     random_crop_d = albu.RandomCrop(width=crop_size_d[0], height=crop_size_d[1])

    #     croped_a = random_crop_a(image=img_a.copy(), mask=mask_a.copy())
    #     croped_b = random_crop_b(image=img_b.copy(), mask=mask_b.copy())
    #     croped_c = random_crop_c(image=img_c.copy(), mask=mask_c.copy())
    #     croped_d = random_crop_d(image=img_d.copy(), mask=mask_d.copy())

    #     img_crop_a, mask_crop_a = croped_a['image'], croped_a['mask']
    #     img_crop_b, mask_crop_b = croped_b['image'], croped_b['mask']
    #     img_crop_c, mask_crop_c = croped_c['image'], croped_c['mask']
    #     img_crop_d, mask_crop_d = croped_d['image'], croped_d['mask']

    #     top = np.concatenate((img_crop_a, img_crop_b), axis=1)
    #     bottom = np.concatenate((img_crop_c, img_crop_d), axis=1)
    #     img = np.concatenate((top, bottom), axis=0)

    #     top_mask = np.concatenate((mask_crop_a, mask_crop_b), axis=1)
    #     bottom_mask = np.concatenate((mask_crop_c, mask_crop_d), axis=1)
    #     mask = np.concatenate((top_mask, bottom_mask), axis=0)
    #     mask = np.ascontiguousarray(mask)
    #     img = np.ascontiguousarray(img)
    #     img = Image.fromarray(img)
    #     mask = Image.fromarray(mask)
    #     # print(img.shape)

    #     return img, mask