import os

import cv2 

import math

import numpy as np

import pandas as pd

import matplotlib.pyplot as plt

from matplotlib.patches import Rectangle

from sklearn.cluster import KMeans

from skimage.morphology import erosion, opening, closing, square, \

                               disk, convex_hull_image, remove_small_holes

from skimage.measure import label, regionprops

from skimage.filters import sobel

from skimage.feature import canny

from scipy import ndimage as ndi

from skimage.segmentation import watershed

SMALL_FONT = 13

MEDIUM_FONT = 15

LARGE_FONT = 18

plt.rc('font', size=SMALL_FONT)          # controls default text sizes

plt.rc('axes', titlesize=SMALL_FONT)     # fontsize of the axes title

plt.rc('axes', labelsize=MEDIUM_FONT)    # fontsize of the x and y labels

plt.rc('xtick', labelsize=SMALL_FONT)    # fontsize of the tick labels

plt.rc('ytick', labelsize=SMALL_FONT)    # fontsize of the tick labels

plt.rc('legend', fontsize=MEDIUM_FONT)   # legend fontsize

plt.rc('figure', titlesize=LARGE_FONT)   # fontsize of the figure title

plt.rcParams["figure.figsize"] = (10, 5)

def readSortedSlices(path):

    slices = []

    for s in os.listdir(path):

        slices.append(path + '/' + s)       

    slices.sort(key = lambda s: int(s[s.find('_') + 1 : s.find('.')]))

    ID = slices[0][slices[0].find('/') + 1 : slices[0].find('_')]

    print('CT scan of Patient %s consists of %d slices.' % (ID, len(slices)))  

    return (slices, ID)

def getSliceImages(slices):

    return list(map(readImg, slices))

def readImg(path, showOutput=0):

    img = cv2.imread(path)

    img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

    if showOutput:

        plt.title('A CT Scan Image Slice')

        plt.imshow(img, cmap='gray')

    return img

def imgKMeans(img, K, showOutput=0, showHistogram=0):

    '''

    Apply KMeans on an image with the number of clusters K

    Input: Image, Number of clusters K

    Output: Dictionary of cluster center labels and points, Output segmented image

    '''

    imgflat = np.reshape(img, img.shape[0] * img.shape[1]).reshape(-1, 1)

    kmeans = KMeans(n_clusters=K, verbose=0)

    kmmodel = kmeans.fit(imgflat)

    labels = kmmodel.labels_

    centers = kmmodel.cluster_centers_

    center_labels = dict(zip(np.arange(K), centers))

    output = np.array([center_labels[label] for label in labels])

    output = output.reshape(img.shape[0], img.shape[1]).astype(int)

    if showOutput:

        fig, axes = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(10, 5))

        axes = axes.ravel()

        axes[0].imshow(img, cmap='gray')

        axes[0].set_title('Original Image')

        axes[1].imshow(output)

        axes[1].set_title('Image after KMeans (K = ' + str(K) + ')')

    return center_labels, output

def preprocessImage(img, showOutput=0):

    '''

    Preprocess the image by applying truncated thresholding using KMeans

    Input: Image

    Output: Preprocessed image

    '''

    centroids, segmented_img = imgKMeans(img, 3, showOutput=0)

    sorted_center_values = sorted([i[0] for i in centroids.values()])

    threshold = (sorted_center_values[-1] + sorted_center_values[-2]) / 2

    retval, procImg = cv2.threshold(img, threshold, 255, cv2.THRESH_TOZERO) 

    if showOutput:

        fig, axes = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(10, 5))

        axes = axes.ravel()

        axes[0].imshow(img, cmap='gray')

        axes[0].set_title('Original Image')

        axes[1].imshow(procImg, cmap='gray')

        axes[1].set_title('Processed Image - After Thresholding')

    return procImg, threshold

