"""

Tools for converting, normalizing, and fixing the brats data.

"""

import glob

import os

import warnings

import shutil

import SimpleITK as sitk

import numpy as np

from nipype.interfaces.ants import N4BiasFieldCorrection

# from brats.train_fetal import config

def append_basename(in_file, append):

    dirname, basename = os.path.split(in_file)

    base, ext = basename.split(".", 1)

    return os.path.join(dirname, base + append + "." + ext)

def get_background_mask(in_folder, out_file, truth_name="GlistrBoost_ManuallyCorrected"):

    """

    This function computes a common background mask for all of the data in a subject folder.

    :param in_folder: a subject folder from the BRATS dataset.

    :param out_file: an image containing a mask that is 1 where the image data for that subject contains the background.

    :param truth_name: how the truth file is labeled int he subject folder

    :return: the path to the out_file

    """

    background_image = None

    for name in config["all_modalities"] + [truth_name]:

        image = sitk.ReadImage(get_image(in_folder, name))

        if background_image:

            if name == truth_name and not (image.GetOrigin() == background_image.GetOrigin()):

                image.SetOrigin(background_image.GetOrigin())

            background_image = sitk.And(image == 0, background_image)

        else:

            background_image = image == 0

    sitk.WriteImage(background_image, out_file)

    return os.path.abspath(out_file)

def convert_image_format(in_file, out_file):

    sitk.WriteImage(sitk.ReadImage(in_file), out_file)

    return out_file

def window_intensities_data(data, min_percent=1, max_percent=99):

    image = sitk.GetImageFromArray(data)

    image = sitk.IntensityWindowing(image,

                                    np.percentile(data, min_percent),

                                    np.percentile(data, max_percent))

    return sitk.GetArrayFromImage(image)

def window_intensities(in_file, out_file, min_percent=1, max_percent=99):

    image = sitk.ReadImage(in_file, sitk.sitkFloat32)

    image_data = sitk.GetArrayFromImage(image)

    out_image = sitk.IntensityWindowing(image, np.percentile(image_data, min_percent),

                                               np.percentile(image_data, max_percent))

    sitk.WriteImage(out_image, out_file)

    return os.path.abspath(out_file)

def correct_bias(in_file, out_file, image_type=sitk.sitkFloat64):

    """

    Corrects the bias using ANTs N4BiasFieldCorrection. If this fails, will then attempt to correct bias using SimpleITK

    :param in_file: input file path

    :param out_file: output file path

    :param image_type:

    :return: file path to the bias corrected image

    """

    correct = N4BiasFieldCorrection()

    correct.inputs.input_image = in_file

    correct.inputs.output_image = out_file

    try:

        done = correct.run()

        return done.outputs.output_image

    except IOError:

        warnings.warn(RuntimeWarning("ANTs N4BIasFieldCorrection could not be found."

                                     "Will try using SimpleITK for bias field correction"

                                     " which will take much longer. To fix this problem, add N4BiasFieldCorrection"

                                     " to your PATH system variable. (example: EXPORT PATH=${PATH}:/path/to/ants/bin)"))

        input_image = sitk.ReadImage(in_file, image_type)

        output_image = sitk.N4BiasFieldCorrection(input_image, input_image > 0)

        sitk.WriteImage(output_image, out_file)

        return os.path.abspath(out_file)

def rescale(in_file, out_file, minimum=0, maximum=20000):

    image = sitk.ReadImage(in_file)

    sitk.WriteImage(sitk.RescaleIntensity(image, minimum, maximum), out_file)

    return os.path.abspath(out_file)

def get_image(subject_folder, name):

    file_card = os.path.join(subject_folder, "*" + name + ".nii")

    try:

        return glob.glob(file_card)[0]

    except IndexError:

        raise RuntimeError("Could not find file matching {}".format(file_card))

def background_to_zero(in_file, background_file, out_file):

    sitk.WriteImage(sitk.Mask(sitk.ReadImage(in_file), sitk.ReadImage(background_file, sitk.sitkUInt8) == 0),

                    out_file)

    return out_file

def check_origin(in_file, in_file2):

    image = sitk.ReadImage(in_file)

    image2 = sitk.ReadImage(in_file2)

    if not image.GetOrigin() == image2.GetOrigin():

        image.SetOrigin(image2.GetOrigin())

        sitk.WriteImage(image, in_file)

def normalize_image(in_file, out_file, bias_correction=True):

    if bias_correction:

        correct_bias(in_file, out_file)

    else:

        shutil.copy(in_file, out_file)

    return out_file

def convert_brats_folder(in_folder, out_folder, truth_name="truth",

                         no_bias_correction_modalities=None):

    for name in config["all_modalities"]:

        image_file = get_image(in_folder, name)

        out_file = os.path.abspath(os.path.join(out_folder, name + ".nii"))

        perform_bias_correction = no_bias_correction_modalities and name not in no_bias_correction_modalities

        normalize_image(image_file, out_file, bias_correction=perform_bias_correction)

    # copy the truth file

    try:

        truth_file = get_image(in_folder, truth_name)

    except RuntimeError:

        truth_file = get_image(in_folder, truth_name.split("_")[0])

    out_file = os.path.abspath(os.path.join(out_folder, "truth.nii.gz"))

    shutil.copy(truth_file, out_file)

    check_origin(out_file, get_image(in_folder, config["all_modalities"][0]))

def convert_brats_data(brats_folder, out_folder, overwrite=False, no_bias_correction_modalities=("flair",)):

    """

    Preprocesses the BRATS data and writes it to a given output folder. Assumes the original folder structure.

    :param brats_folder: folder containing the original brats data

    :param out_folder: output folder to which the preprocessed data will be written

    :param overwrite: set to True in order to redo all the preprocessing

    :param no_bias_correction_modalities: performing bias correction could reduce the signal of certain modalities. If

    concerned about a reduction in signal for a specific modality, specify by including the given modality in a list

    or tuple.

    :return:

    """

    for subject_folder in glob.glob(os.path.join(brats_folder, "*")):

        if os.path.isdir(subject_folder):

            subject = os.path.basename(subject_folder)

            new_subject_folder = os.path.join(out_folder, os.path.basename(os.path.dirname(subject_folder)),

                                              subject)

            if not os.path.exists(new_subject_folder) or overwrite:

                if not os.path.exists(new_subject_folder):

                    os.makedirs(new_subject_folder)

                print('subject_folder: ' + subject_folder)

                convert_brats_folder(subject_folder, new_subject_folder,

                                     no_bias_correction_modalities=no_bias_correction_modalities)