from __future__ import print_function

import numpy as np

import os

from skimage.io import imread

from skimage.transform import rescale

from skimage.transform import rotate

image_rows = 256

image_cols = 256

channels = 3    # refers to neighboring slices; if set to 3, takes previous and next slice as additional channels

modalities = 1  # refers to pre, flair and post modalities; if set to 3, uses all and if set to 1, only flair

def load_data(path):

    """

    Assumes filenames in given path to be in the following format as defined in `preprocessing3D.m`:

    for images: <case_id>_<slice_number>.tif

    for masks: <case_id>_<slice_number>_mask.tif

        Args:

            path: string to the folder with images

        Returns:

            np.ndarray: array of images

            np.ndarray: array of masks

            np.chararray: array of corresponding images' filenames without extensions

    """

    images_list = os.listdir(path)

    total_count = len(images_list) / 2

    images = np.ndarray((total_count, image_rows, image_cols,

                         channels * modalities), dtype=np.uint8)

    masks = np.ndarray((total_count, image_rows, image_cols), dtype=np.uint8)

    names = np.chararray(total_count, itemsize=64)

    i = 0

    for image_name in images_list:

        if 'mask' in image_name:

            continue

        names[i] = image_name.split('.')[0]

        slice_number = int(names[i].split('_')[-1])

        patient_id = '_'.join(names[i].split('_')[:-1])

        image_mask_name = image_name.split('.')[0] + '_mask.tif'

        img = imread(os.path.join(path, image_name), as_grey=(modalities == 1))

        img_mask = imread(os.path.join(path, image_mask_name), as_grey=True)

        if channels > 1:

            img_prev = read_slice(path, patient_id, slice_number - 1)

            img_next = read_slice(path, patient_id, slice_number + 1)

            img = np.dstack((img_prev, img[..., np.newaxis], img_next))

        elif modalities == 1:

            img = np.array([img])

        img_mask = np.array([img_mask])

        images[i] = img

        masks[i] = img_mask

        i += 1

    images = images.astype('float32')

    masks = masks[..., np.newaxis]

    masks = masks.astype('float32')

    masks /= 255.

    return images, masks, names

def read_slice(path, patient_id, slice):

    img = np.zeros((image_rows, image_cols))

    img_name = patient_id + '_' + str(slice) + '.tif'

    img_path = os.path.join(path, img_name)

    try:

        img = imread(img_path, as_grey=(modalities == 1))

    except Exception:

        pass

    return img[..., np.newaxis]