##############################################

#                                            #

# Main program (later this will be changed)  #

# for simplicities sake this will only call  # 

# function from algorithms/<>.py files       # 

#                                            #

# Author: Amine Neggazi                      #

# Email: neggazimedlamine@gmail/com          #

# Nick: nemo256                              #

#                                            #

# Please read bc-count/LICENSE               #

#                                            #

##############################################

import glob

import os

import cv2

import numpy as np

import tensorflow as tf

import matplotlib.pyplot as plt

from scipy import ndimage

from skimage.segmentation import watershed

from skimage.feature import peak_local_max

# custom imports

from config import *

import data

from model import do_unet, segnet, get_callbacks

def train(model_name='mse', epochs=50):

    '''

    This is the train function, so that we can train multiple models

    according to blood cell types and multiple input shapes aswell.

    :param model_name --> model weights that we want saved

    :param epochs --> how many epochs we want the model to be trained

    :return --> saves the model weights under <model_name>.h5 with

    its respective history file <model_name>_history.npy

    '''

    train_img_list = sorted(glob.glob(f'data/{cell_type}/train/image/*.jpg'))

    test_img_list = sorted(glob.glob(f'data/{cell_type}/test/image/*.jpg'))

    train_mask_list = sorted(glob.glob(f'data/{cell_type}/train/mask/*.jpg'))

    test_mask_list = sorted(glob.glob(f'data/{cell_type}/test/mask/*.jpg'))

    if cell_type == 'rbc':

        train_edge_list = sorted(glob.glob(f'data/{cell_type}/train/edge/*.jpg'))

        test_edge_list = sorted(glob.glob(f'data/{cell_type}/test/edge/*.jpg'))

    elif cell_type == 'wbc' or cell_type == 'plt':

        train_edge_list = None

        test_edge_list = None

    else:

        print('Invalid blood cell type!\n')

        return

    # loading train dataset and test datasets

    train_dataset = data.generator(

        train_img_list,

        train_mask_list,

        train_edge_list,

        type='train'

    )

    test_dataset = data.generator(

        test_img_list,

        test_mask_list,

        test_edge_list,

        type='test'

    )

    # initializing the do_unet model

    if model_type == 'do_unet':

        model = do_unet()

    else:

        model = segnet()

    # create models directory if it does not exist

    if not os.path.exists('models/'):

        os.makedirs('models/')

    # Check for existing weights

    if os.path.exists(f'models/{model_name}.h5'):

        model.load_weights(f'models/{model_name}.h5')

    # fitting the model

    history = model.fit(

        train_dataset.batch(8),

        validation_data=test_dataset.batch(8),

        epochs=epochs,

        steps_per_epoch=125,

        max_queue_size=16,

        use_multiprocessing=True,

        workers=8,

        verbose=1,

        callbacks=get_callbacks(model_name)

    )

    # save the history

    np.save(f'models/{model_name}_history.npy', history.history)

def normalize(img):

    '''

    Normalizes an image

    :param img --> an input image that we want normalized

    :return np.array --> an output image normalized (as a numpy array)

    '''

    return np.array((img - np.min(img)) / (np.max(img) - np.min(img)))

def get_sizes(img,

              padding=padding[1],

              input=input_shape[0],

              output=output_shape[0]):

    '''

    Get full image sizes (x, y) to rebuilt the full image output

    :param img --> an input image we want to get its dimensions

    :param padding --> the default padding used on the test dataset

    :param input --> the input shape of the image (param: img)

    :param output --> the output shape of the image (param: img)

    :return couple --> a couple which contains the image dimensions as in (x, y)

    '''

    offset = padding + (output / 2)

    return [(len(np.arange(offset, img[0].shape[0] - input / 2, output)), len(np.arange(offset, img[0].shape[1] - input / 2, output)))]

def reshape(img,

            size_x,

            size_y):

    '''

    Reshape the full image output using the original sizes (x, y)

    :param img --> an input image we want to reshape

    :param size_x --> the x axis (length) of the input image (param: img)

    :param size_y --> the y axis (length) of the input image (param: img)

    :return img (numpy array) --> the output image reshaped according to the provided dimensions (size_x, size_y)

    '''

    return img.reshape(size_x, size_y, output_shape[0], output_shape[0], 1)

def concat(imgs):

    '''

    Concatenate all the output image chips to rebuild the full image

    :param imgs --> the images that we want to concatenate

    :return full_image --> the concatenation of all the provided images (param: imgs)

    '''

    return cv2.vconcat([cv2.hconcat(im_list) for im_list in imgs[:,:,:,:]])

