# %% Importing packages

import numpy as np

import cv2 as cv

import matplotlib.pyplot as plt

import os

from joblib import Parallel, delayed

# %% Defining Functions

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

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

def load_image_names(folder):

    # This function reads the images within a folder while filtering 

    # out the weird invisible files that macos includes in their folders

    file_list = []

    for file_name in os.listdir(folder):

        # check if the first character of the name is a '.', skip if so

        if file_name[0] != '.': 

            file_list.append(file_name)

    return(file_list)

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

def basic_image_seg_visualization(image_name):

    # this function shows both the RGB image and its corresponding segmentation

    # next to each other

    image = cv.imread(image_name,cv.IMREAD_UNCHANGED)

    # change the color order, because openCV2 reads color in as BGR, not RGB

    color_image = cv.cvtColor(image[:,:,0:3],cv.COLOR_BGR2RGB)

    # create our subplots

    fig, (ax1, ax2) = plt.subplots(1,2)

    # show both the images

    ax1.imshow(color_image)

    ax2.imshow(image[:,:,3], vmin=0,vmax=5)

    plt.show()

    return()

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

def includes_segmentation(image_name, class_id):

    # this function is meant to be run in parallel, but can be run individually.

    # it receives an image name and a class_id, and determins whether and how 

    # much of a certain class is contained in the image.

    image = cv.imread(image_name,cv.IMREAD_UNCHANGED)

    segmentation = image[:,:,3]

    # how many pixels are a part of the class?

    seg_sum = np.sum((segmentation==class_id))

    # total number of pixels?

    total_sum = segmentation.shape[0]*segmentation.shape[1]

    # if the class is in the image, return true and the percentage

    if seg_sum>0:

        return(image_name,True,seg_sum/total_sum)

    else:

        return(image_name,False,0)

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

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

dataset_directory = '/media/briancottle/3a7b7bdc-6753-4423-b5ac-ff074ad75013/sub_sampled_20220526'

os.chdir(dataset_directory)

file_names = load_image_names('.')

# %%

random_index = int(np.random.random()*len(file_names))

basic_image_seg_visualization(file_names[random_index])

# %%

class_id = 5

# check for vasculature

contains_names_vascular = Parallel(n_jobs=20, verbose=1)(delayed(includes_segmentation) \

                                    (name,5) for name in file_names)

# check for neural tissue 

contains_names_neural = Parallel(n_jobs=20, verbose=1)(delayed(includes_segmentation) \

                                    (name,4) for name in file_names)

# %%

# get only the names and percentages for images that contain vasculature

vascular_images = []

vascular_percentages = []

for evaluation in contains_names_vascular:

    name = evaluation[0]

    present = evaluation[1]

    percentage = evaluation[2]

    if present:

        vascular_images.append(name)

        vascular_percentages.append(percentage)

# get only the names and percentages for images that contain vasculature

neural_images = []

neural_percentages = []

for evaluation in contains_names_neural:

    name = evaluation[0]

    present = evaluation[1]

    percentage = evaluation[2]

    if present:

        neural_images.append(name)

        neural_percentages.append(percentage)

# %% 

# reporting on valuse found for vasculature

print(f'the number of images containing vasculature is: {len(vascular_images)}')

print(f'the largest percentage was found in file {vascular_images[np.argmax(vascular_percentages)]}, and was {np.max(vascular_percentages)}')

basic_image_seg_visualization(vascular_images[np.argmax(vascular_percentages)])

plt.hist(vascular_percentages)

plt.show()

# %%

random_index = int(np.random.random()*len(vascular_images))

basic_image_seg_visualization(vascular_images[random_index])

# %%