"""Functions for reading MSISBI2015 NRRD data."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import glob

import math

import os.path

import pickle

import matplotlib.pyplot

from imageio import imwrite

from scipy.ndimage import zoom

from six.moves import xrange

from skimage.measure import label, regionprops

from utils.NII import *

from utils.image_utils import crop, crop_center

from utils.tfrecord_utils import *

matplotlib.pyplot.ion()

class MSISBI2015(object):

    PROTOCOL_MAPPINGS = {'FLAIR': ['flair'], 'MPRAGE': ['mprage'], 'PD': ['pd'], 'T2': ['t2']}

    SET_TYPES = ['TRAIN', 'VAL', 'TEST']

    class Options(object):

        def __init__(self):

            self.dir = os.path.dirname(os.path.realpath(__file__))

            self.numSamples = -1

            self.partition = {'TRAIN': 0.7, 'VAL': 0.2, 'TEST': 0.1}

            self.useCrops = False

            self.cropType = 'random'  # random or center

            self.numRandomCropsPerSlice = 5

            self.onlyPatchesWithLesions = False

            self.rotations = 0

            self.cropWidth = 128

            self.cropHeight = 128

            self.cache = False

            self.sliceResolution = None  # format: HxW

            self.addInstanceNoise = False  # Affects only the batch sampling. If True, a tiny bit of noise will be added to every batch

            self.filterProtocol = None  # FLAIR, T1, T2

            self.filterType = "train"  # train or test

            self.axis = 'axial'  # saggital, coronal or axial

            self.debug = False

            self.normalizationMethod = 'standardization'

            self.sliceStart = 0

            self.sliceEnd = 155

            self.format = "raw"  # raw or aligned; If aligned, nii-files will be crawled and loaded

            self.skullStripping = True

            self.viewMapping = {'saggital': 2, 'coronal': 1, 'axial': 0}

    def __init__(self, options=Options()):

        self.options = options

        if options.cache and os.path.isfile(self.pckl_name()):

            f = open(self.pckl_name(), 'rb')

            tmp = pickle.load(f)

            f.close()

            self._epochs_completed = tmp._epochs_completed

            self._index_in_epoch = tmp._index_in_epoch

            self.patientsSplit = tmp.patients_split

            self.patients = tmp.patients

            self._images, self._labels, self._sets = read_tf_record(self.tfrecord_name())

            self._epochs_completed = {'TRAIN': 0, 'VAL': 0, 'TEST': 0}

            self._index_in_epoch = {'TRAIN': 0, 'VAL': 0, 'TEST': 0}

        else:

            # Collect all patients

            self.patients = self._get_patients()

            self.patientsSplit = {}

            if not os.path.isfile(self.split_name()):

                _numPatients = len(self.patients)

                _ridx = numpy.random.permutation(_numPatients)

                _already_taken = 0

                for split in self.options.partition.keys():

                    if self.options.partition[split] <= 1.0:

                        num_patients_for_current_split = math.floor(self.options.partition[split] * _numPatients)

                    else:

                        num_patients_for_current_split = self.options.partition[split]

                    if num_patients_for_current_split > (_numPatients - _already_taken):

                        num_patients_for_current_split = _numPatients - _already_taken

                    self.patientsSplit[split] = _ridx[_already_taken:_already_taken + num_patients_for_current_split]

                    _already_taken += num_patients_for_current_split

                f = open(self.split_name(), 'wb')

                pickle.dump(self.patientsSplit, f)

                f.close()

            else:

                f = open(self.split_name(), 'rb')

                self.patientsSplit = pickle.load(f)

                f.close()

            self._create_numpy_arrays()

            self._epochs_completed = {'TRAIN': 0, 'VAL': 0, 'TEST': 0}

            self._index_in_epoch = {'TRAIN': 0, 'VAL': 0, 'TEST': 0}

            if self.options.cache:

                write_tf_record(self._images, self._labels, self._sets, self.tfrecord_name())

                tmp = copy.copy(self)

                tmp._images = None

                tmp._labels = None

                tmp._sets = None

                f = open(self.pckl_name(), 'wb')

                pickle.dump(tmp, f)

                f.close()

    def _create_numpy_arrays(self):

        # Iterate over all patients and extract slices

        _images = []

        _labels = []

        _sets = []

        for p, patient in enumerate(self.patients):

            if p in self.patientsSplit['TRAIN']:

                _set_of_current_patient = MSISBI2015.SET_TYPES.index('TRAIN')

            elif p in self.patientsSplit['VAL']:

                _set_of_current_patient = MSISBI2015.SET_TYPES.index('VAL')

            elif p in self.patientsSplit['TEST']:

                _set_of_current_patient = MSISBI2015.SET_TYPES.index('TEST')

            for n, nii_filename in enumerate(patient['filtered_files']):

                # try:

                _images_tmp, _labels_tmp = self.gather_data(patient, nii_filename)

                _images += _images_tmp

                _labels += _labels_tmp

                _sets += [_set_of_current_patient] * len(_images_tmp)

        self._images = numpy.array(_images).astype(numpy.float32)

        self._labels = numpy.array(_labels).astype(numpy.float32)

        if self._images.ndim < 4:

            self._images = numpy.expand_dims(self._images, 3)

        self._sets = numpy.array(_sets).astype(numpy.int32)

    def gather_data(self, patient, nii_filename):

        _images = []

        _labels = []

        nii, nii_seg, nii_skullmap = self.load_volume_and_groundtruth(nii_filename, patient)

        # Iterate over all slices and collect them

        # We only want to select in the range from 15 to 125 (in axial view)

        for s in xrange(self.options.sliceStart, min(self.options.sliceEnd, nii.num_slices_along_axis(self.options.axis))):

            if 0 < self.options.numSamples < len(_images):

                break

            slice_data = nii.get_slice(s, self.options.axis)

            slice_seg = nii_seg.get_slice(s, self.options.axis)

            # Skip the slice if it is "empty"

            if numpy.percentile(slice_data, 90) < 0.2:

                continue

            if self.options.sliceResolution is not None:

                # Pad withzeros to top and bottom, if the image is too small

                if slice_data.shape[0] < self.options.sliceResolution[0]:

                    before_y = math.floor((self.options.sliceResolution[0] - slice_data.shape[0]) / 2.0)

                    after_y = math.ceil((self.options.sliceResolution[0] - slice_data.shape[0]) / 2.0)

                if slice_data.shape[1] < self.options.sliceResolution[1]:

                    before_x = math.floor((self.options.sliceResolution[1] - slice_data.shape[1]) / 2.0)

                    after_x = math.ceil((self.options.sliceResolution[1] - slice_data.shape[1]) / 2.0)

                if slice_data.shape[0] < self.options.sliceResolution[0] or slice_data.shape[1] < self.options.sliceResolution[1]:

                    slice_data = np.pad(slice_data, ((before_y, after_y), (before_x, after_x)), 'constant', constant_values=(0, 0))

                    slice_seg = np.pad(slice_seg, ((before_y, after_y), (before_x, after_x)), 'constant', constant_values=(0, 0))

                slice_data = zoom(slice_data, float(self.options.sliceResolution[0]) / float(

                    slice_data.shape[0]))

                slice_seg = zoom(

                    slice_seg, float(self.options.sliceResolution[0]) / float(slice_seg.shape[0]), mode="nearest"

                )

                slice_seg[slice_seg < 0.9] = 0.0

                slice_seg[slice_seg >= 0.9] = 1.0

                # assert numpy.max(slice_data) <= 1.0, "Resized slice range is outside [0; 1]!"

            # Either collect crops

            if self.options.useCrops:

                if self.options.cropType == 'random':

                    rx = numpy.random.randint(0, high=(slice_data.shape[1] - self.options.cropWidth),

                                              size=self.options.numRandomCropsPerSlice)

                    ry = numpy.random.randint(0, high=(slice_data.shape[0] - self.options.cropHeight),

                                              size=self.options.numRandomCropsPerSlice)

                    for r in range(self.options.numRandomCropsPerSlice):

                        _images.append(crop(slice_data, ry(r), rx(r), self.options.cropHeight, self.options.cropWidth))

                        _labels.append(crop(slice_data, ry(r), rx(r), self.options.cropHeight, self.options.cropWidth))

                elif self.options.cropType == 'center':

                    slice_data_cropped = crop_center(slice_data, self.options.cropWidth, self.options.cropHeight)

                    slice_seg_cropped = crop_center(slice_seg, self.options.cropWidth, self.options.cropHeight)

                    _images.append(slice_data_cropped)

                    _labels.append(slice_seg_cropped)

                elif self.options.cropType == 'lesions':

                    cc_slice = label(slice_seg)

                    props = regionprops(cc_slice)

                    if len(props) > 0:

                        for prop in props:

                            cx = prop['centroid'][1]

                            cy = prop['centroid'][0]

                            if cy < self.options.cropHeight // 2:

                                cy = self.options.cropHeight // 2

                            if cy > (slice_data.shape[0] - (self.options.cropHeight // 2)):

                                cy = (slice_data.shape[0] - (self.options.cropHeight // 2))

                            if cx < self.options.cropWidth // 2:

                                cx = self.options.cropWidth // 2

                            if cx > (slice_data.shape[1] - (self.options.cropWidth // 2)):

                                cx = (slice_data.shape[1] - (self.options.cropWidth // 2))

                            image_crop = crop(slice_data, int(cy) - (self.options.cropHeight // 2), int(cx) - (self.options.cropWidth // 2),

                                              self.options.cropHeight, self.options.cropWidth)

                            seg_crop = crop(slice_seg, int(cy) - (self.options.cropHeight // 2), int(cx) - (self.options.cropWidth // 2),

                                            self.options.cropHeight, self.options.cropWidth)

                            if image_crop.shape[0] != self.options.cropHeight or image_crop.shape[1] != self.options.cropWidth:

                                continue

                            _images.append(image_crop)

                            _labels.append(seg_crop)

            # Or whole slices

            else:

                _images.append(slice_data)

                _labels.append(slice_seg)

        return _images, _labels

    def load_volume_and_groundtruth(self, nii_filename, patient):

        # Load the nrrd

        try:

            nii = NII(nii_filename)

            nii_groundtruth = NII(patient['groundtruth'])

            nii.denoise()

            nii.set_view_mapping(self.options.viewMapping)

        except:

            print('MSISBI2015: Failed to open file ' + nii_filename)

        # Make sure ground-truth is binary and nrrd doesnt have NaNs

        nii.data[np.isnan(nii.data)] = 0.0

        nii_groundtruth.data[nii_groundtruth.data < 0.9] = 0.0

        nii_groundtruth.data[nii_groundtruth.data >= 0.9] = 1.0

        # Do skull-stripping, if desired

        if self.options.skullStripping:

            try:

                nii_skullmap = NII(patient['skullmap'])

                nii_skullmap.set_view_mapping(self.options.viewMapping)

                nii.apply_skullmap(nii_skullmap)

            except:

                print('MSISBI2015: Failed to open file ' + patient['skullmap'] + ', skipping skullremoval')

        # In-place normalize the loaded volume

        nii.normalize(method=self.options.normalizationMethod, lowerpercentile=0, upperpercentile=99.8)

        # nii_skullmap.data = nii_skullmap.data > 0.0

        return nii, nii_groundtruth, nii_skullmap

    # Hidden helper function, not supposed to be called from outside!

    def _get_patients(self):

        return MSISBI2015.get_patients(self.options)

    @staticmethod

    def get_patients(options):

        folders = ["training01", "training02", "training03", "training04", "training05"]

        # Iterate over all folders in folders and collect all patients

        patients = []

        for f, folder in enumerate(folders):

            # Get all files that can be used for training and validation

            _patients = glob.glob(os.path.join(options.dir, folder, "preprocessed", folder + "_*_flair_pp.nii"))

            for p, pname in enumerate(_patients):

                patient = {}

                _tmp = os.path.normpath(pname).split(os.path.sep)

                patient['name'] = _tmp[-1].replace("_flair_pp.nii", "")

                patient['fullpath'] = os.path.join(options.dir, folder, "preprocessed")

                patient["filtered_files"] = []

                for protocol, protocol_array in MSISBI2015.PROTOCOL_MAPPINGS.items():

                    if len(options.filterProtocols) > 0 and protocol not in options.filterProtocols:

                        continue

                    else:

                        if options.format == "raw":

                            patient[protocol] = os.path.join(options.dir, folder, "preprocessed", patient['name'] + '_' + protocol_array[0] + '_pp.nii')

                        elif options.format == "aligned":

                            patient[protocol] = os.path.join(options.dir, folder, "preprocessed",

                                                             patient['name'] + '_' + protocol_array[0] + '.aligned.nii.gz')

                    if len(options.filterProtocols) > 0 and protocol not in options.filterProtocols:

                        continue

                    else:

                        if options.format == "raw":

                            patient["filtered_files"] += [

                                os.path.join(options.dir, folder, "preprocessed", patient['name'] + '_' + protocol_array[0] + '_pp.nii')]

                        elif options.format == "aligned":

                            patient["filtered_files"] += [

                                os.path.join(options.dir, folder, "preprocessed", patient['name'] + '_' + protocol_array[0] + '.aligned.nii.gz')]

                if options.format == "raw":

                    patient['groundtruth'] = os.path.join(options.dir, folder, "masks", patient['name'] + "_mask1.nii")

                    patient['skullmap'] = os.path.join(options.dir, folder, "preprocessed", patient['name'] + "_skullmap.nii.gz")

                elif options.format == "aligned":

                    patient['groundtruth'] = os.path.join(options.dir, folder, "preprocessed", patient['name'] + "_mask1.aligned.nii.gz")

                    patient['skullmap'] = os.path.join(options.dir, folder, "preprocessed", patient['name'] + "_skullmap.aligned.nii.gz")

                # Append to the list of all patients

                patients.append(patient)

        return patients

    # Returns the indices of patients which belong to either TRAIN, VAL or TEST. Your choice

    def get_patient_idx(self, split='TRAIN'):

        return self.patientsSplit[split]

    def get_patient_split(self):

        return self.patientsSplit

    @property

    def images(self):

        return self._images

    def get_images(self, set=None):

        _setIdx = MSISBI2015.SET_TYPES.index(set)

        images_in_set = numpy.where(self._sets == _setIdx)[0]

        return self._images[images_in_set]

    def get_image(self, i):

        return self._images[i, :, :, :]

    def get_label(self, i):

        return self._labels[i, :, :, :]

    def get_patient(self, i):

        return self.patients[i]

    @property

    def labels(self):

        return self._labels

    @property

    def sets(self):

        return self._sets

    @property

    def meta(self):

        return self._meta

    @property

    def num_examples(self):

        return self._images.shape[0]

    @property

    def width(self):

        return self._images.shape[2]

    @property

    def height(self):

        return self._images.shape[1]

    @property

    def num_channels(self):

        return self._images.shape[3]

    @property

    def epochs_completed(self):

        return self._epochs_completed

    def name(self):

        _name = "MSISBI2015"

        if self.options.numSamples > 0:

            _name += '_n{}'.format(self.options.numSamples)

        _name += "_p{}-{}".format(self.options.partition['TRAIN'], self.options.partition['VAL'])

        if self.options.useCrops:

            _name += "_{}crops{}x{}".format(self.options.cropType, self.options.cropWidth, self.options.cropHeight)

            if self.options.cropType == "random":

                _name += "_{}cropsPerSlice".format(self.options.numRandomCropsPerSlice)

        if self.options.sliceResolution is not None:

            _name += "_res{}x{}".format(self.options.sliceResolution[0], self.options.sliceResolution[1])

        _name += "_{}".format(self.options.format)

        return _name

    def split_name(self):

        return os.path.join(self.dir(), 'split-{}-{}.pckl'.format(self.options.partition['TRAIN'], self.options.partition['VAL']))

    def pckl_name(self):

        return os.path.join(self.dir(), self.name() + ".pckl")

    def tfrecord_name(self):

        return os.path.join(self.dir(), self.name() + ".tfrecord")

    def dir(self):

        return self.options.dir

    def export_slices(self, dir):

        for i in range(self.num_examples):

            imwrite(os.path.join(dir, '{}.png'.format(i)), np.squeeze(self.get_image(i) * 255).astype('uint8'))

    def visualize(self, pause=1):

        f, (ax1, ax2) = matplotlib.pyplot.subplots(1, 2)

        images_tmp, labels_tmp, _ = self.next_batch(10)

        for i in range(images_tmp.shape[0]):

            img = numpy.squeeze(images_tmp[i])

            lbl = numpy.squeeze(labels_tmp[i])

            ax1.imshow(img)

            ax1.set_title('Patch')

            ax2.imshow(lbl)

            ax2.set_title('Groundtruth')

            matplotlib.pyplot.pause(pause)

    def num_batches(self, batchsize, set='TRAIN'):

        _setIdx = MSISBI2015.SET_TYPES.index(set)

        images_in_set = numpy.where(self._sets == _setIdx)[0]

        return len(images_in_set) // batchsize

    def next_batch(self, batch_size, shuffle=True, set='TRAIN', return_brainmask=True):

        """Return the next `batch_size` examples from this data set."""

        _setIdx = MSISBI2015.SET_TYPES.index(set)

        images_in_set = numpy.where(self._sets == _setIdx)[0]

        samples_in_set = len(images_in_set)

        start = self._index_in_epoch[set]

        # Shuffle for the first epoch

        if self._epochs_completed[set] == 0 and start == 0 and shuffle:

            perm0 = numpy.arange(samples_in_set)

            numpy.random.shuffle(perm0)

            self._images[images_in_set] = self.images[images_in_set[perm0]]

            self._labels[images_in_set] = self.labels[images_in_set[perm0]]

            self._sets[images_in_set] = self.sets[images_in_set[perm0]]

        # Go to the next epoch

        if start + batch_size > samples_in_set:

            # Finished epoch

            self._epochs_completed[set] += 1

            # Get the rest examples in this epoch

            rest_num_examples = samples_in_set - start

            images_rest_part = self._images[images_in_set[start:samples_in_set]]

            labels_rest_part = self._labels[images_in_set[start:samples_in_set]]

            # Shuffle the data

            if shuffle:

                perm = numpy.arange(samples_in_set)

                numpy.random.shuffle(perm)

                self._images[images_in_set] = self.images[images_in_set[perm]]

                self._labels[images_in_set] = self.labels[images_in_set[perm]]

                self._sets[images_in_set] = self.sets[images_in_set[perm]]

            # Start next epoch

            start = 0

            self._index_in_epoch[set] = batch_size - rest_num_examples

            end = self._index_in_epoch[set]

            images_new_part = self._images[images_in_set[start:end]]

            labels_new_part = self._labels[images_in_set[start:end]]

            images_tmp = numpy.concatenate((images_rest_part, images_new_part), axis=0)

            labels_tmp = numpy.concatenate((labels_rest_part, labels_new_part), axis=0)

        else:

            self._index_in_epoch[set] += batch_size

            end = self._index_in_epoch[set]

            images_tmp = self._images[images_in_set[start:end]]

            labels_tmp = self._labels[images_in_set[start:end]]

        if self.options.addInstanceNoise:

            noise = numpy.random.normal(0, 0.01, images_tmp.shape)

            images_tmp += noise

        # Check the batch

        assert images_tmp.size, "The batch is empty!"

        assert labels_tmp.size, "The labels of the current batch are empty!"

        if return_brainmask:

            brainmasks = images_tmp > 0.05

        else:

            brainmasks = None

        return images_tmp, labels_tmp, brainmasks