import random

import torch

from torch.utils.data import Dataset, DataLoader

import numpy as np

from sklearn.model_selection import train_test_split, KFold

from utils import get_paths

import albumentations as A

class Covid19Dataset(Dataset):

    def __init__(self, paths, transform=None):

        self.paths = paths

        self.transform = transform

        self._len = len(self.paths)

    def __len__(self):

        return self._len

    def __getitem__(self, index):

        path = self.paths[index]

        loaded = np.load(path)

        image = loaded['image']

        mask = loaded['mask']

        if self.transform:

            transformed = self.transform(image=image, mask=mask)

            image = transformed['image']

            mask = transformed['mask']

        image = torch.from_numpy(np.array([image], dtype=np.float))

        image = image.type(torch.FloatTensor)

        mask = torch.from_numpy(np.array([mask], dtype=np.uint8))

        return image, mask

def data_generator(cfg):

    image_paths = get_paths(cfg)

    image_paths = np.asarray(image_paths)

    train_paths, val_paths = [], []

    if not cfg.kfold:

        _train_paths, _val_paths = train_test_split(image_paths, test_size=cfg.val_size, random_state=cfg.seed)

    else:

        kf = KFold(n_splits=cfg.n_splits)

        for i, (train_index, val_index) in enumerate(kf.split(image_paths)):

            if i + 1 == cfg.fold_number:

                _train_paths = image_paths[train_index]

                _val_paths = image_paths[val_index]

    for paths in _train_paths:

        train_paths.extend(paths)

    for paths in _val_paths:

        val_paths.extend(paths)

    random.shuffle(train_paths)

    random.shuffle(val_paths)

    return train_paths, val_paths

def get_transforms(cfg):

    # getting transforms from albumentations

    pre_transforms = [getattr(A, item["name"])(**item["params"]) for item in cfg.pre_transforms]

    augmentations = [getattr(A, item["name"])(**item["params"]) for item in cfg.augmentations]

    post_transforms = [getattr(A, item["name"])(**item["params"]) for item in cfg.post_transforms]

    # concatenate transforms

    train = A.Compose(pre_transforms + augmentations + post_transforms)

    test = A.Compose(pre_transforms + post_transforms)

    return train, test

def get_loaders(cfg):

    train_transforms, test_transforms = get_transforms(cfg)

    train_paths, val_paths = data_generator(cfg)

    train_ds = Covid19Dataset(train_paths, transform=train_transforms)

    val_ds = Covid19Dataset(val_paths, transform=train_transforms)

    train_dl = DataLoader(train_ds, batch_size=cfg.batch_size, drop_last=True)

    val_dl = DataLoader(val_ds, batch_size=cfg.batch_size, drop_last=True)

    return train_dl, val_dl