from __future__ import print_function

import os

import numpy as np
from skimage.io import imread
from skimage.transform import rescale
from skimage.transform import rotate

image_rows = 256
image_cols = 256

channels = 1    # refers to neighboring slices; if set to 3, takes previous and next slice as additional channels
modalities = 3  # 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 oversample(images, masks, augment=False):
    """
    Repeats 2 times every slice with nonzero mask.

        Args:
            np.ndarray: array of images
            np.ndarray: array of masks

        Returns:
            np.ndarray: array of oversampled images
            np.ndarray: array of oversampled masks
    """
    images_o = []
    masks_o = []
    for i in range(len(masks)):
        if np.max(masks[i]) < 1:
            continue

        if augment:
            image_a, mask_a = augmentation_rotate(images[i], masks[i])
            images_o.append(image_a)
            masks_o.append(mask_a)
            image_a, mask_a = augmentation_scale(images[i], masks[i])
            images_o.append(image_a)
            masks_o.append(mask_a)
            continue

        for _ in range(2):
            images_o.append(images[i])
            masks_o.append(masks[i])

    images_o = np.array(images_o)
    masks_o = np.array(masks_o)

    return np.vstack((images, images_o)), np.vstack((masks, masks_o))


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]


def augmentation_rotate(img, img_mask):
    angle = np.random.uniform(5.0, 15.0) * np.random.choice([-1.0, 1.0], 1)[0]

    img = rotate(img, angle, resize=False, order=3, preserve_range=True)
    img_mask = rotate(img_mask, angle, resize=False, order=0, preserve_range=True)

    return img, img_mask


def augmentation_scale(img, img_mask):
    scale = 1.0 + np.random.uniform(0.04, 0.08) * np.random.choice([-1.0, 1.0], 1)[0]

    img = rescale(img, scale, order=3, preserve_range=True)
    img_mask = rescale(img_mask, scale, order=0, preserve_range=True)
    if scale > 1:
        img = center_crop(img, image_rows, image_cols)
        img_mask = center_crop(img_mask, image_rows, image_cols)
    else:
        img = zeros_pad(img, image_rows)
        img_mask = zeros_pad(img_mask, image_rows)

    return img, img_mask


def center_crop(img, cropx, cropy):
    startx = img.shape[1] // 2 - (cropx // 2)
    starty = img.shape[0] // 2 - (cropy // 2)
    return img[starty : starty + cropy, startx : startx + cropx]


def zeros_pad(img, size):
    pad_before = int(round(((size - img.shape[0]) / 2.0)))
    pad_after = size - img.shape[0] - pad_before
    if len(img.shape) > 2:
        return np.pad(
            img,
            ((pad_before, pad_after), (pad_before, pad_after), (0, 0)),
            mode="constant",
        )
    return np.pad(img, (pad_before, pad_after), mode="constant")