def getForegroundMask(img, fg_threshold, showOutput=0):

    retval, init_fg_mask = cv2.threshold(img, fg_threshold, 255, cv2.THRESH_BINARY)

    # Morphological operations to clean the mask

    fg_mask_opened = opening(init_fg_mask, square(3))

    fg_mask_opened2 = opening(fg_mask_opened, disk(4))

    # Perform edge-based segmentation of the foreground...

    # Detect contours that delineate the foreground with the Canny edge detector

    edges = canny(fg_mask_opened2) # Background is uniform - edges are on the boundary/inside ROI

    # Fill the inner part of the boundary using morphology ops

    fg_mask = ndi.binary_fill_holes(fg_mask_opened2)

    # Plot all steps

    if showOutput:

        fig, axes = plt.subplots(2, 3, sharex=True, sharey=True, figsize=(16, 10))

        axes = axes.ravel()

        axes[0].set_title('Original Image')

        axes[0].imshow(img, cmap='gray')

        axes[1].set_title('On Performing Thresholding\'s')

        axes[1].imshow(init_fg_mask, cmap='gray')

        axes[2].set_title('On Opening with Square SE (3)')

        axes[2].imshow(fg_mask_opened, cmap='gray')

        axes[3].set_title('On Opening with Disk SE (4)')

        axes[3].imshow(fg_mask_opened2, cmap='gray')

        axes[4].set_title('Outer Boundary Delineation with Canny\'s')

        axes[4].imshow(edges, cmap='gray')

        axes[5].set_title('Foreground Mask')

        axes[5].imshow(fg_mask, cmap='gray')

    return fg_mask

def getLungTracheaMasks(img, fg_mask, fg_threshold, showOutput=0):

    # Distinguish black pixels of the background from those of the lungs

    enhanced = img.copy()

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

        for j in range(img.shape[1]):

            if fg_mask[i][j] == 0:

                enhanced[i][j] = 255

    # Extract lungs from the foreground mask

    retval, initial_lung_mask = cv2.threshold(enhanced, fg_threshold, 255, cv2.THRESH_BINARY_INV) 

    # Clean up the lung mask with morphological operations

    lung_mask_op = opening(initial_lung_mask, square(2))

    lung_mask_opcl = closing(lung_mask_op, disk(6))

    lung_mask_opclrm = ndi.binary_fill_holes(lung_mask_opcl)

    # Get connected components of the segmented image and label them

    label_img = label(lung_mask_opclrm)

    lung_regions = regionprops(label_img)

    # Upon experimentation: areas of regions < 1500 are wind pipe structures

    trachea_labels = []

    for i in lung_regions:

        if i.area < 1500:

            trachea_labels.append(i.label)

    # Create trachea mask as a summation of all those regions

    trachea_mask = np.zeros(img.shape, dtype=np.uint8)

    for row in range(label_img.shape[0]):

        for col in range(label_img.shape[1]):

            if label_img[row][col] in trachea_labels:

                trachea_mask[row][col] = 255

    # Lung mask is made of all the other regions

    lung_mask = lung_mask_opclrm * np.invert(trachea_mask) 

    # Lung mask is all black? Convex hull set to 0, since convex hull op on empty img errors out

    if sum(sum(lung_mask)) > 0:

        ch_lung_mask = convex_hull_image(lung_mask)

    else:

        ch_lung_mask = lung_mask.copy()

    initial_int_heart_mask = ch_lung_mask * np.invert(lung_mask) * np.invert(trachea_mask)

    int_heart_mask_op1 = opening(initial_int_heart_mask, square(5))

    int_heart_mask_op2 = opening(int_heart_mask_op1, disk(4))

    heart_label_img = label(int_heart_mask_op2)

    heart_regions = regionprops(heart_label_img)

    areas = {}

    for i in heart_regions:

        areas[i.label] = i.area

    if areas:

        heart_label = max(areas, key=areas.get)

        int_heart_mask = np.where(heart_label_img==heart_label, np.uint8(255), np.uint8(0))

    else:

        int_heart_mask = np.zeros(img.shape, dtype=np.uint8)

    if showOutput:

        fig, axes = plt.subplots(4, 3, sharex=True, sharey=True, figsize=(20, 20))

        axes = axes.ravel()

        axes[0].set_title('Original Image')

        axes[0].imshow(img, cmap='gray')

        axes[1].set_title('Enhanced Image')

        axes[1].imshow(enhanced, cmap='gray')

        axes[2].set_title('Initial Lung Mask')

        axes[2].imshow(initial_lung_mask, cmap='gray')

        axes[3].set_title('On Opening with Square SE (2)')

        axes[3].imshow(lung_mask_op, cmap='gray')

        axes[4].set_title('On Closing with Disk SE (6)')

        axes[4].imshow(lung_mask_opcl, cmap='gray')

        axes[5].set_title('On Filling Regions')

        axes[5].imshow(lung_mask_opclrm, cmap='gray')

        axes[6].set_title('Trachea Mask/Primary Bronchi')

        axes[6].imshow(trachea_mask, cmap='gray')

        axes[7].set_title('Lung Mask')

        axes[7].imshow(lung_mask, cmap='gray')

        axes[8].set_title('Convex Hull of Lung Mask')

        axes[8].imshow(ch_lung_mask, cmap='gray')

        axes[9].set_title('Initial Intermediate Heart Mask')

        axes[9].imshow(initial_int_heart_mask, cmap='gray')

        axes[10].set_title('On Opening with Square SE (3)')

        axes[10].imshow(int_heart_mask_op1, cmap='gray')

        axes[11].set_title('Intermediate heart mask')

        axes[11].imshow(int_heart_mask, cmap='gray')

    return trachea_mask, lung_mask, ch_lung_mask, int_heart_mask

