import glob

import matplotlib.patches as patches

import json

import numpy as np

from matplotlib.path import Path

import dicom

from sklearn.utils import shuffle

import cv2

def get_roi(image, contour, shape_out = 32):

    """

    Create a binary mask with ROI from contour. 

    Extract the maximum square around the contour.

    :param image: input image (needed for shape only)

    :param contour: numpy array contour (d, 2)

    :return: numpy array mask ROI (shape_out, shape_out)

    """

    X_min, Y_min = contour[:,0].min(), contour[:,1].min()

    X_max, Y_max = contour[:,0].max(), contour[:,1].max()  

    w = X_max - X_min

    h = Y_max - Y_min

    mask_roi = np.zeros(image.shape)

    if w > h :

        mask_roi[int(Y_min - (w -h)/2):int(Y_max + (w -h)/2), int(X_min):int(X_max)] = 1.0

    else :

        mask_roi[int(Y_min):int(Y_max), int(X_min - (h-w)/2):int(X_max + (h -w)/2)] = 1.0

    return cv2.resize(mask_roi, (shape_out, shape_out), interpolation = cv2.INTER_NEAREST)

def create_dataset(image_shape=64, n_set='train', original_image_shape=256, 

                   roi_shape=32, data_path='./Data/'):

    """

    Creating the dataset from the images and the contour for the CNN.

    :param image_shape: image dataset desired size

    :param original_image_shape: original image size

    :param roi_shape: binary ROI mask shape

    :param data_path: path for the dataset

    :return: correct size image dataset, full size image dataset, label (contours) dataset

    """

    if n_set == 'train':

        number_set = 3

        name_set = 'Training'

    elif n_set == 'test':

        number_set = 1

        name_set = 'Online'         

    # Create dataset

    series = json.load(open('series_case.json'))[n_set]

    images, images_fullsize, contours, contour_mask = [], [], [], []

    # Loop over the series

    for case, serie in series.items():

        image_path_base = data_path + 'challenge_%s/%s/IM-%s' % (name_set.lower(),case, serie)

        contour_path_base = data_path + 'Sunnybrook Cardiac MR Database ContoursPart%s/\

%sDataContours/%s/contours-manual/IRCCI-expert/' % (number_set, name_set, case)

        contours_list = glob.glob(contour_path_base + '*')

        contours_list_series = [k.split('/')[7].split('-')[2] for k in contours_list]

        # Loop over the contours/images

        for c in contours_list_series:

            # Get contours and images path

            idx_contour = contours_list_series.index(c)

            image_path = image_path_base + '-%s.dcm' % c

            contour_path = contours_list[idx_contour]

            # open image as numpy array and resize to (image_shape, image_shape)

            image_part = dicom.read_file(image_path).pixel_array  

            # open contours as numpy array

            contour = []

            file = open(contour_path, 'r') 

            for line in file: 

                contour.append(tuple(map(float, line.split())))

            contour = np.array(contour)

            # append binary ROI mask 

            contours.append(get_roi(image_part, contour))

            # create mask contour with experts contours

            x, y = np.meshgrid(np.arange(256), np.arange(256)) # make a canvas with coordinates

            x, y = x.flatten(), y.flatten()

            points = np.vstack((x,y)).T 

            p = Path(contour) # make a polygon

            grid = p.contains_points(points)

            mask_contour = grid.reshape(256,256)

            mask_contour=mask_contour*1

            contour_mask.append(mask_contour)

            # Open image and resize it 

            images.append(cv2.resize(image_part, (image_shape, image_shape)))

            images_fullsize.append(cv2.resize(image_part, (original_image_shape, original_image_shape)))

    X_fullsize = np.array(images_fullsize)

    X = np.reshape(np.array(images), [len(images), image_shape, image_shape, 1])

    Y = np.reshape(np.array(contours), [len(contours), 1, roi_shape, roi_shape])

    print('Dataset shape :', X.shape, Y.shape)

    return shuffle(X, X_fullsize, Y, contour_mask, random_state=0)