def denoise(img):

    '''

    Remove noise from an image

    :param img --> the input image that we want to denoise (remove the noise)

    :return image --> the denoised output image

    '''

    # read the image

    img = cv2.imread(img)

    # return the denoised image

    # if cell_type == 'plt':

    #     return cv2.fastNlMeansDenoising(img, 19, 19, 7, 21)

    # return cv2.fastNlMeansDenoising(img, 23, 23, 7, 21)

    return img

def predict(imgName='Im037_0'):

    '''

    Predict (segment) blood cell images using the trained model (do_unet)

    :param img --> the image we want to predict (from the test/ directory)

    :return --> saves the predicted (segmented blood cell image) under the folder output/

    '''

    # Check for existing predictions

    if not os.path.exists(f'{output_directory}/{imgName}'):

        os.makedirs(f'{output_directory}/{imgName}', exist_ok=True)

    else:

        print('Prediction already exists!')

        return

    test_img = sorted(glob.glob(f'data/ALL-IDB1-FULL/{imgName}.jpg'))

    # initializing the do_unet model

    if model_type == 'do_unet':

        model = do_unet()

    else:

        model = segnet()

    # Check for existing weights

    if os.path.exists(f'models/{model_name}.h5'):

        model.load_weights(f'models/{model_name}.h5')

    # load test data

    img = data.load_image(test_img, padding=padding[0])

    img_chips = data.slice_image(

        img,

        padding=padding[1],

        input_size=input_shape[0],

        output_size=output_shape[0],

    )

    # segment all image chips

    output = model.predict(img_chips)

    if cell_type == 'rbc':

        new_mask_chips = np.array(output[0])

        new_edge_chips = np.array(output[1])

    elif cell_type == 'wbc' or cell_type == 'plt':

        new_mask_chips = np.array(output)

    # get image dimensions

    dimensions = [get_sizes(img)[0][0], get_sizes(img)[0][1]]

    # reshape chips arrays to be concatenated

    new_mask_chips = reshape(new_mask_chips, dimensions[0], dimensions[1])

    if cell_type == 'rbc':

        new_edge_chips = reshape(new_edge_chips, dimensions[0], dimensions[1])

    # get rid of none necessary dimension

    new_mask_chips = np.squeeze(new_mask_chips)

    if cell_type == 'rbc':

        new_edge_chips = np.squeeze(new_edge_chips)

    # concatenate chips into a single image (mask and edge)

    new_mask = concat(new_mask_chips)

    if cell_type == 'rbc':

        new_edge = concat(new_edge_chips)

    # save predicted mask and edge

    plt.imsave(f'{output_directory}/{imgName}/mask.png', new_mask)

    if cell_type == 'rbc':

        plt.imsave(f'{output_directory}/{imgName}/edge.png', new_edge)

        plt.imsave(f'{output_directory}/{imgName}/edge_mask.png', new_mask - new_edge)

def denoise_full_image(imgName='Im037_0'):

    # denoise all the output images

    new_mask  = denoise(f'{output_directory}/{imgName}/mask.png')

    if cell_type == 'rbc':

        new_edge  = denoise(f'{output_directory}/{imgName}/edge.png')

        edge_mask = denoise(f'{output_directory}/{imgName}/edge_mask.png')

    # save predicted mask and edge after denoising

    plt.imsave(f'{output_directory}/{imgName}/denoise.png', new_mask)

    if cell_type == 'rbc':

        plt.imsave(f'{output_directory}/{imgName}/edge.png', new_edge)

        plt.imsave(f'{output_directory}/{imgName}/edge_mask.png', edge_mask)

