import torch
import torchio
from torchio import ImagesDataset, Image, ZNormalization, Compose
from config.config import *
from config.paths import train_images_folder, train_labels_folder, train_images, train_labels
from semseg.data_loader_torchio import get_pad_3d_image
from augm.lambda_channel import LambdaChannel

transforms_dict = {
    ZNormalization(): 1.0,
    # LambdaChannel(z_score_normalization, types_to_apply=torchio.INTENSITY): 1.0,
    # RandomAffine(): 0.25,
    # RandomElasticDeformation(max_displacement=3): 0.25,
    LambdaChannel(get_pad_3d_image(pad_ref=(48, 64, 48), zero_pad=False)): 1.0,
}

transform = Compose(transforms_dict)

subject_list = list()

idx = 0
for idx,(train_image, train_label) in enumerate(zip(train_images, train_labels)):
    image_path = os.path.join(train_images_folder, train_image)
    label_path = os.path.join(train_labels_folder, train_label)

    s1 = torchio.Subject(
        t1    = Image(type=torchio.INTENSITY, path=image_path),
        label = Image(type=torchio.LABEL, path=label_path),
    )

    subject_list.append(s1)

subjects_dataset = ImagesDataset(subject_list, transform=transform)
subject_sample = subjects_dataset[0]

for idx in range(0,len(train_images[:10])):
    subject_sample = subjects_dataset[idx]
    print("Iter {} on {}".format(idx+1,len(train_images)))
    print("t1.shape       = {}".format(subject_sample.t1.shape))
    print("label.shape    = {}".format(subject_sample.label.shape))
    print("t1 [min - max] = [{:.1f} : {:.1f}]".format(subject_sample.t1.data.min(),subject_sample.t1.data.max()))
    print("label.unique   = {}".format(subject_sample.label.data.unique()))


config = SemSegMRIConfig()
train_data = torch.utils.data.DataLoader(subjects_dataset, batch_size=config.batch_size,
                                         shuffle=False, num_workers=config.num_workers)


iterable_data_loader = iter(train_data)
el = next(iterable_data_loader)
inputs = el['t1']['data']
labels = el['label']['data']
print("Shape of Batch: [input {}] [label {}]".format(inputs.shape, labels.shape))

for i, el in enumerate(train_data):
    print("Iteration {} on {}".format(i+1,len(train_data)))
    inputs = el['t1']['data']
    labels = el['label']['data']
    print("Shape of Batch: [input {}] [label {}]".format(inputs.shape, labels.shape))
    print("Range Inputs  : [{:.2f} %% {:.2f}]".format(inputs.min(),inputs.max()))