import matplotlib.pyplot as plt

import nrrd

import numpy as np

# Class for working with a NII file in the context of machine learning

class NRRD:

    VIEW_MAPPING = {'saggital': 0, 'coronal': 1, 'axial': 2}

    def __init__(self, filename):

        self.data, self.info = nrrd.read(filename)

    @property

    def num_saggital_slices(self):

        return self.data.shape[NRRD.VIEW_MAPPING['saggital']]

    @property

    def num_coronal_slices(self):

        return self.data.shape[NRRD.VIEW_MAPPING['coronal']]

    @property

    def num_axial_slices(self):

        return self.data.shape[NRRD.VIEW_MAPPING['axial']]

    @staticmethod

    def set_view_mapping(mapping):

        NRRD.VIEW_MAPPING = mapping

    def shape(self):

        return self.data.shape

    @staticmethod

    def get_axis_index(axis):

        return NRRD.VIEW_MAPPING[axis]

    def num_slices_along_axis(self, axis):

        return self.data.shape[NRRD.VIEW_MAPPING[axis]]

    def normalize(self, method='scaling', lowerpercentile=None, upperpercentile=None):

        # Convert the attribute "data" to float()

        self.data = self.data.astype(np.float32)

        if lowerpercentile is not None:

            qlow = np.percentile(self.data, lowerpercentile)

        if upperpercentile is not None:

            qup = np.percentile(self.data, upperpercentile)

        if lowerpercentile is not None:

            self.data[self.data < qlow] = qlow

        if upperpercentile is not None:

            self.data[self.data > qup] = qup

        if method == 'scaling':

            # Divide "data" by its maximum value

            self.data -= self.data.min()

            self.data = np.multiply(self.data, 1.0 / self.data.max())

        elif method == 'standardization':

            self.data = self.data - np.mean(self.data)

            self.data = self.data / np.std(self.data)

    def get_slice(self, the_slice, axis='axial'):

        indices = [slice(None)] * self.data.ndim

        indices[NRRD.VIEW_MAPPING[axis]] = the_slice

        return self.data[indices]

    def get_data(self):

        return self.data

    def set_to_zero(self):

        self.data.fill(0.0)

    def visualize(self, axis='axial', pause=0.2):

        for i in range(self.data.shape[NRRD.VIEW_MAPPING[axis]]):

            img = self.get_slice(i, axis=axis)

            plt.imshow(img)

            plt.pause(pause)