def evaluate(model_name='mse'):

    '''

    Evaluate an already trained model

    :param model_name --> the model weights that we want to evaluate

    :return --> output the evaluated model weights directly to the screen.

    '''

    test_img_list = sorted(glob.glob(f'data/{cell_type}/test/image/*.jpg'))

    test_mask_list = sorted(glob.glob(f'data/{cell_type}/test/mask/*.jpg'))

    if cell_type == 'rbc':

        test_edge_list = sorted(glob.glob(f'data/{cell_type}/test/edge/*.jpg'))

    elif cell_type == 'wbc' or cell_type == 'plt':

        test_edge_list = None

    else:

        print('Invalid blood cell type!\n')

        return

    # initializing the do_unet model

    if model_type == 'do_unet':

        model = do_unet()

    else:

        model = segnet()

    # load weights

    if os.path.exists(f'models/{model_name}.h5'):

        model.load_weights(f'models/{model_name}.h5')

    else:

        train(model_name)

    # load test data

    if cell_type == 'rbc':

        img, mask, edge = data.load_data(test_img_list, test_mask_list, test_edge_list, padding=padding[0])

        img_chips, mask_chips, edge_chips = data.slice(

            img,

            mask,

            edge,

            padding=padding[1],

            input_size=input_shape[0],

            output_size=output_shape[0]

        )

    elif cell_type == 'wbc' or cell_type == 'plt':

        img, mask = data.load_data(test_img_list, test_mask_list, padding=padding[0])

        img_chips, mask_chips = data.slice(

            img,

            mask,

            padding=padding[1],

            input_size=input_shape[0],

            output_size=output_shape[0]

        )

    # print the evaluated accuracies

    if cell_type == 'rbc':

        print(model.evaluate(img_chips, (mask_chips, edge_chips)))

    else:

        print(model.evaluate(img_chips, (mask_chips)))

def threshold(img='edge.png', imgName='Im037_0'):

    '''

    This is the threshold function, which applied an otsu threshold

    to the input image (param: img)

    :param img --> the image we want to threshold

    :return --> saves the output thresholded image under the folder output/<cell_type>/threshold_<img>.png

    '''

    if not os.path.exists(f'{output_directory}/{imgName}/{img}'):

        print('Image does not exist!')

        return

    # substract if img is edge_mask

    if img == 'edge_mask.png':

        mask = cv2.imread(f'{output_directory}/{imgName}/threshold_mask.png')

        edge = cv2.imread(f'{output_directory}/{imgName}/threshold_edge.png')

        # substract mask - edge

        image = mask - edge

    else:

        # getting the input image

        image = cv2.imread(f'{output_directory}/{imgName}/{img}')

        # convert to grayscale and apply otsu's thresholding

        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        if cell_type == 'plt':

            threshold, image = cv2.threshold(image, 67, 255, cv2.THRESH_BINARY)

        elif cell_type == 'wbc':

            threshold, image = cv2.threshold(image, 127, 255, cv2.THRESH_BINARY)

        else:

            threshold, image = cv2.threshold(image, 0, 255, cv2.THRESH_OTSU)

    # save the resulting thresholded image

    plt.imsave(f'{output_directory}/{imgName}/threshold_{img}', image, cmap='gray')

def hough_transform(img='edge.png', imgName='Im037_0'):

    '''

    This is the Circle Hough Transform function (CHT), which counts the

    circles from an input image.

    :param img --> the input image that we want to count circles from.

    :return --> saves the output image under the folder output/<cell_type>/hough_transform.png

    '''

    if not os.path.exists(f'{output_directory}/{imgName}/{img}'):

        print('Image does not exist!')

        return

    # getting the input image

    image = cv2.imread(f'{output_directory}/{imgName}/{img}')

    # convert to grayscale

    img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    if cell_type == 'wbc':

        # apply surface filter

        img, ret_count = surfaceFilter(img, min_size=2000)

        img = ((img > 0) * 255.).astype(np.uint8)

    # apply hough circles

    if cell_type == 'rbc':

        if model_type == 'segnet':

            circles = cv2.HoughCircles(img, cv2.HOUGH_GRADIENT, 1, minDist=33, maxRadius=55, minRadius=28, param1=30, param2=20)

        else:

            circles = cv2.HoughCircles(img, cv2.HOUGH_GRADIENT, 1, minDist=33, maxRadius=56, minRadius=29, param1=28, param2=20)

    elif cell_type == 'wbc':

        if model_type == 'do_unet':

            circles = cv2.HoughCircles(img, cv2.HOUGH_GRADIENT, 1, minDist=51, maxRadius=120, minRadius=48, param1=70, param2=20)

        else:

            circles = cv2.HoughCircles(img, cv2.HOUGH_GRADIENT, 1, minDist=70, maxRadius=120, minRadius=33, param1=62, param2=12)

    elif cell_type == 'plt':

        circles = cv2.HoughCircles(img, cv2.HOUGH_GRADIENT, 1.3, minDist=16, maxRadius=25, minRadius=1, param1=13, param2=11)

    output = img.copy()

    # ensure at least some circles were found

    if circles is not None:

        # convert the (x, y) coordinates and radius of the circles to integers

        circles = np.round(circles[0, :]).astype("int")

        # loop over the (x, y) coordinates and radius of the circles

        for (x, y, r) in circles:

            # draw the circle in the output image, then draw a rectangle

            # corresponding to the center of the circle

            cv2.circle(output, (x, y), r, (0, 0, 255), 2)

            cv2.rectangle(output, (x - 5, y - 5), (x + 5, y + 5), (0, 0, 255), -1)

        # save the output image

        plt.imsave(f'{output_directory}/{imgName}/hough_transform.png',

                   np.hstack([img, output]))

        # show the hough_transform results

        print(f'Hough Transform: {len(circles)}')

        if len(circles) == None:

            return 0

        else:

            return len(circles)

    else:

        return 0

