# 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 albumentations.pytorch.transforms import ToTensorV2
from torch.utils.data import Dataset, DataLoader
import cv2 as cv
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 *

# %% ../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 = cv.cvtColor(cv.imread(str(self.image_fns[i])), cv.COLOR_BGR2RGB)
        if self.is_org:
            # return empty mask if original (not annotated) image
            mask = np.zeros(image.shape)
        else:
            mask = cv.imread(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)
