# -*- coding: utf-8 -*-
"""
Created on Thu May 12 06:55:20 2022

@author: petra

script to ilustarte algorithm step by step
images that are made in this script are also used in diploma thesis
"""


globals().clear()


# IMPORTS

import matplotlib.pyplot as plt
import numpy as np
import copy
import scipy
import math
import skimage
import scipy.io
import scipy.ndimage
import scipy.interpolate
import skimage.segmentation
import skimage.color
from scipy.ndimage import zoom
from numpy import matlib


# FUNCTIONS
def normaliseImg(img):
    '''
    img: input image
    function to normalise the image so that the values are in between 0 and 1
    output: normalised image
    '''
    imgMax = np.max(img)
    imgMin = np.min(img)

    for i in range(np.shape(img)[0]):
        for j in range(np.shape(img)[1]):
            img[i, j] = (img[i, j] - imgMin) / (imgMax - imgMin)

    return img

def HomomorphicFilter(img, cutoffFreq, a, b):
    '''
    homomorphic filtration
    
    img: input image.
    cutoffFreq: cut-off frequency.
    a: constant, added in filtration to the image.
    b: constant, multiply the impact of filter on the image.

    output: filtered image.

    '''

    img[img == 0] = np.min(img[img > 0])  # to make sure, there will be no 0 when applying log
    img = np.log(np.array(img, dtype="float"))  # applying the log

    img = np.fft.fft2(img)  # transfer into fft domain

    imgShape = img.shape  # shape of image

    x = img.shape[0] / 2
    y = img.shape[1] / 2

    H = np.zeros(imgShape)  # preparation for filter

    xx, yy = np.meshgrid(range(imgShape[0]), range(imgShape[1]))  # make grid for future filter

    H = np.exp((-((xx - x) ** 2 + (yy - y) ** 2) / (2 * (cutoffFreq) ** 2)))  # turn meshgrid into -96:96, then put it into the Gauss formula
    H = 1 - H  # inverse high and low pass

    H = np.fft.fftshift(H)  # fft

    I = (a + b * H) * img  # filtration - multiplication of filter and image i n fft domain

    I = np.fft.ifft2(I)  # transfer into image domain
    I = np.real(I)
    I = np.exp(I)  # exponent

    return I

def histogramEqualisation(im,nbr_bins):
    
    '''
    function to equalise the histogram of the data
    im: image to equalise
    nbr_bins: number of bins from which the histogram will be created (and then transformed)
    
    output: equalised image

    '''
    imhist,bins = np.histogram(im.flatten(), nbr_bins, normed=True, density=True)
    cdf = imhist.cumsum() # cumulative distribution function
    cdf = np.max(im) * cdf / cdf[-1] # normalize
 
    # use linear interpolation of cdf to find new pixel values
    im2 = np.interp(im.flatten(), bins[:-1], cdf)
 
    return im2.reshape(im.shape) #, cdf

def curvefit2(x, y):
    '''
    transformation of coordinates  to higher polynom values
    '''
    a = x**2
    b = y**2
    c = (x * y) **1
    d = x ** 1
    e = y ** 1
    f = 1

    number = 6

    returning = np.array([])
    returning = np.append(returning, a)
    returning = np.append(returning, b)
    returning = np.append(returning, c)
    returning = np.append(returning, d)
    returning = np.append(returning, e)
    returning = np.append(returning, f)
    returning = returning.reshape(1, number)

    return returning, number

def imgradient_fcn(img, operator):
    '''
    to calculate magnitude of image
    img: input image
    operator: operator to choose, between Prewitt and Sobel
    
    output: magnitude and direction of image
    '''

    if (operator == 'prewitt'):
        h = np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])  # prewitt operator'

    elif (operator == 'sobel'):
        h = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])  # sobel operator

    else:
        h = np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])

    Gx = scipy.ndimage.correlate(img, h.transpose(), mode='nearest')
    Gy = scipy.ndimage.correlate(img, h, mode='nearest')

    Gmag = np.sqrt((np.abs(Gx) ** 2) + (np.abs(Gy) ** 2))
    Gdir = np.arctan2(-Gy, Gx) * (180 / np.pi)

    return Gmag, Gdir

