# -*- coding: utf-8 -*-
import numpy as np 
import pandas as pd
import dicom
import os
import scipy.ndimage as ndimage
import matplotlib.pyplot as plt
import scipy.ndimage # added for scaling

import cv2
import time
import glob

from skimage import measure, morphology, segmentation
import SimpleITK as sitk


# make sure I'm in the cwd I think I'm in
print("curr dir: ")
print(os.getcwd())


"""
SCRIPT INFO


This is the first script to run if the patients' images are
in dicom format. If they are in .mhd format then you'll need
to find the script that processes those.

This python script will take in all patients' (or a single
patient's) dicom image slices and export them to segmented
files for the next stage of processing. The next stage will
be nodule identification. This is just preprocessing before 
that.

Explanation of where this file goes in the directory structure:
*only showing dirs necessary for this script to function*
    
    
    

<PROJECT ROOT>
|
|__py_script (name this whatever)
|   |
|   |__<THIS FILE YOU'RE READING.py>
|
|__input
    |
    |__stage1
    |   |
    |   |__<patient one dir name>
    |   |   |
    |   |   |__<image id>.dcm (~200 of these per patient)
    |   |   |__<image id>.dcm
    |   |  
    |   |
    |   |__<patient two dir name>
    |   |   |   
    |   |   |__<image id>.dcm (~200 of these per patient)
    |   |   |__<image id>.dcm
    |   |
    |   |
    |   |__<patient 'n' dir name>
    |   |   |   
    |   |   |__<image id>.dcm (~200 of these per patient)
    |   |   |__<image id>.dcm
    |   |  
    |   |
    |__stage1_2x2x2 (create this, leave it empty)
    |   |
    |   |
    |   |   (these files generated by this script, one .npz per patient)
    |   |__<patient one id>.npz 
    |   |__<patient two id>.npz
    |   |__<patient n id>.npz
    |
    |
    |__stage1_segmented_2x2x2 (create this, leave it empty) 
    |   |
    |   |__<patient one id>.npz (this script generates these)
    |   |__<patient two id>.npz
    |   |__<patient n id>.npz 
    |
    |
    |__stage1_segmented_2x2x2_crop (create this, leave empty)
        |
        |__<patient one id>.npz (this script generates these)
        |__<patient two id>.npz
        |__<patient n id>.npz 



Based on the above info:
Path to dicom files is:
"../input/stage1/<patient id dir name>/<image id>.dcm"


"""




# CONSTANTS ------------------------------------------------------------------------------------------------

RESIZE_SPACING = [2,2,2]  # z, y, x  (x & y MUST be the same)
RESOLUTION_STR = "2x2x2"

img_rows = 448 
img_cols = 448 # global values

DO_NOT_USE_SEGMENTED = True

#STAGE = "stage1"

STAGE_DIR_BASE = "../input/%s/"  # on one cluster we had input_shared

LUNA_MASKS_DIR = "../luna/data/original_lung_masks/"
luna_subset = 0       # initial 
LUNA_BASE_DIR = "../luna/data/original_lungs/subset%s/"  # added on AWS; data as well 
LUNA_DIR = LUNA_BASE_DIR % luna_subset

CSVFILES = "../luna/data/original_lungs/CSVFILES/%s"
LUNA_ANNOTATIONS = CSVFILES % "annotations.csv"
LUNA_CANDIDATES =  CSVFILES % "candidates.csv"

MARKER_INTERNAL_THRESH = -400  
MARKER_FRAME_WIDTH = 9      # 9 seems OK for the half special case ...

BINARY_CLOSING_SIZE = 7  #was 7 before final;   5 for disk seems sufficient - for safety let's go with 6 or even 7



# Function defs --------------------------------------------------------------------------------------------------