def compute_roi_pred(X_fullsize, y_pred, contour_mask, idx, roi_shape=32):

    """

    Computing and cropping a ROI from the original image for further processing in the next stage

    :param X_fullsize: full size training set (256x256)

    :param y_pred: predictions

    :param contour_mask label: (contours) dataset

    :param idx: desired image prediction index

    :param roi_shape: shape of the binary mask

    """

    # up sampling from 32x32 to original MR size

    pred = cv2.resize(y_pred[idx].reshape((roi_shape, roi_shape)), (256,256), cv2.INTER_NEAREST)

    # select the non null pixels

    pos_pred = np.array(np.where(pred > 0.5))

    # get the center of the mask

    X_min, Y_min = pos_pred[0, :].min(), pos_pred[1, :].min()

    X_max, Y_max = pos_pred[0, :].max(), pos_pred[1, :].max()  

    X_middle = X_min + (X_max - X_min) / 2

    Y_middle = Y_min + (Y_max - Y_min) / 2

    # Find ROI coordinates

    X_top = int(X_middle - 50)

    Y_top = int(Y_middle - 50)

    X_down = int(X_middle + 50)

    Y_down = int(Y_middle + 50)

    # crop ROI of size 100x100

    mask_roi = np.zeros((256, 256))

    mask_roi = cv2.rectangle(mask_roi, (X_top, Y_top), (X_down, Y_down), 1, -1)*255

    return X_fullsize[idx][X_top:X_down, Y_top:Y_down], mask_roi, contour_mask[idx][X_top:X_down, Y_top:Y_down]

def prediction_plot(X, model, idx=None):

    """

    Compute the Inferred shape binary mask using the trained stacked AE model

    :param X: dataset to predict

    :param model: trained AE model

    :param idx: index of the particular picture to return

    :return: inferred shape binary mask, infered shape on the MR image

    """

    if not idx:

        idx= np.random.randint(len(X))

    contours = model.predict(X)

    contour = contours[idx].reshape((64,64))

    # thresholding

    binary = cv2.threshold(contour, 0, 1, cv2.INTERSECT_NONE)

    return binary[1], binary[1]*X[idx].reshape(64,64), idx

def dice_metric(X, Y):

    """

    Dice metric for measuring the contour overlap 

    :param X,Y: 2D numpy arrays

    :return: metric scalar

    """

    return np.sum(X[Y==1])*2.0 / (np.sum(X) + np.sum(Y))

def conformity_coefficient(X, Y):

    """

    Conformity coefficient for measuring the  ratio of the number of mis-segmented pixels  

    :param X,Y: 2D numpy arrays

    :return: metric scalar

    """

    if dice_metric(X,Y) !=0:

        return (3*dice_metric(X,Y)-2)/dice_metric(X,Y)

def stats_results(Y_true, Y_pred, idx=None, print_=False):

    """

    DM and CC for the whole dataset. If one index is passed, return DM and CC for this index

    :param Y_true: True labels (Y_train)

    :param Y_pred: Prediction (binarys)

    :param print_: Boolean to print stats

    :return: DM and CC arrays

    """

    dm_tot = np.array([dice_metric(Y_true[k].reshape((64,64)), Y_pred[k]) for k in range(len(Y_true))])

    cc_tot = np.array([conformity_coefficient(Y_true[k].reshape((64,64)), Y_pred[k]) for k in range(len(Y_true))])

    cc_tot = cc_tot[cc_tot != None]

    #return dm_tot, cc_tot

    if idx:

        print('For image %s :\nDice Metric : %.2f\nConformity Coefficient : %.2f\n' % (

                idx, dice_metric(Y_true[idx].reshape((64,64)), Y_pred[idx]),

                conformity_coefficient(Y_true[idx].reshape((64,64)), Y_pred[idx])))

    if print_:

        print('For the full dataset :\nDice Metric : %.2f\nConformity Coefficient : %.2f' % (

                                            dm_tot.mean(),cc_tot.mean()))

    return dm_tot, cc_tot