def imgradientxy_fcn(img, operator):
    '''
    to calculate gradients of image
    img: input image
    operator: operator to choose, between Prewitt and Sobel
    
    output: gradients in both directions of image
    '''
    if (operator == 'prewitt'):
        h = np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])  # prewitt operator'

    elif (operator == 'sobel'):
        h = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])  # sobel operator

    else:
        h = np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])

    Gx = scipy.ndimage.correlate(img, h.transpose(), mode='nearest')
    Gy = scipy.ndimage.correlate(img, h, mode='nearest')

    # mode{‘reflect’, ‘constant’, ‘nearest’, ‘mirror’, ‘wrap’}

    return Gx, Gy

def move_to_end(points):
    '''
    move the first point to last position
    '''    
    p = points[0,:].reshape(1,2)
    points = points[1:,:]
    points = np.append(points,p,axis=0)
    return points

def CreateMatrixCoeff(M, a, b, c):
    '''
    function to create a coeficient matrix - internal forces influencing the contour
    M: lenght of contour
    a-b-c: calculated coefficients from set parameters
    
    output: coefficient matrix
    '''   
    A = np.diag(c)
    A = A + np.diag(b[0:M - 1], k=1)
    A = A + np.diag(a[2:M], k=-2)
    A = A + np.diag(b[1:M], k=-1)
    A = A + np.diag(a[0:M - 2], k=2)
    A[0, -1] = b[0]
    A[0, -2] = a[0]
    A[1, -1] = a[1]
    A[-1, 0] = b[-1]
    A[-2, 0] = a[-2]
    A[-1, 1] = a[-1]

    return A

def distances(points):
    '''
    function to calculate distances between points
    
    '''    
    d = np.diff(points, n=1, axis=0)
    # pom = d[0,:].reshape(1,2)
    # d = d[1:,:]
    # d = np.append(d, pom, axis=0)
    d2 = (np.array([points[-1, 0] - points[0, 0], points[-1, 1] - points[0, 1]])).reshape(1, 2)
    d = np.concatenate((d2, d), axis=0)
    d = d ** 2
    d = np.sum(d, axis=1)
    d = np.sqrt(d)

    return d

def PolyArea(x,y):
    '''
    function to calculate area inside the polygon
    '''   
    
    p = 0.5*np.abs(np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1)))
    return p

def compute_baloon_force(points):
    '''
    calculation of balloon force (normal force to contour)
    '''   
    
    d = np.diff(points, n=1, axis=0)
    # pom = d[0,:].reshape(1,2)
    # d = d[1:,:]
    # d = np.append(d, pom, axis=0)
    d2 = (np.array([points[-1, 0] - points[0, 0], -points[-1, 1] + points[0, 1]])).reshape(1, 2)
    d = np.concatenate((d, d2), axis=0)
    d[:, 0] = -d[:, 0]
    d = np.flip(d, axis=1)

    d_new = copy.copy(d)

    for i in range(np.shape(d)[0] - 1):
        d_new[i, :] = d[i, :] + d[i + 1, :]

    d_new[-1, :] = d_new[-1, :] + d[0, :]
    d = d_new

    # normalise the vector
    Area_old = PolyArea(points[:, 0], points[:, 1]) + 50
    norms = np.sqrt(np.sum(d ** 2, axis=1))
    d = (d / (matlib.repmat(norms, 2, 1).transpose())) / Area_old
    d = np.roll(d, 1, 0)

    return d


plt.close('all')

# %%

img = scipy.io.loadmat('ukazkovy_jeden_obr_14_4.mat')
img = img['mydata']

gt = scipy.io.loadmat('ukazkova_jedna_maska_14_4.mat')
gt = gt['mydata']


row = 93
col = 97


# plt.figure(num=1)
# plt.imshow(img, cmap='gray')
# plt.axis('off')
# plt.title('14th measure, 4th slice', fontsize=15)
# plt.suptitle('    Original measured image', fontsize=20)
# plt.show()



plt.figure(num=13)
plt.subplot(2,4,1)
plt.imshow(img, cmap='gray')
plt.axis('off')
plt.title('Original measured image', fontsize=15)
plt.suptitle('14th measure, 4th slice', fontsize=20)
plt.show()

# %% INHOMOGENITY ILLUMINATION SOLUTION

# correspond to preprocessing_image function
# make same energy for all pxs, normalise the image and filter

# img_norm = img / np.linalg.norm(img, ord=None)
img_norm = normaliseImg(img)
img_norm = scipy.ndimage.gaussian_filter(img_norm, sigma=1.5)
img_norm[img_norm < 0.5 * np.mean(img_norm)] = 0