# Load the scans in given folder path (loads the most recent acquisition)
def load_scan(path):
    slices = [dicom.read_file(path + '/' + s) for s in os.listdir(path)]
    #slices.sort(key = lambda x: int(x.InstanceNumber))
       
    acquisitions = [x.AcquisitionNumber for x in slices]
    
    vals, counts = np.unique(acquisitions, return_counts=True)
    vals = vals[::-1]  # reverse order so the later acquisitions are first (the np.uniques seems to always return the ordered 1 2 etc.
    counts = counts[::-1]
    
    ## take the acquistions that has more entries; if these are identical take the later  entrye
    acq_val_sel = vals[np.argmax(counts)]
  

    ##acquisitions = sorted(np.unique(acquisitions), reverse=True)
    
    if len(vals) > 1:
        print ("WARNING ##########: MULTIPLE acquisitions & counts, acq_val_sel, path: ", vals, counts, acq_val_sel, path)
    slices2= [x for x in slices if x.AcquisitionNumber == acq_val_sel]
    
    slices = slices2
    
   
    ## ONE path includes 2 acquisitions (2 sets), take the latter acquiisiton only whihch cyupically is better than the first/previous ones.
    ## example of the     '../input/stage1/b8bb02d229361a623a4dc57aa0e5c485'
    
    #slices.sort(key = lambda x: int(x.ImagePositionPatient[2]))  # from v 8, BUG should be float
    slices.sort(key = lambda x: float(x.ImagePositionPatient[2]))  # from v 9
    try:
        slice_thickness = np.abs(slices[0].ImagePositionPatient[2] - slices[1].ImagePositionPatient[2])
    except:
        slice_thickness = np.abs(slices[0].SliceLocation - slices[1].SliceLocation)
        
    for s in slices:
        s.SliceThickness = slice_thickness
        
    return slices



def get_3d_data_slices(slices):  # get data in Hunsfield Units
    slices.sort(key = lambda x: float(x.ImagePositionPatient[2]))  # from v 9
    
    image = np.stack([s.pixel_array for s in slices])
    image = image.astype(np.int16)  # ensure int16 (it may be here uint16 for some images )
    image[image == -2000] = 0   #correcting cyindrical bound entrioes to 0
    
    # Convert to Hounsfield units (HU)
    # The intercept is usually -1024
    for slice_number in range(len(slices)):  # from v 8
        intercept = slices[slice_number].RescaleIntercept
        slope = slices[slice_number].RescaleSlope
        
        if slope != 1:  # added 16 Jan 2016, evening
            image[slice_number] = slope * image[slice_number].astype(np.float64)
            image[slice_number] = image[slice_number].astype(np.int16)
            
        image[slice_number] += np.int16(intercept)
    
    return np.array(image, dtype=np.int16)




def resample(image, scan, new_spacing=[1,1,1]):
    # Determine current pixel spacing
    spacing = map(float, ([scan[0].SliceThickness] + scan[0].PixelSpacing))
    spacing = np.array(list(spacing))
    
    #scan[2].SliceThickness


    resize_factor = spacing / new_spacing
    new_real_shape = image.shape * resize_factor
    new_shape = np.round(new_real_shape)
    real_resize_factor = new_shape / image.shape
    new_spacing = spacing / real_resize_factor
    
    image = scipy.ndimage.interpolation.zoom(image, real_resize_factor, mode='nearest')  ### early orig modified 
    
    return image, new_spacing