def component_labeling(img='edge.png', imgName='Im037_0'):

    '''

    This is the Connected Component Labeling (CCL), which labels all the connected objects from an input image

    :param img --> the input image that we want to apply CCL to.

    :return --> saves the output image under the folder output/<cell_type>/component_labeling.png

    '''

    if not os.path.exists(f'{output_directory}/{imgName}/{img}'):

        print('Image does not exist!')

        return

    # getting the input image

    image = cv2.imread(f'{output_directory}/{imgName}/{img}')

    # convert to grayscale

    img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # # converting those pixels with values 1-127 to 0 and others to 1

    # img = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)[1]

    # applying surfaceFilter

    if cell_type == 'wbc':

        result_image, ret_count = surfaceFilter(img, min_size=1400)

    elif cell_type == 'rbc':

        result_image, ret_count = surfaceFilter(img, min_size=200)

    else:

        result_image, ret_count = surfaceFilter(img)

    # saving image after Component Labeling

    plt.imsave(f'{output_directory}/{imgName}/component_labeling.png',

               np.hstack([img, result_image]))

    # show number of labels detected

    print(f'Connected Component Labeling: {ret_count - 1}')

    return ret_count - 1

def distance_transform(img='threshold_edge_mask.png', imgName='Im037_0'):

    '''

    This is the Euclidean Distance Transform function (EDT), which applied the distance transform algorithm to an input image>

    :param img --> the input image that we want to apply EDT to.

    :return --> saves the output image under the folder output/<cell_type>/distance_transform.png

    '''

    if not os.path.exists(f'{output_directory}/{imgName}/{img}'):

        print('Image does not exist!')

        return

    # getting the input image

    image = cv2.imread(f'{output_directory}/{imgName}/{img}')

    # convert to numpy array

    img = np.asarray(image)

    # convert to grayscale

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

    img = ndimage.distance_transform_edt(img)

    # saving image after Component Labeling

    plt.imsave(f'{output_directory}/{imgName}/distance_transform.png', img, cmap='gray')

def surfaceFilter(image, min_size = None, max_size = None):

    img = image.copy()

    ret, labels = cv2.connectedComponents(img)

    label_codes = np.unique(labels)

    result_image = labels

    if 9999 in result_image:

        print("error the image contains the null number 9999")

    i = 0

    background_index = 0

    max = 0

    for label in label_codes:

        count = (labels == label).sum()

        #find the background index

        if count > max:

            max = count

            background_index = i

        if min_size is not None and (count < min_size):

            result_image[labels == label] = 9999

            ret = ret - 1

        if max_size is not None and (count > max_size):

            result_image[labels == label] = 9999

            ret = ret - 1

        i = i + 1

    result_image[result_image == 9999] = label_codes[background_index]

    return result_image, ret

def count(img='threshold_mask.png', imgName='Im037_0'):

    if not os.path.exists(f'{output_directory}/{imgName}/{img}'):

        print('Image does not exist!')

        return

    # getting the input image

    image = cv2.imread(f'{output_directory}/{imgName}/{img}')

    # convert to numpy array

    img = np.asarray(image)

    # convert to grayscale

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

    if cell_type == 'rbc':

        min_distance = 40

        threshold_abs = None

        exclude_border = True

    elif cell_type == 'wbc':

        min_distance = 51

        threshold_abs = 24

        exclude_border = False

    elif cell_type == 'plt':

        min_distance = 20

        img = ndimage.binary_dilation(img)

        threshold_abs = 4

        exclude_border = False

    edt = ndimage.distance_transform_edt(img)

    coords = peak_local_max(edt, 

                            indices=True,

                            num_peaks=2000,

                            min_distance=min_distance, 

                            threshold_abs=threshold_abs,

                            exclude_border=exclude_border,

                            labels=img)

    # print(coords[:, 1])

    canvas = np.ones(img.shape + (3,), dtype=np.uint8) * 255

    i = 255

    for c in coords:

        o_c = (int(c[1]), int(c[0]))

        cv2.circle(canvas, o_c, 20, (i, 0, 0), -1)

        i = i - 1

    # saving image after counting

    plt.imsave(f'{output_directory}/{imgName}/output.png', canvas, cmap='gray')

    print(f'Euclidean Distance Transform: {len(coords)}')

    return len(coords)