# plt.figure(num=2)
# plt.imshow(img_norm, cmap='gray')
# plt.axis('off')
# plt.title('14th measure, 4th slice', fontsize=15)
# plt.suptitle('Image after preprocessing', fontsize=20)
# plt.show()


plt.figure(num=13)
plt.subplot(2,4,2)
plt.imshow(img_norm, cmap='gray')
plt.axis('off')
plt.title('Image after preprocessing', fontsize=15)
plt.show()

# 
# correspond to make_plane function
# plain that equalize the inhomogeneity in image

X, Y = np.meshgrid(range(np.shape(img_norm)[0]), range(np.shape(img_norm)[1]))

X = X.flatten()
X = X + 1e-5
X = X.astype(float)
Y = Y.flatten()
Y = Y + 1e-5
Y = Y.astype(float)
Z = img_norm.flatten()
Z = Z
Z = Z.astype(float)

x = []
y = []
z = []    

for i in range(len(X)):
    value = Z[i]
    if value > np.mean(img_norm):
        x = np.append(x, X[i])
        y = np.append(y, Y[i])
        z = np.append(z, Z[i])


for i in range(len(x)):
    arr, cols = curvefit2(x[i], y[i])
    if (i == 0):
        A_matrix = np.zeros((len(x), cols))

    arr = arr.reshape(1, cols)
    # ### save recalculated coefficients in matrix
    A_matrix[i, :] = arr

coeff, r, rank, s = np.linalg.lstsq(A_matrix, z, rcond=None)


# ### same process with full image
for i in range(len(X)):
    arr, cols = curvefit2(X[i], Y[i])
    if (i == 0):
        B_matrix = np.zeros((len(X), cols))

    arr = arr.reshape(1, cols)
    B_matrix[i, :] = arr

# ### calculate plane by matrix multiplication
plane = B_matrix @ coeff


plane = plane.reshape(np.shape(img_norm)[0], np.shape(img_norm)[1])

# plt.figure(num=3)
# plt.imshow(plane)
# plt.colorbar()
# plt.axis('off')
# plt.title('14th measure, 4th slice', fontsize=15)
# plt.suptitle('Image inhomogeneities equalisation plane', fontsize=20)
# plt.show()


# # plt.figure(num=14)
# # plt.subplot(2,2,1)
# # plt.imshow(plane)
# # plt.colorbar()
# # plt.rc('ytick', labelsize=10) # labelsize=20
# # plt.axis('off')
# # plt.title('Image inhomogeneities equalisation plane', fontsize=20)
# # plt.suptitle('Visualisation of ', fontsize=20)
# # plt.show()





# correspond to image_calibrations funtion
# calibrate image for segmentation

# ### image division by plane
image_after = img_norm / plane

image_after[image_after < 0] = 0
image_after[image_after > 10] = 0

# plt.figure(num=4)
# plt.imshow(image_after, cmap='gray')
# plt.axis('off')
# plt.title('14th measure, 4th slice', fontsize=15)
# plt.suptitle('Image after inhomogenities equalisation', fontsize=20)
# plt.show()


plt.figure(num=13)
plt.subplot(2,4,3)
plt.imshow(image_after, cmap='gray')
plt.axis('off')
plt.title('After inhomogenities equalisation', fontsize=15)
plt.show()


img_fit = copy.copy(image_after)
img_fit = HomomorphicFilter(img_fit, cutoffFreq=8, a=1, b=2)
img_fit = img_fit ** 0.5
img_fit[img_fit > 10 * np.mean(img_fit)] = np.mean(img_fit)
img_fit[img_fit < 2 * np.median(img_fit)] = 0

# %% IMAGE PREPROCESSING

# plt.figure(num=5)
# plt.hist(img_fit, bins=10, histtype='bar', log=True)
# plt.title('Image histogram before equalisation (log scale)', fontsize=20) # fontsize=40
# plt.rc('xtick', labelsize=10) # labelsize=20
# plt.rc('ytick', labelsize=10) # labelsize=20
# plt.xlabel('Pixel intensities', fontsize=15) # fontsize=30
# plt.ylabel('Frequency of intensity occurence', fontsize=15) # fontsize=30
# plt.show()