def generate_markers(image):
    #Creation of the internal Marker
    
    useTestPlot = False
    if useTestPlot:
        timg = image
        plt.imshow(timg, cmap='gray')
        plt.show()

    add_frame_vertical = True  
    if add_frame_vertical:   # add frame for potentially closing the lungs that touch the edge, but only vertically
                        
        fw = MARKER_FRAME_WIDTH  # frame width (it looks that 2 is the minimum width for the algorithms implemented here, namely the first 2 operations for the marker_internal)
   
        xdim = image.shape[1]
        #ydim = image.shape[0]
        img2 = np.copy(image)
        #y3 = ydim // 3
        
        img2 [:, 0]    = -1024
        img2 [:, 1:fw] = 0

        img2 [:, xdim-1:xdim]    = -1024
        img2 [:, xdim-fw:xdim-1] = 0
               
        marker_internal = img2 < MARKER_INTERNAL_THRESH  
    else:
        marker_internal = image < MARKER_INTERNAL_THRESH  # was -400
        
    
    useTestPlot = False
    if useTestPlot:
        timg = marker_internal
        plt.imshow(timg, cmap='gray')
        plt.show()
       
    
    correct_edges2 = False  ## NOT a good idea - no added value
    if correct_edges2:   
        marker_internal[0,:]   = 0
        marker_internal[:,0]   = 0
        #marker_internal[:,1]   = True
        #marker_internal[:,2]   = True
        marker_internal[511,:]   = 0
        marker_internal[:,511]   = 0
    
    marker_internal = segmentation.clear_border(marker_internal, buffer_size=0)
    marker_internal_labels = measure.label(marker_internal)
    areas = [r.area for r in measure.regionprops(marker_internal_labels)]
    areas.sort()
    if len(areas) > 2:
        for region in measure.regionprops(marker_internal_labels):
            if region.area < areas[-2]:
                for coordinates in region.coords:                
                       marker_internal_labels[coordinates[0], coordinates[1]] = 0
    marker_internal = marker_internal_labels > 0
    #Creation of the external Marker
    external_a = ndimage.binary_dilation(marker_internal, iterations=10)  # was 10
    external_b = ndimage.binary_dilation(marker_internal, iterations=55)   # was 55 
    marker_external = external_b ^ external_a
    #Creation of the Watershed Marker matrix
    #marker_watershed = np.zeros((512, 512), dtype=np.int)  # origi
    marker_watershed = np.zeros((marker_external.shape), dtype=np.int)
    
    marker_watershed += marker_internal * 255
    marker_watershed += marker_external * 128
    
    return marker_internal, marker_external, marker_watershed




def rescale_n(n,reduce_factor):
    return max( 1, int(round(n / reduce_factor)))




