""" Utility functions. """

import numpy as np

import os

import random

import tensorflow as tf

from tensorflow.contrib.layers.python import layers as tf_layers

from tensorflow.python.platform import flags

import SimpleITK as sitk

from scipy import ndimage

import itertools

from tensorflow.contrib import slim

from scipy.ndimage import _ni_support

from scipy.ndimage.morphology import distance_transform_edt, binary_erosion,\

    generate_binary_structure

FLAGS = flags.FLAGS

## Image reader

def get_images(paths, labels, nb_samples=None, shuffle=True):

    if nb_samples is not None:

        sampler = lambda x: random.sample(x, nb_samples)

    else:

        sampler = lambda x: x

    images = [(i, os.path.join(path, image)) \

        for i, path in zip(labels, paths) \

        for image in sampler(os.listdir(path))]

    if shuffle:

        random.shuffle(images)

    return images

## Loss functions

def mse(pred, label):

    pred = tf.reshape(pred, [-1])

    label = tf.reshape(label, [-1])

    return tf.reduce_mean(tf.square(pred-label))

def xent(pred, label):

    return tf.nn.softmax_cross_entropy_with_logits_v2(logits=pred, labels=label)

def kd(data1, label1, data2, label2, bool_indicator, n_class=7, temperature=2.0):

    kd_loss = 0.0

    eps = 1e-16

    prob1s = []

    prob2s = []

    for cls in range(n_class):

        mask1 = tf.tile(tf.expand_dims(label1[:, cls], -1), [1, n_class])

        logits_sum1 = tf.reduce_sum(tf.multiply(data1, mask1), axis=0)

        num1 = tf.reduce_sum(label1[:, cls])

        activations1 = logits_sum1 * 1.0 / (num1 + eps) # add eps for prevent un-sampled class resulting in NAN

        prob1 = tf.nn.softmax(activations1 / temperature)

        prob1 = tf.clip_by_value(prob1, clip_value_min=1e-8, clip_value_max=1.0)  # for preventing prob=0 resulting in NAN

        mask2 = tf.tile(tf.expand_dims(label2[:, cls], -1), [1, n_class])

        logits_sum2 = tf.reduce_sum(tf.multiply(data2, mask2), axis=0)

        num2 = tf.reduce_sum(label2[:, cls])

        activations2 = logits_sum2 * 1.0 / (num2 + eps)

        prob2 = tf.nn.softmax(activations2 / temperature)

        prob2 = tf.clip_by_value(prob2, clip_value_min=1e-8, clip_value_max=1.0)

        KL_div = (tf.reduce_sum(prob1 * tf.log(prob1 / prob2)) + tf.reduce_sum(prob2 * tf.log(prob2 / prob1))) / 2.0

        kd_loss += KL_div * bool_indicator[cls]

        prob1s.append(prob1)

        prob2s.append(prob2)

    kd_loss = kd_loss / n_class

    return kd_loss, prob1s, prob2s

def JS(data1, label1, data2, label2, bool_indicator, n_class=7, temperature=2.0):

    kd_loss = 0.0

    eps = 1e-16

    prob1s = []

    prob2s = []

    for cls in range(n_class):

        mask1 = tf.tile(tf.expand_dims(label1[:, cls], -1), [1, n_class])

        logits_sum1 = tf.reduce_sum(tf.multiply(data1, mask1), axis=0)

        num1 = tf.reduce_sum(label1[:, cls])

        activations1 = logits_sum1 * 1.0 / (num1 + eps) # add eps for prevent un-sampled class resulting in NAN

        prob1 = tf.nn.softmax(activations1 / temperature)

        prob1 = tf.clip_by_value(prob1, clip_value_min=1e-8, clip_value_max=1.0)  # for preventing prob=0 resulting in NAN

        mask2 = tf.tile(tf.expand_dims(label2[:, cls], -1), [1, n_class])

        logits_sum2 = tf.reduce_sum(tf.multiply(data2, mask2), axis=0)

        num2 = tf.reduce_sum(label2[:, cls])

        activations2 = logits_sum2 * 1.0 / (num2 + eps)

        prob2 = tf.nn.softmax(activations2 / temperature)

        prob2 = tf.clip_by_value(prob2, clip_value_min=1e-8, clip_value_max=1.0)

        mean_prob = (prob1 + prob2) / 2

        JS_div = (tf.reduce_sum(prob1 * tf.log(prob1 / mean_prob)) + tf.reduce_sum(prob2 * tf.log(prob2 / mean_prob))) / 2.0

        kd_loss += JS_div * bool_indicator[cls]

        prob1s.append(prob1)

        prob2s.append(prob2)

    kd_loss = kd_loss / n_class

    return kd_loss, prob1s, prob2s

