# Copyright 2019 Population Health Sciences and Image Analysis, German Center for Neurodegenerative Diseases(DZNE)

#

#    Licensed under the Apache License, Version 2.0 (the "License");

#    you may not use this file except in compliance with the License.

#    You may obtain a copy of the License at

#

#        http://www.apache.org/licenses/LICENSE-2.0

#

#    Unless required by applicable law or agreed to in writing, software

#    distributed under the License is distributed on an "AS IS" BASIS,

#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

#    See the License for the specific language governing permissions and

#    limitations under the License.

import numpy as np

from skimage.measure import label

def largets_connected_componets(labels):

    """Calculate the largest connected component, all the labels are unified to one

    Args:

        labels: ndarray (int or float) label image or volume

        neighbors : {4, 8}, int, optional

        Whether to use 4- or 8-“connectivity”. In 3D, 4-“connectivity” means connected pixels have to share face, whereas with 8-“connectivity”,

        they have to share only edge or vertex. Deprecated, use ``connectivity`` instead.

    Returns:

        out :ndarray, the input array only with the largest connected component

"""

    mask = np.copy(labels)

    mask[labels > 0] = 1

    connected_labels, num = label(mask,connectivity=3,background=0, return_num=True)

    #0 is background data, so check with out zero

    #largestCC = np.argmax(np.bincount(connected_labels.flat)[1:])

    if num !=1 :

        unique, counts = np.unique(connected_labels, return_counts=True)

        largest=np.argmax(counts[1:]) + 1 #0 is background data, so check with out zero

        mask[connected_labels != largest] = 0

    mask = np.array(mask, dtype=np.int8)

    return labels * mask

def swap_axes(data,plane):

    if plane == 'axial':

        return  data

    elif plane == 'frontal':

        data = np.swapaxes(data, 1, 0)

        return data

    elif plane == 'sagital':

        data = np.swapaxes(data, 2, 0)

        return data

def check_size(data,patch_size):

    x_low=int(np.floor(-1*(data.shape[1]-patch_size[0])/2))

    x_high=int(np.ceil(-1*(data.shape[1]-patch_size[0])/2))

    y_low=int(np.floor(-1*(data.shape[2]-patch_size[1])/2))

    y_high=int(np.ceil(-1*(data.shape[2]-patch_size[1])/2))

    new_arr=np.zeros((data.shape[0],patch_size[0],patch_size[1]))

    new_arr[:,x_low:patch_size[0]-x_high,y_low:patch_size[1]-y_high]=data[:,:,:]

    return new_arr

def change_data_plane(arr, plane='axial',return_index=False):

    if plane == 'axial':

        if return_index:

            return arr,0, arr.shape[0]

        else:

            return arr

    elif plane == 'frontal' or plane == 'coronal':

            if len(arr.shape) == 4:

                new_arr = np.zeros((arr.shape[1], arr.shape[1], arr.shape[2],arr.shape[3]))

                for slice in range(arr.shape[3]):

                    aux_arr=arr[:,:,:,slice]

                    aux_arr = np.swapaxes(aux_arr, 1, 0)

                    idx_low = int((new_arr.shape[1] / 2) - (aux_arr.shape[1] / 2))

                    idx_high = int((new_arr.shape[1] / 2) + (aux_arr.shape[1] / 2))

                    new_arr[:, idx_low:idx_high, :,slice] = aux_arr

                if return_index:

                    return new_arr,idx_low,idx_high

                else:

                    return new_arr

            else:

                new_arr = np.zeros((arr.shape[1], arr.shape[1], arr.shape[2]))

                arr = np.swapaxes(arr, 1, 0)

                idx_low = int((new_arr.shape[1] / 2) - (arr.shape[1] / 2))

                idx_high = int((new_arr.shape[1] / 2) + (arr.shape[1] / 2))

                new_arr[:, idx_low:idx_high, :] = arr

                if return_index:

                    return new_arr,idx_low,idx_high

                else:

                    return new_arr

    elif plane == 'sagital' or plane == 'sagittal':

            if len(arr.shape)== 4:

                new_arr = np.zeros((arr.shape[2], arr.shape[1], arr.shape[2],arr.shape[3]))

                for slice in range(arr.shape[3]):

                    aux_arr=arr[:,:,:,slice]

                    aux_arr = np.swapaxes(aux_arr, 2, 0)

                    idx_low = int((new_arr.shape[2] / 2) - (aux_arr.shape[2] / 2))

                    idx_high = int((new_arr.shape[2] / 2) + (aux_arr.shape[2] / 2))

                    new_arr[:, :, idx_low:idx_high,slice] = aux_arr[:]

                if return_index:

                    return new_arr,idx_low,idx_high

                else:

                    return new_arr

            else:

                new_arr = np.zeros((arr.shape[2], arr.shape[1], arr.shape[2]))

                arr = np.swapaxes(arr, 2, 0)

                idx_low = int((new_arr.shape[2] / 2) - (arr.shape[2] / 2))

                idx_high = int((new_arr.shape[2] / 2) + (arr.shape[2] / 2))

                new_arr[:, :, idx_low:idx_high] = arr

                if return_index:

                    return new_arr,idx_low,idx_high

                else:

                    return new_arr

def find_labels(arr):

    idx=(np.where(arr > 0))

    min_idx=np.min(idx[0])

    max_idx=np.max(idx[0])

    return max_idx,min_idx