def seperate_lungs_cv2(image):      # for increased speed
    #Creation of the markers as shown above:
    marker_internal, marker_external, marker_watershed = generate_markers(image)
    
    #image_size = image.shape[0]
    reduce_factor = 512 / image.shape[0]
    
    #Creation of the Sobel-Gradient
    sobel_filtered_dx = ndimage.sobel(image, 1)
    sobel_filtered_dy = ndimage.sobel(image, 0)
    sobel_gradient = np.hypot(sobel_filtered_dx, sobel_filtered_dy)
    sobel_gradient *= 255.0 / np.max(sobel_gradient)
    
    useTestPlot = False
    if useTestPlot:
        timg = sobel_gradient
        plt.imshow(timg, cmap='gray')
        plt.show()
        
    #Watershed algorithm
    watershed = morphology.watershed(sobel_gradient, marker_watershed)
    
    if useTestPlot:
        timg = marker_external
        plt.imshow(timg, cmap='gray')
        plt.show()    
    
    #Reducing the image created by the Watershed algorithm to its outline
    #wsize = rescale_n(3,reduce_factor)  # THIS IS TOO SMALL, dynamically adjusting the size for the watersehed algorithm
    outline = ndimage.morphological_gradient(watershed, size=(3,3))   # original (3,3), (wsize, wsize) is too small to create an outline
    outline = outline.astype(bool)
    outline_u = outline.astype(np.uint8)  #added
    
    #Performing Black-Tophat Morphology for reinclusion
    #Creation of the disk-kernel and increasing its size a bit
    blackhat_struct = [[0, 0, 1, 1, 1, 0, 0],
                       [0, 1, 1, 1, 1, 1, 0],
                       [1, 1, 1, 1, 1, 1, 1],
                       [1, 1, 1, 1, 1, 1, 1],
                       [1, 1, 1, 1, 1, 1, 1],
                       [0, 1, 1, 1, 1, 1, 0],
                       [0, 0, 1, 1, 1, 0, 0]]
    

    use_reduce_factor = True
    if use_reduce_factor:
        blackhat_struct = ndimage.iterate_structure(blackhat_struct, rescale_n(8,reduce_factor))  # dyanmically adjust the number of iterattions; original was 8
    else:
        blackhat_struct = ndimage.iterate_structure(blackhat_struct, 8) 
    
    blackhat_struct_cv2 = blackhat_struct.astype(np.uint8)
    #Perform the Black-Hat
    
    #outline += ndimage.black_tophat(outline, structure=blackhat_struct)  # original slow 

    #outline1 = outline + (cv2.morphologyEx(outline_u, cv2.MORPH_BLACKHAT, kernel=blackhat_struct_cv2)).astype(np.bool)
    #outline2 = outline + ndimage.black_tophat(outline, structure=blackhat_struct)
    #np.array_equal(outline1,outline2)  # True

    outline += (cv2.morphologyEx(outline_u, cv2.MORPH_BLACKHAT, kernel=blackhat_struct_cv2)).astype(np.bool)  # fats


    if useTestPlot:
        timg = outline
        plt.imshow(timg, cmap='gray')
        plt.show()
        

    #Use the internal marker and the Outline that was just created to generate the lungfilter
    lungfilter = np.bitwise_or(marker_internal, outline)
    
    if useTestPlot:
        timg = lungfilter
        plt.imshow(timg, cmap='gray')
        plt.show()
    
    #Close holes in the lungfilter
    #fill_holes is not used here, since in some slices the heart would be reincluded by accident
    ##structure = np.ones((BINARY_CLOSING_SIZE,BINARY_CLOSING_SIZE)) # 5 is not enough, 7 is
    structure2 = morphology.disk(2)  # used to fill the gaos/holes close to the border (otherwise the large sttructure would create a gap by the edge)
    if use_reduce_factor:
        structure3 = morphology.disk(rescale_n(BINARY_CLOSING_SIZE,reduce_factor)) # dynanically adjust; better , 5 seems sufficient, we use 7 for safety/just in case
    else:
        structure3 = morphology.disk(BINARY_CLOSING_SIZE) # dynanically adjust; better , 5 seems sufficient, we use 7 for safety/just in case
    
    
    ##lungfilter = ndimage.morphology.binary_closing(lungfilter, structure=structure, iterations=3) #, ORIGINAL iterations=3)  # was structure=np.ones((5,5))
    lungfilter2 = ndimage.morphology.binary_closing(lungfilter, structure=structure2, iterations=3)  # ADDED
    lungfilter3 = ndimage.morphology.binary_closing(lungfilter, structure=structure3, iterations=3)
    lungfilter = np.bitwise_or(lungfilter2, lungfilter3)
    
    ### NOTE if no iterattions, i.e. default 1 we get holes within lungs for the disk(5) and perhaps more
    
    #Apply the lungfilter (note the filtered areas being assigned -2000 HU)
    #image.shape
    #segmented = np.where(lungfilter == 1, image, -2000*np.ones((512, 512)).astype(np.int16)) # was -2000 someone suggested 30
    segmented = np.where(lungfilter == 1, image, -2000*np.ones(image.shape).astype(np.int16)) # was -2000 someone suggested 30
    
    return segmented, lungfilter, outline, watershed, sobel_gradient, marker_internal, marker_external, marker_watershed