# plt.figure(num=14)
# plt.subplot(2,2,3)
# plt.hist(img_fit, bins=10, histtype='bar', log=True)
# plt.title('Image histogram before equalisation (log scale)', fontsize=20) # fontsize=40
# plt.rc('xtick', labelsize=10) # labelsize=20
# plt.rc('ytick', labelsize=10) # labelsize=20
# plt.xlabel('Pixel intensities', fontsize=15) # fontsize=30
# plt.ylabel('Frequency of intensity occurence', fontsize=15) # fontsize=30
# plt.show()


img_fit = histogramEqualisation(img_fit, nbr_bins=64000)

# plt.figure(num=6)
# plt.hist(img_fit, bins=10, histtype='bar', log=True)
# plt.title('Image histogram after equalisation (log scale)', fontsize=20) # fontsize=40
# plt.rc('xtick', labelsize=10) # labelsize=20
# plt.rc('ytick', labelsize=10) # labelsize=20
# plt.xlabel('Pixel intensities', fontsize=15) # fontsize=30
# plt.ylabel('Frequency of intensity occurence', fontsize=15) # fontsize=30
# plt.show()

# plt.figure(num=14)
# plt.subplot(2,2,4)
# plt.hist(img_fit, bins=10, histtype='bar', log=True)
# plt.title('Image histogram after equalisation (log scale)', fontsize=20) # fontsize=40
# plt.rc('xtick', labelsize=10) # labelsize=20
# plt.rc('ytick', labelsize=10) # labelsize=20
# plt.xlabel('Pixel intensities', fontsize=15) # fontsize=30
# plt.ylabel('Frequency of intensity occurence', fontsize=15) # fontsize=30
# plt.show()


# plt.figure(num=7)
# plt.imshow(img_fit, cmap='gray')
# plt.axis('off')
# plt.title('14th measure, 4th slice', fontsize=15)
# plt.suptitle('Image after histogram equalisation', fontsize=20)
# plt.show()


plt.figure(num=13)
plt.subplot(2,4,4)
plt.imshow(img_fit, cmap='gray')
plt.axis('off')
plt.title('After histogram equalisation', fontsize=15)
plt.show()



img_fit = normaliseImg(img_fit)
img_fit[img_fit < 1 * np.mean(img_fit)] = 0 # 2* np.mean pro diastoly, 1* np.mean pro systoly
img_fit = np.max(img_fit) - img_fit
img_fit = skimage.morphology.area_closing(img_fit)
img_fit = np.max(img_fit) - img_fit
image_after = img_fit

# plt.figure(num=8)
# plt.imshow(img_fit, cmap='gray')
# plt.axis('off')
# plt.title('14th measure, 4th slice', fontsize=15)
# plt.suptitle('Image prepared for segmentation', fontsize=20)
# plt.show()


plt.figure(num=13)
plt.subplot(2,4,5)
plt.imshow(img_fit, cmap='gray')
plt.axis('off')
plt.title('Image prepared for segmentation', fontsize=15)
plt.show()


# %% PREPARATIONS FOR SEGMENTATION


# correspond to segmentation_alg function
# part of code, where segmentation takes place

Gmag, Gdir = imgradient_fcn(image_after, operator='prewitt')
Gx, Gy = imgradientxy_fcn(Gmag, operator='prewitt')    
Gx = Gx
Gy = Gy

# plt.figure(num=9)
# plt.imshow(Gmag, cmap='gray')
# plt.axis('off')
# plt.title('14th measure, 4th slice', fontsize=15)
# plt.suptitle('Edge map', fontsize=20)
# plt.show()


plt.figure(num=13)
plt.subplot(2,4,6)
plt.imshow(Gmag, cmap='gray')
plt.axis('off')
plt.title('Edge map', fontsize=15)
plt.show()

Xx = np.array(range(img.shape[1]))
Yy = np.array(range(img.shape[0]))
Gx_inter = scipy.interpolate.RegularGridInterpolator((Xx,Yy), Gx.transpose())
Gy_inter = scipy.interpolate.RegularGridInterpolator((Xx,Yy), Gy.transpose())


W = 1 # edge force
W_ballon = 400 #baloon force
alpha = 0.5 # contraction force
beta = 6 # sturdiness
gamma = 0.2 # step size
threshold = 1e-2 # change in area treshold



npoints = 200
center = np.array([col, row])
radius = 1
N = np.zeros((npoints, 2))

for i in range(npoints):
    N[i,0] = radius * math.cos(i / (npoints) * 2 * math.pi) + center[1]
    N[i,1] = radius * math.sin(i / (npoints) * 2 * math.pi) + center[0]