def chullSpineMask(img, int_heart_mask, showOutput=0):

    # If no heart

    if not int_heart_mask.any():

        return int_heart_mask, int_heart_mask

    int_heart_pixel = img.copy()

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

        for j in range(img.shape[1]):

            if int_heart_mask[i][j] == 0:

                int_heart_pixel[i][j] = 0

    centroids, segmented_heart_img = imgKMeans(int_heart_pixel, 3, showOutput=0)

    spine_threshold = (max(centroids.values()))[0]

    retval, initial_spine_mask = cv2.threshold(int_heart_pixel, spine_threshold, 255, cv2.THRESH_BINARY) 

    bone_mask = closing(initial_spine_mask, disk(20))  

    label_spine = label(bone_mask)

    spine_regions = regionprops(label_spine)

    # Assumption: The spine's area is greater than that of any calcium deposits

    labels = []

    areas = {}

    geometric_measures = {}

    for i in spine_regions:

        labels.append(i.label)

        areas[i.label] = i.area

        geometric_measures[i.label] = [i.centroid, i.orientation, i.axis_major_length]

    spine_label = max(areas, key=areas.get)

    labels.remove(spine_label)

    spine_mask = np.where(label_spine==spine_label, np.uint8(255), np.uint8(0))

#     if labels:

#         calcium_deposit_mask = np.where(label_spine==labels[0], np.uint8(255), np.uint8(0))

#     else:

#         calcium_deposit_mask = np.zeros(img.shape, dtype=np.uint8)

    label_heart = label(int_heart_mask)

    heart_regions = regionprops(label_heart)

    heart_region_area = heart_regions[0].area

    spine_region_area = areas[spine_label] 

    frac_heart = (heart_region_area - spine_region_area)/heart_region_area

    if frac_heart < 0.5:

        heart_mask = np.zeros(img.shape, dtype=np.uint8)

        make_spine_mask = 0

    else:

        make_spine_mask = 1

        # Center point of the spine - get the centroid 

        y0, x0 = geometric_measures[spine_label][0]

        orientation = geometric_measures[spine_label][1]

        # Top-most point of the spine

        # top_most_coordinate = centroid - (slightly_more_than_half * axis_major_length * sin(angle))

        x2 = x0 - math.sin(orientation) * 0.6 * geometric_measures[spine_label][2]

        y2 = y0 - math.cos(orientation) * 0.6 * geometric_measures[spine_label][2]

        chull_spine_mask = spine_mask.copy()

        # Vertical axis

        for i in range(math.ceil(y2), img.shape[1]):

            if i > math.ceil(y0):

                # Horizontal axis

                for j in range(img.shape[0]):

                    chull_spine_mask[i][j] = 255

            else:

                # Horizontal axis

                for j in range(math.ceil(x0)):

                    chull_spine_mask[i][j] = 255

        heart_mask = int_heart_mask * np.invert(chull_spine_mask)

    if showOutput:

        fig, axes = plt.subplots(3, 3, sharex=True, sharey=True, figsize=(15, 15))

        axes = axes.ravel()

        axes[0].set_title('Intermediate Heart Mask')

        axes[0].imshow(int_heart_mask, cmap='gray')

        axes[1].set_title('Intermediate Heart Segment')

        axes[1].imshow(int_heart_pixel, cmap='gray')

        axes[2].set_title('Intermediate Heart Segment on K-Means (K = 3)')

        axes[2].imshow(segmented_heart_img)

        axes[3].set_title('Spine Mask')

        axes[3].imshow(initial_spine_mask, cmap='gray')

        axes[4].set_title('On Closing with Disk SE (20)')

        axes[4].imshow(spine_mask, cmap='gray')

        axes[5].set_title('On Opening with Square SE (4)')

        axes[5].imshow(spine_mask, cmap='gray')

        axes[6].set_title('Centroid and uppermost point')

        axes[6].imshow(spine_mask, cmap='gray')

        if make_spine_mask:

            axes[6].plot((x0, x2), (y0, y2), '-r', linewidth=1.5)

            axes[6].plot(x0, y0, '.g', markersize=5)

            axes[6].plot(x2, y2, '.b', markersize=5)

            axes[7].set_title('Convex Hull of Spine Mask')

            axes[7].imshow(chull_spine_mask, cmap='gray')

        else:

            axes[7].set_title('Convex Hull of Spine Mask')

            axes[7].imshow(chull_spine_mask, cmap='gray')

        axes[8].set_title('Heart Mask')

        axes[8].imshow(heart_mask, cmap='gray')

    return spine_mask, heart_mask