def contrastive(feature1, label1, feature2, label2, bool_indicator=None, margin=50):

    l1 = tf.argmax(label1, axis=1)

    l2 = tf.argmax(label2, axis=1)

    pair = tf.to_float(tf.equal(l1,l2))

    delta = tf.reduce_sum(tf.square(feature1-feature2), 1) + 1e-10

    match_loss = delta

    delta_sqrt = tf.sqrt(delta + 1e-10)

    mismatch_loss = tf.square(tf.nn.relu(margin - delta_sqrt))

    if bool_indicator is None:

        loss = tf.reduce_mean(0.5 * (pair * match_loss + (1-pair) * mismatch_loss))

    else:

        loss = 0.5 * tf.reduce_sum(match_loss*pair)/tf.reduce_sum(pair)

    debug_dist_positive = tf.reduce_sum(delta_sqrt * pair)/tf.reduce_sum(pair)

    debug_dist_negative = tf.reduce_sum(delta_sqrt * (1-pair))/tf.reduce_sum(1-pair)

    return loss, pair, delta, debug_dist_positive, debug_dist_negative

def compute_distance(feature1, label1, feature2, label2):

    l1 = tf.argmax(label1, axis=1)

    l2 = tf.argmax(label2, axis=1)

    pair = tf.to_float(tf.equal(l1,l2))

    delta = tf.reduce_sum(tf.square(feature1-feature2), 1)

    delta_sqrt = tf.sqrt(delta + 1e-16)

    dist_positive_pair = tf.reduce_sum(delta_sqrt * pair)/tf.reduce_sum(pair)

    dist_negative_pair = tf.reduce_sum(delta_sqrt * (1-pair))/tf.reduce_sum(1-pair)

    return dist_positive_pair, dist_negative_pair

def _get_segmentation_cost(softmaxpred, seg_gt, n_class=2):

    """

    calculate the loss for segmentation prediction

    :param seg_logits: probability segmentation from the segmentation network

    :param seg_gt: ground truth segmentaiton mask

    :return: segmentation loss, according to the cost_kwards setting, cross-entropy weighted loss and dice loss

    """

    dice = 0

    for i in xrange(n_class):

        #inse = tf.reduce_sum(softmaxpred[:, :, :, i]*seg_gt[:, :, :, i])

        inse = tf.reduce_sum(softmaxpred[:, :, :, i]*seg_gt[:, :, :, i])

        l = tf.reduce_sum(softmaxpred[:, :, :, i])

        r = tf.reduce_sum(seg_gt[:, :, :, i])

        dice += 2.0 * inse/(l+r+1e-7) # here 1e-7 is relaxation eps

    dice_loss = 1 - 1.0 * dice / n_class

    # ce_weighted = 0

    # for i in xrange(n_class):

    #     gti = seg_gt[:,:,:,i]

    #     predi = softmaxpred[:,:,:,i]

    #     ce_weighted += -1.0 * gti * tf.log(tf.clip_by_value(predi, 0.005, 1))

    # ce_weighted_loss = tf.reduce_mean(ce_weighted)

    # total_loss =  dice_loss 

    return dice_loss#, dice_loss, ce_weighted_loss

