# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_data.ipynb.

# %% auto 0
__all__ = ['get_dataset_fns', 'get_preprocessing', 'get_augmentation', 'ATNCDataset', 'DataLoaders', 'get_datasets',
           'get_dataloaders', 'preproc_img']

# %% ../nbs/02_data.ipynb 3
from pathlib import Path
import random
import albumentations as album
from torch.utils.data import Dataset, DataLoader
import numpy as np
from .utility import to_tensor, mask_map_classes, mask_code_to_continues_indices, load_cls_labels, mask_continues_indices_to_codes, mask_to_rgb
from .visualization import *
from .model import *
from PIL import Image

# %% ../nbs/02_data.ipynb 5
def get_dataset_fns(path, data_dirs, img_dirs, seed=42, ext='png', dist=[0.8,0.1,0.1], img_count=-1):
    """
    Returns 3 randomized lists of filenames for train, val, and test datasets
    """
    all_fns = []
    for data_dir in data_dirs:
        for img_dir in img_dirs:
            img_path = path/data_dir/img_dir
            all_fns += sorted(list(img_path.glob('*.'+ext)))

    random.Random(seed).shuffle(all_fns)
    if img_count == -1:
        img_count = len(all_fns)
    train_count = round(img_count * dist[0])
    val_count = round(img_count * dist[1])
    test_count = img_count - train_count - val_count
        
    train_fns = all_fns[0:train_count]
    val_fns = all_fns[train_count:val_count+train_count]
    test_fns = all_fns[train_count + val_count:train_count + val_count + test_count]
    
    return train_fns, val_fns, test_fns
                


# %% ../nbs/02_data.ipynb 8
def get_preprocessing(preprocessing_fn=None):
    """Construct preprocessing transform    
    Args:
        preprocessing_fn (callable): data normalization function 
            (can be specific for each pretrained neural network)
    Return:
        transform: albumentations.Compose
    """ 
    _transform = []
    if preprocessing_fn:
        _transform.append(album.Lambda(image=preprocessing_fn))
    _transform.append(album.Lambda(image=to_tensor, mask=to_tensor))
    
    return album.Compose(_transform)

# %% ../nbs/02_data.ipynb 10
def get_augmentation(transforms):
    """Construct augmentation transform
    Args:
        list of augmentation transforms
    Return:
        transform: albumentations.Compose
    """
    return album.Compose(transforms)

# %% ../nbs/02_data.ipynb 13
class ATNCDataset(Dataset):
    """Dataset. Read images, apply augmentation and preprocessing transformations.
    
    Args:
        image_fns ([Paths]): paths to images folder
        mask_path (str): path to segmentation masks folder
        cls_codes (list): class codes of select classes to extract from segmentation mask
        augmentation (albumentations.Compose): data transfromation pipeline 
            (e.g. flip, scale, etc.)
        preprocessing (albumentations.Compose): data preprocessing 
            (e.g. noralization, shape manipulation, etc.)
    
    """
    
    def __init__(
            self,
            image_fns,
            mask_path,
            cls_codes=None,
            class_mappings=None,
            augmentation=None,
            preprocessing=None,
            is_org=False
            ):
        self.image_fns = image_fns
        self.mask_fns = self.get_mask_fns(image_fns, mask_path)
        self.cls_codes = cls_codes
        self.augmentation = augmentation
        self.preprocessing = preprocessing
        self.is_org=is_org
        self.class_mappings = class_mappings
    
    def __getitem__(self, i):
        image = np.array(Image.open(str(self.image_fns[i])))
        if self.is_org:
            # return empty mask if original (not annotated) image
            mask = np.zeros(image.shape)
        else:
            mask = np.array(Image.open(str(self.mask_fns[i])))#[:,:,0]

        if self.class_mappings is not None:
            mask = mask_map_classes(mask, self.class_mappings)
        if self.cls_codes is not None:
           mask = mask_code_to_continues_indices(mask, self.cls_codes)
        
        if self.augmentation:
            sample = self.augmentation(image=image, mask=mask)
            image, mask = sample['image'], sample['mask']
            
        if self.preprocessing:
            sample = self.preprocessing(image=image, mask=mask)
            image, mask = sample['image'], sample['mask']
        
        return image, mask.astype('int')
    
    def __len__(self):
        return len(self.image_fns)

    #def get_mask_fns(self, image_fns, mask_path):
    #    return [o.parent.parent/mask_path/('seg_'+o.name.split('_',1)[-1]) for o in image_fns]
    
    def get_mask_fns(self, image_fns, mask_path):
        mask_fns = []
        for o in image_fns:
            if 'img' in o.name:
                source_name = 'img'
            else:
                source_name = 'noisyImg'
            
            mask_fn = o.parent.parent/mask_path/(o.name.replace(source_name,'seg'))
            if mask_fn.exists() == False:
                mask_name = mask_fn.name.split('.png')[0] + '.tif'
                mask_fn = mask_fn.parent/mask_name
            mask_fns.append(mask_fn)    
        return mask_fns #[o.parent.parent/mask_path/(o.name.replace('img','seg')) for o in image_fns]

# %% ../nbs/02_data.ipynb 15
class DataLoaders():
    def __init__(self, train_dl, valid_dl):
        self.train = train_dl
        self.valid= valid_dl

def get_datasets(train_fns, val_fns, masks_folder, cls_codes, class_mappings, train_augmentation, valid_augmentation, model_def):
    preprocessing = get_preprocessing(get_preprocessiong_fn(model_def))
    train_ds = ATNCDataset(image_fns=train_fns,
                           mask_path=masks_folder,
                           cls_codes=cls_codes,
                           class_mappings=class_mappings,
                           augmentation=train_augmentation,
                           preprocessing=preprocessing,
                           is_org=False)
    valid_ds = ATNCDataset(image_fns=val_fns,
                           mask_path=masks_folder,
                           cls_codes=cls_codes,
                           class_mappings=class_mappings,
                           augmentation=valid_augmentation,
                           preprocessing=preprocessing,
                           is_org=False)
    return train_ds, valid_ds

def get_dataloaders(train_ds, valid_ds, train_bs=4,valid_bs = 4, num_workers = 4):
    train_dl = DataLoader(train_ds, batch_size=train_bs, shuffle=True, num_workers=num_workers)
    valid_dl = DataLoader(valid_ds, batch_size=valid_bs, shuffle=False, num_workers=num_workers)
    return DataLoaders(train_dl, valid_dl)

# %% ../nbs/02_data.ipynb 17
def preproc_img(img, model_def, aug_transform=[]):
    preproc = get_preprocessiong_fn(model_def)
    transform = [
        album.PadIfNeeded(min_height = None, min_width = None, pad_height_divisor=32, pad_width_divisor=32, always_apply=True),
            ] + aug_transform
    aug = get_augmentation(transform)
    tmp_img = aug(image = img)['image']
    tmp_img = preproc(tmp_img)
    
    return to_tensor(tmp_img).unsqueeze(dim=0)