def segmentHeart(img, heart_mask, showOutput=0):

    seg_heart = cv2.bitwise_and(img, img, mask=heart_mask)

    if showOutput:

        plt.figure(figsize=(10, 5))

        plt.title('Segmented Heart')

        plt.imshow(seg_heart, cmap='gray')

    return seg_heart

def segmentHeartLungsTrachea(img, heart_mask, lung_mask, trachea_mask, showOutput=0):

    seg_heart = cv2.bitwise_and(img, img, mask=heart_mask)

    seg_lungs = cv2.bitwise_and(img, img, mask=lung_mask)

    seg_trachea = cv2.bitwise_and(img, img, mask=trachea_mask)

    if showOutput:

        fig, [ax1, ax2, ax3] = plt.subplots(1, 3, figsize=(12, 6), sharex=True, sharey=False)

        ax1.set_title('Segmented Heart')

        ax1.imshow(seg_heart, cmap='gray')

        ax2.set_title('Segmented Lungs')

        ax2.imshow(seg_lungs, cmap='gray')

        ax3.set_title('Segmented Trachea')

        ax3.imshow(seg_trachea, cmap='gray')

    return seg_heart, seg_lungs, seg_trachea

def applyMaskColor(mask, mask_color):

    masked = np.concatenate(([mask[ ... , np.newaxis] * color for color in mask_color]), axis=2)

    # Matplotlib expects color intensities to range from 0 to 1 if a float

    maxValue = np.amax(masked)

    minValue = np.amin(masked)

    # Therefore, scale the color image accordingly

    if maxValue - minValue == 0:

        return masked

    else:

        masked = masked / (maxValue - minValue)

    return masked

def getColoredMasks(img, heart_mask, lung_mask, trachea_mask, showOutput=0):

    heart_mask_color = np.array([256, 0, 0])

    lung_mask_color = np.array([0, 256, 0])

    trachea_mask_color = np.array([0, 0, 256])

    heart_colored = applyMaskColor(heart_mask, heart_mask_color)

    lung_colored = applyMaskColor(lung_mask, lung_mask_color)

    trachea_colored = applyMaskColor(trachea_mask, trachea_mask_color)

    colored_masks = heart_colored + lung_colored + trachea_colored

    if showOutput:

        fig, axes = plt.subplots(2, 2, figsize=(10, 10))

        ax = axes.ravel()

        ax[0].set_title("Original Image")

        ax[0].imshow(img, cmap='gray')

        ax[1].set_title("Heart Mask")

        ax[1].imshow(heart_colored)

        ax[2].set_title("Lung Mask")

        ax[2].imshow(lung_colored)

        ax[3].set_title("Masks")

        ax[3].imshow(colored_masks)

    return heart_colored, lung_colored, trachea_colored, colored_masks