def _get_compactness_cost(y_pred, y_true): 

    """

    y_pred: BxHxWxC

    """

    """

    lenth term

    """

    # y_pred = tf.one_hot(y_pred, depth=2)

    # print (y_true.shape)

    # print (y_pred.shape)

    y_pred = y_pred[..., 1]

    y_true = y_pred[..., 1]

    x = y_pred[:,1:,:] - y_pred[:,:-1,:] # horizontal and vertical directions 

    y = y_pred[:,:,1:] - y_pred[:,:,:-1]

    delta_x = x[:,:,1:]**2

    delta_y = y[:,1:,:]**2

    delta_u = tf.abs(delta_x + delta_y) 

    epsilon = 0.00000001 # where is a parameter to avoid square root is zero in practice.

    w = 0.01

    length = w * tf.reduce_sum(tf.sqrt(delta_u + epsilon), [1, 2])

    area = tf.reduce_sum(y_pred, [1,2])

    compactness_loss = tf.reduce_sum(length ** 2 / (area * 4 * 3.1415926))

    return compactness_loss, tf.reduce_sum(length), tf.reduce_sum(area), delta_u

# def _get_sample_masf(y_true):

#     """

#     y_pred: BxHxWx2

#     """

#     positive_mask = np.expand_dims(y_true[..., 1], axis=3)

#     metrix_label_group = np.expand_dims(np.array([1, 0, 1, 1, 0]), axis = 1)

#     # print (positive_mask.shape)

#     coutour_group = np.zeros(positive_mask.shape)

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

#         slice_i = positive_mask[i]

#         if metrix_label_group[i] == 1:

#             sample = (slice_i == 1)

#         elif metrix_label_group[i] == 0:

#             sample = (slice_i == 0)

#         coutour_group[i] = sample

#     return coutour_group, metrix_label_group

def _get_coutour_sample(y_true):

    """

    y_true: BxHxWx2

    """

    positive_mask = np.expand_dims(y_true[..., 1], axis=3)

    metrix_label_group = np.expand_dims(np.array([1, 0, 1, 1, 0]), axis = 1)

    coutour_group = np.zeros(positive_mask.shape)

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

        slice_i = positive_mask[i]

        if metrix_label_group[i] == 1:

            # generate coutour mask

            erosion = ndimage.binary_erosion(slice_i[..., 0], iterations=1).astype(slice_i.dtype)

            sample = np.expand_dims(slice_i[..., 0] - erosion, axis = 2)

        elif metrix_label_group[i] == 0:

            # generate background mask

            dilation = ndimage.binary_dilation(slice_i, iterations=5).astype(slice_i.dtype)

            sample = dilation - slice_i 

        coutour_group[i] = sample

    return coutour_group, metrix_label_group

# def _get_negative(y_true):

def _get_boundary_cost(y_pred, y_true): 

    """

    y_pred: BxHxWxC

    """

    """

    lenth term

    """

    # y_pred = tf.one_hot(y_pred, depth=2)

    # print (y_true.shape)

    # print (y_pred.shape)

    y_pred = y_pred[..., 1]

    y_true = y_pred[..., 1]

    x = y_pred[:,1:,:] - y_pred[:,:-1,:] # horizontal and vertical directions 

    y = y_pred[:,:,1:] - y_pred[:,:,:-1]

    delta_x = x[:,:,1:]**2

    delta_y = y[:,1:,:]**2

    delta_u = tf.abs(delta_x + delta_y) 

    epsilon = 0.00000001 # where is a parameter to avoid square root is zero in practice.

    w = 0.01

    length = w * tf.reduce_sum(tf.sqrt(delta_u + epsilon), [1, 2]) # equ.(11) in the paper

    area = tf.reduce_sum(y_pred, [1,2])

    compactness_loss = tf.reduce_sum(length ** 2 / (area * 4 * 3.1415926))

    return compactness_loss, tf.reduce_sum(length), tf.reduce_sum(area)

