import torch
import os
import torchio
from torchio import Image, ImagesDataset, SubjectsDataset


class SemSegConfig():
    train_images = None
    train_labels = None
    val_images   = None
    val_labels   = None
    do_normalize = True
    augmentation = None
    zero_pad     = True
    pad_ref      = (64,64,64)
    batch_size   = 4
    num_workers  = 0


def TorchIODataLoader3DTraining(config: SemSegConfig) -> torch.utils.data.DataLoader:
    print('Building TorchIO Training Set Loader...')
    subject_list = list()
    for idx, (image_path, label_path) in enumerate(zip(config.train_images, config.train_labels)):
        s1 = torchio.Subject(
            t1=Image(type=torchio.INTENSITY, path=image_path),
            label=Image(type=torchio.LABEL, path=label_path),
        )

        subject_list.append(s1)

    # Deprecated
    # subjects_dataset = ImagesDataset(subject_list, transform=config.transform_train)
    subjects_dataset = SubjectsDataset(subject_list, transform=config.transform_train)
    train_data = torch.utils.data.DataLoader(subjects_dataset, batch_size=config.batch_size,
                                             shuffle=True, num_workers=config.num_workers)
    print('TorchIO Training Loader built!')
    return train_data


def TorchIODataLoader3DValidation(config: SemSegConfig) -> torch.utils.data.DataLoader:
    print('Building TorchIO Validation Set Loader...')
    subject_list = list()
    for idx, (image_path, label_path) in enumerate(zip(config.val_images, config.val_labels)):
        s1 = torchio.Subject(
            t1=Image(type=torchio.INTENSITY, path=image_path),
            label=Image(type=torchio.LABEL, path=label_path),
        )

        subject_list.append(s1)

    # Deprecated
    # subjects_dataset = ImagesDataset(subject_list, transform=config.transform_val)
    subjects_dataset = SubjectsDataset(subject_list, transform=config.transform_val)
    val_data = torch.utils.data.DataLoader(subjects_dataset, batch_size=config.batch_size,
                                           shuffle=False, num_workers=config.num_workers)
    print('TorchIO Validation Loader built!')
    return val_data


def get_pad_3d_image(pad_ref: tuple = (64, 64, 64), zero_pad: bool = True):
    def pad_3d_image(image):
        if zero_pad:
            value_to_pad = 0
        else:
            value_to_pad = image.min()
        pad_ref_channels = (image.shape[0], *pad_ref)
        # print("image.shape = {}".format(image.shape))
        if value_to_pad == 0:
            image_padded = torch.zeros(pad_ref_channels)
        else:
            image_padded = value_to_pad * torch.ones(pad_ref_channels)
        image_padded[:,:image.shape[1],:image.shape[2],:image.shape[3]] = image
        # print("image_padded.shape = {}".format(image_padded.shape))
        return image_padded
    return pad_3d_image


def z_score_normalization(inputs):
    input_mean = torch.mean(inputs)
    input_std = torch.std(inputs)
    return (inputs - input_mean)/input_std