import numpy as np
from scipy.spatial import distance
from scipy import ndimage
from skimage import measure


def define_parameter():
    
    params = {}
    
    
    
    params['lungMinValue']      = -1024
    params['lungMaxValue']      = -800
    params['lungThreshold']     = -900
    
    
    
    params['xRangeRatio1']      = 0.4
    params['xRangeRatio2']      = 0.75
    params['zRangeRatio1']      = 0.9
    params['zRangeRatio2']      = 1.0
    
    
    params['airwayRadiusMask']  = 15  
    params['airwayRadiusX']     = 8   
    params['airwayRadiusZ']     = 15  
    params['super2infer']       = 0  
    
    return params



def generate_structure_trachea(Radius, RadiusX, RadiusZ):
    
    struct_trachea = np.zeros([2*Radius+1,2*Radius+1,RadiusZ])
    for i in range(0,2*Radius+1):
        for j in range(0,2*Radius+1):
            if distance.euclidean([Radius+1,Radius+1],[i,j]) < RadiusX:
                struct_trachea[i,j,:] = 1
            else:
                struct_trachea[i,j,:] = 0
    
    return struct_trachea


def bbox2_3D(img,label,margin,limit):
    
    imgtmp = np.zeros(img.shape)
    imgtmp[img == label] = 1
    
    x = np.any(imgtmp, axis=(1, 2))
    y = np.any(imgtmp, axis=(0, 2))
    z = np.any(imgtmp, axis=(0, 1))

    xmin, xmax = np.where(x)[0][[0, -1]]
    ymin, ymax = np.where(y)[0][[0, -1]]
    zmin, zmax = np.where(z)[0][[0, -1]]

    xmin = xmin - margin - 1
    xmin = max(0,xmin)
    ymin = ymin - margin - 1
    ymin = max(0,ymin)
    zmin = zmin - margin - 1
    zmin = max(0,zmin)        
    xmax = xmax + margin + 1
    xmax = min(xmax,limit[0])
    ymax = ymax + margin + 1
    ymax = min(ymax,limit[1])
    zmax = zmax + margin + 1
    zmax = min(zmax,limit[2])
        
    return xmin, xmax, ymin, ymax, zmin, zmax

	
def generate_initLoc(params, I, Mlung, Radius, RadiusZ, struct_trachea):
	
	mind           = np.argwhere(Mlung == 1)
	initLoc        = [0,0,0]
	minDiff        = float('inf')
	
	if params['super2infer']:
	    slice_no  = np.min(mind[:,2])
	    Itmp      = I[:,:,slice_no:slice_no+RadiusZ]
	else:
	    slice_no  = np.max(mind[:,2])
	    Itmp      = I[:,:,slice_no-RadiusZ:slice_no]

	Mtmp = np.ones(Itmp.shape);
	Mtmp[Itmp < params['lungMinValue']] = 0
	Mtmp[Itmp > params['lungMaxValue']] = 0
	Itmp = Mtmp;
	Mtmp = np.sum(Mtmp, axis = 2)

	for i in range(Radius, Itmp.shape[0]-Radius):
	    for j in range(Radius, Itmp.shape[1]-Radius):
	        if Mtmp[i,j] > 0:   
	            struct_Itmp = Itmp[i-Radius:i+Radius+1,j-Radius:j+Radius+1,:]
	            currVal     = struct_Itmp - struct_trachea
	            currVal     = np.sum(np.square(currVal))

	            if currVal  < minDiff:
	                initLoc = [i,j,slice_no]
	                minDiff = currVal

	print( 'initial location = '+str(initLoc))
	
	return slice_no, initLoc


def close_space_dilation(params, I, Mlung, Radius, RadiusX, RadiusZ, struct_s, slice_no, initLoc):
	
	iterNoPerSlice = RadiusX
	maxFactor      = RadiusX/2
	maxChange      = RadiusX*RadiusX*RadiusX*50
	totalChange    = 1
	tempCheck      = 0	
	[m,n,p]        = I.shape
	Mtmp           = np.zeros([m,n,p])
	struct_m       = ndimage.iterate_structure(struct_s, 2)
	
	if params['super2infer']:
	    Mtmp[initLoc[0]-Radius:initLoc[0]+Radius+1,
	         initLoc[1]-Radius:initLoc[1]+Radius+1,
	         0:slice_no+RadiusZ] = 1
	else:
	    Mtmp[initLoc[0]-Radius:initLoc[0]+Radius+1,
	         initLoc[1]-Radius:initLoc[1]+Radius+1,
	         slice_no-RadiusZ:p-1] = 1
	Mtmp  = np.multiply(Mtmp, Mlung)
	Minit = ndimage.binary_closing(Mtmp, structure = struct_s, iterations = 1)
	Minit = np.int8(Minit)
	Minit[Minit > 0] = 2 

	while totalChange > 0:

	    maxSegmentChange = 0;
	    tempCheck        = tempCheck + 1     
	    L                = measure.label(np.floor(Minit/2))
	    Minit[Minit > 1] = 1  

	    for label in np.unique(L[:]):

	        if label != 0 and np.sum(L[:] == label) > 10:


	            xmin, xmax, ymin, ymax, zmin, zmax = bbox2_3D(L,label,iterNoPerSlice,[m,n,p])                                       
	            Mtmp                = Minit[xmin:xmax,ymin:ymax,zmin:zmax]
	            Itmp                = I[xmin:xmax,ymin:ymax,zmin:zmax]
	            Ltmp                = L[xmin:xmax,ymin:ymax,zmin:zmax]
	            Ltmp[Ltmp != label] = 0
	            Ltmp[Ltmp > 0]      = 1;

	            for iterCount in range(0, iterNoPerSlice):
	                Ltmp = ndimage.binary_dilation(Ltmp, structure = struct_s, iterations = 1)
	                Ltmp = np.int8(Ltmp)
	                Ltmp[Itmp > params['lungThreshold']] = 0                

	            Ltmp = ndimage.binary_closing(Ltmp, structure = struct_s, iterations = 1)
	            Ltmp = np.int8(Ltmp)
	            Ltmp[Mtmp > 0] = 0
	            Ltmp[Ltmp > 0] = 2
	            Ltmp = Ltmp + Mtmp

	            segmentChange = np.sum(Ltmp[:]>1)        
	            if segmentChange < maxChange or tempCheck < 10:
	                Minit[xmin:xmax,ymin:ymax,zmin:zmax] = Ltmp
	                if segmentChange > maxSegmentChange:
	                    maxSegmentChange = segmentChange

	    if tempCheck < 10:
	        maxChange = max(maxFactor*maxSegmentChange,maxChange)
	    else:        
	        maxChange = min(maxFactor*maxSegmentChange,maxChange)

	    totalChange = np.sum(Minit[:]>1)

	    print ('iter = '+str(tempCheck)+' airway sum = '+str(np.sum(Minit[:]>0))\
	                        +' airway change = '+str(totalChange))
	
	Minit[Minit > 0] = 1
	Minit = ndimage.binary_opening(Minit, structure = struct_s, iterations = 1)
	Minit = ndimage.binary_dilation(Minit, structure = struct_m, iterations = 1)
	Maw = np.int8(Minit)
							
	return Maw