def check_folder(log_dir):

    if not os.path.exists(log_dir):

        print ("Allocating '{:}'".format(log_dir))

        os.makedirs(log_dir)

    return log_dir

def _eval_dice(gt_y, pred_y, detail=False):

    class_map = {  # a map used for mapping label value to its name, used for output

        "0": "bg",

        "1": "CZ",

        "2": "prostate"

    }

    dice = []

    for cls in xrange(1,2):

        gt = np.zeros(gt_y.shape)

        pred = np.zeros(pred_y.shape)

        gt[gt_y == cls] = 1

        pred[pred_y == cls] = 1

        dice_this = 2*np.sum(gt*pred)/(np.sum(gt)+np.sum(pred))

        dice.append(dice_this)

        if detail is True:

            #print ("class {}, dice is {:2f}".format(class_map[str(cls)], dice_this))

            logging.info("class {}, dice is {:2f}".format(class_map[str(cls)], dice_this))

    return dice

def __surface_distances(result, reference, voxelspacing=None, connectivity=1):

    """

    The distances between the surface voxel of binary objects in result and their

    nearest partner surface voxel of a binary object in reference.

    """

    result = np.atleast_1d(result.astype(np.bool))

    reference = np.atleast_1d(reference.astype(np.bool))

    if voxelspacing is not None:

        voxelspacing = _ni_support._normalize_sequence(voxelspacing, result.ndim)

        voxelspacing = np.asarray(voxelspacing, dtype=np.float64)

        if not voxelspacing.flags.contiguous:

            voxelspacing = voxelspacing.copy()

    # binary structure

    footprint = generate_binary_structure(result.ndim, connectivity)

    # test for emptiness

    if 0 == np.count_nonzero(result): 

        raise RuntimeError('The first supplied array does not contain any binary object.')

    if 0 == np.count_nonzero(reference): 

        raise RuntimeError('The second supplied array does not contain any binary object.')    

    # extract only 1-pixel border line of objects

    result_border = result ^ binary_erosion(result, structure=footprint, iterations=1)

    reference_border = reference ^ binary_erosion(reference, structure=footprint, iterations=1)

    # compute average surface distance        

    # Note: scipys distance transform is calculated only inside the borders of the

    #       foreground objects, therefore the input has to be reversed

    dt = distance_transform_edt(~reference_border, sampling=voxelspacing)

    sds = dt[result_border]

    return sds

def asd(result, reference, voxelspacing=None, connectivity=1):

    sds = __surface_distances(result, reference, voxelspacing, connectivity)

    asd = sds.mean()

    return asd

def calculate_hausdorff(lP,lT,spacing):

    return asd(lP, lT, spacing)

def _eval_haus(pred, gt, spacing, detail=False):

    '''

    :param pred: whole brain prediction

    :param gt: whole

    :param detail:

    :return: a list, indicating Dice of each class for one case

    '''

    haus = []

    for cls in range(1,2):

        pred_i = np.zeros(pred.shape)

        pred_i[pred == cls] = 1

        gt_i = np.zeros(gt.shape)

        gt_i[gt == cls] = 1

        # hausdorff_distance_filter = sitk.HausdorffDistanceImageFilter()

        # hausdorff_distance_filter.Execute(gt_i, pred_i)

        haus_cls = calculate_hausdorff(gt_i, (pred_i), spacing)

        haus.append(haus_cls)

        if detail is True:

            logging.info("class {}, haus is {:4f}".format(class_map[str(cls)], haus_cls))

    # logging.info("4 class average haus is {:4f}".format(np.mean(haus)))

    return haus

def _connectivity_region_analysis(mask):

    s = [[0,1,0],

         [1,1,1],

         [0,1,0]]

    label_im, nb_labels = ndimage.label(mask)#, structure=s)

    sizes = ndimage.sum(mask, label_im, range(nb_labels + 1))

    # plt.imshow(label_im)        

    label_im[label_im != np.argmax(sizes)] = 0

    label_im[label_im == np.argmax(sizes)] = 1

    return label_im

