# -*- coding: utf-8 -*-
"""
Created on Sat May  7 17:37:05 2022

@author: petra

take data from folder Finalni
"""


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

    '''
    # get image histogram
    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
    '''
    # funkce na vypocet gradientu Gmag, Gdir
    # vstup je obrazek, vystup je Magnitude a Direction

    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
    '''
    # funkce na vypocet gradientu ve smeru x, y
    # vstup je obrazek, vystup je Gx, Gy

    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)
    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)
    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

def make_plane(img):
    # mozna spis nazev neco na tvorbu te plochy, kterou to prokladam
    
    
    # create meshgrid
    X, Y = np.meshgrid(range(np.shape(img)[0]), range(np.shape(img)[1]))
    
    # make vector from array, then pick only values bigger than...
    X = X.flatten()
    X = X + 1e-5
    X = X.astype(float)
    Y = Y.flatten()
    Y = Y + 1e-5
    Y = Y.astype(float)
    Z = img.flatten()
    Z = Z
    Z = Z.astype(float)
    
    x = []
    y = []
    z = []    
    
    for i in range(len(X)):
        value = Z[i]
        if value > np.mean(img):
            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)[0], np.shape(img)[1])

    return plane

def make_circle(npoints, center, radius):
    '''
    function to create circle
    npoints: number of points from which the contour will be created
    center: center of circle
    radius: radius of created circle
    '''
    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)

    return N

def segmentation_alg(img, W, W_ballon, alpha, beta, gamma, threshold, N, num_iter):
    
    '''
    segmentation algorithm
    img: input image, preprocessed
    W: edge force, defining strength of attraction of contour to edge
    W_ballon: force defining ballon force, how much the contour will be inflated
    alpha: contraction force, how much will the contour collapse into itsels
    beta: sturdiness of conotur, how much will the contour be deformable
    gamma: step of changes
    threshold: defines the minimal difference in areas, so that the algorithm continues 
    N: predefined contour
    num_iter: number of maximal iteration for segmentation
    
    output: contour
    '''
    
    # 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
    
    Gmag, Gdir = imgradient_fcn(img, operator='prewitt')
    Gx, Gy = imgradientxy_fcn(Gmag, operator='prewitt')        
    
    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())



    for i in range(num_iter):
        # ### calculate the distances between the points
        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

    return N

def create_mask(N):
    '''
    function to crete mask from calculated contour
    N: given contour
    
    output: BW image, 1 == segmented area, 0 == area outside
    '''
    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)

    return mask_fin
    
def image_calibrations(img, cutoffFreq, a, b):
    '''
    function performing image calibrations, preparing image to segmentation
    img: input image
    cutoffFreq: cutoff frequency for homomorphic filter
    a,b: coefficients for homomorphic filter
    '''
    img[img < 0] = 0
    img[img > 10] = 0
        
    img = HomomorphicFilter(img, cutoffFreq=cutoffFreq, a=a, b=b)
    img = img ** 0.5
    img[img > 10 * np.mean(img)] = np.mean(img)
    img[img < 2 * np.median(img)] = 0
    img = histogramEqualisation(img, nbr_bins=64000)
    img = normaliseImg(img)
    img[img < 1 * np.mean(img)] = 0 
    img = np.max(img) - img
    img = skimage.morphology.area_closing(img)
    img = np.max(img) - img

    return img    
    
def preprocessing_image(img):
    '''
    function to preprocess image - energy normalisation, normalisation and Gaussian filtration
    '''
    # img = img / np.linalg.norm(img, ord=None)
    img = normaliseImg(img)
    img = scipy.ndimage.gaussian_filter(img, sigma=1.5)
    
    img[img < 0.5 * np.mean(img)] = 0
    
    return img    


# PREPARE THE SCRIPT

plt.close('all')


# %%

images = scipy.io.loadmat('images_sys.mat')
images = images['mydata']
# original images, representing one measure serie, in diastola

masks_sys = scipy.io.loadmat('masks_sys_gt.mat')
masks_sys = masks_sys['mydata']
# manually segmented masks

centers = scipy.io.loadmat('centers.mat')
centers = centers['mydata']
# centers of ventricles, diastolas, systolas

# %% defining centers of ventricles


data_range = np.shape(images)[0]

print('Do you want to use predefined ventricle centers, or do you prefer define your own ventricle centers? ')
print('If you prefer using predefined centers, type: yes. If you prefer to define them on your own, type: no')
ventricle_center = input('Predefined centers: yes, own centers: no   ')
ventricle_center = str.lower(ventricle_center)

loop_counter = 0

while (loop_counter < 3):
    loop_counter = loop_counter + 1

    if (ventricle_center == 'no'):
        for i in range(data_range):
            plt.figure(num=1)
            plt.subplot(2,3,i+1)
            plt.imshow(images[i,:,:], cmap='gray')
            plt.axis('off')
        plt.suptitle('Select center for each systola', fontsize=20)
        points = plt.ginput(n=5, timeout=0)
        plt.show()
        
        plt.close(1)
        
        loop_counter = 5
    
    
    elif (ventricle_center == 'yes'):
        points = []
        for i in range(data_range):
            points.append(centers[1,:])
            
        loop_counter = 5
        
    else:
        print('You may did some typo.') 
        ventricle_center = input('If you want to use predefined centers, type: yes. Otherwise, if you prefer define them, type: no.   ')
        ventricle_center = str.lower(ventricle_center)

if (loop_counter == 3):
    print('You may did some typos, I will use predefined centers for segmentation.    ')    
    points = []
    for i in range(data_range):
        points.append(centers[0,:])

# %% segmentation algorithm

segmented_masks_sys = []
contours_sys = []
dice_sys = []


for i in range(data_range):
# i = 0

    ventricle_c = points[i]
    col = ventricle_c[1]
    row = ventricle_c[0]

    # select image from dataset
    img = images[i,:,:]
    
    # preprocessing the image
    img_fit = preprocessing_image(img)
    
    # create plane to fit to the image
    plane = make_plane(img_fit)
    
    # image division by created plane
    img_fit = img_fit / plane
    
    # image calibrations 
    img_fit = image_calibrations(img_fit, 8, 1, 2)
    
    # image is preprocessed
    # part of segmentation algorithm
    
    # create a starting polygon / circle
    # N = make_circle(npoints=200, center=np.array([col, row]), radius=1)
    N = make_circle(200, np.array([col, row]), 7)
    
    # segmentation algorithm
    # N = segmentation_alg(img_fit, W=1, W_ballon=400, alpha=0.5, beta=6, gamma=0.2, threshold=1e-2, N=N, num_iter=1000)
    N = segmentation_alg(img_fit, 1, 0, 2, 1.5, 0.2, 1e-2, N, num_iter=200)
    
    # create mask from segmented contour
    mask_fin = create_mask(N)
    
    # evaluation
    ground_truth = masks_sys[i,:,:]
    segmented = mask_fin
    dice = np.sum(segmented[ground_truth == 1]) * 2.0 / (np.sum(segmented) + np.sum(ground_truth))
    
    # save results
    segmented_masks_sys.append(segmented)
    contours_sys.append(N)
    dice_sys.append(dice)
    
    print(' ')
    print('For systola segmentation, in loop number: ' + str(i) + ', the dice coefficient is: ' + str(round(dice,2)))


scipy.io.savemat('segmented_sys_masks.mat', {'mydata': segmented_masks_sys})
# scipy.io.savemat('contours_sys.mat', {'mydata': contours_sys})
# scipy.io.savemat('dice_sys.mat', {'mydata': dice_sys})

# %% data visualisation


for i in range(data_range):
    contour = contours_sys[i]
    plt.figure(num=1)
    plt.subplot(2,3,i+1)
    plt.imshow(images[i,:,:], cmap='gray')
    plt.plot(contour[:,0], contour[:,1])
    plt.axis('off')
plt.suptitle('Segmented left ventricles during diastolas', fontsize=20)
plt.show()
    
   