def segment_all(stage, part=0, processors=1, showSummaryPlot=True):  # stage added to simplify the stage1 and stage2 calculations

    count = 0
    STAGE_DIR = STAGE_DIR_BASE % stage
    folders = glob.glob(''.join([STAGE_DIR,'*']))
    if len(folders) == 0:
        print ("ERROR, check directory, no folders found in: ", STAGE_DIR )
        
        
    # file paths are set up correctly and working
        
    for folder in folders:
        
        count += 1
        if count % processors == part:    # do this part in this process, otherwise skip
            path = folder
            
            slices = load_scan(path)
            image_slices = get_3d_data_slices(slices)
            #mid = len(image_slices) // 2
            #img_sel = mid
            
            useTestPlot = False
            if useTestPlot:
                print("Shape before segmenting\t", image_slices.shape)
                plt.hist(image_slices.flatten(), bins=80, color='c')
                plt.xlabel("Hounsfield Units (HU)")
                plt.ylabel("Frequency")
                plt.show()


            start = time.time()
            resampleImages = True
            if resampleImages:
                image_resampled, spacing = resample(image_slices, slices, RESIZE_SPACING)   # let's start wkith this small resolutuion for workign our the system (then perhaps 2, 0.667, 0.667)
                print("Shape_before_&_after_resampling\t", image_slices.shape,image_resampled.shape)
                if useTestPlot:
                    plt.imshow(image_slices[image_slices.shape[0]//2], cmap=plt.cm.bone)
                    plt.show() 
                    plt.imshow(image_resampled[image_resampled.shape[0]//2], cmap=plt.cm.bone)
                    np.max(image_slices)     
                    np.max(image_resampled)
                    
                    np.min(image_slices)     
                    np.min(image_resampled)   
                    plt.show() 
                image_slices = image_resampled
                        
            shape = image_slices.shape
            l_segmented = np.zeros(shape).astype(np.int16)
            l_lungfilter = np.zeros(shape).astype(np.bool)
            l_outline = np.zeros(shape).astype(np.bool)
            l_watershed = np.zeros(shape).astype(np.int16)
            l_sobel_gradient = np.zeros(shape).astype(np.float32)
            l_marker_internal = np.zeros(shape).astype(np.bool)
            l_marker_external = np.zeros(shape).astype(np.bool)
            l_marker_watershed = np.zeros(shape).astype(np.int16)  
            
            # start = time.time()
            i=0
            for i in range(shape[0]):
                l_segmented[i], l_lungfilter[i], l_outline[i], l_watershed[i], l_sobel_gradient[i], l_marker_internal[i], l_marker_external[i], l_marker_watershed[i] = seperate_lungs_cv2(image_slices[i])
            print("Rescale & Seg time, and path: ", ((time.time() - start)), path )
                
            
            if useTestPlot:
                plt.hist(image_slices.flatten(), bins=80, color='c')
                plt.xlabel("Hounsfield Units (HU)")
                plt.ylabel("Frequency")
                plt.show()
            
            
                plt.hist(l_segmented.flatten(), bins=80, color='c')
                plt.xlabel("Hounsfield Units (HU)")
                plt.ylabel("Frequency")
                plt.show()
                
                img_sel_i = shape[0] // 2
                # Show some slice in the middle
                plt.imshow(image_slices[img_sel_i], cmap=plt.cm.gray)
                plt.show()
                
                # Show some slice in the middle
                plt.imshow(l_segmented[img_sel_i], cmap='gray')
                plt.show()
            
            path_rescaled = path.replace(stage, ''.join([stage, "_", RESOLUTION_STR]), 1)
            path_segmented = path.replace(stage, ''.join([stage, "_segmented_", RESOLUTION_STR]), 1)
            path_segmented_crop = path.replace(stage, ''.join([stage, "_segmented_", RESOLUTION_STR, "_crop"]), 1)
         
            # currently failing: No directory: '../input/stage1_2x2x2
            np.savez_compressed (path_rescaled, image_slices)
            np.savez_compressed (path_segmented, l_segmented)
            
            mask = l_lungfilter.astype(np.int8)
            
            regions = measure.regionprops(mask)  # this measures the largest region and is a bug when the mask is not the largest region !!!

     
            bb = regions[0].bbox
            #print(bb)
            zlen = bb[3] - bb[0]
            ylen = bb[4] - bb[1]
            xlen = bb[5] - bb[2]
            
            dx = 0  # could  be reduced
            ## have to reduce dx as for istance at least image the lungs stretch right to the border evebn without cropping 
            ## namely for '../input/stage1/be57c648eb683a31e8499e278a89c5a0'
            
            crop_max_ratio_z = 0.6  # 0.8 is to big    make_submit2(45, 1)
            crop_max_ratio_y = 0.4
            crop_max_ratio_x = 0.6
            
            bxy_min = np.min(bb[1:3]) 
            bxy_max = np.max(bb[4:6])
            mask_shape= mask.shape
            image_shape = l_segmented.shape
            
            mask_volume = zlen*ylen*zlen /(mask_shape[0] * mask_shape[1] * mask_shape[2])
            mask_volume_thresh = 0.08  # anything below is too small (maybe just one half of the lung or something very small0)
            mask_volume_check =   mask_volume >   mask_volume_thresh
            # print ("Mask Volume: ", mask_volume )    
            
            ### DO NOT allow the mask to touch x & y ---> if it does it is likely a wrong one as for:
            ## folders[3] , path = '../input/stage1/9ba5fbcccfbc9e08edcfe2258ddf7
            
            # maskOK = False
            if bxy_min >0 and bxy_max < 512 and mask_volume_check and zlen/mask_shape[0] > crop_max_ratio_z and ylen/mask_shape[1] > crop_max_ratio_y and xlen/mask_shape[2]  > crop_max_ratio_x:
                
                ## square crop and at least dx elements on both sides on x & y
                bxy_min = np.min(bb[1:3]) 
                bxy_max = np.max(bb[4:6])
                
                if bxy_min == 0 or bxy_max == 512:
                    # Mask to bigg, auto-correct
                    print("The following mask likely too big, autoreducing by:", dx)
                    
                    bxy_min = np.max((bxy_min, dx)) 
                    bxy_max = np.min ((bxy_max, mask_shape[1] - dx))
                
                image = l_segmented[bb[0]:bb[3], bxy_min:bxy_max, bxy_min:bxy_max]
                mask =   mask[bb[0]:bb[3], bxy_min:bxy_max, bxy_min:bxy_max]
                #maskOK = True
         
                print ("Shape, cropped, bbox ", mask_shape, mask.shape, bb)
        
            elif bxy_min> 0 and bxy_max < 512 and mask_volume_check and zlen/mask.shape[0] > crop_max_ratio_z:
                ## cut on z at least
            
                image = l_segmented[bb[0]:bb[3], dx: image_shape[1] - dx, dx: image_shape[2] - dx]
                #mask =   mask[bb[0]:bb[3], dx: mask_shape[1] - dx, dx: mask_shape[2] - dx]
                print("Mask too small, NOT auto-cropping x-y: shape, cropped, bbox, ratios, violume:", mask_shape, image.shape, bb, path, zlen/mask_shape[0], ylen/mask_shape[1], xlen/mask_shape[2], mask_volume)
        
        
            else:
                image = l_segmented[0:mask_shape[0], dx: image_shape[1] - dx, dx: image_shape[2] - dx]
                #mask =   mask[0:mask_shape[0], dx: mask_shape[1] - dx, dx: mask_shape[2] - dx]
                print("Mask wrong, NOT auto-cropping: shape, cropped, bbox, ratios, volume:", mask_shape, image.shape, bb, path, zlen/mask_shape[0], ylen/mask_shape[1], xlen/mask_shape[2], mask_volume)
                
            
            if showSummaryPlot:
                img_sel_i = shape[0] // 2
                # Show some slice in the middle
                useSeparatePlots = False
                if useSeparatePlots:
                    plt.imshow(image_slices[img_sel_i], cmap=plt.cm.gray)
                    plt.show()
                    
                    # Show some slice in the middle
                    plt.imshow(l_segmented[img_sel_i], cmap='gray')
                    plt.show()
                else:                     
                    f, ax = plt.subplots(1, 2, figsize=(6,3))
                    ax[0].imshow(image_slices[img_sel_i],cmap=plt.cm.bone)
                    ax[1].imshow(l_segmented[img_sel_i],cmap=plt.cm.bone)
                    plt.show()   
                # Show some slice in the middle
                #plt.imshow(image[image.shape[0] // 2], cmap='gray')  # don't show it for simpler review 
                #plt.show()
                
            
            np.savez_compressed(path_segmented_crop, image)
            #print("Mask count: ", count)
            #print ("Shape: ", image.shape)
    return part, processors, count






# PART 2 - prepare competitions data, for stage 1 and stage 2  ------------------------------------------------------------
part = 0
processors = 1          # you may run several of these jobs; define processors to say 4, and start 4 separate jobs with part = 0, 1, 2, 3 respectively
showSummaryPlot = True
# for stage in ["stage1", "stage2"]:  
stage = "stage1"
        



start_time = time.time()
print ("Starting segmentation, stage, part of a set/processors: ", stage, part, processors)
part, processors, count = segment_all(stage, part, processors, showSummaryPlot)
print ("Completed, part, processors,total count tried, total time: ", stage, part, processors, count, time.time()-start_time)
