import os
import json
import torch
import glob

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

from imgaug import augmenters as iaa
from imgaug.augmentables.segmaps import SegmentationMapsOnImage

import numpy as np

def create_nested_dir(log_path):
    # Create the experiment directory if not present
    if not os.path.isdir(log_path):
        os.makedirs(log_path)
        os.makedirs(os.path.join(log_path, 'checkpoint'))

def load_dataset_dist():
  
    with open(os.path.join('configuration', 'cases_division.json'), 'r') as f:
        dataset = json.load(f)

    return dataset       

def get_data_loaders(data_aug, cases, dataset_dir, batch_size):
    dataloaders = {}

    dataloaders['Train'] = get_dataset(
        dataset_dir, data_aug, cases=cases['train'], balanced_filelist=None, batch_size=batch_size)

    dataloaders['Valid'] = get_dataset(
        dataset_dir, 'none', cases=cases['valid'], batch_size=batch_size)

    return dataloaders 

def get_dataset(data_dir, data_aug, cases=[], balanced_filelist=None, imageFolder='Images', maskFolder='Masks', batch_size=4):

    data_transforms = {
        'Train': transforms.Compose([ToTensor()]),
        'Test': transforms.Compose([ToTensor()]),
    }

    image_dataset = SegNumpyDataset(
        data_aug=data_aug, root_dir=data_dir, cases=cases, transform=data_transforms['Train'], maskFolder=maskFolder, imageFolder=imageFolder, balanced_filelist=balanced_filelist)

    dataloader = DataLoader(
        image_dataset, batch_size=batch_size, drop_last=True, shuffle=True, num_workers=4)

    return dataloader    

class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""

    def __call__(self, sample, maskresize=None, imageresize=None):
        image, mask = sample['image'], sample['mask']
        if len(mask.shape) == 2:
            mask = mask.reshape((1,)+mask.shape)
        if len(image.shape) == 2:
            image = image.reshape((1,)+image.shape)
        return {'image': torch.from_numpy(image).float(),
                'mask': torch.from_numpy(mask).float()}

class SegNumpyDataset(Dataset):
    """Segmentation Dataset"""

    def __init__(self, root_dir, cases, imageFolder, maskFolder, data_aug, cases_number_format=False, transform=None, balanced_filelist=None):
        self.in_channels  = 3
        self.root_dir = root_dir
        self.transform = transform
        self.data_aug = data_aug

        if cases_number_format:
            cases_names = ["case_{:05d}".format(i) for i in cases]
        else:
            cases_names = cases

        image_names = []
        mask_names = []

        if balanced_filelist is None:
            for case in cases_names:
                image_names.extend(glob.glob(os.path.join(
                    self.root_dir, case, imageFolder, '*')))
                mask_names.extend(glob.glob(os.path.join(
                    self.root_dir, case, maskFolder, '*')))
        else:
            # Essa condição é necessária, pois no data aug offline o nome dos arquivos muda.
            if data_aug != 'offline':
                for case in cases_names:
                    image_list = set(os.listdir(os.path.join(
                        self.root_dir, case, imageFolder)))
                    set_balanced = set(balanced_filelist)

                    image_list = list(set_balanced.intersection(image_list))
                    fullpath_image_list = [os.path.join(self.root_dir, case, imageFolder, x)
                                           for x in image_list]
                    fullpath_mask_list = [os.path.join(self.root_dir, case, maskFolder, "masc_"+str(x))
                                          for x in image_list]

                    image_names.extend(fullpath_image_list)
                    mask_names.extend(fullpath_mask_list)
            else:
                for case in cases_names:
                    image_list = set(os.listdir(os.path.join(
                        self.root_dir, case, imageFolder)))

                    balanced_filelist_aug = []
                    # adiciona os data aug manualmente

                    for fl in balanced_filelist:
                        for i in range(0, 5):
                            # case_00000-0-aug-0
                            balanced_filelist_aug.append(
                                "{}-aug-{}.npz".format(fl.replace(".npz", ""), i))

                    set_balanced = set(balanced_filelist_aug)

                    image_list = list(set_balanced.intersection(image_list))
                    fullpath_image_list = [os.path.join(self.root_dir, case, imageFolder, x)
                                           for x in image_list]
                    fullpath_mask_list = [os.path.join(self.root_dir, case, maskFolder, "masc_"+str(x))
                                          for x in image_list]

                    image_names.extend(fullpath_image_list)
                    mask_names.extend(fullpath_mask_list)

        self.image_names = sorted(image_names)
        self.mask_names = sorted(mask_names)

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

    def __getitem__(self, idx):

        image = np.load(self.image_names[idx])
        mask = np.load(self.mask_names[idx])

        __, file_extension = os.path.splitext(self.image_names[idx])

        if file_extension == '.npz':
            image = image['arr_0']
            mask = mask['arr_0']


        if self.in_channels == 1:
            image = image[1]
        
        if self.data_aug == 'online':

            segmap = SegmentationMapsOnImage(mask, shape=(256, 256))

            seq = iaa.Sequential([
                
                iaa.Affine(
                    scale=(0.5, 1.2),
                    rotate=(-15, 15)
                ),  # rotate the image
                iaa.Flipud(0.5),
                iaa.PiecewiseAffine(scale=(0.01, 0.05)),
                iaa.Sometimes(
                    0.1,
                    iaa.GaussianBlur((0.1, 1.5)),
                ),
                iaa.Sometimes(
                    0.1,
                    iaa.LinearContrast((0.5, 2.0), per_channel=0.5),
                )
            ])

            image = image.transpose(1, 2, 0)
            # Apply augmentations for image and mask
            image, mask = seq(image=image, segmentation_maps=segmap)
            image = image.copy()
            mask = mask.copy()
            image = image.transpose(2, 0, 1)
            mask = mask.get_arr()

        sample = {'image': image, 'mask': mask}

        if self.transform:
            sample = self.transform(sample)

        return sample