def accuracy(real, predicted):

    acc = (1 - (np.absolute(int(predicted) - int(real)) / int(real))) * 100

    if int(real) == 0.:

        if int(predicted) == 0.:

            return 100.

        else:

            return 0.

    if acc <= 100. and acc > 0.:

        return acc

    elif acc < 0.:

        return np.absolute(acc / 100.)

    else:

        return 0.

def predict_all_idb():

    image_list = sorted(glob.glob('data/ALL-IDB1/*'))

    if not os.path.exists(output_directory):

        os.makedirs(output_directory, exist_ok=True)

    real_count = []

    with open(f'data/{cell_type}_count.txt', 'r+') as rc:

        file = rc.read().splitlines()

        for line in file:

            real_count += [line.split(' ')[-1]]

    i = 0

    cht_accuracy = []

    ccl_accuracy = []

    edt_accuracy = []

    with open(f'{output_directory}/{cell_type}_results.txt', 'a+') as r:

        r.write('Image Real_Count CHT CCL EDT CHT_Accuracy CCL_Accuracy EDT_Accuracy\n')

        for image in image_list:

            img = image.split('/')[-1].split('.')[0]

            print(f'--------------------------------------------------')

            predict(img)

            # denoise_full_image(img)

            threshold('mask.png', img)

            print(f'Image <-- {img} -->')

            print(f'Real Count: {real_count[i]}')

            if cell_type == 'rbc':

                threshold('edge.png', img)

                threshold('edge_mask.png', img)

                distance_transform('threshold_edge_mask.png', img)

                cht_count = hough_transform('edge.png', img)

            else:

                distance_transform('threshold_mask.png', img)

                cht_count = hough_transform('threshold_mask.png', img)

            edt_count = count('threshold_mask.png', img)

            ccl_count = component_labeling('threshold_mask.png', img)

            cht_accuracy += [accuracy(real_count[i], cht_count)]

            ccl_accuracy += [accuracy(real_count[i], ccl_count)]

            edt_accuracy += [accuracy(real_count[i], edt_count)]

            # accuracy = np.mean([cht_accuracy, ccl_accuracy])

            r.write(f'{img} {real_count[i]} {cht_count} {ccl_count} {edt_count} {np.round(cht_accuracy[i], 2)} {np.round(ccl_accuracy[i], 2)} {np.round(edt_accuracy[i], 2)}\n')

            i = i + 1

        r.write(f'Total -1 -1 -1 -1 {np.round(np.mean(cht_accuracy), 2)} {np.round(np.mean(ccl_accuracy), 2)} {np.round(np.mean(edt_accuracy), 2)}\n')

        if cell_type == 'rbc':

            print(f'Accuracy: {np.round(np.mean(cht_accuracy), 2)}%')

        elif cell_type == 'wbc':

            print(f'Accuracy: {np.round(np.mean(edt_accuracy), 2)}%')

        else:

            print(f'Accuracy: {np.round(np.mean(ccl_accuracy), 2)}%')

if __name__ == '__main__':

    '''

    The main function, which handles all the function call

    (later on, this will dynamically call functions according user input)

    '''

    # train('plt_segnet', epochs=12)

    # evaluate(model_name='plt_segnet')

    image = 'Im037_0'

    predict(imgName=image)

    # denoise_full_image(imgName=image)

    threshold('mask.png', image)

    if cell_type == 'rbc':

        threshold('edge.png', image)

        threshold('edge_mask.png', image)

        distance_transform('threshold_edge_mask.png', image)

        hough_transform('edge.png', image)

    else:

        distance_transform('threshold_mask.png', image)

        hough_transform('threshold_mask.png', image)

    count('threshold_mask.png', image)

    component_labeling('threshold_mask.png', image)

    # predict_all_idb()