N = move_to_end(N)


# %%  SEGMENTATION ALGORITHM

for i in range(200):
    
    # distance calculation
    h = distances(N)
    
    # ### initialisation for matrix of coefficients
    a = beta / (h**4)
    b = -4*beta / (h**4) - alpha / (h**2)
    c = 6*beta / (h**4) + 2*alpha / (h**2)
        
    A = CreateMatrixCoeff(len(h), a, b, c)
      
    # ### calculate current extern forces
    FextX = []
    FextY = [] 
    for j in range(len(N)):
        gx = Gx_inter([N[j,0],N[j,1]])
        gy = Gy_inter([N[j,0],N[j,1]])
        
        FextX.append(gx)
        FextY.append(gy)
    
    F_ext = np.array([FextX, FextY]).T
    F_shape = np.shape(F_ext)
    F_ext = F_ext.reshape(F_shape[1], F_shape[2])
    
    # ### actualise extern forces, depends also on ballon force
    F_ext = F_ext + W_ballon * compute_baloon_force(N)
    
    # ### calculate the area of the current segment
    area_old = PolyArea(N[:,0],N[:,1])
    
    # ### recalculate the new points of the snake
    N = np.abs((np.linalg.inv(np.eye(np.shape(A)[0]) + gamma * A)) @ (N + W * gamma * F_ext))
    
    # ### calculate new area of the segment
    area_new = PolyArea(N[:,0],N[:,1])
    
    # ### stop condition
    if (np.abs(area_old - area_new) < threshold):
        # print('Terminated at iteration ' + str(i))
        break



# %% CREATE MASK FROM CONTOUR

fig = plt.figure(num=100, figsize=(5,5))
plt.fill(N[:,0], N[:,1])
plt.xlim([0,191])
plt.ylim([191,0])
plt.show()


fig.canvas.draw()
mask = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
mask = mask.reshape(fig.canvas.get_width_height()[::-1] + (3,))
mask = skimage.color.rgb2gray(mask)
mask = mask[61:444, :]
mask = mask[:, 65:448]
mask_under = zoom(mask,0.5) #downsample the upsampled image
mask_under[0,:] = 1
mask_under[mask_under>=0.999999] = 1
mask_fin = np.floor(mask_under)
mask_fin = 1 - mask_fin

plt.close(100)


# plt.figure(num=10)
# plt.imshow(mask_fin, cmap='gray')
# plt.axis('off')
# plt.title('14th measure, 4th slice', fontsize=15)
# plt.suptitle('Created mask', fontsize=20)
# plt.show()

# plt.figure(num=14)
# plt.subplot(2,2,2)
# plt.imshow(mask_fin, cmap='gray')
# plt.axis('off')
# plt.title('Created mask', fontsize=20)
# plt.show()
# %%  EVALUATION OF SEGMENTATION


seg = mask_fin
dice = np.sum(seg[gt == 1]) * 2.0 / (np.sum(seg) + np.sum(gt))

print(' ')
print('Dice coefficient for current slice (14th measure, 4th slice) in diastola is: ' + str(round(dice, 2)))
    

# plt.figure(num=11)
# plt.imshow(img_fit, cmap='gray')
# plt.plot(N[:,0], N[:,1], 'r')
# plt.axis('off')
# plt.title('14th measure, 4th slice', fontsize=15)
# plt.suptitle('Preprocessed image with final contour', fontsize=20)
# plt.show()

# plt.figure(num=12)
# plt.imshow(img, cmap='gray')
# plt.plot(N[:,0], N[:,1], 'r')
# plt.axis('off')
# plt.title('14th measure, 4th slice', fontsize=15)
# plt.suptitle('Original image with final contour', fontsize=20)
# plt.show()
    
 
plt.figure(num=13)
plt.subplot(2,4,7)
plt.imshow(img_fit, cmap='gray')
plt.plot(N[:,0], N[:,1], 'r')
plt.axis('off')
plt.title('Preprocessed image with final contour', fontsize=15)
plt.show()

plt.figure(num=13)
plt.subplot(2,4,8)
plt.imshow(img, cmap='gray')
plt.plot(N[:,0], N[:,1], 'r')
plt.axis('off')
plt.title('Original image with final contour', fontsize=15)
plt.show()
    
 
    
 
    
 
    
 
    
 
    