def _crop_object_region(mask, prediction):

    limX, limY, limZ = np.where(mask>0)

    min_z = np.min(limZ)

    max_z = np.max(limZ)

    prediction[..., :np.min(limZ)] = 0

    prediction[..., np.max(limZ)+1:] = 0

    return prediction

def parse_fn(data_path):

    '''

    :param image_path: path to a folder of a patient

    :return: normalized entire image with its corresponding label

    In an image, the air region is 0, so we only calculate the mean and std within the brain area

    For any image-level normalization, do it here

    '''

    path = data_path.split(",")

    image_path = path[0]

    label_path = path[1]

    #itk_image = zoom2shape(image_path, [512,512])#os.path.join(image_path, 'T1_unbiased_brain_rigid_to_mni.nii.gz'))

    #itk_mask = zoom2shape(label_path, [512,512], label=True)#os.path.join(image_path, 'T1_brain_seg_rigid_to_mni.nii.gz'))

    itk_image = sitk.ReadImage(image_path)#os.path.join(image_path, 'T1_unbiased_brain_rigid_to_mni.nii.gz'))

    itk_mask = sitk.ReadImage(label_path)#os.path.join(image_path, 'T1_brain_seg_rigid_to_mni.nii.gz'))

    # itk_image = sitk.ReadImage(os.path.join(image_path, 'T2_FLAIR_unbiased_brain_rigid_to_mni.nii.gz'))

    image = sitk.GetArrayFromImage(itk_image)

    mask = sitk.GetArrayFromImage(itk_mask)

    #image[image >= 1000] = 1000

    binary_mask = np.ones(mask.shape)

    mean = np.sum(image * binary_mask) / np.sum(binary_mask)

    std = np.sqrt(np.sum(np.square(image - mean) * binary_mask) / np.sum(binary_mask))

    image = (image - mean) / std  # normalize per image, using statistics within the brain, but apply to whole image

    mask[mask==2] = 1

    return image.transpose([1,2,0]), mask.transpose([1,2,0]) # transpose the orientation of the

def parse_fn_haus(data_path):

    '''

    :param image_path: path to a folder of a patient

    :return: normalized entire image with its corresponding label

    In an image, the air region is 0, so we only calculate the mean and std within the brain area

    For any image-level normalization, do it here

    '''

    path = data_path.split(",")

    image_path = path[0]

    label_path = path[1]

    #itk_image = zoom2shape(image_path, [512,512])#os.path.join(image_path, 'T1_unbiased_brain_rigid_to_mni.nii.gz'))

    #itk_mask = zoom2shape(label_path, [512,512], label=True)#os.path.join(image_path, 'T1_brain_seg_rigid_to_mni.nii.gz'))

    itk_image = sitk.ReadImage(image_path)#os.path.join(image_path, 'T1_unbiased_brain_rigid_to_mni.nii.gz'))

    itk_mask = sitk.ReadImage(label_path)#os.path.join(image_path, 'T1_brain_seg_rigid_to_mni.nii.gz'))

    # itk_image = sitk.ReadImage(os.path.join(image_path, 'T2_FLAIR_unbiased_brain_rigid_to_mni.nii.gz'))

    spacing = itk_mask.GetSpacing()

    image = sitk.GetArrayFromImage(itk_image)

    mask = sitk.GetArrayFromImage(itk_mask)

    #image[image >= 1000] = 1000

    binary_mask = np.ones(mask.shape)

    mean = np.sum(image * binary_mask) / np.sum(binary_mask)

    std = np.sqrt(np.sum(np.square(image - mean) * binary_mask) / np.sum(binary_mask))

    image = (image - mean) / std  # normalize per image, using statistics within the brain, but apply to whole image

    mask[mask==2] = 1

    return image.transpose([1,2,0]), mask.transpose([1,2,0]), spacing

def show_all_variables():

    model_vars = tf.trainable_variables()

    slim.model_analyzer.analyze_vars(model_vars, print_info=True)